diff --git a/atom/config.py b/atom/config.py index 90ce774f91..96bc8866b7 100644 --- a/atom/config.py +++ b/atom/config.py @@ -639,6 +639,18 @@ def glm5_kpool_block_size(index_kpool: int) -> int: return index_kpool * _MQA_LOGITS_PRESHUFFLE_ROWS +def _glm5_next_unsupported_features(config: "Config") -> list[str]: + """Return parallel modes that do not yet preserve GLM-5.3 k-pool state.""" + unsupported = [] + if config.prefill_context_parallel_size > 1: + unsupported.append("PCP") + if config.decode_context_parallel_size > 1: + unsupported.append("DCP") + if config.enable_tbo or config.enable_tbo_decode: + unsupported.append("TBO") + return unsupported + + _CONFIG_REGISTRY: dict[str, str] = { "deepseek_v32": "deepseek_v3", "deepseek_v4": "deepseek_v3", # V4 reuses V3 schema; V4-specific fields @@ -1134,6 +1146,8 @@ class SpeculativeConfig: "qwen3_5_moe_text": "qwen3_5_mtp", "mimo_v2": "mimo_v2_mtp", "mimo_v2_flash": "mimo_v2_mtp", + "glm5_next": "glm5_next_mtp", + "glm5_next_text": "glm5_next_mtp", } # mtp_model_type → (n_predict_attr, architecture) @@ -1142,6 +1156,7 @@ class SpeculativeConfig: "deepseek_v4_mtp": ("num_nextn_predict_layers", "DeepseekV4MTPModel"), "qwen3_next_mtp": ("num_nextn_predict_layers", "Qwen3NextMTPModel"), "qwen3_5_mtp": ("mtp_num_hidden_layers", "Qwen3_5MTPModel"), + "glm5_next_mtp": ("num_nextn_predict_layers", "Glm5NextMTPModel"), } def use_dspark(self) -> bool: @@ -2079,15 +2094,7 @@ def __post_init__(self): # set here and the attention builder sizes the index cache from it. is_glm5_next = any("Glm5Next" in str(a) for a in arches) if is_glm5_next: - unsupported_features = [] - if self.prefill_context_parallel_size > 1: - unsupported_features.append("PCP") - if self.decode_context_parallel_size > 1: - unsupported_features.append("DCP") - if self.speculative_config is not None: - unsupported_features.append("speculative decoding") - if self.enable_tbo or self.enable_tbo_decode: - unsupported_features.append("TBO") + unsupported_features = _glm5_next_unsupported_features(self) if unsupported_features: raise ValueError( "GLM-5.3-Flash text serving does not yet support " diff --git a/atom/model_ops/attentions/kimi_mla_gdn_attn.py b/atom/model_ops/attentions/kimi_mla_gdn_attn.py index b351b2a55d..3624b52594 100644 --- a/atom/model_ops/attentions/kimi_mla_gdn_attn.py +++ b/atom/model_ops/attentions/kimi_mla_gdn_attn.py @@ -197,6 +197,16 @@ def _index_rows_per_block(self) -> int: ) return rows + def _kpool_history_size(self) -> int: + """Return ring rows needed to survive speculative rejection.""" + pool = self._kpool_size() + spec = self.model_runner.config.speculative_config + if spec is None: + return pool + # Two speculative windows plus the incomplete pool survive rejection. + window = spec.num_speculative_tokens + 1 + return 1 << (pool + 2 * window - 1).bit_length() + def _kpool_tail_bytes(self) -> int: """Per-request tail bytes across every indexer-owning layer.""" kpool = self._kpool_size() @@ -204,7 +214,9 @@ def _kpool_tail_bytes(self) -> int: return 0 hf = self.model_runner.config.hf_config index_cache_layer_ids, _ = self._index_cache_layout() - per_layer = 2 * kpool * hf.index_head_dim * torch.bfloat16.itemsize + per_layer = ( + 2 * self._kpool_history_size() * hf.index_head_dim * torch.bfloat16.itemsize + ) return len(index_cache_layer_ids) * per_layer def _kpool_tail_plane_shape(self) -> tuple[int, int] | None: @@ -263,7 +275,7 @@ def allocate_per_req_cache(self, entries: dict[str, int]) -> dict: len(index_cache_layer_ids), entries.get(STATE_SLOT_CLASS, 0), 2, # 0 = K, 1 = gate score - self._kpool_size(), + self._kpool_history_size(), hf.index_head_dim, ), dtype=torch.bfloat16, diff --git a/atom/model_ops/glm5_next/indexer.py b/atom/model_ops/glm5_next/indexer.py index 8ff954aa2b..2cfb4fe8a2 100644 --- a/atom/model_ops/glm5_next/indexer.py +++ b/atom/model_ops/glm5_next/indexer.py @@ -141,8 +141,36 @@ def _sparse_attn_indexer_kpool( raise NotImplementedError( "GLM-5.3 kpool does not support DCP/PCP; use dcp=pcp=1" ) - if not context.is_prefill and attn_metadata.max_seqlen_q > 1: - raise NotImplementedError("GLM-5.3 kpool does not support speculative decode") + # Prefer explicit scheduler metadata. The ragged-query condition covers + # metadata implementations that do not expose ``num_spec_decodes``. + is_speculative_verify = getattr(attn_metadata, "num_spec_decodes", 0) > 0 or ( + not context.is_prefill and attn_metadata.max_seqlen_q > 1 + ) + if is_speculative_verify: + from .speculative import run_speculative_kpool_indexer + + run_speculative_kpool_indexer( + attn_metadata, + kv_cache, + q_fp8, + k, + gate_score, + weights, + compress_ape, + tail_cache, + state_slot_idx_in, + state_slot_idx, + positions, + sparse_kv_indices_buffer, + index_kpool, + topk_tokens, + topk_out_width, + get_current_atom_config().kv_cache_block_size, + max_model_len, + scale_fmt, + stable_topk, + ) + return result device = hidden_states.device block_size = get_current_atom_config().kv_cache_block_size diff --git a/atom/model_ops/glm5_next/speculative.py b/atom/model_ops/glm5_next/speculative.py new file mode 100644 index 0000000000..6fdbb8ecb9 --- /dev/null +++ b/atom/model_ops/glm5_next/speculative.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Position-addressed k-pool history for speculative verification. + +Keep one position-addressed ring of raw keys/gates. A rejected suffix may +overwrite future residues, but the next verification supplies those positions +as fresh rows before they are consumed. Pooled entries beyond the committed +position stay invisible; closing the same pool again overwrites its speculative +entry before scoring. +""" + +import torch +import triton +import triton.language as tl + +from . import kpool + + +@triton.jit +def _update_kpool_history_kernel( + keys, + gates, + positions, + cu_seqlens_q, + source_slots, + destination_slots, + history, + HISTORY_SIZE: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + request_idx = tl.program_id(0) + residue = tl.program_id(1) + source_slot = tl.load(source_slots + request_idx) + destination_slot = tl.load(destination_slots + request_idx) + query_start = tl.load(cu_seqlens_q + request_idx) + query_end = tl.load(cu_seqlens_q + request_idx + 1) + if destination_slot < 0 or query_end <= query_start: + return + last_position = tl.load(positions + query_end - 1) + position = last_position - (last_position - residue) % HISTORY_SIZE + first_position = tl.load(positions + query_start) + query_row = query_start + position - first_position + is_fresh = ( + (position >= first_position) + & (query_row >= query_start) + & (query_row < query_end) + ) + offsets = tl.arange(0, HEAD_DIM) + for plane in tl.static_range(2): + history_offset = ( + (tl.maximum(source_slot, 0) * 2 + plane) * HISTORY_SIZE + residue + ) * HEAD_DIM + previous = tl.load(history + history_offset + offsets) + source = keys if plane == 0 else gates + value = tl.load( + source + query_row * HEAD_DIM + offsets, + mask=is_fresh, + other=0, + ) + value = tl.where(is_fresh, value, previous) + destination_offset = ( + (destination_slot * 2 + plane) * HISTORY_SIZE + residue + ) * HEAD_DIM + tl.store(history + destination_offset + offsets, value) + + +def get_query_request_indices( + cu_seqlens_q: torch.Tensor, + num_query_tokens: int, +) -> torch.Tensor: + """Map each packed query row to its request index.""" + # Unlike repeat_interleave, this has no data-dependent output allocation. + return torch.bucketize( + torch.arange(num_query_tokens, device=cu_seqlens_q.device), + cu_seqlens_q[1:], + right=True, + ).clamp_max(cu_seqlens_q.numel() - 2) + + +def build_speculative_pool_candidates( + history: torch.Tensor, + keys: torch.Tensor, + gates: torch.Tensor, + positions: torch.Tensor, + cu_seqlens_q: torch.Tensor, + source_slots: torch.Tensor, + pool_bias: torch.Tensor, + pool_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Form the pool ending at each query, before any history is overwritten.""" + num_query_tokens = keys.shape[0] + request_indices = get_query_request_indices(cu_seqlens_q, num_query_tokens) + request_start_rows = cu_seqlens_q[:-1].to(torch.int64)[request_indices] + request_start_positions = positions[request_start_rows] + pool_positions = ( + positions.to(torch.int64)[:, None] + - pool_size + + 1 + + torch.arange(pool_size, device=keys.device) + ) + query_rows = ( + request_start_rows[:, None] + pool_positions - request_start_positions[:, None] + ) + fresh_mask = pool_positions >= request_start_positions[:, None] + history_slots = source_slots.to(torch.int64)[request_indices].clamp_min(0) + history_positions = pool_positions % history.shape[2] + previous_keys = history[history_slots[:, None], 0, history_positions] + previous_gates = history[history_slots[:, None], 1, history_positions] + query_rows = query_rows.clamp(0, num_query_tokens - 1) + candidate_keys = torch.where(fresh_mask[..., None], keys[query_rows], previous_keys) + candidate_gates = torch.where( + fresh_mask[..., None], gates[query_rows], previous_gates + ) + pooled_keys = kpool.pool_and_rotate( + candidate_keys, + candidate_gates, + pool_bias, + ) + return pooled_keys, request_indices + + +def update_speculative_kpool_history( + history: torch.Tensor, + keys: torch.Tensor, + gates: torch.Tensor, + positions: torch.Tensor, + cu_seqlens_q: torch.Tensor, + source_slots: torch.Tensor, + destination_slots: torch.Tensor, +) -> None: + """Copy committed history and overlay fresh verification rows.""" + _update_kpool_history_kernel[(cu_seqlens_q.numel() - 1, history.shape[2])]( + keys, + gates, + positions, + cu_seqlens_q, + source_slots, + destination_slots, + history, + HISTORY_SIZE=history.shape[2], + HEAD_DIM=keys.shape[1], + ) + + +@triton.jit +def _map_token_indices_to_slots_kernel( + token_indices, + request_indices, + block_tables, + output_indptr, + output, + INPUT_WIDTH: tl.constexpr, + BLOCK_TABLE_STRIDE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + OUTPUT_SIZE: tl.constexpr, + BLOCK_WIDTH: tl.constexpr, +): + row_index = tl.program_id(0) + column_indices = tl.program_id(1) * BLOCK_WIDTH + tl.arange(0, BLOCK_WIDTH) + output_start = tl.load(output_indptr + row_index) + output_end = tl.load(output_indptr + row_index + 1) + request_index = tl.load(request_indices + row_index) + token_index = tl.load( + token_indices + row_index * INPUT_WIDTH + column_indices, + mask=column_indices < INPUT_WIDTH, + other=-1, + ) + is_valid = (token_index >= 0) & (token_index // BLOCK_SIZE < BLOCK_TABLE_STRIDE) + block_index = tl.load( + block_tables + request_index * BLOCK_TABLE_STRIDE + token_index // BLOCK_SIZE, + mask=is_valid, + other=0, + ) + cache_slot = tl.where( + is_valid, + block_index * BLOCK_SIZE + token_index % BLOCK_SIZE, + 0, + ) + tl.store( + output + output_start + column_indices, + cache_slot, + mask=(column_indices < INPUT_WIDTH) + & (output_start + column_indices < output_end) + & (output_start + column_indices < OUTPUT_SIZE), + ) + + +def map_token_indices_to_slots( + token_indices: torch.Tensor, + request_indices: torch.Tensor, + block_tables: torch.Tensor, + output_indptr: torch.Tensor, + output: torch.Tensor, + block_size: int, +) -> None: + """Convert request-local token indices to physical cache slots.""" + grid = ( + token_indices.shape[0], + triton.cdiv(token_indices.shape[1], 128), + ) + _map_token_indices_to_slots_kernel[grid]( + token_indices, + request_indices, + block_tables, + output_indptr, + output, + INPUT_WIDTH=token_indices.shape[1], + BLOCK_TABLE_STRIDE=block_tables.stride(0), + BLOCK_SIZE=block_size, + OUTPUT_SIZE=output.numel(), + BLOCK_WIDTH=128, + ) + + +def run_speculative_kpool_indexer( + metadata, + kv_cache: torch.Tensor, + queries: torch.Tensor, + keys: torch.Tensor, + gates: torch.Tensor, + weights: torch.Tensor, + pool_bias: torch.Tensor, + history: torch.Tensor, + source_slots: torch.Tensor, + destination_slots: torch.Tensor, + positions: torch.Tensor, + sparse_kv_indices: torch.Tensor, + pool_size: int, + topk_tokens: int, + output_width: int, + block_size: int, + max_model_len: int, + scale_fmt: str, + stable_topk: bool, +) -> None: + """Run the pooled indexer for a speculative verification batch.""" + from aiter.ops.cache import indexer_k_quant_and_cache + from aiter.ops.topk import top_k_per_row_decode + from aiter.ops.triton.attention.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits + + cu_seqlens_q = metadata.cu_seqlens_q + num_query_tokens, num_heads, head_dim = queries.shape + pooled_keys, request_indices = build_speculative_pool_candidates( + history, + keys, + gates, + positions, + cu_seqlens_q, + source_slots, + pool_bias, + pool_size, + ) + closes_pool = (positions % pool_size == pool_size - 1) & ( + destination_slots[request_indices] >= 0 + ) + pool_ids = torch.where( + closes_pool, + positions // pool_size, + -1, + ).to(torch.int64) + pool_rows_per_block = block_size // pool_size + cache_slots = kpool.pool_slot_mapping( + metadata.block_tables, + pool_ids, + request_indices, + pool_rows_per_block, + ) + pooled_kv_cache = kv_cache.view( + -1, + pool_rows_per_block, + kv_cache.shape[-1], + ) + indexer_k_quant_and_cache( + pooled_keys, + pooled_kv_cache, + cache_slots, + head_dim, + scale_fmt, + preshuffle=True, + ) + update_speculative_kpool_history( + history, + keys, + gates, + positions, + cu_seqlens_q, + source_slots, + destination_slots, + ) + if metadata.max_seqlen_k <= topk_tokens: + return + # A per-token paged scoring path also handles ragged verification. Chunking + # bounds scratch memory during long prefills; optimize only after parity. + selected = torch.full( + (num_query_tokens, output_width), + -1, + device=keys.device, + dtype=torch.int32, + ) + max_pools = triton.cdiv(max_model_len, pool_size) + for begin in range(0, num_query_tokens, 128): + end = min(num_query_tokens, begin + 128) + count = end - begin + sequence_lengths = (positions[begin:end] + 1).to(torch.int32) + pool_lengths = (sequence_lengths // pool_size).contiguous() + logits = torch.empty( + (count, max_pools), + device=keys.device, + dtype=torch.float32, + ) + block_tables = metadata.block_tables[request_indices[begin:end]].contiguous() + deepgemm_fp8_paged_mqa_logits( + queries[begin:end].view(count, 1, num_heads, head_dim), + pooled_kv_cache.unsqueeze(-2), + weights[begin:end], + logits, + pool_lengths, + block_tables, + max_pools, + KVBlockSize=pool_rows_per_block, + Preshuffle=True, + ) + selected_pools = torch.empty( + (count, topk_tokens // pool_size), + device=keys.device, + dtype=torch.int32, + ) + top_k_per_row_decode( + logits, + 1, + pool_lengths, + selected_pools, + count, + logits.stride(0), + logits.stride(1), + k=topk_tokens // pool_size, + stable=stable_topk, + ) + kpool.expand_pools_and_append_tail( + selected_pools, + sequence_lengths, + pool_size, + out=selected[begin:end], + ) + map_token_indices_to_slots( + selected, + request_indices, + metadata.block_tables, + metadata.sparse_kv_indptr, + sparse_kv_indices, + block_size, + ) diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index 166fbc36f6..0ac6515700 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -66,6 +66,23 @@ def __init__( class DeepSeekMultiTokenPredictorLayer(nn.Module): + @staticmethod + def build_mtp_block( + atom_config: Config, + prefix: str, + layer_idx: int, + alt_stream: torch.cuda.Stream | None, + ) -> nn.Module: + return DeepseekV2DecoderLayer( + prefix=prefix, + config=atom_config.hf_config, + cache_config=atom_config.kv_cache_dtype, + quant_config=atom_config.quant_config, + layer_num=layer_idx, + is_mtp_block=True, + alt_stream=alt_stream, + ) + def __init__( self, atom_config: Config, @@ -92,16 +109,11 @@ def __init__( config=config, prefix=prefix, quant_config=atom_config.quant_config ) - quant_config = atom_config.quant_config - - self.mtp_block = DeepseekV2DecoderLayer( - prefix=prefix, - config=self.config, - cache_config=atom_config.kv_cache_dtype, - quant_config=quant_config, - layer_num=layer_idx, - is_mtp_block=True, - alt_stream=alt_stream, + self.mtp_block = self.build_mtp_block( + atom_config, + prefix, + layer_idx, + alt_stream, ) def forward( @@ -169,6 +181,9 @@ def __init__( *, atom_config: Config, prefix: str = "", + layer_cls: type[DeepSeekMultiTokenPredictorLayer] = ( + DeepSeekMultiTokenPredictorLayer + ), ): super().__init__() config = atom_config.hf_config @@ -183,7 +198,7 @@ def __init__( # to map the exact layer index from weights self.layers = torch.nn.ModuleDict( { - str(idx): DeepSeekMultiTokenPredictorLayer( + str(idx): layer_cls( atom_config, f"{prefix}.layers.{idx}", layer_idx=idx, @@ -334,6 +349,10 @@ def compact_topk_indices(self, slot_ids: torch.Tensor) -> None: @support_torch_compile class DeepSeekMTP(nn.Module): + predictor_layer_cls = DeepSeekMultiTokenPredictorLayer + packed_modules_mapping_override: dict[str, tuple[str, int]] | None = None + supports_indexer_projection_fusion = True + def __init__(self, atom_config: Config, prefix: str = ""): super().__init__() self.config = atom_config.hf_config @@ -352,7 +371,11 @@ def __init__(self, atom_config: Config, prefix: str = ""): ): atom_config.quant_config.apply_default_exclude_layers(["*.eh_proj"]) - if hasattr(self.config, "q_lora_rank") and self.config.q_lora_rank is not None: + if self.packed_modules_mapping_override is not None: + self.packed_modules_mapping = dict(self.packed_modules_mapping_override) + elif ( + hasattr(self.config, "q_lora_rank") and self.config.q_lora_rank is not None + ): self.packed_modules_mapping = { "q_a_proj": ("fused_qkv_a_proj", 0), "kv_a_proj_with_mqa": ("fused_qkv_a_proj", 1), @@ -366,7 +389,9 @@ def __init__(self, atom_config: Config, prefix: str = ""): } model_prefix = maybe_prefix(prefix, "model") - if hasattr(self.config, "index_topk"): + if self.supports_indexer_projection_fusion and hasattr( + self.config, "index_topk" + ): indexer_prefixes = [ f"{model_prefix}.layers.{idx}.self_attn.indexer" for idx in range( @@ -390,6 +415,7 @@ def __init__(self, atom_config: Config, prefix: str = ""): self.model = DeepSeekMultiTokenPredictor( atom_config=atom_config, prefix=model_prefix, + layer_cls=self.predictor_layer_cls, ) def remap_mtp_weight_name(self, name: str) -> str | None: diff --git a/atom/models/glm5_next.py b/atom/models/glm5_next.py index 41eaac9c0f..f1908dd198 100644 --- a/atom/models/glm5_next.py +++ b/atom/models/glm5_next.py @@ -44,9 +44,9 @@ ``KimiKDAAttention`` -- and all of its state-cache, TP and CUDA-graph integration -- reusable unchanged. -Not yet wired: the MTP draft layer (checkpoint layer 45) or multimodal input. -The checkpoint's unreachable vision tower is skipped on the text-only path. -See ``recipes/GLM-5.3-Flash.md``. +The checkpoint's layer 45 is wired separately as the reusable NextN draft block +in ``glm5_next_mtp.py``. Multimodal input remains outside this text-only path, +which skips the checkpoint's vision tower. See ``recipes/GLM-5.3-Flash.md``. """ from itertools import islice @@ -393,6 +393,7 @@ def __init__( config, quant_config: QuantizationConfig | None = None, prefix: str = "", + reduce_results: bool = True, ) -> None: super().__init__() self.tp_size = get_tp_group().world_size @@ -400,6 +401,7 @@ def __init__( self.n_routed_experts = int(config.n_routed_experts) self.n_shared_experts = int(config.n_shared_experts or 0) self.swiglu_limit = float(getattr(config, "swiglu_limit", 10.0)) + self.reduce_results = reduce_results ep_group = get_ep_group().device_group self.ep_size = ep_group.size() @@ -468,7 +470,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: out = out * self.routed_scaling_factor if shared_output is not None: out = out + shared_output - if self.tp_size > 1: + if self.tp_size > 1 and self.reduce_results: out = tensor_model_parallel_all_reduce(out) return out.view(num_tokens, hidden_dim) @@ -581,6 +583,7 @@ def __init__( cache_config, use_wk_weights_proj_fusion: bool = False, prefix: str = "", + is_mtp: bool = False, ) -> None: super().__init__( atom_config, @@ -593,6 +596,7 @@ def __init__( prefix, ) self.index_kpool = int(getattr(config, "index_kpool", 1) or 1) + self.is_mtp = is_mtp always_select_tail = bool( getattr(config, "index_kpool_always_select_tail", True) ) @@ -692,6 +696,11 @@ def forward_impl( state_slot_idx_in = ( None if gdn is None else gdn.non_spec_state_indices_in_tensor ) + if state_slot_idx is None and gdn is not None: + # ReplaySSM retains one committed request slot during verify. + state_slot_idx = gdn.slot_idx + if state_slot_idx is None and gdn.spec_state_indices_tensor is not None: + state_slot_idx = gdn.spec_state_indices_tensor[:, 0] if state_slot_idx is None: # The profile/warmup forward carries no block tables, so the # builder leaves gdn_metadata unset. That forward is a dummy @@ -724,7 +733,7 @@ def forward_impl( tail_cache, state_slot_idx_in, state_slot_idx, - positions, + positions - int(self.is_mtp), self.sparse_kv_indices_buffer, self.topk_tokens, self.index_kpool, @@ -813,6 +822,7 @@ def __init__( layer_num: int, rotary_emb: nn.Module, prefix: str = "", + is_mtp: bool = False, ) -> None: super().__init__() config = _text_config(atom_config.hf_config) @@ -894,6 +904,7 @@ def __init__( atom_config.kv_cache_dtype, False, # GLM-5.3 ships wk / weights_proj unfused f"{prefix}.indexer", + is_mtp=is_mtp, ) self.rotary_emb = rotary_emb @@ -909,6 +920,9 @@ def __init__( # that `is_sparse=True` then makes MLA read, so leaving it uncalled # gives the sparse path an empty selection buffer rather than an error. self.run_indexer = not force_dense + # NextN step 0 populates the draft layer's sparse selection; later + # speculative steps reuse it when index_share_for_mtp_iteration=true. + self.skip_topk = False mla_modules = MLAModules( q_lora_rank=self.q_lora_rank, @@ -960,7 +974,7 @@ def forward( ) # Drive the indexer before MLA: it writes this layer's pooled index # keys and the selected KV slots that the sparse MLA path then reads. - if self.run_indexer: + if self.run_indexer and not self.skip_topk: self.indexer(hidden_states, q_c, None, positions, self.rotary_emb) return self.mla_attn(q_c, kv_c, k_pe, positions) diff --git a/atom/models/glm5_next_mtp.py b/atom/models/glm5_next_mtp.py new file mode 100644 index 0000000000..21973018f5 --- /dev/null +++ b/atom/models/glm5_next_mtp.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""GLM-5.3-Flash NextN support for ATOM speculative decoding.""" + +import copy +from typing import ClassVar + +import torch +from torch import nn + +from atom.config import Config +from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.rotary_embedding import NoPositionalRotaryEmbedding +from atom.models.deepseek_mtp import ( + DeepSeekMTP, + DeepSeekMultiTokenPredictorLayer, +) +from atom.models.deepseek_v2 import ENABLE_ALLREDUCE_RMSNORM_FUSION +from atom.models.glm5_next import ( + _ROPE_PAD, + Glm5NextMLAAttention, + Glm5NextMoE, + _normalize_glm5_next_config, +) + + +def _add_mtp_quant_excludes(atom_config: Config) -> None: + """Mirror layer-45 BF16 exclusions onto the runtime ``mtp_block`` path.""" + quant_config = atom_config.quant_config + if quant_config is None: + return + layer_prefix = f"model.layers.{atom_config.hf_config.num_hidden_layers}." + block_prefixes = ( + "input_layernorm", + "post_attention_layernorm", + "self_attn", + "mlp", + ) + for attr in ("exclude_layers", "online_exclude_layers"): + excludes = getattr(quant_config, attr, None) + if not excludes: + continue + additions = [] + for name in excludes: + if not name.startswith(layer_prefix): + continue + suffix = name[len(layer_prefix) :] + if suffix.startswith(block_prefixes): + additions.append(f"{layer_prefix}mtp_block.{suffix}") + excludes.extend(name for name in additions if name not in excludes) + + +def _prepare_mtp_config(atom_config: Config) -> Config: + """Return an isolated config for the checkpoint's NextN layer.""" + mtp_config = copy.copy(atom_config) + draft_config = getattr( + getattr(atom_config, "speculative_config", None), + "draft_model_hf_config", + None, + ) + if draft_config is not None: + mtp_config.hf_config = copy.copy(draft_config) + + if atom_config.quant_config is not None: + mtp_config.quant_config = copy.copy(atom_config.quant_config) + for attr in ("exclude_layers", "online_exclude_layers"): + excludes = getattr(mtp_config.quant_config, attr, None) + if excludes is not None: + setattr(mtp_config.quant_config, attr, list(excludes)) + + _normalize_glm5_next_config(mtp_config.hf_config) + _add_mtp_quant_excludes(mtp_config) + return mtp_config + + +class Glm5NextMTPDecoderLayer(nn.Module): + """Checkpoint layer 45, which is a regular residual MLA+MoE block. + + Unlike backbone layers, the NextN layer has no mHC parameters. It still + needs GLM's k-pool indexer and clamped SwiGLU experts, so using the generic + DeepSeek decoder would silently change both attention selection and MoE + numerics. + """ + + def __init__( + self, + atom_config: Config, + prefix: str, + layer_num: int, + ) -> None: + super().__init__() + config = atom_config.hf_config + rotary_emb = NoPositionalRotaryEmbedding( + head_size=_ROPE_PAD, + rotary_dim=_ROPE_PAD, + max_position_embeddings=int(config.max_position_embeddings), + base=10000.0, + is_neox_style=True, + dtype=torch.bfloat16, + ) + self.self_attn = Glm5NextMLAAttention( + atom_config, + layer_num, + rotary_emb, + prefix=f"{prefix}.self_attn", + is_mtp=True, + ) + self.mlp = Glm5NextMoE( + config, + atom_config.quant_config, + prefix=f"{prefix}.mlp", + reduce_results=not ENABLE_ALLREDUCE_RMSNORM_FUSION, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self.self_attn(hidden_states, positions) + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + return self.mlp(hidden_states), residual + + +class Glm5NextMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayer): + """DeepSeek-compatible predictor prologue with a GLM NextN block.""" + + @staticmethod + def build_mtp_block( + atom_config: Config, + prefix: str, + layer_idx: int, + alt_stream: torch.cuda.Stream | None, + ) -> nn.Module: + del alt_stream + return Glm5NextMTPDecoderLayer( + atom_config=atom_config, + prefix=prefix, + layer_num=layer_idx, + ) + + +class Glm5NextMTP(DeepSeekMTP): + """Load checkpoint layer 45 through ATOM's existing NextN runtime.""" + + predictor_layer_cls: ClassVar[type[DeepSeekMultiTokenPredictorLayer]] = ( + Glm5NextMultiTokenPredictorLayer + ) + packed_modules_mapping_override: ClassVar[dict[str, tuple[str, int]]] = { + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + supports_indexer_projection_fusion: ClassVar[bool] = False + weights_mapping: ClassVar[dict[str, str]] = { + "index_kpool_compress_gate": "index_kpool_compress_gate.weight", + } + + def __init__(self, atom_config: Config, prefix: str = ""): + super().__init__(atom_config=_prepare_mtp_config(atom_config), prefix=prefix) + + def remap_mtp_weight_name(self, name: str) -> str | None: + name = name.replace("model.language_model.", "model.") + return super().remap_mtp_weight_name(name) diff --git a/atom/spec_decode/drafter.py b/atom/spec_decode/drafter.py index de9b1a8f0c..c4b37cc725 100644 --- a/atom/spec_decode/drafter.py +++ b/atom/spec_decode/drafter.py @@ -120,6 +120,7 @@ def _identity_extract(output: Any, _module: nn.Module) -> torch.Tensor: "MiMoV2MTPModel": "atom.models.mimo_v2_mtp.MiMoV2MTP", "MiMoV2FlashMTPModel": "atom.models.mimo_v2_mtp.MiMoV2MTP", "Qwen3_5MTPModel": "atom.models.qwen3_5_mtp.Qwen3_5MTP", + "Glm5NextMTPModel": "atom.models.glm5_next_mtp.Glm5NextMTP", "Eagle3LlamaModel": "atom.models.eagle3_llama.Eagle3LlamaModel", "Eagle3DeepseekMLAModel": "atom.models.eagle3_deepseek_mla.Eagle3DeepseekMLAModel", "K3DSparkModel": "atom.models.kimi_k3_dspark.KimiK3DSpark", @@ -220,7 +221,12 @@ def warmup_draft_graphs(self, build_context, stream) -> None: all_gather that goes INTO the recording, so a synthetic context left describing the target bakes a collective the pass never runs at. """ - if not self.draft_graphs: + from atom.utils import envs + + # The forced-rejection diagnostic returns sentinel draft ids before any + # draft forward. Its declared graphs are therefore unreachable, and + # warming them would try to compile an otherwise-uninitialized wrapper. + if envs.ATOM_DEBUG_FORCE_SKIP_DRAFT_MODEL or not self.draft_graphs: return runner = self.runner capture_sizes = sorted(runner.capture_sizes) # capture leaves it descending diff --git a/atom/spec_decode/eagle_proposer.py b/atom/spec_decode/eagle_proposer.py index 913607d904..1e6477fa69 100644 --- a/atom/spec_decode/eagle_proposer.py +++ b/atom/spec_decode/eagle_proposer.py @@ -138,9 +138,13 @@ def _declare_draft_graphs(self): return () draft_hf = self.speculative_config.draft_model_hf_config # DeepSeek-V4 carries the mHC residual, so its hidden is [N, hc, dim] - # rather than [N, dim]. `hc_mult` is absent on every architecture that - # does not, which is exactly the two-dimensional case. - hc = getattr(draft_hf, "hc_mult", None) + # rather than [N, dim]. GLM inherits hc_mult from its backbone, but + # its NextN block carries an ordinary two-dimensional residual. + hc = ( + None + if draft_hf.architectures[0] == "Glm5NextMTPModel" + else getattr(draft_hf, "hc_mult", None) + ) inputs = { # int64, not the int32 of the token buffer step 0 reads: a mid-step's # ids come from `compute_draft_ids`, which is an argmax. The loop @@ -157,11 +161,14 @@ def _declare_draft_graphs(self): dtype=self.dtype, ), } - # Keep this capability at the non-compiled call site. DeepSeekMTPModel - # has the two-dimensional hidden-state and shared-head contracts needed - # to feed its fixed graph inputs directly; other draft architectures + # Keep this capability at the non-compiled call site. These MTP models + # have the two-dimensional hidden-state and shared-head contracts needed + # to feed their fixed graph inputs directly; other draft architectures # remain on the owned-output path. - self._reuse_step_buffers = draft_hf.architectures[0] == "DeepSeekMTPModel" + self._reuse_step_buffers = draft_hf.architectures[0] in { + "DeepSeekMTPModel", + "Glm5NextMTPModel", + } self.step = DraftGraph( forward=self._step_forward, epilogue=self._step_head, @@ -552,6 +559,7 @@ def propose( ) if envs.ATOM_DEBUG_FORCE_SKIP_DRAFT_MODEL: draft_token_ids.fill_(-1) + return draft_token_ids var = self.runner.forward_vars # Eaale3 only support mha currently draft_uses_mha = hasattr(self.runner, "draft_kv_builder") diff --git a/recipes/GLM-5.3-Flash.md b/recipes/GLM-5.3-Flash.md index 81ff52e954..009254b1b5 100644 --- a/recipes/GLM-5.3-Flash.md +++ b/recipes/GLM-5.3-Flash.md @@ -5,8 +5,8 @@ > vision tower is skipped) and scores **gsm8k 0.9682 / 0.9689** at 3-shot on the > dense path and **0.9659 / 0.9666** at 16-shot on the pooled k-pool path, over > all 1319 questions (chat, TP8, bf16 KV; see §7). Per-layer hidden states match -> the transformers reference to cosine ≥ 0.9997 at all 45 layers. Not yet done: -> the MTP draft layer or multimodal serving. See §8. +> the transformers reference to cosine ≥ 0.9997 at all 45 layers. The GLM-5.3 +> MTP draft layer passes TP8 functional and throughput validation. See §8. ```bash python -m atom.examples.simple_inference --model /models/GLM-5.3-Flash -tp 4 \ @@ -91,9 +91,9 @@ as the checkpoint names them (`hc_attn_fn`, ...), keeps `q/k/v_conv1d` separate like the checkpoint (Kimi-K3 already does), and folds the low-rank KDA output gate after load, so its whole `weights_mapping` is one rule: `"model.language_model." -> "model."`. `model.visual.*` is dropped via -`skip_weight_prefixes`, and checkpoint layer 45 (MTP) is dropped automatically by -the loader's past-last-layer filter. Only the expert and q/k/v fusions go through -`packed_modules_mapping`. +`skip_weight_prefixes`. Checkpoint layer 45 is dropped from the target by the +loader's past-last-layer filter and loaded separately by the MTP wrapper. Only +the expert and q/k/v fusions go through `packed_modules_mapping`. ## 3. What is validated @@ -330,7 +330,9 @@ backends forward it into their fused activation path. The dense layers use The model now carries `@support_torch_compile`; TP8 level-3 compilation and whole-forward CUDA graph capture are smoke-tested. Sharing the identity RoPE cache across all 11 MLA layers reduced measured `peak_torch` from 42.71 GiB to -41.47 GiB per TP rank. MTP remains unsupported. +41.47 GiB per TP rank. On MI308X, MTP3 raised single-request output throughput +from 104.67 to 181.54 tok/s (+73.44%), while MTP1 raised concurrency-8 output +throughput from 409.96 to 559.98 tok/s (+36.59%). ## 7. Measured serving accuracy @@ -384,17 +386,20 @@ Three things about scoring this model that will otherwise waste a run: `model_ops/glm5_next/{indexer,kpool}.py` implements the paged/ragged pooled indexer; measured at 16-shot in §7 and checked directly against its torch kernel references. -2. **MTP draft layer** (checkpoint layer 45: `eh_proj` / `enorm` / `hnorm` / - `shared_head.norm`, plus its own indexer). `index_share_for_mtp_iteration` - means it reuses the main model's top-k. +2. ~~**MTP draft layer.**~~ Done — the GLM-specific checkpoint layer 45 + (`eh_proj` / `enorm` / `hnorm` / `shared_head.norm`, k-pool MLA and clamped + SwiGLU MoE) runs through the existing NextN runtime. TP8 MTP1 and MTP3 passed + text requests, natural acceptance lengths 0 through 3, forced rejection, + concurrent short/long requests, and an 8/8 quality smoke suite. 3. **Parallel feature coverage.** PCP, DCP and TBO remain explicitly rejected until their pooled-index metadata/state layouts have dedicated tests. 4. **Multimodal serving.** Land the image processor, input builder, tower and packed-weight tests together. `Glm5NextProcessor` only exists in transformers >= 5.16 while ATOM pins 5.12.1; video additionally needs frame sampling. Until then `model.visual.*` is skipped so text serving does not pay its VRAM. -5. **Performance**: tune the newly-enabled compiled path, add MTP speculative - decoding, and drop the `hc` torch fallback once the fused path is trusted. +5. **Performance**: MTP3 is the measured single-request latency/throughput + choice; MTP1 is the concurrency-8 throughput choice. Add a sustained soak + and drop the `hc` torch fallback once the fused path is trusted. 6. **Upstream the transformers FP8 bug** (§4a) and the gfx950 Triton failure (§4b). Upstream, for reference: sglang PR #36507 (16.6k lines, 144 files) and vLLM PR diff --git a/tests/model_ops/test_glm5_kpool_speculative.py b/tests/model_ops/test_glm5_kpool_speculative.py new file mode 100644 index 0000000000..18aa35e526 --- /dev/null +++ b/tests/model_ops/test_glm5_kpool_speculative.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest + + +def _get_query_request_indices(): + try: + from atom.model_ops.glm5_next.speculative import get_query_request_indices + except (ModuleNotFoundError, RuntimeError) as error: + missing_aiter = ( + isinstance(error, ModuleNotFoundError) + and error.name is not None + and (error.name == "aiter" or error.name.startswith("aiter.")) + ) + missing_device = isinstance(error, RuntimeError) and "rocminfo" in str(error) + if not (missing_aiter or missing_device): + raise + pytest.skip("AITER imports require a visible ROCm device") + return get_query_request_indices + + +def test_get_query_request_indices_maps_ragged_verification_tokens_to_requests(): + import torch + + get_query_request_indices = _get_query_request_indices() + cu = torch.tensor([0, 3, 3, 7, 9], dtype=torch.int32) + + rows = get_query_request_indices(cu, 9) + + # The second request is empty; no token may be assigned to it. + assert rows.tolist() == [0, 0, 0, 2, 2, 2, 2, 3, 3] + + +def test_get_query_request_indices_preserves_single_request_positions(): + import torch + + get_query_request_indices = _get_query_request_indices() + cu = torch.tensor([0, 4], dtype=torch.int32) + + assert get_query_request_indices(cu, 4).tolist() == [0, 0, 0, 0] + + +def test_pool_candidates_mix_history_with_fresh_rows(): + import torch + + if not torch.cuda.is_available(): + pytest.skip("requires a ROCm GPU") + from atom.model_ops.glm5_next import kpool + from atom.model_ops.glm5_next.speculative import build_speculative_pool_candidates + + torch.manual_seed(19) + pool = 4 + dim = 128 + history = torch.randn(2, 2, pool, dim, device="cuda", dtype=torch.bfloat16) + keys = torch.randn(5, dim, device="cuda", dtype=torch.bfloat16) + gates = torch.randn_like(keys) + positions = torch.tensor([10, 11, 20, 21, 22], device="cuda") + cu = torch.tensor([0, 2, 5], device="cuda", dtype=torch.int32) + slots = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + ape = torch.randn(pool, dim, device="cuda") + + got, requests = build_speculative_pool_candidates( + history, keys, gates, positions, cu, slots, ape, pool + ) + + expected_keys = [] + expected_gates = [] + for row, (request, start) in enumerate(zip([0, 0, 1, 1, 1], [0, 0, 2, 2, 2])): + start_position = int(positions[start]) + query_position = int(positions[row]) + row_keys = [] + row_gates = [] + for position in range(query_position - pool + 1, query_position + 1): + if position < start_position: + row_keys.append(history[request, 0, position % pool]) + row_gates.append(history[request, 1, position % pool]) + else: + fresh_row = start + position - start_position + row_keys.append(keys[fresh_row]) + row_gates.append(gates[fresh_row]) + expected_keys.append(torch.stack(row_keys)) + expected_gates.append(torch.stack(row_gates)) + expected = kpool.pool_compress_ref( + torch.stack(expected_keys), torch.stack(expected_gates), ape + ) + expected = kpool.hadamard128_ref(expected.to(torch.bfloat16).float()).to( + torch.bfloat16 + ) + + assert requests.tolist() == [0, 0, 1, 1, 1] + torch.testing.assert_close(got, expected, rtol=2e-2, atol=2e-2) + + +def test_history_update_preserves_unwritten_residues(): + import torch + + if not torch.cuda.is_available(): + pytest.skip("requires a ROCm GPU") + from atom.model_ops.glm5_next.speculative import update_speculative_kpool_history + + pool = 4 + dim = 128 + history = ( + torch.arange(4 * 2 * pool * dim, device="cuda", dtype=torch.float32) + .reshape(4, 2, pool, dim) + .to(torch.bfloat16) + ) + before = history.clone() + keys = torch.full((5, dim), 101, device="cuda", dtype=torch.bfloat16) + gates = torch.full((5, dim), 202, device="cuda", dtype=torch.bfloat16) + positions = torch.tensor([10, 11, 20, 21, 22], device="cuda") + cu = torch.tensor([0, 2, 5], device="cuda", dtype=torch.int32) + slots_in = torch.tensor([0, 1], device="cuda", dtype=torch.int32) + slots_out = torch.tensor([2, 3], device="cuda", dtype=torch.int32) + + update_speculative_kpool_history( + history, keys, gates, positions, cu, slots_in, slots_out + ) + torch.cuda.synchronize() + + # Request 0 writes absolute positions 10/11 (ring residues 2/3), copying + # residues 0/1 from its source slot. Request 1 writes 20/21/22 (0/1/2). + torch.testing.assert_close(history[2, :, :2], before[0, :, :2]) + assert torch.all(history[2, 0, 2:] == 101) + assert torch.all(history[2, 1, 2:] == 202) + assert torch.all(history[3, 0, :3] == 101) + assert torch.all(history[3, 1, :3] == 202) + torch.testing.assert_close(history[3, :, 3], before[1, :, 3]) + + +def test_token_indices_map_to_request_cache_slots(): + import torch + + if not torch.cuda.is_available(): + pytest.skip("requires a ROCm GPU") + from atom.model_ops.glm5_next.speculative import map_token_indices_to_slots + + token_indices = torch.tensor( + [[0, 15, 16, 31], [0, 1, 16, 32]], + device="cuda", + dtype=torch.int32, + ) + request_indices = torch.tensor([0, 1], device="cuda", dtype=torch.int64) + block_tables = torch.tensor( + [[10, 20], [30, 40]], + device="cuda", + dtype=torch.int32, + ) + output_indptr = torch.tensor([0, 4, 8], device="cuda", dtype=torch.int32) + output = torch.empty(8, device="cuda", dtype=torch.int32) + + map_token_indices_to_slots( + token_indices, + request_indices, + block_tables, + output_indptr, + output, + block_size=16, + ) + torch.cuda.synchronize() + + assert output.tolist() == [160, 175, 320, 335, 480, 481, 640, 0] diff --git a/tests/test_glm5_next_mtp_routing.py b/tests/test_glm5_next_mtp_routing.py new file mode 100644 index 0000000000..f0d9ac7c89 --- /dev/null +++ b/tests/test_glm5_next_mtp_routing.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +from types import SimpleNamespace + +import pytest + +from atom.config import SpeculativeConfig, _glm5_next_unsupported_features + + +def _aiter_unavailable(error): + missing_aiter = ( + isinstance(error, ModuleNotFoundError) + and error.name is not None + and (error.name == "aiter" or error.name.startswith("aiter.")) + ) + missing_device = isinstance(error, RuntimeError) and "rocminfo" in str(error) + return missing_aiter or missing_device + + +def _mtp_symbols(): + try: + from atom.models.glm5_next_mtp import ( + Glm5NextMTP, + _add_mtp_quant_excludes, + ) + except (ModuleNotFoundError, RuntimeError) as error: + if not _aiter_unavailable(error): + raise + pytest.skip("AITER imports require a visible ROCm device") + return Glm5NextMTP, _add_mtp_quant_excludes + + +def _eagle_proposer(): + try: + from atom.spec_decode.eagle_proposer import EagleProposer + except (ModuleNotFoundError, RuntimeError) as error: + if not _aiter_unavailable(error): + raise + pytest.skip("AITER imports require a visible ROCm device") + return EagleProposer + + +def test_glm5_next_text_routes_to_glm_mtp_model(): + config = SimpleNamespace( + model_type="glm5_next_text", + architectures=["Glm5NextForConditionalGeneration"], + num_nextn_predict_layers=1, + ) + config.update = lambda values: [ + setattr(config, key, value) for key, value in values.items() + ] + + SpeculativeConfig.hf_config_override(config) + + assert config.model_type == "glm5_next_mtp" + assert config.architectures == ["Glm5NextMTPModel"] + assert config.n_predict == 1 + + +def test_glm5_next_allows_mtp_but_rejects_unimplemented_parallel_modes(): + config = SimpleNamespace( + speculative_config=object(), + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + enable_tbo=False, + enable_tbo_decode=False, + ) + + assert _glm5_next_unsupported_features(config) == [] + + config.prefill_context_parallel_size = 2 + config.decode_context_parallel_size = 2 + config.enable_tbo = True + assert _glm5_next_unsupported_features(config) == ["PCP", "DCP", "TBO"] + + +def test_glm5_next_mtp_remaps_language_model_checkpoint_layer(): + Glm5NextMTP, _ = _mtp_symbols() + model = object.__new__(Glm5NextMTP) + model.config = SimpleNamespace( + num_hidden_layers=45, + num_nextn_predict_layers=1, + ) + + assert ( + model.remap_mtp_weight_name( + "model.language_model.layers.45.self_attn.q_a_proj.weight" + ) + == "model.layers.45.mtp_block.self_attn.q_a_proj.weight" + ) + assert ( + model.remap_mtp_weight_name("model.language_model.layers.45.eh_proj.weight") + == "model.layers.45.eh_proj.weight" + ) + assert ( + model.remap_mtp_weight_name( + "model.language_model.layers.45.shared_head.norm.weight" + ) + == "model.layers.45.shared_head.norm.weight" + ) + assert ( + model.remap_mtp_weight_name( + "model.language_model.layers.44.self_attn.q_a_proj.weight" + ) + is None + ) + assert Glm5NextMTP.weights_mapping == { + "index_kpool_compress_gate": "index_kpool_compress_gate.weight" + } + + +def test_glm5_next_mtp_mirrors_block_quant_excludes(): + _, _add_mtp_quant_excludes = _mtp_symbols() + quant_config = SimpleNamespace( + exclude_layers=[ + "model.layers.45.self_attn.q_a_proj", + "model.layers.45.input_layernorm", + "model.layers.45.eh_proj", + "model.layers.44.self_attn.q_a_proj", + ], + online_exclude_layers=["model.layers.45.mlp.gate"], + ) + atom_config = SimpleNamespace( + hf_config=SimpleNamespace(num_hidden_layers=45), + quant_config=quant_config, + ) + + _add_mtp_quant_excludes(atom_config) + + assert "model.layers.45.mtp_block.self_attn.q_a_proj" in ( + quant_config.exclude_layers + ) + assert "model.layers.45.mtp_block.input_layernorm" in (quant_config.exclude_layers) + assert "model.layers.45.mtp_block.mlp.gate" in (quant_config.online_exclude_layers) + assert "model.layers.45.mtp_block.eh_proj" not in quant_config.exclude_layers + + +@pytest.mark.parametrize( + "architecture, expected", + [ + ("Glm5NextMTPModel", (4096,)), + ("DeepSeekV4MTPModel", (4, 4096)), + ("DeepSeekMTPModel", (4096,)), + ], +) +def test_mtp_graph_stages_actual_residual_shape(architecture, expected): + import torch + + EagleProposer = _eagle_proposer() + + hf = SimpleNamespace(architectures=[architecture], hidden_size=4096) + if architecture != "DeepSeekMTPModel": + hf.hc_mult = 4 + proposer = SimpleNamespace( + runner=SimpleNamespace(use_mrope=False), + mtp_k=3, + speculative_config=SimpleNamespace(draft_model_hf_config=hf), + dtype=torch.bfloat16, + _step_forward=lambda *a, **kw: None, + _step_head=lambda *a, **kw: None, + _step_warmup_inputs=lambda *a, **kw: None, + ) + (graph,) = EagleProposer._declare_draft_graphs(proposer) + assert graph.inputs["hidden_states"].shape == expected + graph.bind(SimpleNamespace(max_num_seqs=8), "cpu") + source = torch.ones((1, *expected), dtype=torch.bfloat16) + staged = graph.stage(4, {"hidden_states": source})["hidden_states"] + assert staged.shape == (4, *expected) + torch.testing.assert_close(staged, source.expand_as(staged))