From 7eee65ea9b458585ba414d4331657e46ece36f60 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Sat, 27 Jun 2026 22:10:29 +0800 Subject: [PATCH 01/17] [GLM5.1/5.2] Fix acc issue for long prompt Signed-off-by: zejunchen-zejun --- atom/model_ops/attention_mla.py | 18 +++++ atom/model_ops/attentions/aiter_mla.py | 97 +++++++++++++++++++++++--- 2 files changed, 107 insertions(+), 8 deletions(-) diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index ee98185f41..0f61552a31 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -874,6 +874,24 @@ def _forward_prefill_mla( sm_scale=self.scale, q_scale=self._q_scale, kv_scale=self._k_scale, + work_meta_data=getattr( + attn_metadata, "sparse_prefill_work_meta_data", None + ), + work_indptr=getattr( + attn_metadata, "sparse_prefill_work_indptr", None + ), + work_info_set=getattr( + attn_metadata, "sparse_prefill_work_info_set", None + ), + reduce_indptr=getattr( + attn_metadata, "sparse_prefill_reduce_indptr", None + ), + reduce_final_map=getattr( + attn_metadata, "sparse_prefill_reduce_final_map", None + ), + reduce_partial_map=getattr( + attn_metadata, "sparse_prefill_reduce_partial_map", None + ), ) else: mla_prefill_fwd( diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 3895026831..b924ea0498 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -198,6 +198,40 @@ def __init__(self, model_runner): dtype=torch.int32, device=self.device, ) + ( + (spp_wmd_size, spp_wmd_type), + (spp_wi_size, spp_wi_type), + (spp_wis_size, spp_wis_type), + (spp_ri_size, spp_ri_type), + (spp_rfm_size, spp_rfm_type), + (spp_rpm_size, spp_rpm_type), + ) = get_mla_metadata_info_v1( + self.max_num_batched_tokens, + 1, # sparse prefill treats each query token as q_len=1 + self.padded_num_attention_heads, + self.dtype_q, + self.dtype_kv, + is_sparse=True, + fast_mode=True, + ) + mla_metadata["sparse_prefill_work_meta_data"] = torch.empty( + spp_wmd_size, dtype=spp_wmd_type, device=self.device + ) + mla_metadata["sparse_prefill_work_indptr"] = torch.empty( + spp_wi_size, dtype=spp_wi_type, device=self.device + ) + mla_metadata["sparse_prefill_work_info_set"] = torch.empty( + spp_wis_size, dtype=spp_wis_type, device=self.device + ) + mla_metadata["sparse_prefill_reduce_indptr"] = torch.empty( + spp_ri_size, dtype=spp_ri_type, device=self.device + ) + mla_metadata["sparse_prefill_reduce_final_map"] = torch.empty( + spp_rfm_size, dtype=spp_rfm_type, device=self.device + ) + mla_metadata["sparse_prefill_reduce_partial_map"] = torch.empty( + spp_rpm_size, dtype=spp_rpm_type, device=self.device + ) if self.is_sparse and max_seqlen_qo > 1: # Allocate a second set of persistent work buffers for sparse MTP @@ -744,14 +778,25 @@ def prepare_prefill(self, batch: ScheduledBatch): self.prepare_block_tables(batch) attn_metadata.block_tables = var["block_tables"].copy_to_gpu(bs) counts = var["cu_seqlens_q"].np[1 : bs + 1] - var["cu_seqlens_q"].np[:bs] + local_offsets = np.concatenate( + [np.arange(s, dtype=np.int32) for s in counts] + ) if attn_metadata.has_cached: - # Full context (cached + new): use cu_seqlens_k for indexer + # Full context (cached + new): each query token can see the cached + # prefix plus previous query tokens in this chunk, not future chunk + # tokens. + seq_starts = var["cu_seqlens_k"].np[:bs] + seq_lens = var["cu_seqlens_k"].np[1 : bs + 1] - seq_starts + cached_lens = seq_lens - counts + repeated_seq_starts = np.repeat(seq_starts, counts) + repeated_cached_lens = np.repeat(cached_lens, counts) var["cu_seqlen_ks"].np[:sum_scheduled_tokens] = np.repeat( - var["cu_seqlens_k"].np[:bs], counts + seq_starts, counts ) - var["cu_seqlen_ke"].np[:sum_scheduled_tokens] = np.repeat( - var["cu_seqlens_k"].np[1 : bs + 1], counts + var["cu_seqlen_ke"].np[:sum_scheduled_tokens] = ( + repeated_seq_starts + repeated_cached_lens + local_offsets + 1 ) + sparse_counts = repeated_cached_lens + local_offsets + 1 else: var["cu_seqlen_ke"].np[:sum_scheduled_tokens] = ( np.arange(sum_scheduled_tokens, dtype=np.int32) + 1 @@ -759,6 +804,7 @@ def prepare_prefill(self, batch: ScheduledBatch): var["cu_seqlen_ks"].np[:sum_scheduled_tokens] = np.repeat( var["cu_seqlens_q"].np[:bs], counts ) + sparse_counts = local_offsets + 1 attn_metadata.cu_seqlen_ks = var["cu_seqlen_ks"].copy_to_gpu( sum_scheduled_tokens ) @@ -780,15 +826,50 @@ def prepare_prefill(self, batch: ScheduledBatch): ) var["sparse_kv_indptr"].np[0] = 0 var["sparse_kv_indptr"].np[1 : sum_scheduled_tokens + 1] = np.cumsum( - np.minimum( - np.concatenate([np.arange(1, s + 1) for s in counts]), - self.index_topk, - ), + np.minimum(sparse_counts, self.index_topk), dtype=np.int32, ) attn_metadata.sparse_kv_indptr = var["sparse_kv_indptr"].copy_to_gpu( sum_scheduled_tokens + 1 ) + get_mla_metadata_v1( + attn_metadata.sparse_cu_seqlens_q, + attn_metadata.sparse_kv_indptr, + attn_metadata.kv_last_page_lens, + self.padded_num_attention_heads, + 1, # nhead_kv + True, + var["sparse_prefill_work_meta_data"], + var["sparse_prefill_work_info_set"], + var["sparse_prefill_work_indptr"], + var["sparse_prefill_reduce_indptr"], + var["sparse_prefill_reduce_final_map"], + var["sparse_prefill_reduce_partial_map"], + page_size=self.block_size, + dtype_q=self.dtype_q, + dtype_kv=self.dtype_kv, + kv_granularity=max(self.block_size, 16), + max_seqlen_qo=1, + uni_seqlen_qo=1, + fast_mode=1, + max_split_per_batch=16, + ) + attn_metadata.sparse_prefill_work_meta_data = var[ + "sparse_prefill_work_meta_data" + ] + attn_metadata.sparse_prefill_work_info_set = var[ + "sparse_prefill_work_info_set" + ] + attn_metadata.sparse_prefill_work_indptr = var["sparse_prefill_work_indptr"] + attn_metadata.sparse_prefill_reduce_indptr = var[ + "sparse_prefill_reduce_indptr" + ] + attn_metadata.sparse_prefill_reduce_final_map = var[ + "sparse_prefill_reduce_final_map" + ] + attn_metadata.sparse_prefill_reduce_partial_map = var[ + "sparse_prefill_reduce_partial_map" + ] if hasattr(self.model_runner, "drafter") or attn_metadata.has_cached: # Populate kv_last_page_lens for full sequence (needed for MLA prefill with From b3107d5c32af94cd9d12411e822d9e50e3c0f681 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Mon, 29 Jun 2026 12:11:38 +0800 Subject: [PATCH 02/17] fix GLM5.2 acc drop Signed-off-by: zejunchen-zejun --- atom/models/deepseek_v2.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index ff47a5fcc1..01d3161799 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -174,6 +174,15 @@ def _supports_fused_indexer_kernel_config(config: PretrainedConfig) -> bool: ) +def _is_neox_rope_style( + config: PretrainedConfig, interleave_attr: str, default_interleave: bool +) -> bool: + interleave = getattr(config, interleave_attr, default_interleave) + if interleave is None: + interleave = default_interleave + return not bool(interleave) + + def _can_fuse_indexer_wk_weights_proj( config: PretrainedConfig, quant_config: Optional[QuantizationConfig], @@ -1595,14 +1604,16 @@ def forward( weights = self.weights_proj(hidden_states) if not self.use_qk_rope_cache_fusion: - q_pe, _ = torch.split( + q_pe, q_nope = torch.split( q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) k = self.k_norm(k) - k_pe, _ = torch.split( + k_pe, k_nope = torch.split( k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) - q_pe, k_pe = rotary_emb(positions, q_pe, k_pe) + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) + q = torch.cat([q_pe, q_nope], dim=-1) + k = torch.cat([k_pe.squeeze(1), k_nope], dim=-1) q = q.view(-1, self.head_dim) q_fp8, q_scale = self.quant_func(q, quant_dtype=dtypes.fp8) @@ -1822,7 +1833,7 @@ def __init__( max_position=max_position_embeddings, base=rope_theta, rope_scaling=rope_scaling, - is_neox_style=False, + is_neox_style=_is_neox_rope_style(config, "rope_interleave", True), ) if rope_scaling: mscale_all_dim = rope_scaling.get("mscale_all_dim", False) From 9fd3b3d3dc02f059649393854c091228ef8e6f23 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Mon, 29 Jun 2026 15:56:40 +0800 Subject: [PATCH 03/17] add potential fix for glm5.2 Signed-off-by: zejunchen-zejun --- atom/model_ops/attention_mla.py | 21 +++++++++++++++++++-- atom/models/deepseek_v2.py | 6 ++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index 0f61552a31..ea846ebed3 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -169,6 +169,13 @@ class MLAModules: kv_b_proj: torch.nn.Module o_proj: torch.nn.Module indexer: Optional[torch.nn.Module] + # Model-level sparse flag. A v3.2 / GLM-5.2 model runs sparse MLA on ALL its + # layers. GLM-5.2 IndexShare "shared" layers carry no indexer module yet must + # still run sparse attention (reusing the prior "full" layer's top-k), so + # sparsity must be derived from the model, not from whether this layer owns + # an indexer. Defaults keep non-sparse models unchanged. + is_sparse: bool = False + topk_tokens: Optional[int] = None def dynamic_per_batched_tensor_quant( @@ -231,10 +238,20 @@ def __init__( self.one_scale = torch.tensor(1.0, dtype=torch.float32) self._k_scale = self.one_scale self._q_scale = self.one_scale - self.is_sparse_mla = mla_modules.indexer is not None + # Derive sparsity from the model-level flag, not from whether THIS layer + # owns an indexer: GLM-5.2 IndexShare "shared" layers have indexer=None + # but must still run sparse MLA, reusing the prior "full" layer's top-k. + # (`mla_modules.is_sparse` defaults False, so non-sparse models and the + # `indexer is not None` fallback keep their previous behavior.) + self.is_sparse_mla = mla_modules.is_sparse or (mla_modules.indexer is not None) self.topk_tokens = ( - mla_modules.indexer.topk_tokens if mla_modules.indexer is not None else None + mla_modules.indexer.topk_tokens + if mla_modules.indexer is not None + else mla_modules.topk_tokens ) + # Shared layers have no indexer buffer at construction; the metadata + # builder rebinds it to the shared `_sparse_kv_indices_gpu` at runtime, + # so the layer reads the prior full layer's selected indices. self.sparse_kv_indices_buffer = ( mla_modules.indexer.sparse_kv_indices_buffer if mla_modules.indexer is not None diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 01d3161799..b802d84d7f 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -1899,6 +1899,12 @@ def __init__( kv_b_proj=self.kv_b_proj, o_proj=self.o_proj, indexer=self.indexer, + # v3.2 / GLM-5.2 runs sparse MLA on every layer. For GLM-5.2 IndexShare + # "shared" layers self.indexer is None, but they must still run sparse + # attention and reuse the prior full layer's top-k, so flag sparsity at + # the model level rather than per-layer. + is_sparse=self.is_v32, + topk_tokens=(config.index_topk if self.is_v32 else None), ) self.mla_attn = Attention( From 61fdbdea18430ada3631a286e776f8a5398d3c1e Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Fri, 3 Jul 2026 10:58:40 +0800 Subject: [PATCH 04/17] retrieve the no-op for ROPE Signed-off-by: zejunchen-zejun --- atom/models/deepseek_v2.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index b802d84d7f..618c5e89c9 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -1604,16 +1604,18 @@ def forward( weights = self.weights_proj(hidden_states) if not self.use_qk_rope_cache_fusion: - q_pe, q_nope = torch.split( + q_pe, _ = torch.split( q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) k = self.k_norm(k) - k_pe, k_nope = torch.split( + k_pe, _ = torch.split( k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) - q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1)) - q = torch.cat([q_pe, q_nope], dim=-1) - k = torch.cat([k_pe.squeeze(1), k_nope], dim=-1) + # rotary_emb runs aiter's in-place rope kernel, which honors the + # (strided) views returned by torch.split and writes the rotated + # values straight back into q/k's own storage. So q/k already carry + # the rope here; no re-concat is needed. + rotary_emb(positions, q_pe, k_pe) q = q.view(-1, self.head_dim) q_fp8, q_scale = self.quant_func(q, quant_dtype=dtypes.fp8) From c78daeed5adac0e7d3130482e543ed9f3917f65a Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Fri, 3 Jul 2026 11:04:37 +0800 Subject: [PATCH 05/17] retrieve outplace return Signed-off-by: zejunchen-zejun --- atom/models/deepseek_v2.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 618c5e89c9..30c369e3bb 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -1611,11 +1611,7 @@ def forward( k_pe, _ = torch.split( k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) - # rotary_emb runs aiter's in-place rope kernel, which honors the - # (strided) views returned by torch.split and writes the rotated - # values straight back into q/k's own storage. So q/k already carry - # the rope here; no re-concat is needed. - rotary_emb(positions, q_pe, k_pe) + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe) q = q.view(-1, self.head_dim) q_fp8, q_scale = self.quant_func(q, quant_dtype=dtypes.fp8) From 165f838ca33cdccddda8dfa85082b6bdb4e2ab46 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Wed, 1 Jul 2026 13:48:54 +0800 Subject: [PATCH 06/17] [GLM5.2][perf] Enable fused indexer qk-rope+quant+cache kernel GLM-5.2 (glm_moe_dsa) was hard-excluded from the fused indexer path added for DeepSeek-V3.2 in #788, forcing the indexer's rope + fp8-quant + kv-cache write to run as 5-6 separate ops per layer per token. GLM's indexer is structurally identical to V3.2 (index_head_dim=128, qk_rope_head_dim=64, per_1x128 fp8 quant, always-neox indexer rope), so the fused indexer_qk_rope_quant_and_cache kernel is math-equivalent to the per-op path. Notably the recent GLM acc fix (da360a21) fixed a bug that only exists in the per-op path (dropping q_nope on recombine); the fused kernel never had it. - Allow glm_moe_dsa through _supports_fused_indexer_kernel_config, gated by a new env ATOM_ENABLE_GLM_FUSED_INDEXER (default on) for easy A/B + rollback. - Keep the wk+weights_proj GEMM merge OFF for GLM (decoupled in _can_fuse_indexer_wk_weights_proj): that fusion requires merging two checkpoint tensors and GLM names weights_proj "indexers_proj" without a pre-mergeable layout. The dominant win (rope+quant+cache) needs no merge. DeepSeek-V3.2 behavior is unchanged. Co-Authored-By: Claude Opus 4.8 --- atom/models/deepseek_v2.py | 19 ++++++++++++++++++- atom/utils/envs.py | 8 ++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 30c369e3bb..ef9c8aadb3 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -127,6 +127,7 @@ ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION = ( envs.ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION ) +ENABLE_GLM_FUSED_INDEXER = envs.ATOM_ENABLE_GLM_FUSED_INDEXER _FP8_DTYPES = tuple( dtype for dtype in ( @@ -166,8 +167,16 @@ def _enable_non_triton_global_mxfp4_input_norm_quant( def _supports_fused_indexer_kernel_config(config: PretrainedConfig) -> bool: if not hasattr(config, "index_topk"): return False + # GLM-5.2 (glm_moe_dsa) shares DeepSeek-V3.2's sparse-MLA indexer: same dims + # (index_head_dim=128, qk_rope_head_dim=64), same per_1x128 fp8 quant, and the + # indexer rope is always neox for both. The fused kernel path is therefore + # math-equivalent to the per-op path, so allow it here (gated by an env flag for + # easy rollback). The GEMM-merge of wk+weights_proj stays off for GLM — see + # _can_fuse_indexer_wk_weights_proj — because GLM checkpoints name the projections + # differently and are not pre-mergeable. if getattr(config, "model_type", None) == "glm_moe_dsa": - return False + if not ENABLE_GLM_FUSED_INDEXER: + return False return ( getattr(config, "index_head_dim", None) == 128 and getattr(config, "qk_rope_head_dim", None) == 64 @@ -192,6 +201,14 @@ def _can_fuse_indexer_wk_weights_proj( return False if not _supports_fused_indexer_kernel_config(config): return False + # GLM-5.2: enable the qk-rope+quant+cache kernel fusion but keep indexer.wk / + # indexer.weights_proj as separate checkpoint tensors. GLM checkpoints name + # weights_proj "indexers_proj" and don't ship a pre-mergeable wk_weights_proj + # layout, so the GEMM-merge (which relies on packed_modules_mapping merging two + # tensors) is left off until the loader is validated for GLM. The dominant win + # (rope+quant+cache fusion) does not require this merge. + if getattr(config, "model_type", None) == "glm_moe_dsa": + return False if quant_config is None: return True diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 3bd2a53abb..77b2989140 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -70,6 +70,14 @@ "ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION": lambda: ( os.getenv("ATOM_ENABLE_DS_INDEXER_QK_ROPE_CACHE_FUSION", "1") == "1" ), + # GLM-5.2 (glm_moe_dsa): enable the fused indexer qk-rope + fp8-quant + kv-cache + # kernel (indexer_qk_rope_quant_and_cache), same path DeepSeek-V3.2 uses. GLM's + # indexer dims (index_head_dim=128, qk_rope_head_dim=64, per_1x128, neox rope) are + # identical to V3.2, so the fusion is math-equivalent to the unfused path. Set to + # "0" to fall back to the per-op Python path if a regression is suspected. + "ATOM_ENABLE_GLM_FUSED_INDEXER": lambda: ( + os.getenv("ATOM_ENABLE_GLM_FUSED_INDEXER", "1") == "1" + ), "ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION": lambda: ( os.getenv("ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION", "1") == "1" ), From 2d2e6056ac311d800e6af03e01b30282e72c9d5b Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Wed, 1 Jul 2026 17:56:27 +0800 Subject: [PATCH 07/17] [GLM5.2][perf] Enable indexer wk+weights_proj GEMM merge for GLM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove the guard that kept the indexer wk+weights_proj GEMM-merge off for glm_moe_dsa. The guard was overly conservative: it assumed GLM checkpoints ship the projection under a non-standard "indexers_proj" name that couldn't feed the packed_modules_mapping merge. In fact "indexers_proj" only appears in the HF quant config's modules_to_not_convert (remapped via quant_exclude_name_mapping); the actual checkpoint tensors use the standard indexer.wk (fp8 block-scale) + indexer.weights_proj (bf16) paths — exactly the layout IndexerWkWeightsProjLinear's fp8-wk load and the merge expect. GLM now takes the same path as DeepSeek-V3.2: wk + weights_proj collapse into a single BF16 GEMM (IndexerWkWeightsProjLinear), on top of the already-enabled qk-rope+quant+cache fusion. Still gated by ATOM_ENABLE_GLM_FUSED_INDEXER for rollback. DeepSeek-V3.2 behavior unchanged. Verify on GPU: weights load without missing indexer.wk / indexer.weights_proj, accuracy unchanged, and the two small indexer GEMMs collapse to one in the decode breakdown. Co-Authored-By: Claude Opus 4.8 --- atom/models/deepseek_v2.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index ef9c8aadb3..3648510797 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -171,9 +171,10 @@ def _supports_fused_indexer_kernel_config(config: PretrainedConfig) -> bool: # (index_head_dim=128, qk_rope_head_dim=64), same per_1x128 fp8 quant, and the # indexer rope is always neox for both. The fused kernel path is therefore # math-equivalent to the per-op path, so allow it here (gated by an env flag for - # easy rollback). The GEMM-merge of wk+weights_proj stays off for GLM — see - # _can_fuse_indexer_wk_weights_proj — because GLM checkpoints name the projections - # differently and are not pre-mergeable. + # easy rollback). This also enables the wk+weights_proj GEMM-merge for GLM — its + # checkpoint uses the standard indexer.wk / indexer.weights_proj tensor names + # (the "indexers_proj" alias only lives in the HF quant config), so the merge + # loads correctly; see _can_fuse_indexer_wk_weights_proj. if getattr(config, "model_type", None) == "glm_moe_dsa": if not ENABLE_GLM_FUSED_INDEXER: return False @@ -201,14 +202,13 @@ def _can_fuse_indexer_wk_weights_proj( return False if not _supports_fused_indexer_kernel_config(config): return False - # GLM-5.2: enable the qk-rope+quant+cache kernel fusion but keep indexer.wk / - # indexer.weights_proj as separate checkpoint tensors. GLM checkpoints name - # weights_proj "indexers_proj" and don't ship a pre-mergeable wk_weights_proj - # layout, so the GEMM-merge (which relies on packed_modules_mapping merging two - # tensors) is left off until the loader is validated for GLM. The dominant win - # (rope+quant+cache fusion) does not require this merge. - if getattr(config, "model_type", None) == "glm_moe_dsa": - return False + # GLM-5.2 (glm_moe_dsa) reuses the same indexer weight layout as DeepSeek-V3.2: + # separate indexer.wk (fp8 block-scale) + indexer.weights_proj (bf16). The + # "indexers_proj" name only appears in the HF quant config's modules_to_not_convert + # list (remapped via quant_exclude_name_mapping); the actual checkpoint tensors use + # the standard indexer.wk / indexer.weights_proj paths, which is exactly what the + # packed_modules_mapping merge and IndexerWkWeightsProjLinear's fp8-wk load expect. + # So GLM takes the same wk+weights_proj GEMM-merge path as V3.2 below. if quant_config is None: return True From 6b8259bdbfda78a13eac0838b972426c790f52e3 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Mon, 6 Jul 2026 09:48:05 +0800 Subject: [PATCH 08/17] add online quant for mxfp8 Signed-off-by: zejunchen-zejun --- atom/quantization/quark/utils.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/atom/quantization/quark/utils.py b/atom/quantization/quark/utils.py index f14410772e..060fd26c0f 100644 --- a/atom/quantization/quark/utils.py +++ b/atom/quantization/quark/utils.py @@ -164,6 +164,8 @@ def quant_weight_online( - MXFP4 (``dtypes.fp4x2``): use the aiter HIP kernel with ``Even`` round mode (:func:`quant_mxfp4_online_even`), matching the offline Quark kernel. + - MXFP8 (``per_1x32`` + ``dtypes.fp8``): fp8 weights with a per-32 block + scale. See the e8m0 note below for why ``scale_type`` must be forced. - FP8 (incl. ptpc_fp8 per-token / per-channel): use the aiter quant function resolved from ``get_hip_quant(online_quant_type)``. @@ -178,4 +180,17 @@ def quant_weight_online( if online_quant_dtype == dtypes.fp4x2: return quant_mxfp4_online_even(weight) quant_func = get_hip_quant(online_quant_type) + if online_quant_type == QuantType.per_1x32 and online_quant_dtype == dtypes.fp8: + # MXFP8: aiter's per_1x32 fp8 quantizer defaults scale_type=fp32 for + # backward compat, but that fp32 scale is NOT what we want here — + # Fp8MoEMethod.create_weights allocates an e8m0 (uint8) scale buffer for + # per_1x32, and the MXFP8 GEMM / flydsl MoE kernels only consume the + # e8m0 byte scale. Force scale_type=fp8_e8m0 so the produced scale + # matches the buffer and the kernel; DO NOT drop this argument or the + # scale silently reverts to fp32 and loading/inference breaks. + return quant_func( + weight, + quant_dtype=online_quant_dtype, + scale_type=dtypes.fp8_e8m0, + ) return quant_func(weight, quant_dtype=online_quant_dtype) From c03472140fcf3c58a8a6afad99582e19b3dc98e6 Mon Sep 17 00:00:00 2001 From: zejunchen-zejun Date: Mon, 6 Jul 2026 10:17:37 +0800 Subject: [PATCH 09/17] update recipe Signed-off-by: zejunchen-zejun --- recipes/GLM-5.md | 55 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/recipes/GLM-5.md b/recipes/GLM-5.md index 4c80151f10..2df5f49162 100644 --- a/recipes/GLM-5.md +++ b/recipes/GLM-5.md @@ -132,3 +132,58 @@ Reference numbers on 8×MI355X (TP8, FP8 weights, bf16 KV cache), using the benc | 8192 | 1024 | 1 | 73 | 669 | 409 | 13.2 | | 8192 | 1024 | 16 | 645 | 5818 | 418 | 23.3 | | 8192 | 1024 | 64 | 1210 | 10853 | 483 | 51.3 | + +## GLM-5.2 FP8 and MXFP4 Server Recipes + +Use the following docker image for these recipes: + +```bash +docker pull docker.io/rocm/atom-dev:nightly_202606301541 +``` + +Inside the container, install ATOM from the `zejun/opt_GLM5.2_0701` branch. AITER can remain unchanged. + +```bash +cd PATH_TO_ATOM +git checkout zejun/opt_GLM5.2_0701 +pip install -e . +``` + +### GLM-5.2 FP8 Server + +```bash +#!/bin/bash + +model_path=/shared/data/amd_int/models/GLM-5.2-FP8 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +TP=4 + +rm -rf /root/.cache/atom/* + +python -m atom.entrypoints.openai_server \ + --model "$model_path" \ + --server-port 8000 \ + --kv_cache_dtype fp8 \ + --no-enable_prefix_caching \ + --online_quant_config '{"layer_quant_config":{"model.layers.*.mlp.experts":"mxfp8"}}' \ + -tp $TP 2>&1 | tee server.log & +``` + +### GLM-5.2 MXFP4 Server + +```bash +#!/bin/bash + +model_path=/shared/data/amd_int/models/GLM-5.2-MXFP4 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +TP=4 + +rm -rf /root/.cache/atom/* + +python -m atom.entrypoints.openai_server \ + --model "$model_path" \ + --server-port 8000 \ + --kv_cache_dtype fp8 \ + --no-enable_prefix_caching \ + -tp $TP 2>&1 | tee server.log & +``` From 7986b0aa4ece8c94fa1ea61e5cda187aa1f74b1e Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Mon, 6 Jul 2026 03:36:24 +0000 Subject: [PATCH 10/17] add mtp recipe Signed-off-by: whx-sjtu --- recipes/GLM-5.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/recipes/GLM-5.md b/recipes/GLM-5.md index 2df5f49162..f8926ee679 100644 --- a/recipes/GLM-5.md +++ b/recipes/GLM-5.md @@ -187,3 +187,23 @@ python -m atom.entrypoints.openai_server \ --no-enable_prefix_caching \ -tp $TP 2>&1 | tee server.log & ``` + +### GLM-5.2 MXFP4 MTP Server +```bash +#!/bin/bash + +model_path=/shared/data/amd_int/models/GLM-5.2-MXFP4 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +TP=4 + +rm -rf /root/.cache/atom/* + +python -m atom.entrypoints.openai_server \ + --model "$model_path" \ + --server-port 8004 \ + --kv_cache_dtype fp8 \ + --no-enable_prefix_caching \ + --num-speculative-tokens 3 \ + --method mtp \ + -tp $TP 2>&1 | tee server_mtp.log +``` From d2c9f437cc5e5c72f69af181c29ba2485526ee2f Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Mon, 6 Jul 2026 09:17:50 +0000 Subject: [PATCH 11/17] fuse mtp mla prepare kernel Signed-off-by: whx-sjtu --- atom/model_ops/attentions/aiter_mla.py | 59 ++++++++++++++++++++++---- atom/utils/block_convert.py | 58 +++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 8 deletions(-) diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index b924ea0498..7569c58d6a 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -8,6 +8,7 @@ import numpy as np import torch +import triton from atom.utils import envs from aiter import ( decode_update_mla_metadata_v1, @@ -20,6 +21,7 @@ from atom.utils import CpuGpuBuffer from atom.utils.block_convert import ( kv_indices_generate_triton, + mtp_prepare_decode_mla_kernel, ) from atom.utils.forward_context import AttentionMetaData, Context @@ -99,6 +101,11 @@ def get_impl_cls() -> Type["MLAAttention"]: class AiterMLAMetadataBuilder(CommonAttentionBuilder): + # EagleProposer folds the per-draft-step position/context bump into + # prepare_mtp_decode's fused kernel when this is set (matches the MHA + # backend). The fused kernel handles both sparse and dense MLA. + fuse_mtp_decode_position_update = True + def __init__(self, model_runner): if envs.ATOM_MLA_PAGE_SIZE > 1: self.block_size = envs.ATOM_MLA_PAGE_SIZE @@ -111,6 +118,10 @@ def __init__(self, model_runner): f"got --block-size {model_runner.block_size}" ) CommonAttentionBuilder.__init__(self, model_runner) + # Single-program block for the fused MTP-decode metadata kernel. Sized + # to the max batch (runtime bs <= max_bs) so one tl.cumsum spans the + # whole batch in a single launch. + self._mtp_fuse_block = triton.next_power_of_2(self.max_bs + 1) config = model_runner.config hf_config = config.hf_config # `self.num_attention_heads` set by CommonAttentionBuilder.__init__. @@ -589,21 +600,53 @@ def prepare_mtp_decode( positions: torch.Tensor, # [total_tokens] int32 only_update: bool = False, num_reject_tokens: torch.Tensor = None, + *, + update_context_lens: bool = False, + positions_out: torch.Tensor | None = None, + last_token_indices: torch.Tensor | None = None, ): + """Per-draft-step MLA metadata update, fused into a single kernel. + + One ``_mtp_prepare_decode_mla_kernel`` launch performs, in place: + - ``kv_indptr += cu_seqlens_q`` (needed by kv_indices + slot_mapping), + - (sparse) per-seq ``min(kv_count, index_topk)`` cumsum -> + ``sparse_kv_indptr``, + - (fused position update) ``positions += 1`` when ``positions_out`` is + given, and ``context_lens += 1`` when ``update_context_lens`` is set. + + ``fuse_mtp_decode_position_update`` makes EagleProposer route the + per-step position/context bumps through here instead of launching them + as separate kernels. ``last_token_indices`` is accepted for signature + parity with the MHA backend but unused (MLA's ``positions`` is already + one entry per sequence at this point). + """ + del last_token_indices # MLA positions are already per-seq (1 per token) var = self.model_runner.forward_vars kv_indptr = var["kv_indptr"].gpu[: bs + 1] + cu_seqlens_q = var["cu_seqlens_q"].gpu[: bs + 1] if self.is_sparse: - # Update dense kv_indptr (needed for kv_indices generation and slot_mapping) - kv_indptr += var["cu_seqlens_q"].gpu[: bs + 1] - # Recompute sparse_kv_indptr: per-seq sparse count = min(dense_kv_count, index_topk) sparse_kv_indptr = var["sparse_kv_indptr"].gpu[: bs + 1] - kv_counts = kv_indptr[1 : bs + 1] - kv_indptr[:bs] - sparse_counts = torch.clamp(kv_counts, max=self.index_topk) - sparse_kv_indptr[0] = 0 - sparse_kv_indptr[1 : bs + 1] = torch.cumsum(sparse_counts, dim=0) else: assert self.block_size == 1 - kv_indptr += var["cu_seqlens_q"].gpu[: bs + 1] + sparse_kv_indptr = None + + update_positions = positions_out is not None + context_lens = var["context_lens"].gpu[:bs] if update_context_lens else None + + mtp_prepare_decode_mla_kernel[(1,)]( + kv_indptr, + cu_seqlens_q, + sparse_kv_indptr if self.is_sparse else kv_indptr, + positions_out if update_positions else kv_indptr, + context_lens if update_context_lens else kv_indptr, + bs, + self.index_topk if self.is_sparse else 0, + positions_out.stride(0) if update_positions else 1, + IS_SPARSE=self.is_sparse, + UPDATE_POSITIONS=update_positions, + UPDATE_CONTEXT_LENS=update_context_lens, + BLOCK=self._mtp_fuse_block, + ) kv_indices_generate_triton( var["block_tables"].gpu[:bs], diff --git a/atom/utils/block_convert.py b/atom/utils/block_convert.py index 7ebfa41f2e..590ac8f69f 100644 --- a/atom/utils/block_convert.py +++ b/atom/utils/block_convert.py @@ -213,6 +213,64 @@ def kv_indices_generate_triton( return kv_indices_convert +@triton.jit +def mtp_prepare_decode_mla_kernel( + kv_indptr_ptr, # int32 [bs+1] in/out: += cu_seqlens_q + cu_seqlens_q_ptr, # int32 [bs+1] in + sparse_kv_indptr_ptr, # int32 [bs+1] out (IS_SPARSE only) + positions_ptr, # int64 [bs] in/out (UPDATE_POSITIONS only) + context_lens_ptr, # int32 [bs] in/out (UPDATE_CONTEXT_LENS only) + bs, + index_topk, + position_stride, + IS_SPARSE: tl.constexpr, + UPDATE_POSITIONS: tl.constexpr, + UPDATE_CONTEXT_LENS: tl.constexpr, + BLOCK: tl.constexpr, +): + """Single-program fused per-draft-step MTP-decode metadata for MLA. + + Collapses the elementwise chain that used to launch ~4 tiny kernels + (``kv_indptr += cu_seqlens_q``; ``counts = diff(kv_indptr)``; + ``clamp(index_topk)``; ``cumsum -> sparse_kv_indptr``) plus eagle's two + per-step bumps (``positions += 1``, ``context_lens += 1``) into one launch. + One program owns the whole batch so the sparse cumsum is a single in-register + ``tl.cumsum``; the caller guarantees ``bs + 1 <= BLOCK`` (bs <= max_bs). All + reads happen before any write, so the in-place ``kv_indptr`` update is + hazard-free. ``sparse_kv_indptr[0]`` is written on-device (no Python-scalar + H2D copy, which is what caused the original hot-path Host-Device sync). + """ + offs = tl.arange(0, BLOCK) + mask_p1 = offs < (bs + 1) # indptr arrays are length bs+1 + mask_bs = offs < bs # per-seq arrays are length bs + + # Load originals first (into registers) so the kv_indptr store below cannot + # clobber the shifted [i+1] read used for per-seq counts. + kvi = tl.load(kv_indptr_ptr + offs, mask=mask_p1, other=0) + cuq = tl.load(cu_seqlens_q_ptr + offs, mask=mask_p1, other=0) + new_kvi = kvi + cuq + + if IS_SPARSE: + kvi_next = tl.load(kv_indptr_ptr + offs + 1, mask=mask_bs, other=0) + cuq_next = tl.load(cu_seqlens_q_ptr + offs + 1, mask=mask_bs, other=0) + counts = tl.where(mask_bs, (kvi_next + cuq_next) - new_kvi, 0) + clamped = tl.minimum(counts, index_topk) + # inclusive scan: csum[i] == sparse_kv_indptr[i + 1] + csum = tl.cumsum(clamped, axis=0) + tl.store(sparse_kv_indptr_ptr + offs + 1, csum, mask=mask_bs) + tl.store(sparse_kv_indptr_ptr + offs, tl.zeros_like(csum), mask=(offs == 0)) + + tl.store(kv_indptr_ptr + offs, new_kvi, mask=mask_p1) + + if UPDATE_POSITIONS: + pos = tl.load(positions_ptr + offs * position_stride, mask=mask_bs, other=0) + tl.store(positions_ptr + offs * position_stride, pos + 1, mask=mask_bs) + + if UPDATE_CONTEXT_LENS: + ctx = tl.load(context_lens_ptr + offs, mask=mask_bs, other=0) + tl.store(context_lens_ptr + offs, ctx + 1, mask=mask_bs) + + if __name__ == "__main__": # Example usage and test From 0089cefb1c9cf135fb757fa12990eb8232c59a89 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Mon, 6 Jul 2026 11:09:44 +0000 Subject: [PATCH 12/17] llm_embd not sharded Signed-off-by: whx-sjtu --- atom/config.py | 7 ++++++ atom/model_ops/embed_head.py | 20 ++++++++++++++- atom/models/deepseek_mtp.py | 28 ++++++++++++++++----- atom/models/deepseek_v2.py | 49 ++++++++++++++++++++++++++++++++---- atom/models/eagle3_llama.py | 4 +-- atom/spec_decode/eagle.py | 16 +++++++++--- atom/utils/envs.py | 17 +++++++++---- 7 files changed, 119 insertions(+), 22 deletions(-) diff --git a/atom/config.py b/atom/config.py index a651dd9299..a8707e91b0 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1288,6 +1288,13 @@ def compute_hash(self) -> str: ), ) ) + # Vocab-embedding replication (ATOM_REPLICATE_VOCAB_EMBED) changes both the + # embed weight shape ([vocab] vs [vocab/tp]) and the embed op (local + # F.embedding vs masked-embedding + all-reduce), so it alters the compiled + # graph and MUST be part of its key. Without this, toggling the flag — or + # deploying it on top of a cache built with the other setting — reuses a + # stale artifact and trips assert_size_stride at runtime. + factors.append(bool(envs.ATOM_REPLICATE_VOCAB_EMBED)) hash_str = hashlib.md5( str(factors).encode(), usedforsecurity=False diff --git a/atom/model_ops/embed_head.py b/atom/model_ops/embed_head.py index 15a8505c97..f0152a432d 100644 --- a/atom/model_ops/embed_head.py +++ b/atom/model_ops/embed_head.py @@ -103,6 +103,24 @@ def masked_embedding( return _masked_embedding_launcher(x, weight, vocab_start_idx, vocab_end_idx) +def _replicated_embedding_fake(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return torch.empty( + x.numel(), + weight.shape[1], + dtype=weight.dtype, + device=weight.device, + ) + + +@torch_compile_guard(gen_fake=_replicated_embedding_fake) +def replicated_embedding(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + # Keep the lookup opaque to torch.compile: inductor otherwise fuses the + # embedding gather into the surrounding graph, which corrupts the MTP draft + # rollout (acceptance collapses ~69%->45%) — the same reason + # VocabParallelEmbedding routes through the masked_embedding custom op. + return F.embedding(x, weight) + + class VocabParallelEmbedding(nn.Module): def __init__( @@ -184,7 +202,7 @@ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): param.data.copy_(loaded_weight) def forward(self, x: torch.Tensor): - return F.embedding(x, self.weight) + return replicated_embedding(x, self.weight) class ParallelLMHead(VocabParallelEmbedding): diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index dff057eb8e..d8efa9246b 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -6,7 +6,11 @@ import torch.nn as nn from aiter.dist.communication_op import tensor_model_parallel_all_reduce from atom.config import Config, QuantizationConfig -from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding +from atom.model_ops.embed_head import ( + ParallelLMHead, + ReplicatedEmbedding, + VocabParallelEmbedding, +) from atom.model_ops.layernorm import RMSNorm from atom.model_ops.linear import ReplicatedLinear from atom.model_ops.moe import FusedMoE @@ -15,7 +19,11 @@ from atom.utils.decorators import support_torch_compile from transformers import DeepseekV2Config, DeepseekV3Config, PretrainedConfig -from .deepseek_v2 import DeepseekV2DecoderLayer, _can_fuse_indexer_wk_weights_proj +from .deepseek_v2 import ( + DeepseekV2DecoderLayer, + _can_fuse_indexer_wk_weights_proj, + use_replicated_vocab_embed, +) from .utils import ckpt_has_tensor_suffix, maybe_prefix @@ -136,10 +144,18 @@ def __init__( ) } ) - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - ) + if use_replicated_vocab_embed(config): + # GLM-5.2 MTP: full table per rank, no post-embedding all-reduce. + # (Shared with the target's replicated embed by EagleProposer at load.) + self.embed_tokens = ReplicatedEmbedding( + config.vocab_size, + config.hidden_size, + ) + else: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) def forward( self, diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 3648510797..81277c22cc 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -60,7 +60,11 @@ triton_gather_kv_indices_sparse, ) from atom.model_ops.base_attention import Attention -from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding +from atom.model_ops.embed_head import ( + ParallelLMHead, + ReplicatedEmbedding, + VocabParallelEmbedding, +) from atom.model_ops.layernorm import LayerNorm, RMSNorm from atom.model_ops.linear import ( ColumnParallelLinear, @@ -90,6 +94,7 @@ from atom.model_ops import module_dispatch_ops as _module_dispatch_ops # noqa: F401 from atom.utils.decorators import mark_trace, support_torch_compile from atom.utils.forward_context import get_forward_context +from atom.plugin import is_plugin_mode from atom.plugin.vllm.attention.layer_sparse_mla import ( IndexerDecoratorForPluginMode, DeepseekV32IndexerCacheDecoratorForPluginMode, @@ -2261,6 +2266,33 @@ def forward( return hidden_states, residual +def use_replicated_vocab_embed(config: PretrainedConfig) -> bool: + """Whether to hold the full vocab embedding on every TP rank (local lookup, + no post-embedding all-reduce) instead of a ``VocabParallelEmbedding`` shard. + + Enabled by default (gated by ``ATOM_REPLICATE_VOCAB_EMBED``) for GLM-5.2 + (``glm_moe_dsa``) — both the main model and its MTP draft — whose embedding is + independent of the still TP-sharded ``lm_head`` (``tie_word_embeddings=False``), + so the lookup is bit-identical to the sharded masked-embedding + all-reduce + path. Skipped in plugin mode (vLLM/SGLang/RTP own the embedding lifecycle) and + whenever the embedding is tied to the sharded head. + + The GLM-5.2 main model keeps ``model_type == "glm_moe_dsa"``; its MTP draft + config has ``model_type`` rewritten to ``"deepseek_mtp"`` (see + ``SpeculativeConfig.hf_config_override``) but still carries the GLM-only + ``index_share_for_mtp_iteration`` flag, so we detect either. + """ + if not envs.ATOM_REPLICATE_VOCAB_EMBED: + return False + if is_plugin_mode(): + return False + if getattr(config, "tie_word_embeddings", False): + return False + return getattr(config, "model_type", None) == "glm_moe_dsa" or bool( + getattr(config, "index_share_for_mtp_iteration", False) + ) + + @support_torch_compile class DeepseekV2Model(nn.Module): def __init__( @@ -2281,10 +2313,17 @@ def __init__( self.is_v32 = hasattr(config, "index_topk") if get_pp_group().is_first_rank: - self.embed_tokens = VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - ) + if use_replicated_vocab_embed(config): + # GLM-5.2: full table per rank, no post-embedding all-reduce. + self.embed_tokens = ReplicatedEmbedding( + config.vocab_size, + config.hidden_size, + ) + else: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) else: self.embed_tokens = PPMissingLayer() diff --git a/atom/models/eagle3_llama.py b/atom/models/eagle3_llama.py index 0aac1f2f67..3f8380fb42 100644 --- a/atom/models/eagle3_llama.py +++ b/atom/models/eagle3_llama.py @@ -340,8 +340,8 @@ def __init__(self, atom_config: Config, prefix: str = "", layer_offset: int = 0) # full on every rank — a local lookup with no post-embedding all-reduce. # Bit-identical to the sharded path; on by default (trades memory for one # fewer collective per draft step). Falls back to the sharded embedding - # when ATOM_EAGLE_REPLICATE_EMBED=0. - if envs.ATOM_EAGLE_REPLICATE_EMBED: + # when ATOM_REPLICATE_VOCAB_EMBED=0. + if envs.ATOM_REPLICATE_VOCAB_EMBED: self.embed_tokens = ReplicatedEmbedding( config.vocab_size, config.hidden_size ) diff --git a/atom/spec_decode/eagle.py b/atom/spec_decode/eagle.py index 7a921de241..70dd580a08 100644 --- a/atom/spec_decode/eagle.py +++ b/atom/spec_decode/eagle.py @@ -324,11 +324,21 @@ def load_model(self, target_model: nn.Module) -> None: self.model.share_with_target(target_base, loaded) return - # Share embed_tokens with the target model + # Share embed_tokens with the target model. Match on the *logical* vocab + # (num_embeddings) and hidden dim rather than the stored weight shape, so a + # replicated target embed ([vocab, hidden], ATOM_REPLICATE_VOCAB_EMBED) is + # still shared onto a TP-sharded draft embed ([vocab/tp, hidden]) — the + # draft then reuses the target's replicated table (no post-embed + # all-reduce). When both are sharded this is identical to the old check. + draft_embed = self.model.model.embed_tokens + target_embed = target_base.model.embed_tokens + draft_vocab = getattr(draft_embed, "num_embeddings", None) + target_vocab = getattr(target_embed, "num_embeddings", None) if ( get_pp_group().world_size == 1 - and self.model.model.embed_tokens.weight.shape - == target_base.model.embed_tokens.weight.shape + and draft_vocab is not None + and draft_vocab == target_vocab + and draft_embed.weight.shape[1] == target_embed.weight.shape[1] ): logger.info( "Assuming the EAGLE head shares the same vocab embedding" diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 77b2989140..2aa52daffe 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -81,11 +81,18 @@ "ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION": lambda: ( os.getenv("ATOM_ENABLE_ALLREDUCE_RMSNORM_FUSION", "1") == "1" ), - # Replicate the EAGLE3 draft vocab embedding on every TP rank (full table per - # rank, local lookup) instead of sharding it — eliminates the post-embedding - # all-reduce. The draft embed is independent of the (sharded) lm_head. - "ATOM_EAGLE_REPLICATE_EMBED": lambda: ( - os.getenv("ATOM_EAGLE_REPLICATE_EMBED", "1") == "1" + # Replicate the vocab embedding on every TP rank (full table per rank, purely + # local lookup) instead of TP-sharding it — eliminates the post-embedding + # all-reduce. Applies to BOTH the main/target model and the speculative draft + # (EAGLE3 head + MTP draft steps), so the collective is dropped on every embed + # on the critical path. Only used where the embedding is independent of the + # still TP-sharded lm_head (EAGLE3 draft; GLM-5.2 target+MTP, whose + # tie_word_embeddings=False), so the lookup is bit-identical to the sharded + # masked-embedding + all-reduce path. Trades (tp-1)/tp of the embedding's + # memory per rank for one fewer collective per embed. Default on; set "0" to + # fall back to the sharded VocabParallelEmbedding. + "ATOM_REPLICATE_VOCAB_EMBED": lambda: ( + os.getenv("ATOM_REPLICATE_VOCAB_EMBED", "1") == "1" ), "ATOM_ENABLE_GDN_DECODE_LOSSY_FAST": lambda: ( os.getenv("ATOM_ENABLE_GDN_DECODE_LOSSY_FAST", "0").lower() == "1" From cd3eaf8c524e8be6c2c045dc579a9563da1e5a6c Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Mon, 6 Jul 2026 12:26:48 +0000 Subject: [PATCH 13/17] argmax opt Signed-off-by: whx-sjtu --- atom/models/deepseek_mtp.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index d8efa9246b..6cae32d4c8 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -186,6 +186,24 @@ def compute_logits( logits = mtp_layer.shared_head.head(mtp_layer.shared_head(hidden_states)) return logits + def compute_draft_token( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + """Greedy draft token via distributed argmax over the TP-sharded vocab — + avoids all-gathering the full [N, vocab] logits every draft step. + + Mirrors compute_logits() (same norm + shared head), but reduces each + rank's logit shard to (max_val, global_idx) and all-gathers only [N, 2] + instead of the O(vocab) logits. Token-identical to + compute_logits(...).argmax(-1). + """ + current_step_idx = spec_step_idx % self.num_mtp_layers + mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)] + normed = mtp_layer.shared_head(hidden_states) + return mtp_layer.shared_head.head.compute_argmax_token(normed) + @support_torch_compile class DeepSeekMTP(nn.Module): @@ -275,6 +293,20 @@ def compute_logits( ) -> torch.Tensor | None: return self.model.compute_logits(hidden_states, spec_step_idx) + def compute_draft_token( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + """Distributed greedy argmax for the MTP draft rollout (GLM-5.2). + + EagleProposer picks this over compute_logits().argmax(-1) when present + (``_draft_argmax_fused``), so the draft never all-gathers the full + [N, vocab] logits — it all-gathers only the packed [N, 2] per-rank + reductions. See DeepSeekMultiTokenPredictor.compute_draft_token. + """ + return self.model.compute_draft_token(hidden_states, spec_step_idx) + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) From b5db6bee4adc8da0ef445b014e30d9eaa88366f0 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Mon, 6 Jul 2026 14:34:50 +0000 Subject: [PATCH 14/17] fuse 2 norm Signed-off-by: whx-sjtu --- atom/model_ops/layernorm.py | 103 ++++++++++++++++++++++++++++++++++++ atom/models/deepseek_mtp.py | 18 ++++--- 2 files changed, 114 insertions(+), 7 deletions(-) diff --git a/atom/model_ops/layernorm.py b/atom/model_ops/layernorm.py index 0367630d0b..1162ed3fc6 100644 --- a/atom/model_ops/layernorm.py +++ b/atom/model_ops/layernorm.py @@ -1054,6 +1054,109 @@ def __call__( ) +# --------------------------------------------------------------------------- +# Fused dual RMSNorm + concat — single Triton launch. +# +# Two independent RMSNorms over the SAME hidden dim (different inputs, different +# weights, shared eps) whose bf16 results are concatenated on the last axis: +# out = cat([rmsnorm(xe, we), rmsnorm(xh, wh)], dim=-1) # [M, 2H] +# One program per row does both norms and writes each into its half of the +# [M, 2H] output, so the concat is free — it eliminates the separate cat +# kernel's read+write of 2*M*H (halving traffic vs two norms + torch.cat) and +# folds three launches (enorm, hnorm, cat) into one. Used by the DeepSeek/GLM +# MTP eh_proj input (enorm(embed) ++ hnorm(prev_hidden)). +# --------------------------------------------------------------------------- + + +@triton.jit +def _fused_dual_rmsnorm_cat_kernel( + xe_ptr, + xh_ptr, + we_ptr, + wh_ptr, + out_ptr, + H, + stride_xe_m, + stride_xh_m, + stride_out_m, + eps, + BLOCK_H: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_H) + mask = cols < H + + # enorm half -> out[row, :H] + xe = tl.load(xe_ptr + row * stride_xe_m + cols, mask=mask, other=0.0).to(tl.float32) + we = tl.load(we_ptr + cols, mask=mask, other=0.0).to(tl.float32) + rstd_e = tl.rsqrt(tl.sum(xe * xe, axis=-1) / H + eps) + out_e = (xe * rstd_e * we).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + row * stride_out_m + cols, out_e, mask=mask) + + # hnorm half -> out[row, H:2H] + xh = tl.load(xh_ptr + row * stride_xh_m + cols, mask=mask, other=0.0).to(tl.float32) + wh = tl.load(wh_ptr + cols, mask=mask, other=0.0).to(tl.float32) + rstd_h = tl.rsqrt(tl.sum(xh * xh, axis=-1) / H + eps) + out_h = (xh * rstd_h * wh).to(out_ptr.dtype.element_ty) + tl.store(out_ptr + row * stride_out_m + H + cols, out_h, mask=mask) + + +def _fused_dual_rmsnorm_cat_fake( + xe: torch.Tensor, + we: torch.Tensor, + xh: torch.Tensor, + wh: torch.Tensor, + eps: float, +) -> torch.Tensor: + M, H = xe.shape + return torch.empty((M, 2 * H), dtype=xe.dtype, device=xe.device) + + +@torch_compile_guard(gen_fake=_fused_dual_rmsnorm_cat_fake) +def fused_dual_rmsnorm_cat( + xe: torch.Tensor, + we: torch.Tensor, + xh: torch.Tensor, + wh: torch.Tensor, + eps: float, +) -> torch.Tensor: + """cat([rmsnorm(xe, we), rmsnorm(xh, wh)], dim=-1) in one Triton launch. + + Args: + xe, xh: (M, H) bf16/fp16 inputs (same shape). + we, wh: (H,) RMSNorm weights (one per input). + eps: shared RMSNorm epsilon. + Returns: + (M, 2H) tensor: [:, :H] = rmsnorm(xe, we), [:, H:] = rmsnorm(xh, wh). + """ + assert xe.shape == xh.shape, f"shape mismatch {xe.shape} vs {xh.shape}" + assert xe.dim() == 2, f"expected 2-D inputs, got {xe.dim()}-D" + xe = xe.contiguous() + xh = xh.contiguous() + M, H = xe.shape + out = torch.empty((M, 2 * H), dtype=xe.dtype, device=xe.device) + if M == 0: + return out + BLOCK_H = triton.next_power_of_2(H) + num_warps = 8 if BLOCK_H >= 4096 else 4 + _fused_dual_rmsnorm_cat_kernel[(M,)]( + xe, + xh, + we, + wh, + out, + H, + xe.stride(0), + xh.stride(0), + out.stride(0), + eps, + BLOCK_H=BLOCK_H, + num_warps=num_warps, + num_stages=2, + ) + return out + + @torch_compile_guard() def layernorm2d_fwd_( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float, dim: int diff --git a/atom/models/deepseek_mtp.py b/atom/models/deepseek_mtp.py index 6cae32d4c8..ff6f993236 100644 --- a/atom/models/deepseek_mtp.py +++ b/atom/models/deepseek_mtp.py @@ -11,7 +11,7 @@ ReplicatedEmbedding, VocabParallelEmbedding, ) -from atom.model_ops.layernorm import RMSNorm +from atom.model_ops.layernorm import RMSNorm, fused_dual_rmsnorm_cat from atom.model_ops.linear import ReplicatedLinear from atom.model_ops.moe import FusedMoE from atom.models.utils import IntermediateTensors @@ -95,13 +95,17 @@ def forward( spec_step_index: int = 0, ) -> torch.Tensor: assert inputs_embeds is not None - masked_inputs_embeds = inputs_embeds - inputs_embeds = self.enorm(masked_inputs_embeds) - previous_hidden_states = self.hnorm(previous_hidden_states) - - hidden_states = self.eh_proj( - torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + # Fused enorm(inputs_embeds) ++ hnorm(previous_hidden_states) in a single + # Triton launch (folds the two RMSNorms + the torch.cat; enorm and hnorm + # share eps=rms_norm_eps). bf16-identical to the separate path. + eh_input = fused_dual_rmsnorm_cat( + inputs_embeds, + self.enorm.weight, + previous_hidden_states, + self.hnorm.weight, + self.enorm.eps, ) + hidden_states = self.eh_proj(eh_input) hidden_states, residual = self.mtp_block( positions=positions, hidden_states=hidden_states, residual=None From 4b96233285b6a8ee072bb5917b8ff741ab7e286e Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Tue, 7 Jul 2026 09:38:33 +0000 Subject: [PATCH 15/17] update recipe to use online quant for mtp Signed-off-by: whx-sjtu --- recipes/GLM-5.md | 1 + 1 file changed, 1 insertion(+) diff --git a/recipes/GLM-5.md b/recipes/GLM-5.md index f8926ee679..77d9576315 100644 --- a/recipes/GLM-5.md +++ b/recipes/GLM-5.md @@ -205,5 +205,6 @@ python -m atom.entrypoints.openai_server \ --no-enable_prefix_caching \ --num-speculative-tokens 3 \ --method mtp \ + --online_quant_config '{"layer_quant_config": {"*layers.78*": "ptpc_fp8"}, "exclude_layer": ["*.gate", "*shared_head*", "*embed*"]}' \ -tp $TP 2>&1 | tee server_mtp.log ``` From c8ad6827e42405d9e26e5a79778ee30b98ad533b Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Wed, 8 Jul 2026 06:27:52 +0000 Subject: [PATCH 16/17] enable pcp for DSA Signed-off-by: whx-sjtu --- atom/config.py | 9 + atom/model_engine/llm_engine.py | 32 +++ atom/model_engine/model_runner.py | 4 + atom/model_ops/attention_mla.py | 91 ++++++- atom/model_ops/attentions/aiter_mla.py | 119 +++++++++ atom/models/deepseek_v2.py | 326 ++++++++++++++++++++++++- atom/spec_decode/eagle.py | 51 +++- recipes/GLM-5.md | 54 ++++ 8 files changed, 671 insertions(+), 15 deletions(-) diff --git a/atom/config.py b/atom/config.py index a8707e91b0..b7db2a6243 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1262,6 +1262,15 @@ def compute_hash(self) -> str: factors.append(vllm_factors) factors.append(self.tensor_parallel_size) + # PCP changes the compiled graph: when pcp>1 the indexer runs through the + # opaque `indexer_with_output` op (whose identity output is fed as the MLA + # query) and the indexer takes the round-robin all-gather / separate-rope + # path. A pcp1 vs pcp2 run over the same model+source otherwise hashes + # identically, so without this factor pcp2 loads pcp1's cached artifact + # (no indexer op) and trips copy_misaligned_inputs / assert_size_stride at + # runtime — the same stale-artifact hazard documented for the vocab-embed + # flag below. + factors.append(self.prefill_context_parallel_size) factors.append(self.enable_dp_attention) text_config = getattr(self.hf_config, "text_config", self.hf_config) factors.append( diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index a20b019dc2..a874db3d69 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -84,6 +84,38 @@ def __init__(self, model, tokenizer=None, **kwargs): "independently (double-split) and corrupt the output. For now, " "disable one of them." ) + # PCP for the native ATOM engine is implemented only on the DeepSeek-V3.2 + # / GLM-5.2 sparse-MLA (DSA) path (indexer + sparse attention). Those + # configs expose `index_topk`. Non-sparse models (dense DeepSeek-V2, + # non-DSA archs) have no PCP wiring, so reject pcp>1 for them with a clear + # message instead of silently running the un-split path. + if config.prefill_context_parallel_size > 1 and not hasattr( + config.hf_config, "index_topk" + ): + raise ValueError( + "prefill_context_parallel_size > 1 (-pcp) is currently only " + "supported for sparse-MLA / DSA models (DeepSeek-V3.2, GLM-5.2 — " + "configs with `index_topk`). The requested model " + f"({getattr(config.hf_config, 'model_type', 'unknown')}) is not a " + "DSA model. Run it with -pcp 1." + ) + # PCP v1 requires chunked prefill OFF: the round-robin token split assumes + # the whole prefill sequence is present in one forward (it pads the global + # token count to a multiple of pcp and needs each query's full history). + if config.prefill_context_parallel_size > 1 and config.enable_chunked_prefill: + raise ValueError( + "prefill_context_parallel_size > 1 (-pcp) requires chunked " + "prefill disabled in this release. Pass " + "--no-enable_chunked_prefill." + ) + # PCP v1 requires prefix caching OFF: mixing a full cached KV prefix with + # 1/pcp new tokens in the indexer gather is not yet handled. + if config.prefill_context_parallel_size > 1 and config.enable_prefix_caching: + raise ValueError( + "prefill_context_parallel_size > 1 (-pcp) requires prefix " + "caching disabled in this release. Pass " + "--no-enable_prefix_caching." + ) self.rquest_ids = set() self.io_processor = InputOutputProcessor( config, self.tokenizer, config.kv_cache_block_size diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 582926f306..2948590ba8 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -739,6 +739,10 @@ def __init__(self, rank: int, config: Config): torch.set_default_device("cpu") torch.set_default_dtype(default_dtype) + # PCP is compile-safe: its runtime-varying branches live inside opaque + # splitting ops (indexer_with_output / unified_attention_with_output_base) + # that run eager, so Dynamo never bakes `_pcp_active()` to its dummy-warmup + # value. No PCP-specific compile guard needed here. if self.config.compilation_config.level == 1: self.model = torch.compile(self.model, fullgraph=True, backend="eager") if hasattr(self, "drafter"): diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index ea846ebed3..de4f9c809d 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -40,6 +40,11 @@ ) from aiter.ops.triton.gather_kv_b_proj import gather_kv_b_proj from atom.config import get_current_atom_config +from atom.distributed.pcp_utils import ( + get_pcp_world_size, + pcp_allgather_rerange, + pcp_is_enabled, +) from atom.model_ops.linear import use_triton_gemm from atom.model_ops.utils import get_and_maybe_dequant_weights from atom.utils import envs @@ -1118,6 +1123,46 @@ def _forward_decode( return self._v_up_proj_and_o_proj(o) + def _pcp_write_full_kv(self, kv_cache, k_nope, k_rope, slot_mapping): + """Write an already-roped full k (kv_lora + rope) into the k-cache. + + Used by the PCP prefill path to materialise the full sequence's KV after + the fused MLA kernel produced q_out on 1/pcp queries. Mirrors the + non-fused k-writes used by the dense (`not use_prefill_mla`) prefill + branch so the physical cache layout matches exactly. `k_rope` must + already be rotary-embedded. + """ + if envs.ATOM_USE_TRITON_MLA and envs.ATOM_USE_TRITON_MLA_SHUFFLE_KV: + shuffled_cache = self._shuffled_kv_view(kv_cache) + triton_cat_and_cache_mla( + k_nope.view(-1, self.num_kv_heads, self.kv_lora_rank), + k_rope.view(-1, self.num_kv_heads, self.qk_rope_head_dim), + shuffled_cache, + slot_mapping.flatten(), + self._k_scale, + apply_scale=True, + shuffled_kv_cache=True, + ) + elif self.use_seg_mla: + kv_cache_seg = self._seg_kv_cache_view(kv_cache) + concat_and_cache_mla_seg( + k_nope, + k_rope.squeeze(1), + kv_cache_seg, + slot_mapping.flatten(), + kv_cache_dtype=self.kv_cache_dtype, + scale=self._k_scale, + ) + else: + concat_and_cache_mla( + k_nope, + k_rope.squeeze(1), + kv_cache, + slot_mapping.flatten(), + kv_cache_dtype=self.kv_cache_dtype, + scale=self._k_scale, + ) + def forward_impl( self, q: torch.Tensor, @@ -1208,6 +1253,30 @@ def forward_impl( else: q_nope, q_rope = self._q_proj_and_k_up_proj(q, x_scale=q_scale) + # ---- Prefill Context Parallel -------------------------------- + # q is this rank's 1/pcp queries, so q_out is naturally 1/pcp. But + # the k-cache must hold the FULL sequence (every rank keeps full KV). + # The fused MLA kernel below couples q_out with the k-write on one + # token count, so under PCP it runs on the owned slots (q_out is + # correct; its k-write is throwaway) and the full k-cache is written + # afterwards from the all-gathered k. Gather the raw (un-roped) k and + # key positions BEFORE the fused kernel ropes k in place. + pcp = ( + pcp_is_enabled() + and context.is_prefill + and not context.is_dummy_run + and use_prefill_mla + ) + if pcp: + pcp_ws = get_pcp_world_size() + n_real = attn_metadata.slot_mapping.shape[0] + k_nope_full = pcp_allgather_rerange(k_nope, pcp_ws)[:n_real] + k_rope_full = pcp_allgather_rerange(k_rope, pcp_ws)[:n_real] + positions_full = pcp_allgather_rerange(positions, pcp_ws)[:n_real] + write_slot_mapping = attn_metadata.slot_mapping_owned + else: + write_slot_mapping = attn_metadata.slot_mapping + if self.use_seg_mla: # Seg path: allocate q_out with a padded last dim so each head row # has a 768-byte stride (required by the gfx1250 decode asm). The @@ -1241,7 +1310,7 @@ def forward_impl( k_nope.view(-1, self.num_kv_heads, self.kv_lora_rank), k_rope.view(-1, self.num_kv_heads, self.qk_rope_head_dim), shuffled_cache, - attn_metadata.slot_mapping, + write_slot_mapping, positions, self.rotary_emb.cos_cache, self.rotary_emb.sin_cache, @@ -1262,7 +1331,7 @@ def forward_impl( # Flat seg layout: [num_blocks, page_size*(kv_lora + pe)]. kv_cache_seg, q_out, - attn_metadata.slot_mapping, + write_slot_mapping, self._k_scale, self._q_scale, positions, @@ -1282,7 +1351,7 @@ def forward_impl( self.kv_lora_rank + self.qk_rope_head_dim, ), q_out, - attn_metadata.slot_mapping, + write_slot_mapping, self._k_scale, self._q_scale, positions, @@ -1293,6 +1362,22 @@ def forward_impl( ) # q_out = self.fused_kv_bmm(q, q_scale, k_nope, k_rope, positions, kv_cache, attn_metadata) + if pcp: + # Complete the full k-cache: rope the gathered full k (in + # place) then write every real slot, overwriting the fused + # kernel's throwaway owned-slot write. The rope kernel is + # 2-component and needs a non-None partner, so pass a + # throwaway query of matching length. + self.rotary_emb( + positions_full, k_rope_full, torch.empty_like(k_rope_full) + ) + self._pcp_write_full_kv( + kv_cache, + k_nope_full, + k_rope_full, + attn_metadata.slot_mapping, + ) + if context.is_prefill: output = self._forward_prefill_mla(q_out, kv_cache, attn_metadata) else: diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 7569c58d6a..054cee8eae 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -16,6 +16,13 @@ get_mla_metadata_info_v1, get_mla_metadata_v1, ) +from atom.distributed.pcp_utils import ( + get_pcp_world_size, + pcp_is_enabled, + pcp_pad_dense, + pcp_pad_len, + pcp_round_robin_query_indices, +) from atom.model_engine.scheduler import ScheduledBatch from atom.model_ops.attention_mla import _MLA_MIN_HEADS, MLAAttention from atom.utils import CpuGpuBuffer @@ -914,6 +921,17 @@ def prepare_prefill(self, batch: ScheduledBatch): "sparse_prefill_reduce_partial_map" ] + # ---- Prefill Context Parallel: shrink per-query sparse metadata -- + # to this rank's 1/pcp round-robin queries. Gate on + # `not batch.is_dummy_run` so the reindex stays in lock-step with the + # model's round-robin token split (ForCausalLM._pcp_active() also + # skips dummy/warmup). Per-sequence + KV-write fields (slot_mapping, + # block_tables, cu_seqlens_q/k) stay FULL — every rank keeps full KV. + if pcp_is_enabled() and not batch.is_dummy_run: + self._apply_pcp_reindex( + attn_metadata, sum_scheduled_tokens, sparse_counts + ) + if hasattr(self.model_runner, "drafter") or attn_metadata.has_cached: # Populate kv_last_page_lens for full sequence (needed for MLA prefill with # prefix cache; decode does the same) @@ -1050,6 +1068,107 @@ def _build_mla_chunk_meta( v_workspace=self.v_chunk_workspace, ) + def _apply_pcp_reindex( + self, + attn_metadata: AttentionMetaData, + sum_scheduled_tokens: int, + sparse_counts: np.ndarray, + ) -> None: + """Reduce the per-query sparse-prefill metadata to this PCP rank's + 1/pcp round-robin queries. + + Prefill Context Parallel round-robin splits the token sequence so each + rank runs the model on 1/pcp of the query tokens while still keeping the + FULL KV. Only *query-indexed* metadata shrinks here; *per-sequence* and + *KV-write* fields (slot_mapping, block_tables, cu_seqlens_q/k) stay full + so the full k-cache is still written and gathered. + + The global token count is padded to a multiple of pcp_size; the extra + (dummy) queries get zero-length KV (they attend nothing and their hidden + output is dropped after the model's final all-gather + unpad). + """ + device = self.device + pcp_ws = get_pcp_world_size() + s_real = int(sum_scheduled_tokens) + padded_total = pcp_pad_len(s_real, pcp_ws) + n_pad = padded_total - s_real + owned_q = pcp_round_robin_query_indices(padded_total, pcp_ws).to(device) + n_owned = int(owned_q.shape[0]) + + # --- dense per-query fields: pad with zeros (dummy query -> 0), select. + # cu_seqlen_ks/ke become 0/0 for dummies == empty logits row. + ks_padded = pcp_pad_dense(attn_metadata.cu_seqlen_ks, n_pad) + attn_metadata.cu_seqlen_ks = ks_padded[owned_q].contiguous() + ke_padded = pcp_pad_dense(attn_metadata.cu_seqlen_ke, n_pad) + attn_metadata.cu_seqlen_ke = ke_padded[owned_q].contiguous() + t2s_padded = pcp_pad_dense(attn_metadata.token_to_seq_idxs, n_pad) + attn_metadata.token_to_seq_idxs = t2s_padded[owned_q].contiguous() + + # --- one query per row (incl dummies) -> sparse_cu_seqlens_q = arange. + attn_metadata.sparse_cu_seqlens_q = torch.arange( + n_owned + 1, dtype=torch.int32, device=device + ) + + # --- sparse_kv_indptr: cumsum of min(sparse_counts, topk); dummy -> 0. + sparse_counts_t = torch.as_tensor(sparse_counts, device=device) + owned_counts = pcp_pad_dense(sparse_counts_t, n_pad)[owned_q].to(torch.int64) + owned_counts = torch.clamp(owned_counts, max=self.index_topk) + indptr_owned = torch.zeros(n_owned + 1, dtype=torch.int32, device=device) + indptr_owned[1:] = torch.cumsum(owned_counts, 0).to(torch.int32) + attn_metadata.sparse_kv_indptr = indptr_owned + + # --- sparse kv_last_page_lens: one page per owned query (all 1s). + attn_metadata.kv_last_page_lens = torch.ones( + n_owned, dtype=torch.int32, device=device + ) + + # --- rebuild the sparse-prefill work buffers for the owned queries. + var = self.model_runner.forward_vars + get_mla_metadata_v1( + attn_metadata.sparse_cu_seqlens_q, + attn_metadata.sparse_kv_indptr, + attn_metadata.kv_last_page_lens, + self.padded_num_attention_heads, + 1, # nhead_kv + True, + var["sparse_prefill_work_meta_data"], + var["sparse_prefill_work_info_set"], + var["sparse_prefill_work_indptr"], + var["sparse_prefill_reduce_indptr"], + var["sparse_prefill_reduce_final_map"], + var["sparse_prefill_reduce_partial_map"], + page_size=self.block_size, + dtype_q=self.dtype_q, + dtype_kv=self.dtype_kv, + kv_granularity=max(self.block_size, 16), + max_seqlen_qo=1, + uni_seqlen_qo=1, + fast_mode=1, + max_split_per_batch=16, + ) + attn_metadata.sparse_prefill_work_meta_data = var[ + "sparse_prefill_work_meta_data" + ] + attn_metadata.sparse_prefill_work_info_set = var["sparse_prefill_work_info_set"] + attn_metadata.sparse_prefill_work_indptr = var["sparse_prefill_work_indptr"] + attn_metadata.sparse_prefill_reduce_indptr = var["sparse_prefill_reduce_indptr"] + attn_metadata.sparse_prefill_reduce_final_map = var[ + "sparse_prefill_reduce_final_map" + ] + attn_metadata.sparse_prefill_reduce_partial_map = var[ + "sparse_prefill_reduce_partial_map" + ] + + # --- owned slot_mapping for the fused q_out kernel in MLAAttention. The + # fused MLA kernel that produces q_out also writes k to these slots; + # that write is throwaway (the full-KV completion write in + # MLAAttention overwrites every real slot). Dummy queries clamp to + # the last real slot so they can never touch an unrelated slot. + owned_clamped = torch.clamp(owned_q, max=max(s_real - 1, 0)) + attn_metadata.slot_mapping_owned = attn_metadata.slot_mapping[ + owned_clamped + ].contiguous() + def prepare_decode(self, batch: ScheduledBatch, bs: int): scheduled_bs = batch.total_seqs_num_decode dropout_p = 0.0 diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 81277c22cc..e8b0c89e42 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -92,6 +92,14 @@ # shared with deepseek_v4. DeepseekV2MoE.forward dispatches via this op when # `_use_dual_stream` is True so torch.compile/Dynamo treats stream code as opaque. from atom.model_ops import module_dispatch_ops as _module_dispatch_ops # noqa: F401 +from atom.distributed.pcp_utils import ( + get_pcp_world_size, + pcp_allgather_rerange, + pcp_is_enabled, + pcp_pad_dense, + pcp_pad_len, + pcp_round_robin_split, +) from atom.utils.decorators import mark_trace, support_torch_compile from atom.utils.forward_context import get_forward_context from atom.plugin import is_plugin_mode @@ -143,6 +151,94 @@ ) +def _pcp_active() -> bool: + """True when Prefill Context Parallel must reshape the current forward. + + PCP only reshapes the *sparse-MLA prefill* path (round-robin query split + + full-KV all-gather). It returns False and the code stays byte-for-byte + identical to the non-PCP path when any of these hold: + * pcp_size == 1; + * decode (PCP keeps decode on the full, already-cached KV); + * dummy/warmup runs (graphs are captured on the full token layout — this + matches the metadata builder, which gates its reindex on + `not batch.is_dummy_run`, so model split and metadata reindex stay in + lock-step); + * short batches that fall back to dense MHA prefill (max_seqlen_k <= + index_topk): only the sparse indexer / sparse-attn path is PCP-wired. + + Every PCP call site (model split/gather, indexer, MLA write, metadata + reindex) keys off the SAME batch-global condition so they agree within a + forward. + """ + if not pcp_is_enabled(): + return False + ctx = get_forward_context() + context = getattr(ctx, "context", None) + attn_metadata = getattr(ctx, "attn_metadata", None) + if context is None or attn_metadata is None: + return False + if not bool(context.is_prefill) or bool(getattr(context, "is_dummy_run", False)): + return False + index_topk = getattr(get_current_atom_config().hf_config, "index_topk", None) + if index_topk is None: + return False + return int(getattr(attn_metadata, "max_seqlen_k", 0)) > int(index_topk) + + +def _install_increment_version_pcp_shim() -> None: + """Make ``torch.autograd.graph.increment_version`` tolerate a torch bug that + only surfaces under PCP + torch.compile. + + Under Prefill Context Parallel the sparse indexer must run through a + Dynamo-opaque custom op (``indexer_with_output``) so its runtime + ``_pcp_active()`` branch is not baked to the warmup value. Inserting that op + reshapes the pre-attention piecewise submodule, and torch's *inference* + runtime wrapper (``keep_input_mutations=True``, hard-coded in + ``torch/_inductor/compile_fx.py``) then resolves that submodule's + ``mutated_graph_handled_indices_seen_by_autograd`` against a runtime arg list + that interleaves SymInt shape symbols. The mutated-input index lands on a + SymInt, so ``increment_version`` is handed a non-tensor: + + RuntimeError: increment_version expects each element ... to be a tensor + IndexError: list index out of range (when the index is past the end) + + The baseline (indexer inlined, no extra op) only dodges this by arg-layout + luck. Crucially the version counter is *autograd-only* metadata: in inference + (no backward) it is never read, and the real in-place mutation is applied by + the compiled kernel regardless of this bump. So it is safe to filter the + iterable to real tensors and swallow the out-of-range walk. This is a strict + no-op for every well-formed call (all-tensor iterables materialize to the + same list), so baseline / non-PCP compilation is byte-for-byte unaffected. + """ + import torch.autograd.graph as _agraph + + if getattr(_agraph.increment_version, "_pcp_shim", False): + return + _orig_increment_version = _agraph.increment_version + + def increment_version(tensor): + if not isinstance(tensor, torch.Tensor): + safe = [] + try: + for t in tensor: + if isinstance(t, torch.Tensor): + safe.append(t) + except IndexError: + # torch handed us `(args[i] for i in )`; stop at the + # first out-of-range index (autograd metadata only — see docstring). + pass + tensor = safe + return _orig_increment_version(tensor) + + increment_version._pcp_shim = True + # runtime_wrappers.py calls this via attribute access on the module, so + # rebinding the module attribute is enough for it to pick up the shim. + _agraph.increment_version = increment_version + + +_install_increment_version_pcp_shim() + + def _enable_non_triton_global_mxfp4_input_norm_quant( config: PretrainedConfig, quant_config: Optional[QuantizationConfig], @@ -1175,6 +1271,22 @@ def sparse_attn_indexer( ) runner_block_size = get_current_atom_config().kv_cache_block_size kv_cache = kv_cache.view(-1, runner_block_size, kv_cache.shape[-1]) + # PCP prefill: `k` (and `positions`) arrive as the full PADDED key set + # [S_pad] produced by an all-gather of the round-robin shards. The KV-cache + # write (driven by slot_mapping) and the gathered-KV sizing (total_kv = + # k.shape[0]) both use the real token count S_real == slot_mapping length, + # so trim the round-robin padding here. Gated on the same condition as the + # caller's PCP path (sparse prefill with max_seqlen_k > topk); no-op + # otherwise. + if ( + pcp_is_enabled() + and context.is_prefill + and attn_metadata.max_seqlen_k > topk_tokens + ): + n_real = slot_mapping.shape[0] + k = k[:n_real] + if positions is not None: + positions = positions[:n_real] if use_qk_rope_cache_fusion: q_bf16 = q_input q_fp8 = torch.empty_like(q_bf16, dtype=dtypes.fp8) @@ -1217,10 +1329,13 @@ def sparse_attn_indexer( return weights prefill_metadata = attn_metadata num_prefills = context.batch_size - total_seq_lens = hidden_states.shape[0] - # When has_cached, gather full KV (cached + new) for indexer top-k + # Size the gathered-KV buffer off the KEY length, not the hidden/query + # length. Under PCP the query side (hidden_states) is 1/pcp while `k` is + # the full all-gathered key set, so `k.shape[0]` is the correct full + # token count; without PCP the two are equal so behaviour is unchanged. + # When has_cached, gather full KV (cached + new) for indexer top-k. total_kv = ( - prefill_metadata.total_kv if prefill_metadata.has_cached else total_seq_lens + prefill_metadata.total_kv if prefill_metadata.has_cached else k.shape[0] ) k_fp8 = torch.empty([total_kv, head_dim], device=k.device, dtype=dtypes.fp8) k_scale = torch.empty([total_kv, 1], device=k.device, dtype=torch.float32) @@ -1521,6 +1636,85 @@ def process_weights_after_loading(self): @IndexerDecoratorForPluginMode +def _indexer_with_output_fake( + hidden_states: torch.Tensor, + qr: torch.Tensor, + qr_scale: Optional[torch.Tensor], + positions: torch.Tensor, + layer_name: str, + sparse_kv_indices_buffer: torch.Tensor, +) -> torch.Tensor: + # Identity-passthrough contract: the op returns a fresh tensor shaped like + # `qr` (see `indexer_with_output`). The caller consumes it as the query into + # `mla_attn`, which keeps the op alive and ordered even independent of the + # declared buffer mutation. + return torch.empty_like(qr) + + +def indexer_with_output( + hidden_states: torch.Tensor, + qr: torch.Tensor, + qr_scale: Optional[torch.Tensor], + positions: torch.Tensor, + layer_name: str, + sparse_kv_indices_buffer: torch.Tensor, +) -> torch.Tensor: + """Dynamo-opaque wrapper around ``Indexer.forward_impl``. + + Registered as a REGULAR custom op (like ``sparse_attn_indexer``), NOT a + splitting op. Opacity — not a graph *split* — is what defeats the bake: + Dynamo treats a custom op as a leaf and never traces the body, so the + runtime ``_pcp_active()`` branch inside the indexer (round-robin k all-gather + + separate q/k rope) is evaluated LIVE every forward instead of being frozen + to its dummy-warmup value (``False``). Prefill is not cudagraph-captured, so + the eager all-gather inside runs safely; decode never takes the PCP branch. + + Why regular and not ``@mark_spliting_op``: adding a second split point + between ``qkv_a_proj`` and ``mla_attn`` shrinks the pre-attention submodule + to a handful of ops, and AOT-autograd then mis-resolves that tiny piece's + mutated-input index (``increment_version`` IndexError). Staying a plain + opaque op keeps the pre-attention submodule identical in shape to the + baseline's (which compiles cleanly). + + Buffer mutation MUST be declared. ``forward_impl`` writes the indexer's + top-k result into ``sparse_kv_indices_buffer`` (via the nested eager + ``sparse_attn_indexer`` op) and ``mla_attn`` reads that same buffer to know + which KV to attend to. In the non-PCP path ``sparse_attn_indexer`` runs + *directly* in the traced graph and declares ``mutates_args= + ["sparse_kv_indices_buffer"]``, so inductor keeps the write ordered before + the MLA read. Nesting it inside this opaque op HIDES that write; with + ``mutates_args=[]`` inductor thinks the buffer is unchanged and the MLA reads + a stale / mis-ordered copy — visible as token-stutter corruption ("errerr", + doubled fragments) even on PCP-noop prompts. So the buffer is threaded + through as an explicit arg and re-declared mutated here, restoring the exact + write→read edge the baseline relies on. (Declaring the mutation re-triggers a + latent AOT-autograd ``increment_version`` bug on graphs with SymInt args, + which ``_install_increment_version_pcp_shim`` neutralizes.) + + The identity ``qr`` return, fed by the caller as the ``mla_attn`` query, is + kept as belt-and-suspenders ordering + DCE protection. + """ + self = get_current_atom_config().compilation_config.static_forward_context[ + layer_name + ] + # Side effect: writes sparse_kv_indices_buffer (via nested eager + # sparse_attn_indexer). `self.sparse_kv_indices_buffer` is the same tensor + # object passed in, so the declared mutation matches the real one. + self.forward_impl(hidden_states, qr, qr_scale, positions) + # Fresh tensor equal to qr; consumed by the caller as the mla_attn query. + # Clone (not a bare return of qr) so the runtime output matches the fake's + # fresh-tensor contract and never aliases an input. + return qr.clone() + + +direct_register_custom_op( + op_name="indexer_with_output", + op_func=indexer_with_output, + mutates_args=["sparse_kv_indices_buffer"], + fake_impl=_indexer_with_output_fake, +) + + class Indexer(nn.Module): def __init__( self, @@ -1602,6 +1796,12 @@ def __init__( self.sparse_kv_indices_buffer = torch.empty(0, dtype=torch.int32, device="cuda") atom_config.compilation_config.static_forward_context[prefix] = self + # Rope module used by `forward_impl` (and the `indexer_with_output` + # splitting op, which can't take a module arg). Bound by the owning + # DeepseekV2MLAAttention right after construction; mirrors V4's + # `self.indexer.rotary_emb = self.rotary_emb`. + self.rotary_emb = None + self.sparse_attn_indexer_impl = torch.ops.aiter.sparse_attn_indexer def forward( @@ -1610,8 +1810,43 @@ def forward( qr: torch.Tensor, qr_scale: Optional[torch.Tensor], positions, - rotary_emb, + rotary_emb=None, + ) -> torch.Tensor: + # Under PCP, route the whole indexer through the Dynamo-opaque + # `indexer_with_output` splitting op so the runtime `_pcp_active()` branch + # (round-robin k all-gather + separate q/k rope) evaluates live instead of + # being baked to its warmup value by torch.compile. `pcp_is_enabled()` is + # a run-level constant (pcp world size is fixed for the process), so this + # guard is compile-safe and a no-op — a direct `forward_impl` call, graph + # unchanged — for non-PCP and plugin (SGLang / vLLM / RTP) backends. + if pcp_is_enabled(): + # Returns `qr` (identity); the caller feeds it to mla_attn so the + # opaque op stays live and ordered. The top-k result travels the + # side-buffer (declared mutated, so its write is ordered before the + # sparse-MLA read), not this return value. + return torch.ops.aiter.indexer_with_output( + hidden_states, + qr, + qr_scale, + positions, + self.prefix, + self.sparse_kv_indices_buffer, + ) + return self.forward_impl(hidden_states, qr, qr_scale, positions, rotary_emb) + + def forward_impl( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + qr_scale: Optional[torch.Tensor], + positions, + rotary_emb=None, ) -> torch.Tensor: + # The opaque `indexer_with_output` op can't pass a module, so it relies on + # the bound `self.rotary_emb`; direct callers (non-PCP / plugins) may still + # pass their own rope explicitly, which takes precedence. + if rotary_emb is None: + rotary_emb = self.rotary_emb q = self.wq_b(qr, qr_scale) q = q.view(-1, self.n_head, self.head_dim) @@ -1625,15 +1860,39 @@ def forward( k = self.wk(hidden_states) weights = self.weights_proj(hidden_states) - if not self.use_qk_rope_cache_fusion: + # Under PCP prefill the fused qk-rope-cache kernel cannot be used: it + # ropes q and writes k in one pass keyed on a single token count, but + # here q/weights are this rank's 1/pcp queries while k must become the + # FULL key set (every rank keeps full KV). So force the unfused path, + # all-gather k (and the key positions) to the full padded sequence, and + # rope q (1/pcp) and k (full) separately. The op then scores 1/pcp + # queries against the gathered full KV and writes the full k-cache. + pcp = _pcp_active() + positions_op = positions + if (not self.use_qk_rope_cache_fusion) or pcp: q_pe, _ = torch.split( q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) + if pcp: + pcp_ws = get_pcp_world_size() + # k is 1/pcp (from 1/pcp hidden); gather to full padded [S_pad]. + k = pcp_allgather_rerange(k, pcp_ws) + positions_op = pcp_allgather_rerange(positions, pcp_ws) k = self.k_norm(k) k_pe, _ = torch.split( k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1 ) - q_pe, k_pe = rotary_emb(positions, q_pe, k_pe) + if pcp: + # Rope q (1/pcp) and k (full) separately: they have different + # token counts under PCP so they can't share one rope call. The + # rope kernel is 2-component (ropes query AND key in place) and + # requires a non-None partner, so pass a throwaway of the + # matching length for the side we don't need. rope is in-place + # on the rotary_dim views (q_pe/k_pe alias q/k). + rotary_emb(positions, q_pe, torch.empty_like(q_pe)) + rotary_emb(positions_op, torch.empty_like(k_pe), k_pe) + else: + q_pe, k_pe = rotary_emb(positions, q_pe, k_pe) q = q.view(-1, self.head_dim) q_fp8, q_scale = self.quant_func(q, quant_dtype=dtypes.fp8) @@ -1663,12 +1922,12 @@ def forward( self.k_norm.weight, self.k_norm.bias, self.k_norm.eps, - positions, + positions_op, rotary_emb.cos_cache.squeeze(-2).squeeze(-2), rotary_emb.sin_cache.squeeze(-2).squeeze(-2), self._weights_scale, rotary_emb.is_neox_style, - self.use_qk_rope_cache_fusion, + self.use_qk_rope_cache_fusion and not pcp, ) @@ -1897,6 +2156,10 @@ def __init__( ), f"{prefix}.indexer", ) + # Bind the indexer's rope so forward_impl (and the opaque + # indexer_with_output splitting op) can rope without receiving a + # module argument. Mirrors deepseek_v4.Attention.__init__. + self.indexer.rotary_emb = self.indexer_rope_emb else: self.indexer_rope_emb = None self.indexer = None @@ -2036,13 +2299,23 @@ def forward( kv_c_normed = self.kv_a_layernorm(kv_c) hidden_states_or_q_c_scale = None if self.is_v32 and self.indexer is not None and not self.skip_topk: - self.indexer( + idx_ret = self.indexer( hidden_states, hidden_states_or_q_c, hidden_states_or_q_c_scale, positions, self.indexer_rope_emb, ) + if pcp_is_enabled(): + # Under PCP the indexer runs through the opaque `indexer_with_output` + # split op, which returns `hidden_states_or_q_c` unchanged + # (identity). Feeding it forward as the mla_attn query is what keeps + # the op live under torch.compile — its real result (top-k) is a + # hidden write to the sparse buffer that mla_attn reads via `self` — + # and orders the write before that read. `pcp_is_enabled()` is a + # run-level constant, so baking this branch at trace time is correct + # (non-PCP / plugins keep discarding the return, graph unchanged). + hidden_states_or_q_c = idx_ret return self.mla_attn( hidden_states_or_q_c, @@ -2497,9 +2770,44 @@ def forward( intermediate_tensors: Optional[IntermediateTensors] = None, inputs_embeds: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, IntermediateTensors]: + # ---- Prefill Context Parallel (PCP) query split ------------------ + # During prefill with pcp_size > 1 the token sequence is round-robin + # split so each PCP rank runs the whole model (embed / norm / q-proj / + # MoE) on only 1/pcp of the tokens. The attention modules re-materialise + # the full KV internally (see DeepseekV2MLAAttention / Indexer / + # MLAAttention), so decode and the cache layout are untouched. When PCP + # is inactive (`pcp=1` or decode) this whole block is skipped and the + # forward is identical to the original path. + pcp = _pcp_active() + n_global = positions.shape[0] + if pcp: + pcp_ws = get_pcp_world_size() + n_pad = pcp_pad_len(n_global, pcp_ws) - n_global + positions = pcp_round_robin_split(pcp_pad_dense(positions, n_pad), pcp_ws) + if input_ids is not None: + input_ids = pcp_round_robin_split( + pcp_pad_dense(input_ids, n_pad), pcp_ws + ) + if inputs_embeds is not None: + inputs_embeds = pcp_round_robin_split( + pcp_pad_dense(inputs_embeds, n_pad), pcp_ws + ) + hidden_states = self.model( input_ids, positions, intermediate_tensors, inputs_embeds ) + + # ---- PCP gather: 1/pcp rows -> full token order, then drop pad ---- + # Only the last PP rank produces real hidden states; earlier stages + # forward IntermediateTensors (already 1/pcp) unchanged. + if pcp and get_pp_group().is_last_rank: + if isinstance(hidden_states, tuple): + hs, aux = hidden_states + hs = pcp_allgather_rerange(hs, pcp_ws)[:n_global] + aux = [pcp_allgather_rerange(a, pcp_ws)[:n_global] for a in aux] + hidden_states = (hs, aux) + elif isinstance(hidden_states, torch.Tensor): + hidden_states = pcp_allgather_rerange(hidden_states, pcp_ws)[:n_global] return hidden_states def compute_logits( diff --git a/atom/spec_decode/eagle.py b/atom/spec_decode/eagle.py index 70dd580a08..762f932cc1 100644 --- a/atom/spec_decode/eagle.py +++ b/atom/spec_decode/eagle.py @@ -8,6 +8,13 @@ from aiter import dtypes from aiter.dist.parallel_state import get_pp_group from atom.config import CompilationLevel, Config, KVCacheTensor +from atom.distributed.pcp_utils import ( + get_pcp_world_size, + pcp_allgather_rerange, + pcp_pad_dense, + pcp_pad_len, + pcp_round_robin_split, +) from atom.model_loader.loader import load_model from atom.utils import CpuGpuBuffer, resolve_obj_by_qualname from atom.utils import envs @@ -447,15 +454,53 @@ def propose( # every iteration (used at i>=1 too, even though i==0 sets it). has_flat_kv = "kv_indices" in var + # Same PCP gate the draft's attention modules use (import lazily to avoid + # any import-order coupling with the model module). + from atom.models.deepseek_v2 import _pcp_active + for i in range(self.mtp_k): with record_function(f"draft[{i}/{self.mtp_k} bs={bs}]"): # Re-sync DP token self._refresh_dp_metadata(forward_context, input_ids.shape[0]) + # ---- Prefill Context Parallel (draft i==0 prefill) -------- + # The draft's first pass is a prefill that reuses the target's + # 1/pcp-reindexed attn_metadata, so it must run on this rank's + # 1/pcp query shard (input_ids / positions / previous hidden) and + # all-gather the draft hidden back to full token order before the + # last-token sampling gather. Later draft steps are decode + # (is_prefill False) and run full — identical to the non-PCP path. + # `input_ids` / `positions` / `hidden_states` themselves stay full + # so the post-i==0 decode-metadata setup (which indexes with the + # full `last_token_indices`) is unchanged. + pcp_draft_prefill = i == 0 and _pcp_active() + if pcp_draft_prefill: + pcp_ws = get_pcp_world_size() + n_global_draft = input_ids.shape[0] + n_pad = pcp_pad_len(n_global_draft, pcp_ws) - n_global_draft + d_input_ids = pcp_round_robin_split( + pcp_pad_dense(input_ids, n_pad), pcp_ws + ) + d_positions = pcp_round_robin_split( + pcp_pad_dense(positions, n_pad), pcp_ws + ) + d_hidden = pcp_round_robin_split( + pcp_pad_dense(hidden_states, n_pad), pcp_ws + ) + else: + d_input_ids, d_positions, d_hidden = ( + input_ids, + positions, + hidden_states, + ) ret_hidden_states = self.model( - input_ids=input_ids, - positions=positions, - hidden_states=hidden_states, + input_ids=d_input_ids, + positions=d_positions, + hidden_states=d_hidden, ) + if pcp_draft_prefill: + ret_hidden_states = pcp_allgather_rerange( + ret_hidden_states, pcp_ws + )[:n_global_draft] sample_hidden_states = ( torch.index_select(ret_hidden_states, 0, last_token_indices) diff --git a/recipes/GLM-5.md b/recipes/GLM-5.md index 77d9576315..0d15b91398 100644 --- a/recipes/GLM-5.md +++ b/recipes/GLM-5.md @@ -208,3 +208,57 @@ python -m atom.entrypoints.openai_server \ --online_quant_config '{"layer_quant_config": {"*layers.78*": "ptpc_fp8"}, "exclude_layer": ["*.gate", "*shared_head*", "*embed*"]}' \ -tp $TP 2>&1 | tee server_mtp.log ``` + +## GLM-5.2 Prefill Context Parallel (PCP) + +Prefill Context Parallel accelerates **long-context prefill** (large ISL) by +round-robin splitting the prompt tokens across an extra parallel dimension +(`world = tp × pcp`). Only the query side is sharded — every rank keeps the +**full KV cache**, so decode, the KV-cache layout, and accuracy are unchanged. +The dominant `O(S²)` DSA indexer scoring and the sparse MLA attention run on +`1/pcp` of the queries per rank, cutting TTFT on long inputs. GLM-5.2 runs on +the DeepSeek-V3.2 sparse-MLA (DSA) path, so PCP applies here just like on +DeepSeek-V4. + +Enable it with `--prefill-context-parallel-size` (`-pcp`). `pcp` is orthogonal +to `-tp`, and the two multiply into the number of GPUs used +(`GPUs = tp × pcp`). MTP speculative decoding is supported — the draft's prefill +pass is split and gathered the same way. + +Constraints in this release (the engine raises a clear error otherwise): +- DSA models only (GLM-5.2 / DeepSeek-V3.2 — configs with `index_topk`). +- `--no-enable_chunked_prefill` (PCP needs the whole prompt in one forward). +- `--no-enable_prefix_caching`. +- Not compatible with `--enable-dp-attention` or `--enable-tbo`. +- `pcp = 1` (the default) is a bit-exact no-op — the non-PCP code path is + unchanged. + +### Serving on 4 GPUs (TP2 × PCP2) + +```bash +#!/bin/bash + +model_path=/shared/data/amd_int/models/GLM-5.2-FP8 + +rm -rf /root/.cache/atom/* + +python -m atom.entrypoints.openai_server \ + --model "$model_path" \ + --server-port 8000 \ + --kv_cache_dtype fp8 \ + --no-enable_prefix_caching \ + --no-enable_chunked_prefill \ + -tp 2 -pcp 2 2>&1 | tee server_pcp.log & +``` + +Pure context parallel (`TP1 × PCP4`) is also valid on 4 GPUs — swap `-tp 2 -pcp 2` +for `-tp 1 -pcp 4`. + +### Validating PCP + +Compare a PCP run against a same-GPU-count `-tp 4 -pcp 1` baseline. Accuracy +should match the baseline (PCP does not change decode), and TTFT should improve +as ISL grows. Use a long-input workload, e.g. `ISL=60000`, `OSL=1024`, +concurrency 8, with the benchmark command from +[Performance baseline](#performance-baseline), and the gsm8k +[Accuracy test](#accuracy-test) for parity. From 659dbd9963ba95456b75380c0e2c1152920655b8 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Wed, 8 Jul 2026 12:43:45 +0000 Subject: [PATCH 17/17] remove guards Signed-off-by: whx-sjtu --- atom/model_engine/llm_engine.py | 32 -------------------------------- 1 file changed, 32 deletions(-) diff --git a/atom/model_engine/llm_engine.py b/atom/model_engine/llm_engine.py index a874db3d69..a20b019dc2 100644 --- a/atom/model_engine/llm_engine.py +++ b/atom/model_engine/llm_engine.py @@ -84,38 +84,6 @@ def __init__(self, model, tokenizer=None, **kwargs): "independently (double-split) and corrupt the output. For now, " "disable one of them." ) - # PCP for the native ATOM engine is implemented only on the DeepSeek-V3.2 - # / GLM-5.2 sparse-MLA (DSA) path (indexer + sparse attention). Those - # configs expose `index_topk`. Non-sparse models (dense DeepSeek-V2, - # non-DSA archs) have no PCP wiring, so reject pcp>1 for them with a clear - # message instead of silently running the un-split path. - if config.prefill_context_parallel_size > 1 and not hasattr( - config.hf_config, "index_topk" - ): - raise ValueError( - "prefill_context_parallel_size > 1 (-pcp) is currently only " - "supported for sparse-MLA / DSA models (DeepSeek-V3.2, GLM-5.2 — " - "configs with `index_topk`). The requested model " - f"({getattr(config.hf_config, 'model_type', 'unknown')}) is not a " - "DSA model. Run it with -pcp 1." - ) - # PCP v1 requires chunked prefill OFF: the round-robin token split assumes - # the whole prefill sequence is present in one forward (it pads the global - # token count to a multiple of pcp and needs each query's full history). - if config.prefill_context_parallel_size > 1 and config.enable_chunked_prefill: - raise ValueError( - "prefill_context_parallel_size > 1 (-pcp) requires chunked " - "prefill disabled in this release. Pass " - "--no-enable_chunked_prefill." - ) - # PCP v1 requires prefix caching OFF: mixing a full cached KV prefix with - # 1/pcp new tokens in the indexer gather is not yet handled. - if config.prefill_context_parallel_size > 1 and config.enable_prefix_caching: - raise ValueError( - "prefill_context_parallel_size > 1 (-pcp) requires prefix " - "caching disabled in this release. Pass " - "--no-enable_prefix_caching." - ) self.rquest_ids = set() self.io_processor = InputOutputProcessor( config, self.tokenizer, config.kv_cache_block_size