From 2ece675665dca8b12fd45217cb83b7a40626ab79 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Wed, 8 Jul 2026 06:27:52 +0000 Subject: [PATCH 1/7] enable pcp for DSA Signed-off-by: whx-sjtu (cherry picked from commit c8ad6827e42405d9e26e5a79778ee30b98ad533b) --- 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 04f169bd23..0187901fd2 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1266,6 +1266,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 e7768679b8..04b3b0e090 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -740,6 +740,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 b924ea0498..47fc9c7528 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -15,6 +15,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 @@ -871,6 +878,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) @@ -1007,6 +1025,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 4ec791444f..d936355afb 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -89,6 +89,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.vllm.attention.layer_sparse_mla import ( @@ -140,6 +148,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], @@ -1186,6 +1282,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) @@ -1228,10 +1340,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) @@ -1573,6 +1688,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, @@ -1654,6 +1848,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( @@ -1662,8 +1862,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) @@ -1677,15 +1912,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) @@ -1715,12 +1974,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, ) @@ -1958,6 +2217,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 @@ -2118,13 +2381,23 @@ def forward( # The indexer's wk/weights_proj GEMMs run in BF16. When input_layernorm # fused the quant it emits a bf16 mirror (indexer_hidden); otherwise # hidden_states is already the bf16 normed activation. - self.indexer( + idx_ret = self.indexer( indexer_hidden if indexer_hidden is not None else 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, @@ -2623,9 +2896,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 7a921de241..7c36bdef5e 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 @@ -437,15 +444,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 35f3ee1ed1..8d53ff0186 100644 --- a/recipes/GLM-5.md +++ b/recipes/GLM-5.md @@ -194,3 +194,57 @@ 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 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 14a6a7e6f3de3fee852f11d24150a0cc47772905 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Wed, 8 Jul 2026 12:43:45 +0000 Subject: [PATCH 2/7] remove guards Signed-off-by: whx-sjtu (cherry picked from commit 659dbd9963ba95456b75380c0e2c1152920655b8) --- 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 From 58de955a164b5f8417b62d66123ae9115fada5ed Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Mon, 13 Jul 2026 05:30:01 +0000 Subject: [PATCH 3/7] update pcp recipe Signed-off-by: whx-sjtu --- recipes/GLM-5.md | 9 --------- 1 file changed, 9 deletions(-) diff --git a/recipes/GLM-5.md b/recipes/GLM-5.md index 8d53ff0186..deca34c70f 100644 --- a/recipes/GLM-5.md +++ b/recipes/GLM-5.md @@ -239,12 +239,3 @@ python -m atom.entrypoints.openai_server \ 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 7a2ff2871891ad95229160bec7a90f912578eae6 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Mon, 13 Jul 2026 05:35:07 +0000 Subject: [PATCH 4/7] remove notes Signed-off-by: whx-sjtu --- atom/model_engine/model_runner.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 04b3b0e090..e7768679b8 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -740,10 +740,6 @@ 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"): From b67d39185fb63c8b9436b72bd1a44c3e44601dc6 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Wed, 15 Jul 2026 07:10:58 +0000 Subject: [PATCH 5/7] vllm attn cp Signed-off-by: whx-sjtu --- atom/config.py | 11 + atom/distributed/pcp_utils.py | 281 +++++++++++++++++- atom/models/deepseek_v2.py | 305 +++++++++++++++++-- atom/plugin/config.py | 2 + atom/plugin/graph_capture_patch.py | 78 +++-- atom/plugin/register.py | 57 ++++ atom/plugin/vllm/attention/layer_mla.py | 60 +++- atom/plugin/vllm/attention/metadata.py | 374 +++++++++++++++++++++++- atom/plugin/vllm/tp_group_reuse.py | 183 +++++++++--- atom/utils/envs.py | 23 ++ 10 files changed, 1259 insertions(+), 115 deletions(-) diff --git a/atom/config.py b/atom/config.py index 0187901fd2..6aeeb13f1d 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1069,6 +1069,12 @@ class Config: # per-DP MoE layout and leave it False. Set by the frontend in # atom/plugin/config.py, not queried via is_vllm() at the call site. moe_ep_flatten_tp_across_dp: bool = False + # vLLM plugin: reuse the TP group as the Context-Parallel group (true + # sequence parallelism over tokens). Set by the frontend from + # ATOM_VLLM_ATTN_CP in atom/plugin/config.py; folded into compute_hash + # because it changes the compiled graph (full-head replicated attention, + # MoE all-gather/reduce-scatter, disabled fused-allreduce norms). + vllm_attn_cp: bool = False torch_dtype: torch.dtype = field(init=False) speculative_config: Optional[SpeculativeConfig] = None kv_transfer_config: dict = field(default_factory=dict) @@ -1275,6 +1281,11 @@ def compute_hash(self) -> str: # runtime — the same stale-artifact hazard documented for the vocab-embed # flag below. factors.append(self.prefill_context_parallel_size) + # Reuse-TP-as-CP (plugin) rebuilds attention as full-head replicated, + # swaps the MoE all-reduce for all-gather+reduce-scatter, and drops the + # fused-allreduce in the norms -- a different graph than the cp-off run + # over the same model+source, so it must key the compiled artifact. + factors.append(self.vllm_attn_cp) factors.append(self.enable_dp_attention) text_config = getattr(self.hf_config, "text_config", self.hf_config) factors.append( diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 79ccf37585..4589fa339b 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -24,6 +24,39 @@ ) +def plugin_attn_cp_enabled() -> bool: + """True when the vLLM plugin reuses the TP group as the Context-Parallel + group (env ``ATOM_VLLM_ATTN_CP``, RFC ROCm/ATOM#196). + + This is distinct from :func:`pcp_is_enabled`, which reflects the aiter + ``_PCP`` group world size and is *also* true for native ATOM's separate + ``-pcp`` dimension. ``plugin_attn_cp_enabled`` specifically flags the + reuse-TP-as-CP true sequence-parallel path, whose construction-time model + changes (full-head *replicated* attention, MoE all-gather/reduce-scatter, + disabled fused-allreduce norms) must fire *before* the aiter ``_PCP`` group + is aliased. It therefore keys off the static env flag, not the group world. + + The round-robin split/gather/reindex machinery is shared between native PCP + and plugin CP and stays gated on :func:`pcp_is_enabled` / ``_pcp_active``. + """ + from atom.utils import envs + + return bool(envs.ATOM_VLLM_ATTN_CP) + + +def plugin_attn_cp_decode_split_enabled() -> bool: + """True when reuse-TP-as-CP additionally round-robin splits the DECODE batch. + + Requires :func:`plugin_attn_cp_enabled` (this flag only extends that path). + When off (default) decode stays full/replicated (every rank runs all tokens + with full heads); when on, a pure-decode batch is split 1/cp exactly like + prefill. See ``ATOM_VLLM_ATTN_CP_DECODE_SPLIT`` in ``envs.py``. + """ + from atom.utils import envs + + return bool(envs.ATOM_VLLM_ATTN_CP_DECODE_SPLIT) and plugin_attn_cp_enabled() + + def get_pcp_world_size() -> int: return get_prefill_context_model_parallel_world_size() @@ -77,7 +110,133 @@ def pcp_round_robin_split( # Divisibility by pcp_size is guaranteed upstream by pcp_pad_len (callers # pad before splitting); the view below would error if violated. rest = tuple(input_.shape[1:]) - return input_.view(-1, pcp_size, *rest)[:, pcp_rank].contiguous() + shard = input_.view(-1, pcp_size, *rest)[:, pcp_rank] + # The round-robin slice has inner stride `pcp_size` (it selects every + # pcp_size-th row). `.contiguous()` normalises that to unit stride by + # copying -- EXCEPT when the shard holds a single element (n_owned == 1, + # e.g. a decode batch that pads to exactly pcp_size): a numel<=1 tensor + # reports is_contiguous()==True regardless of stride, so `.contiguous()` + # is a no-op and the stride-`pcp_size` view leaks downstream. That breaks + # kernels asserting unit inner stride -- notably aiter rope + # (`positions.stride(1) == 1`) on the 1/pcp query positions. `clone` with + # an explicit contiguous format always allocates standard (unit-inner) + # strides, so it fixes n_owned==1 while costing no more than the copy + # `.contiguous()` already performs for the (always strided) n_owned>1 case. + return shard.clone(memory_format=torch.contiguous_format) + + +def _pcp_ca_comm(group): + """Return the group's custom-all-reduce communicator, or None. + + The custom AR comm (``ca_comm``) owns the capture-safe collective kernels + (pre-registered IPC pool, ``_IS_CAPTURING`` handling). It exists on the + reuse-TP-as-CP group because TP all-reduce already uses it. + """ + dc = getattr(group, "device_communicator", None) + ca = getattr(dc, "ca_comm", None) if dc is not None else None + if ca is None or getattr(ca, "disabled", True): + return None + return ca + + +def _align_pad_rows_for_custom_ag(inp: torch.Tensor) -> tuple[torch.Tensor, int]: + """Append zero rows so the tensor's total byte size is a multiple of 16. + + aiter's capture-safe custom all-gather (``CustomAllreduce.should_custom_ag``) + only accepts 16-byte-aligned tensors. Row-aligned 2-D gathers (hidden/k, whose + row bytes are already a multiple of 16) never pad; only tiny 1-D int gathers + (a single int64 query id = 8 B on an ``n_owned == 1`` decode shard) do. The + pad rows sort to the tail after round-robin rerange, so the caller drops them + by slicing with the original (unpadded) length. + """ + import math + + row = inp.element_size() + for s in inp.shape[1:]: + row *= int(s) + if row == 0: + return inp, 0 + mult = 16 // math.gcd(16, row) # rows needed for a 16 B-aligned total + rem = inp.shape[0] % mult + if rem == 0: + return inp, 0 + pad = mult - rem + pad_block = inp.new_zeros((pad,) + tuple(inp.shape[1:])) + return torch.cat([inp, pad_block], dim=0), pad + + +def _custom_all_gather_dim0(group, x: torch.Tensor) -> torch.Tensor: + """aiter custom (capture-safe) all-gather along dim 0. + + aiter's ``custom_all_gather`` maps int types to a same-width float for its + memcpy kernel, but int64 -> float64 is rejected by the kernel + (``Unsupported dtype: torch.float64``). 8-byte dtypes (int64/float64, e.g. + ``input_ids`` / ``positions``) are reinterpreted as int32 pairs -- which the + kernel supports (int32 -> float32) -- and restored afterwards. This is a pure + bitcast (``view``), so values are exact and rank-major order is preserved + (each rank's row stays a contiguous 2xint32 that views back to one int64). + """ + if x.dtype in (torch.int64, torch.float64): + g32 = group.all_gather(x.view(torch.int32), use_custom=True, dim=0) + return g32.view(x.dtype) + return group.all_gather(x, use_custom=True, dim=0) + + +def pcp_dim0_all_gather(group, inp: torch.Tensor) -> tuple[torch.Tensor, int]: + """Rank-major all-gather along dim 0, preferring the CAPTURE-SAFE custom path. + + The default aiter ``all_gather(use_custom=False)`` lowers to raw + ``torch.distributed.all_gather_into_tensor`` (RCCL), whose host-side lazy + init/copies invalidate HIP stream capture -- a full CUDA graph over a split + decode batch dies with ``ncclUnhandledCudaError``. The custom all-gather + (``ca_comm``, pre-registered IPC pool, int-dtype aware) IS capturable, so we + take it whenever the tensor fits ``should_custom_ag``; otherwise (large + prefill gathers, which run eager/piecewise and never enter a full graph) we + fall back to RCCL. + + Returns ``(gathered_rank_major, pad_rows)``. When ``pad_rows > 0`` the extra + rows are this rank's tail, so after rerange the real tokens are the first + ``pcp_size * original_len`` entries. + """ + inp = inp.contiguous() + ca = _pcp_ca_comm(group) + if ca is not None: + cand, pad = _align_pad_rows_for_custom_ag(inp) + if ca.should_custom_ag(cand): + return _custom_all_gather_dim0(group, cand), pad + return group.all_gather(inp, dim=0), 0 + + +def pcp_tp_all_gather_dim0(input_: torch.Tensor) -> torch.Tensor: + """Capture-safe rank-major all-gather (dim 0) for the reuse-TP-as-CP MoE + path (all-gather -> experts -> reduce-scatter). + + Mirrors ``tensor_model_parallel_all_gather(x, dim=0)`` but (a) prefers the + custom capturable collective (see :func:`pcp_dim0_all_gather`) so the MoE + gather is legal inside a full CUDA graph over a split decode batch, and (b) + routes through the DEDICATED CP group (``get_pcp_group()``) rather than TP, + so it shares CP's isolated ca_comm / graph-buffer slot allocator with the + embed & indexer gathers instead of TP all-reduce's. Hidden-state rows are + always 16 B-aligned, so no padding is ever applied here. + """ + gathered, _pad = pcp_dim0_all_gather(get_pcp_group(), input_) + return gathered + + +def pcp_reduce_scatter_dim0(input_: torch.Tensor) -> torch.Tensor: + """Capture-safe rank-major reduce-scatter (dim 0) over the dedicated CP group + for the reuse-TP-as-CP MoE path (all-gather -> experts -> reduce-scatter). + + Inverse of :func:`pcp_tp_all_gather_dim0`: sums each rank's partial expert + outputs over the full token set and scatters dim 0 back to this rank's 1/cp + shard. Uses the custom (``use_custom=True``) capturable reduce-scatter on the + CP group's own ca_comm so it is legal inside a full CUDA graph and shares the + CP slot allocator with the matching all-gather (NOT TP all-reduce's). + """ + group = get_pcp_group() + if group.world_size == 1: + return input_ + return group.reduce_scatter_tensor(input_, use_custom=True, dim=0) def pcp_allgather_rerange( @@ -90,25 +249,29 @@ def pcp_allgather_rerange( interleave is restored by `view(pcp, L, *rest).transpose(0, 1)` so that output[t] == global token t. - Mirrors SGLang `cp_all_gather_rerange_output` (round-robin branch). + Mirrors SGLang `cp_all_gather_rerange_output` (round-robin branch). Uses the + capture-safe custom all-gather so this op is legal inside a full CUDA graph + (see :func:`pcp_dim0_all_gather`). """ if pcp_size is None: pcp_size = get_pcp_world_size() if pcp_size <= 1: return input_ group = get_pcp_group() - # aiter all_gather(dim=0) returns rank-major concat: [pcp*L, *rest]. - gathered = group.all_gather(input_.contiguous(), dim=0) local_len = input_.shape[0] rest = tuple(input_.shape[1:]) - # rank-major [pcp, L, *rest] -> transpose -> token-major [L, pcp, *rest] - # -> flatten to [L*pcp, *rest] == original global order. + # rank-major concat [pcp*(L+pad), *rest]; pad rows are each rank's tail. + gathered, pad = pcp_dim0_all_gather(group, input_) + padded_len = local_len + pad + # rank-major [pcp, L+pad, *rest] -> transpose -> token-major + # [L+pad, pcp, *rest] -> flatten. Real global tokens occupy the first + # pcp*L rows (row index l*pcp+r, l=L) sort to the tail. out = ( - gathered.view(pcp_size, local_len, *rest) + gathered.view(pcp_size, padded_len, *rest) .transpose(0, 1) - .reshape(pcp_size * local_len, *rest) + .reshape(pcp_size * padded_len, *rest) ) - return out + return out[: pcp_size * local_len] def pcp_round_robin_query_indices( @@ -168,6 +331,106 @@ def pcp_pad_dense(t: torch.Tensor, n_pad: int) -> torch.Tensor: return torch.cat([t, t.new_zeros(n_pad, *t.shape[1:])], dim=0) +def pcp_sparse_prefill_reindex( + sparse_seqlen: torch.Tensor, # [T] per-query selected-KV length (pre-clamp) + req_id_per_token: torch.Tensor, # [T] per-query request id + slot_mapping: torch.Tensor, # [T] full per-query slot mapping (KV write) + index_topk: int, + pcp_size: Optional[int] = None, + pcp_rank: Optional[int] = None, +) -> dict: + """Reduce the plugin sparse-MLA per-query metadata to this rank's 1/W queries. + + Plugin reuse-TP-as-CP mirror of the native model_runner + ``AiterMLAImpl._apply_pcp_reindex`` (``aiter_mla.py``), adapted to the plugin + ``AiterMlaSparseMetadataForVllm`` field names. Only *query-indexed* fields + shrink to the round-robin owned subset; *per-sequence* / *KV-write* fields + (``slot_mapping``, ``block_table``, ``seq_lens``) stay FULL so the full KV is + still written and gathered. The global token count is padded to a multiple of + ``pcp_size``; padded (dummy) queries get zero-length KV (they attend nothing + and their output is dropped after the model's exit all-gather + unpad). + + Pure tensor math (no kernels), so it is unit-testable on CPU against a + full-batch reference. Returns a dict of the owned-query tensors plus + ``owned_q`` / ``n_owned`` so the builder can rebuild its work buffers. + """ + device = sparse_seqlen.device + s_real = int(sparse_seqlen.shape[0]) + padded_total = pcp_pad_len(s_real, pcp_size) + n_pad = padded_total - s_real + owned_q = pcp_round_robin_query_indices(padded_total, pcp_size, pcp_rank).to(device) + n_owned = int(owned_q.shape[0]) + + # sparse_kv_indptr <- cumsum of min(sparse_seqlen, topk) over owned queries; + # dummy queries padded to 0 => zero-length KV segment. + owned_counts = pcp_pad_dense(sparse_seqlen, n_pad)[owned_q].to(torch.int64) + owned_counts = torch.clamp(owned_counts, max=int(index_topk)) + paged_kv_indptr = torch.zeros(n_owned + 1, dtype=torch.int32, device=device) + paged_kv_indptr[1:] = torch.cumsum(owned_counts, 0).to(torch.int32) + + # one query per row (incl dummies) => qo_indptr = arange, last_page_len = 1s. + qo_indptr = torch.arange(n_owned + 1, dtype=torch.int32, device=device) + paged_kv_last_page_len = torch.ones(n_owned, dtype=torch.int32, device=device) + + # per-query request id, padded (dummy -> 0) then owned-selected. + req_id_owned = ( + pcp_pad_dense(req_id_per_token, n_pad)[owned_q].to(torch.int32).contiguous() + ) + + # owned slot_mapping for any fused q_out kernel; the real full-KV completion + # write in the sparse layer overwrites every real slot, and 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)) + slot_mapping_owned = slot_mapping[owned_clamped].contiguous() + + return { + "owned_q": owned_q, + "n_owned": n_owned, + "paged_kv_indptr": paged_kv_indptr, + "qo_indptr": qo_indptr, + "paged_kv_last_page_len": paged_kv_last_page_len, + "req_id_per_token": req_id_owned, + "slot_mapping_owned": slot_mapping_owned, + } + + +def pcp_reindex_decode( + seq_lens: torch.Tensor, # [num_decode_tokens] per-query cached KV length + block_table: torch.Tensor, # [num_decode_tokens, max_blocks] per-query blocks + num_decode_tokens: int, + pcp_size: Optional[int] = None, + pcp_rank: Optional[int] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Reduce plain-decode per-query metadata to this rank's round-robin queries. + + Plain decode has exactly one query per request, so ``num_decode_tokens == + num_decodes`` and owning a round-robin subset of decode QUERIES is the same + as owning that subset of requests. The global decode-token count is padded to + a multiple of ``pcp_size`` with dummy (``seq_len == 0``, ``block_table == 0``) + queries so every rank runs the same ``n_owned == padded / pcp_size`` count; + the dummies attend nothing (their sparse-MLA KV segment is zero-length) and + their hidden output is dropped after the model's exit all-gather. + + Only these per-owned-query rows shrink here; the sparse-MLA ``slot_mapping`` + (KV write) stays FULL, mirroring the prefill reindex, so every rank still + writes the full replicated KV for the new decode tokens. + + Returns ``(seq_lens_owned, block_table_owned, owned_q, n_owned)``. + """ + device = seq_lens.device + padded_total = pcp_pad_len(num_decode_tokens, pcp_size) + n_pad = padded_total - num_decode_tokens + owned_q = pcp_round_robin_query_indices(padded_total, pcp_size, pcp_rank).to(device) + n_owned = int(owned_q.shape[0]) + seq_lens_owned = pcp_pad_dense(seq_lens[:num_decode_tokens], n_pad)[ + owned_q + ].contiguous() + block_table_owned = pcp_pad_dense(block_table[:num_decode_tokens], n_pad)[ + owned_q + ].contiguous() + return seq_lens_owned, block_table_owned, owned_q, n_owned + + def pcp_reindex_ragged( kv_indptr: torch.Tensor, # [T_global + 1] int32 — global per-query prefix sum kv_indices: torch.Tensor, # [kv_indptr[-1]] — ragged packed values diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index d936355afb..51d9cec0e8 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -39,7 +39,9 @@ top_k_per_row_decode, top_k_per_row_prefill, ) -from aiter.dist.communication_op import tensor_model_parallel_all_reduce +from aiter.dist.communication_op import ( + tensor_model_parallel_all_reduce, +) from aiter.dist.parallel_state import get_pp_group, get_tensor_model_parallel_world_size from aiter.jit.utils.torch_guard import torch_compile_guard from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits @@ -95,7 +97,10 @@ pcp_is_enabled, pcp_pad_dense, pcp_pad_len, + pcp_reduce_scatter_dim0, pcp_round_robin_split, + pcp_tp_all_gather_dim0, + plugin_attn_cp_enabled, ) from atom.utils.decorators import mark_trace, support_torch_compile from atom.utils.forward_context import get_forward_context @@ -169,12 +174,25 @@ def _pcp_active() -> bool: """ if not pcp_is_enabled(): return False + if plugin_attn_cp_enabled(): + # vLLM plugin reuse-TP-as-CP: ATOM's own forward context is never + # populated in plugin mode (the plugin drives vLLM's forward context), + # so the native branch below would always read context=None. Defer to + # the builder-stamped decision instead. + return _plugin_pcp_active() 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)): + if bool(getattr(context, "is_dummy_run", False)): + return False + # Native PCP reshapes only sparse *prefill* (decode stays on the full cached + # KV). The plugin reuse-TP-as-CP path returned above via _plugin_pcp_active(); + # it likewise splits prefill-only (see _plugin_cp_should_split), so this + # prefill gate is the shared contract. The `max_seqlen_k > index_topk` check + # below still excludes the short dense-MHA fallback (not CP-wired). + if not bool(context.is_prefill): return False index_topk = getattr(get_current_atom_config().hf_config, "index_topk", None) if index_topk is None: @@ -182,6 +200,42 @@ def _pcp_active() -> bool: return int(getattr(attn_metadata, "max_seqlen_k", 0)) > int(index_topk) +def _plugin_pcp_active() -> bool: + """Plugin reuse-TP-as-CP round-robin split decision. + + In vLLM plugin mode ATOM's own forward context is never populated, so the + native ``_pcp_active`` gate (which reads ``ctx.context`` / ``ctx.attn_metadata``) + can't be used. The plugin sparse-MLA metadata builders are the single source + of truth: they evaluate the SAME batch-global gate once per forward (real + build, ``max_seq_len > index_topk``) and stamp ``pcp_split`` on the metadata + they emit. Every model-side PCP call site (entry split, indexer full-KV + all-gather, FFN all-gather/reduce-scatter, exit gather) reads that stamp here + so they stay in lock-step with the metadata reindex within a forward. + + Reading is via vLLM's forward context because that is where the built + metadata lives in plugin mode. The metadata is a dict keyed by layer name; + the split decision is batch-global, so any stamped sparse layer settles it. + """ + try: + from vllm.forward_context import ( + get_forward_context as _get_vllm_forward_context, + is_forward_context_available as _vllm_forward_context_available, + ) + except Exception: + return False + if not _vllm_forward_context_available(): + return False + md = _get_vllm_forward_context().attn_metadata + if md is None: + return False + if isinstance(md, dict): + for layer_md in md.values(): + if bool(getattr(layer_md, "pcp_split", False)): + return True + return False + return bool(getattr(md, "pcp_split", False)) + + def _install_increment_version_pcp_shim() -> None: """Make ``torch.autograd.graph.increment_version`` tolerate a torch bug that only surfaces under PCP + torch.compile. @@ -1687,7 +1741,6 @@ def process_weights_after_loading(self): super().process_weights_after_loading() -@IndexerDecoratorForPluginMode def _indexer_with_output_fake( hidden_states: torch.Tensor, qr: torch.Tensor, @@ -1767,6 +1820,7 @@ def indexer_with_output( ) +@IndexerDecoratorForPluginMode class Indexer(nn.Module): def __init__( self, @@ -2022,7 +2076,15 @@ def __init__( self.num_heads = num_heads tp_size = get_tensor_model_parallel_world_size() assert num_heads % tp_size == 0 - self.num_local_heads = num_heads // tp_size + # RFC ROCm/ATOM#196: when the vLLM plugin reuses the TP group as the CP + # group, attention is token-parallel with *full-head replicated* weights + # (each rank runs ALL heads on its 1/cp query shard against the full, + # gathered KV), so heads are NOT sharded across the group. The q/kv/o + # projections then load the full tensor and o_proj skips its all-reduce. + self._attn_cp_full_head = plugin_attn_cp_enabled() and tp_size > 1 + self.num_local_heads = ( + num_heads if self._attn_cp_full_head else num_heads // tp_size + ) self.scaling = self.qk_head_dim**-0.5 self.max_position_embeddings = max_position_embeddings @@ -2068,6 +2130,11 @@ def __init__( else: base_quant_config = quant_config + # Head-parallel projections (q_b/q/kv_b) shard the head dim across TP by + # default; under reuse-TP-as-CP they become replicated (full heads). + HeadColParallel = ( + ReplicatedLinear if self._attn_cp_full_head else ColumnParallelLinear + ) if self.q_lora_rank is not None: # self.q_a_proj = ReplicatedLinear(self.hidden_size, # self.q_lora_rank, @@ -2088,7 +2155,7 @@ def __init__( # honors this flag in LinearBase.process_weights_after_loading. self.fused_qkv_a_proj.needs_preshuffled_weight = True self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) - self.q_b_proj = ColumnParallelLinear( + self.q_b_proj = HeadColParallel( q_lora_rank, self.num_heads * self.qk_head_dim, bias=False, @@ -2097,7 +2164,7 @@ def __init__( source_quant_dtype=source_quant_dtype, ) else: - self.q_proj = ColumnParallelLinear( + self.q_proj = HeadColParallel( self.hidden_size, self.num_heads * self.qk_head_dim, bias=False, @@ -2115,7 +2182,7 @@ def __init__( source_quant_dtype=source_quant_dtype, ) self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) - self.kv_b_proj = ColumnParallelLinear( + self.kv_b_proj = HeadColParallel( self.kv_lora_rank, self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, @@ -2127,15 +2194,28 @@ def __init__( source_quant_dtype if is_rocm_aiter_fp4bmm_enabled() else None ), ) - self.o_proj = RowParallelLinear( - self.num_heads * self.v_head_dim, - self.hidden_size, - bias=False, - quant_config=base_quant_config, - reduce_results=not ENABLE_ALLREDUCE_RMSNORM_FUSION, - prefix=f"{prefix}.o_proj", - source_quant_dtype=None, - ) + if self._attn_cp_full_head: + # Full-head replicated output projection: each rank already holds the + # complete [num_heads * v_head_dim] attention output for its query + # shard, so o_proj is a local GEMM with no cross-rank all-reduce. + self.o_proj = ReplicatedLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=base_quant_config, + prefix=f"{prefix}.o_proj", + source_quant_dtype=None, + ) + else: + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=base_quant_config, + reduce_results=not ENABLE_ALLREDUCE_RMSNORM_FUSION, + prefix=f"{prefix}.o_proj", + source_quant_dtype=None, + ) rope_params = config.rope_parameters rope_theta = rope_params.get("rope_theta") or 10000 @@ -2408,6 +2488,82 @@ def forward( ) +# Reuse-TP-as-CP FFN registry: layer_name -> DeepseekV2DecoderLayer. The opaque +# `cp_ffn` custom op below recovers the weight-bearing layer from here (custom +# ops can only take tensors / str / scalars, not module refs), mirroring how +# `indexer_with_output` recovers its module from static_forward_context. +_CP_FFN_MODULES: dict[str, nn.Module] = {} + + +def cp_ffn(hidden_states: torch.Tensor, layer_name: str) -> torch.Tensor: + """Dynamo-opaque wrapper around the reuse-TP-as-CP FFN collective. + + SHAPE-PRESERVING for BOTH runtime branches (input [N] -> output [N]), so it + is a clean leaf op for torch.compile: + * split ([T/cp] shard): all-gather -> mlp -> reduce-scatter -> [T/cp] + * unsplit ([T] full): mlp -> all-reduce -> [T] + Opacity — like `indexer_with_output` — is what defeats the bake: Dynamo never + traces the body, so the runtime `_pcp_active()` branch inside is evaluated + LIVE every forward instead of being frozen to its dummy-warmup value. Being a + leaf (not a graph *break*) it does not fragment the graph, so a pure-decode + batch (unsplit, no split anywhere) still captures as one FULL cudagraph. + Prefill is not cudagraph-captured, so the eager collectives inside run safely. + """ + self = _CP_FFN_MODULES[layer_name] + return self._cp_ffn_impl(hidden_states) + + +def _cp_ffn_fake(hidden_states: torch.Tensor, layer_name: str) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +direct_register_custom_op( + op_name="deepseek_v2_cp_ffn", + op_func=cp_ffn, + mutates_args=[], + fake_impl=_cp_ffn_fake, +) + + +# Reuse-TP-as-CP embedding registry: model_name -> DeepseekV2Model. +_CP_EMBED_MODULES: dict[str, nn.Module] = {} + + +def cp_embed(input_ids: torch.Tensor, model_name: str) -> torch.Tensor: + """Dynamo-opaque reuse-TP-as-CP token embedding. + + LEADING-DIM-PRESERVING (input_ids [N] -> embeddings [N, H]), so the outer + compiled graph sees a deterministic shape and the runtime `_pcp_active()` + branch inside never bakes: + * split ([T/cp] owned ids): all-gather the owned ids back to the full + replicated token set, run the vocab-parallel embedding there (so its + cross-rank all-reduce combines the SAME tokens on every rank), then + round-robin re-split to this rank's [T/cp] embeddings. + * unsplit ([T] ids, decode/mixed/dummy): plain vocab-parallel embedding. + Passing input_ids (not a pre-embedded inputs_embeds) keeps the arg pattern + identical to the non-CP path, so decode still embeds from its input_ids + cudagraph buffer and the single shared compiled graph serves both. + """ + self = _CP_EMBED_MODULES[model_name] + return self._cp_embed_impl(input_ids) + + +def _cp_embed_fake(input_ids: torch.Tensor, model_name: str) -> torch.Tensor: + self = _CP_EMBED_MODULES[model_name] + w = self.embed_tokens.weight + return torch.empty( + (input_ids.shape[0], w.shape[1]), dtype=w.dtype, device=input_ids.device + ) + + +direct_register_custom_op( + op_name="deepseek_v2_cp_embed", + op_func=cp_embed, + mutates_args=[], + fake_impl=_cp_embed_fake, +) + + class DeepseekV2DecoderLayer(nn.Module): def __init__( self, @@ -2494,6 +2650,20 @@ def uses_quantized_attn_input(layer_quant_config): self.input_norm_quant_type = attn_input_quant_config.quant_type.value self.fuse_input_norm_quant = False self.fuse_ar_input_norm = ENABLE_ALLREDUCE_RMSNORM_FUSION + # RFC ROCm/ATOM#196 (reuse-TP-as-CP): the residual stream is sequence + # parallel [T/cp]. Attention is full-head replicated (o_proj has no + # cross-rank partial) and the FFN is wrapped in all-gather/reduce-scatter + # at this layer (see forward), so NOTHING produces a plain TP partial to + # fold into a fused all-reduce norm. Disable AR fusion so the norms run + # as ordinary (local) RMSNorms on the [T/cp] shard. + self._attn_cp = ( + plugin_attn_cp_enabled() and get_tensor_model_parallel_world_size() > 1 + ) + self._cp_ffn_layer_name = prefix + if self._attn_cp: + self.fuse_ar_input_norm = False + # Register for the opaque `cp_ffn` custom op (see module top). + _CP_FFN_MODULES[prefix] = self # DSA models (e.g., GLM-5/DeepSeek-V3.2): the indexer's wk/weights_proj GEMMs # run in BF16 and consume the same normed activation, so the RMSNorm(+quant) # must also emit the pre-quant bf16 mirror. Gate the mirror on this layer @@ -2540,7 +2710,9 @@ def uses_quantized_attn_input(layer_quant_config): self.mlp = DeepseekV2MoE( config=config, quant_config=quant_config, - reduce_results=not self.fuse_ar_input_norm, + # Under CP the FFN must emit a TP *partial* so this layer's + # reduce-scatter can sum-and-scatter it back to [T/cp]. + reduce_results=(not self.fuse_ar_input_norm) and not self._attn_cp, prefix=f"{prefix}.mlp", alt_stream=alt_stream, ) @@ -2550,7 +2722,7 @@ def uses_quantized_attn_input(layer_quant_config): intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, - reduce_results=not self.fuse_ar_input_norm, + reduce_results=(not self.fuse_ar_input_norm) and not self._attn_cp, prefix=f"{prefix}.mlp", ) # Fuse activation quant into the AR+RMSNorm when the attention input @@ -2583,10 +2755,41 @@ def uses_quantized_attn_input(layer_quant_config): self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps, - fused_allreduce=ENABLE_ALLREDUCE_RMSNORM_FUSION, + fused_allreduce=ENABLE_ALLREDUCE_RMSNORM_FUSION and not self._attn_cp, ) self.routed_scaling_factor = config.routed_scaling_factor + def _cp_ffn_impl(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Reuse-TP-as-CP FFN reduce, run inside the opaque `cp_ffn` op. + + The FFN emits a TP *partial* here (mlp built with reduce_results=False + under CP); how it is reduced depends on whether THIS forward was + round-robin split at model entry (`_pcp_active()` — evaluated LIVE here + because the enclosing op is Dynamo-opaque): + * split ([T/cp] shard): all-gather to the full [T] token set so the + TP-sharded experts see every token, run the FFN, then reduce-scatter + the partial back to this rank's [T/cp] shard. all-gather is rank-major + and reduce-scatter is its exact inverse, so each rank recovers its own + round-robin shard, fully reduced (the FFN is per-token). + * not split (decode / short-dense / mixed / dummy stay on the full + replicated [T] tokens): plain all-reduce, exactly like the non-CP path. + Output shape always equals input shape (net shape-preserving), which is + what lets the wrapper op be a clean leaf for torch.compile. + """ + if _pcp_active(): + # Capture-safe all-gather + reduce-scatter over the DEDICATED CP group + # (its own ca_comm / graph-buffer slot allocator), so the split-decode + # MoE round-trip is legal AND correct inside a full CUDA graph. Both + # legs share the CP slot allocator (isolated from TP all-reduce), + # which the registered custom all-gather needs to replay correctly. + hidden_states = pcp_tp_all_gather_dim0(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = pcp_reduce_scatter_dim0(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + hidden_states = tensor_model_parallel_all_reduce(hidden_states) + return hidden_states + def forward( self, positions: torch.Tensor, @@ -2686,7 +2889,15 @@ def forward( # Fully Connected hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) - hidden_states = self.mlp(hidden_states) + if self._attn_cp: + # Reuse-TP-as-CP FFN boundary, routed through a Dynamo-opaque op so + # the runtime split branch (see `cp_ffn` / `_cp_ffn_impl`) is never + # baked to the dummy-warmup value under torch.compile. + hidden_states = torch.ops.aiter.deepseek_v2_cp_ffn( + hidden_states, self._cp_ffn_layer_name + ) + else: + hidden_states = self.mlp(hidden_states) if isinstance(self.mlp, DeepseekV2MLP) and hidden_states.dtype == torch.float16: # Fix FP16 overflow @@ -2726,6 +2937,18 @@ def __init__( else: self.embed_tokens = PPMissingLayer() + # Reuse-TP-as-CP: route embedding through the opaque `cp_embed` op so a + # round-robin-split rank still embeds its shard correctly (all-gather -> + # full vocab-parallel embed -> re-split). Process-constant gate. + self._cp_embed = ( + plugin_attn_cp_enabled() + and get_tensor_model_parallel_world_size() > 1 + and get_pp_group().is_first_rank + ) + self._cp_embed_name = prefix or "deepseek_v2_model" + if self._cp_embed: + _CP_EMBED_MODULES[self._cp_embed_name] = self + self.alt_stream: Optional[torch.cuda.Stream] = None if getattr(config, "n_shared_experts", None) is not None: self.alt_stream = torch.cuda.Stream() @@ -2761,7 +2984,22 @@ def __init__( ["hidden_states", "residual"], config.hidden_size ) + def _cp_embed_impl(self, input_ids: torch.Tensor) -> torch.Tensor: + """Reuse-TP-as-CP embedding, run inside the opaque `cp_embed` op.""" + if _pcp_active(): + ws = get_pcp_world_size() + n_owned = input_ids.shape[0] + # Owned [T/cp] ids -> full replicated [T_pad] (exact inverse of the + # model-entry round-robin split), embed there so the vocab-parallel + # all-reduce sees the same tokens on every rank, then re-split. + full_ids = pcp_allgather_rerange(input_ids, ws) + full_emb = self.embed_tokens(full_ids) + return pcp_round_robin_split(full_emb, ws)[:n_owned] + return self.embed_tokens(input_ids) + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + if self._cp_embed: + return torch.ops.aiter.deepseek_v2_cp_embed(input_ids, self._cp_embed_name) return self.embed_tokens(input_ids) def forward( @@ -2897,18 +3135,31 @@ def forward( 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. + # 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 the + # cache layout is untouched. Under ATOM_VLLM_ATTN_CP this fires for EVERY + # batch type -- prefill, pure plain-decode, and mixed -- via one shared + # whole-batch round-robin partition (the metadata builders stamp pcp_split, + # which _pcp_active() reads). It is skipped (identical to the original path) + # only when PCP is inactive (`pcp=1`) or the batch is left unstamped (a + # short-context dense-fallback prefill or a spec/MTP verify batch). 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 + # Round-robin split the RAW inputs (input_ids / positions), keeping the + # arg pattern identical to the non-split (decode) forward: input_ids is + # still input_ids, inputs_embeds still None. This is what lets the + # single compiled DeepseekV2Model graph serve both prefill (split) and + # decode (unsplit) — it just processes N tokens either way. The + # VocabParallelEmbedding all-reduce is made correct for this rank's + # DIFFERENT token shard by the Dynamo-opaque `cp_embed` op inside the + # model (all-gather owned ids -> embed the full replicated set -> + # re-split), never by embedding here (which would break decode's + # input_ids cudagraph and the shared-graph arg contract). 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( diff --git a/atom/plugin/config.py b/atom/plugin/config.py index 5280b95e95..98902c552e 100644 --- a/atom/plugin/config.py +++ b/atom/plugin/config.py @@ -208,6 +208,8 @@ def _generate_atom_config_from_vllm_config(config: Any) -> PluginConfig: # vLLM EP shards MoE across the flattened DP x TP device space (and # therefore disables fused shared experts); native uses per-DP MoE. moe_ep_flatten_tp_across_dp=vllm_parallel_config.enable_expert_parallel, + # Reuse the TP group as the Context-Parallel group (RFC ROCm/ATOM#196). + vllm_attn_cp=envs.ATOM_VLLM_ATTN_CP, enable_tbo=vllm_enable_dbo, enable_tbo_decode=vllm_enable_dbo, plugin_config=plugin_config, diff --git a/atom/plugin/graph_capture_patch.py b/atom/plugin/graph_capture_patch.py index 956ec5fd57..ae35ca1320 100644 --- a/atom/plugin/graph_capture_patch.py +++ b/atom/plugin/graph_capture_patch.py @@ -13,36 +13,72 @@ import functools import logging -from contextlib import contextmanager, nullcontext +from contextlib import ExitStack, contextmanager logger = logging.getLogger("atom") -def _get_aiter_ca_capture_context(): - """Lazily get aiter's ca_comm.capture() context, or nullcontext if unavailable.""" +def _ca_comm_of(group): + """Return a group's capturable aiter ca_comm, or None if unavailable.""" + if group is None: + return None + device_communicator = getattr(group, "device_communicator", None) + if device_communicator is None: + return None + ca_comm = getattr(device_communicator, "ca_comm", None) + if ca_comm is None or getattr(ca_comm, "disabled", True): + return None + if getattr(ca_comm, "capture", None) is None: + return None + return ca_comm + + +def _iter_aiter_capture_ca_comms(): + """Yield the aiter ca_comm(s) that must enter capture mode during graph + capture: the TP group, plus the dedicated Context-Parallel (_PCP) group when + it has its OWN ca_comm (reuse-TP-as-CP with a dedicated CP group). + + The model's custom collectives run on aiter's groups (TP all-reduce on TP; + CP all-gather / reduce-scatter on _PCP). Each ca_comm must be in capture + mode so its graph buffers register (flush_graph_buffers on exit); otherwise + its collectives take the non-registered / hipMemcpyAsync path. Deduped by + identity, so the aliased case (_PCP is TP) enters exactly once. + """ + getters = [] try: from aiter.dist.parallel_state import get_tp_group - aiter_tp = get_tp_group() + getters.append(get_tp_group) except Exception: - return nullcontext() - - if aiter_tp is None: - return nullcontext() - - device_communicator = getattr(aiter_tp, "device_communicator", None) - if device_communicator is None: - return nullcontext() - - aiter_ca_comm = getattr(device_communicator, "ca_comm", None) - if aiter_ca_comm is None or getattr(aiter_ca_comm, "disabled", True): - return nullcontext() - - capture_method = getattr(aiter_ca_comm, "capture", None) - if capture_method is None: - return nullcontext() + pass + try: + from aiter.dist.parallel_state import get_pcp_group - return capture_method() + getters.append(get_pcp_group) + except Exception: + pass + + seen = set() + for getter in getters: + try: + group = getter() + except Exception: + # get_pcp_group() asserts when _PCP is unset (non-CP runs): skip. + continue + ca_comm = _ca_comm_of(group) + if ca_comm is not None and id(ca_comm) not in seen: + seen.add(id(ca_comm)) + yield ca_comm + + +@contextmanager +def _get_aiter_ca_capture_context(): + """Nest ca_comm.capture() for every aiter group that needs it (TP + CP).""" + ca_comms = list(_iter_aiter_capture_ca_comms()) + with ExitStack() as stack: + for ca_comm in ca_comms: + stack.enter_context(ca_comm.capture()) + yield def _patched_graph_capture(original_graph_capture): diff --git a/atom/plugin/register.py b/atom/plugin/register.py index 3471b4155c..7932b1e0ca 100644 --- a/atom/plugin/register.py +++ b/atom/plugin/register.py @@ -722,6 +722,51 @@ def set_attn_cls() -> None: logger.info("Use Attention dispatcher for rtp-llm") +def _validate_plugin_attn_cp( + config: Config, tensor_parallel_size: int, use_vllm_atom_owned_ep: bool +) -> None: + """Fail fast on unsupported ATOM_VLLM_ATTN_CP (reuse-TP-as-CP) setups. + + The reuse path (RFC ROCm/ATOM#196) repurposes the whole TP group as the CP + group, so it needs tp>1, a DSA (sparse-MLA) model, and must not collide with + a separate `-pcp` dimension, DP-attention, or the DP+EP aiter-owned-groups + path. `pcp=1` / cp-off stays a bit-exact no-op and never reaches here. + """ + if not config.plugin_config.is_vllm: + raise ValueError( + "ATOM_VLLM_ATTN_CP is a vLLM-plugin-only feature; it is set but the " + "current plugin backend is not vLLM." + ) + if tensor_parallel_size <= 1: + raise ValueError( + "ATOM_VLLM_ATTN_CP requires tensor_parallel_size > 1 (it reuses the " + f"TP group as the CP group), got tp={tensor_parallel_size}." + ) + if getattr(config.hf_config, "index_topk", None) is None: + raise ValueError( + "ATOM_VLLM_ATTN_CP only supports DSA sparse-MLA models (configs with " + "`index_topk`, e.g. GLM-5.2 / DeepSeek-V3.2)." + ) + if getattr(config, "prefill_context_parallel_size", 1) > 1: + raise ValueError( + "ATOM_VLLM_ATTN_CP reuses the TP group as the CP group and is " + "mutually exclusive with a separate `-pcp` " + f"(prefill_context_parallel_size={config.prefill_context_parallel_size})." + ) + if getattr(config, "enable_dp_attention", False): + raise ValueError("ATOM_VLLM_ATTN_CP is not compatible with DP-attention.") + if use_vllm_atom_owned_ep: + raise ValueError( + "ATOM_VLLM_ATTN_CP is not compatible with the DP+EP aiter-owned-" + "groups path (it needs the reused vLLM TP group)." + ) + logger.info( + "ATOM_VLLM_ATTN_CP validated: reuse-TP-as-CP over %d ranks for a DSA " + "sparse-MLA model.", + tensor_parallel_size, + ) + + def init_aiter_dist(config: Config) -> None: """ Initialize aiter dist for using aiter custom collective op. @@ -756,11 +801,23 @@ def init_aiter_dist(config: Config) -> None: "Skip vLLM TP reuse for OOT DP+EP so aiter owns TP/PP/DP/EP groups." ) + if getattr(config, "vllm_attn_cp", False): + _validate_plugin_attn_cp(config, tensor_parallel_size, use_vllm_atom_owned_ep) + if config.plugin_config.is_vllm and not use_vllm_atom_owned_ep: from atom.plugin.vllm.tp_group_reuse import init_aiter_dist_from_vllm if init_aiter_dist_from_vllm(tensor_parallel_size): return + if getattr(config, "vllm_attn_cp", False): + # The reuse-TP-as-CP path aliases aiter `_PCP` inside + # init_aiter_dist_from_vllm; the init_dist_env fallback below does + # not, so CP would silently no-op. Fail loudly instead. + raise RuntimeError( + "ATOM_VLLM_ATTN_CP requires reusing vLLM's TP group, but " + "init_aiter_dist_from_vllm() failed. Check that vLLM's TP group " + "and aiter ca_comm are available." + ) # Fallback: create aiter's own groups (vLLM reuse failed or non-vLLM plugin) from aiter import init_dist_env diff --git a/atom/plugin/vllm/attention/layer_mla.py b/atom/plugin/vllm/attention/layer_mla.py index ec838070a5..44a10f2668 100644 --- a/atom/plugin/vllm/attention/layer_mla.py +++ b/atom/plugin/vllm/attention/layer_mla.py @@ -5,7 +5,7 @@ import aiter import torch -from aiter import dtypes, fused_qk_rope_concat_and_cache_mla +from aiter import concat_and_cache_mla, dtypes, fused_qk_rope_concat_and_cache_mla from aiter.mla import mla_decode_fwd from aiter.ops.triton import ( batched_gemm_a16wfp4 as _fp4_bmm_module, @@ -1250,6 +1250,10 @@ def forward_impl_sparse( sparse_meta = attn_metadata num_actual_toks = sparse_meta.num_actual_tokens + # Plugin reuse-TP-as-CP: this rank holds 1/cp round-robin owned queries + # (num_actual_toks == n_owned) but still writes the FULL replicated KV. + # Mirrors native MLAAttention.forward_impl PCP block. + pcp = bool(getattr(sparse_meta, "pcp_split", False)) # Inputs and outputs may be padded for CUDA graphs output_padded = output @@ -1270,7 +1274,35 @@ def forward_impl_sparse( "positions" ] - positions = positions[:num_actual_toks] + positions_full = None + k_c_normed_full = None + k_pe_full = None + if pcp: + from atom.distributed.pcp_utils import ( + get_pcp_world_size, + pcp_allgather_rerange, + pcp_pad_dense, + pcp_pad_len, + pcp_round_robin_split, + ) + + pcp_ws = get_pcp_world_size() + n_real = sparse_meta.slot_mapping.shape[0] + positions_full = positions[:n_real] + n_pad = pcp_pad_len(n_real, pcp_ws) - n_real + # Owned (round-robin) positions rope the owned q and the throwaway + # owned-slot k write; positions_full ropes the full completion write. + positions = pcp_round_robin_split( + pcp_pad_dense(positions_full, n_pad), pcp_ws + )[:num_actual_toks] + # Gather raw (un-roped) owned k to full BEFORE the fused kernel ropes + # k_pe in place (k_pe is [n_owned, 1, rope]; k_c_normed [n_owned, lora]). + k_c_normed_full = pcp_allgather_rerange(k_c_normed, pcp_ws)[:n_real] + k_pe_full = pcp_allgather_rerange(k_pe, pcp_ws)[:n_real] + write_slot_mapping = sparse_meta.slot_mapping_owned + else: + positions = positions[:num_actual_toks] + write_slot_mapping = sparse_meta.slot_mapping fp8_attention = self.kv_cache_dtype.startswith("fp8") if fp8_attention: from vllm.platforms import current_platform @@ -1319,16 +1351,17 @@ def forward_impl_sparse( device=ql_nope.device, ) if kv_cache.numel() > 0: + kv_cache_view = kv_cache.view( + kv_cache.shape[0], -1, self.kv_lora_rank + self.qk_rope_head_dim + ) fused_qk_rope_concat_and_cache_mla( ql_nope, q_pe, k_c_normed, k_pe.squeeze(1), - kv_cache.view( - kv_cache.shape[0], -1, self.kv_lora_rank + self.qk_rope_head_dim - ), + kv_cache_view, q_out, - sparse_meta.slot_mapping, + write_slot_mapping, self._k_scale, self._q_scale, positions, @@ -1337,6 +1370,21 @@ def forward_impl_sparse( is_neox=self.rotary_emb.is_neox_style, is_nope_first=True, ) + if pcp: + # Complete the FULL replicated k-cache: rope the gathered full k + # in place with full positions, then overwrite every real slot + # (the fused kernel's owned-slot write above is throwaway). The + # rope kernel is 2-component and needs a non-None partner, so pass + # a throwaway of matching length. + self.rotary_emb(positions_full, k_pe_full, torch.empty_like(k_pe_full)) + concat_and_cache_mla( + k_c_normed_full, + k_pe_full.squeeze(1), + kv_cache, + sparse_meta.slot_mapping.flatten(), + kv_cache_dtype=self.kv_cache_dtype, + scale=self._k_scale, + ) if self.head_repeat_factor > 1: q_out = q_out.repeat_interleave(self.head_repeat_factor, dim=1) diff --git a/atom/plugin/vllm/attention/metadata.py b/atom/plugin/vllm/attention/metadata.py index 1339638839..8e74c854f0 100644 --- a/atom/plugin/vllm/attention/metadata.py +++ b/atom/plugin/vllm/attention/metadata.py @@ -9,6 +9,17 @@ from aiter.dist.parallel_state import get_dp_group, get_tp_group from aiter.jit.utils.chip_info import get_gfx from atom.config import get_current_atom_config +from atom.distributed.pcp_utils import ( + get_pcp_rank, + get_pcp_world_size, + pcp_is_enabled, + pcp_pad_dense, + pcp_pad_len, + pcp_reindex_decode, + pcp_round_robin_query_indices, + pcp_sparse_prefill_reindex, + plugin_attn_cp_enabled, +) from atom.model_ops.attention_mla import _MLA_MIN_HEADS from atom.plugin.vllm.attention.layer_mla import ( disabled_mla_persistent_metadata, @@ -289,6 +300,12 @@ class AiterMlaSparseIndexerMetadataForVllm: decode: AiterMlaSparseIndexerDecodeMetadataForVllm | None = None prefill: AiterMlaSparseIndexerPrefillMetadataForVllm | None = None + # See AiterMlaSparseMetadataForVllm.pcp_split. Stamped True when the indexer + # prefill chunks' query-indexed fields (cu_seqlen_ks/ke, token_to_seq, the + # token_start/token_end ranges) have been reduced to this rank's owned queries + # while the KV-gather fields (block_table, cu_seq_lens) stay FULL. + pcp_split: bool = False + # TODO (zyongye) optimize this, this is now vibe coded def kv_spans_from_batches( @@ -398,6 +415,14 @@ class AiterMlaSparseMetadataForVllm: reduce_final_map: torch.Tensor | None = None reduce_partial_map: torch.Tensor | None = None + # Plugin reuse-TP-as-CP (ATOM_VLLM_ATTN_CP): stamped True by the builder when + # the query-indexed fields above have been reduced to this rank's round-robin + # 1/cp owned queries. deepseek_v2._plugin_pcp_active() reads this so the model + # split/gather stays in lock-step with the metadata reindex. slot_mapping stays + # FULL (KV write); slot_mapping_owned is the owned-query subset for fused kernels. + pcp_split: bool = False + slot_mapping_owned: torch.Tensor | None = None + @dataclass class MinimaxM3SparsePrefillMetadata: @@ -1627,6 +1652,69 @@ def build( return attn_metadata +def _plugin_cp_should_split(common_attn_metadata, topk_tokens, decode_threshold=1): + """Batch-global split gate shared by the plugin sparse builders (main + sparse-MLA and sparse indexer), mirrored by deepseek_v2._plugin_pcp_active(). + + UNIVERSAL query-dim split: when reuse-TP-as-CP is enabled (``ATOM_VLLM_ATTN_CP``) + and the aiter _PCP group is aliased to the TP group (world > 1), split EVERY + batch type -- pure prefill, pure plain-decode, and MIXED -- over the query dim + using ONE shared whole-batch round-robin partition. Each rank runs the whole + model forward on its ``T/cp`` query shard with full (replicated) KV; the single + exit all-gather restores the full token order before lm_head. This removes the + old replicated-attention fallback (full 64-head attention on every rank). + + Correctness constraints: + * Any decode region must be single-token plain decode built by a non-spec + builder (``decode_threshold == 1``) so token index == request index for the + owned-decode slice. Spec/MTP verify batches (``decode_threshold > 1``) keep + decode replicated. + * PREFILL has NO seq_len gate. The plugin indexer has no dense-MHA fallback + (it always runs the sparse top-k, selecting all keys for short rows), and + every plugin CP call site keys off the pcp_split stamp with a length- + agnostic K all-gather + owned reindex, so short-context prefill splits + correctly. (The native path DOES have a dense fallback and keeps the gate; + this plugin gate no longer does.) + * PLAIN-DECODE has NO seq_len gate. The round-robin split is over the QUERY + dim with full replicated KV (mathematically correct at any seq_len) and the + decode path always runs the sparse indexer with a ``min(seq_len, topk)`` + clamp (no dense fallback to desync). Matches native deepseek_v4._pcp_active. + """ + if not (plugin_attn_cp_enabled() and pcp_is_enabled()): + return False + num_tokens = int(common_attn_metadata.num_actual_tokens) + if num_tokens <= 0: + return False + from vllm.v1.attention.backends.utils import split_decodes_and_prefills + + _, _, num_decode_tokens, num_prefill_tokens = split_decodes_and_prefills( + common_attn_metadata, decode_threshold=decode_threshold + ) + # Any decode region present must be plain single-token decode built by a + # non-spec builder (decode_threshold == 1) so token index == request index for + # the owned-decode slice; spec/MTP verify batches -- whose multi-token-per- + # request decode layout is not wired for the round-robin split -- stay + # replicated (return False for pure- and mixed-decode spec batches). + if int(num_decode_tokens) > 0 and int(decode_threshold) != 1: + return False + if int(num_prefill_tokens) > 0: + # PREFILL-containing (pure prefill or MIXED): ALWAYS split over the query + # dim, at any seq_len. The plugin indexer (sparse_attn_indexer_plugin_mode) + # has NO dense-MHA fallback -- unlike the native indexer it never early- + # returns on max_seqlen_k <= index_topk; it always runs the sparse top-k, + # which selects all keys when the row is shorter than topk (dense- + # equivalent). Every plugin CP call site (entry split, indexer K all-gather + # + separate q/k rope, exit gather, sparse-MLA) keys off the pcp_split + # stamp, not seq_len, and the K all-gather + owned per-query reindex fire + # for any length, so short-context prefill splits correctly (topk_tokens is + # therefore no longer consulted here). + return True + # PURE plain-decode: always split under CP, no seq_len gate. The round-robin + # split is over the query dim with full replicated KV (correct at any seq_len) + # and the decode indexer clamps min(seq_len, topk) with no dense fallback. + return int(common_attn_metadata.max_query_len) == 1 + + class AiterMlaSparseMetadataBuilder(AttentionMetadataBuilder): """vLLM-only metadata builder for sparse MLA main attention.""" @@ -1745,6 +1833,15 @@ def __init__( self._prev_indices_extent = 0 self._prev_metadata_key = None + # Persistent owned-slot buffer for the PCP decode-split path (see + # _build_pcp_split). FULL cudagraph bakes the address of every tensor a + # captured kernel reads; the reindex helper returns fresh per-step + # allocations, so the owned slot_mapping must be copied into this stable + # buffer that build() refreshes in place on every replay. + self._pcp_slot_mapping_owned = torch.empty( + (max_num_batched_tokens,), dtype=torch.int64, device=device + ) + def _bind_shared_sparse_kv_indices(self, layer_names, config, device, numel): # Resolve and bind a single shared paged_kv_indices buffer. # Reuse the buffer the other ubatch already bound if it exists, otherwise @@ -1810,6 +1907,108 @@ def _resolve_indexer(layer_name): ) return shared_buffer + def _pcp_should_split(self, common_attn_metadata, num_tokens) -> bool: + return _plugin_cp_should_split( + common_attn_metadata, + self.topk_tokens, + decode_threshold=self.reorder_batch_threshold, + ) + + def _build_pcp_split(self, common_attn_metadata, num_tokens, sparse_seqlen): + """Owned-query metadata for the plugin reuse-TP-as-CP sparse path. + + Mirrors native ``AiterMLAImpl._apply_pcp_reindex``: query-indexed fields + (qo_indptr, paged_kv_indptr, paged_kv_last_page_len, req_id_per_token) are + reduced to the round-robin owned subset and the work-split schedule is + rebuilt for them; slot_mapping / block_table / seq_lens stay FULL for the + replicated full-KV write, with slot_mapping_owned carrying the owned subset + for the fused q_out kernel. + """ + reindex = pcp_sparse_prefill_reindex( + sparse_seqlen, + self.req_id_per_token_buffer[:num_tokens], + common_attn_metadata.slot_mapping, + self.topk_tokens, + ) + n_owned = reindex["n_owned"] + # FULL cudagraph safety: route the owned-query metadata through this + # builder's persistent buffers so captured kernels read stable device + # addresses that build() refreshes in place on every replay. The reindex + # helper returns freshly-allocated per-step tensors (it mirrors the native + # eager prefill path); handing those straight to a captured decode kernel + # freezes it on the capture-time allocation and replays read stale memory. + # qo_indptr is arange and paged_kv_last_page_len is all-ones by + # construction, so the persistent buffers already hold those exact values. + qo_indptr = self.qo_indptr[: n_owned + 1] + paged_kv_last_page_len = self.paged_kv_last_page_len[:n_owned] + self.paged_kv_indptr[: n_owned + 1].copy_( + reindex["paged_kv_indptr"], non_blocking=True + ) + paged_kv_indptr = self.paged_kv_indptr[: n_owned + 1] + self.req_id_per_token_buffer[:n_owned].copy_( + reindex["req_id_per_token"], non_blocking=True + ) + req_id_per_token = self.req_id_per_token_buffer[:n_owned] + self._pcp_slot_mapping_owned[:n_owned].copy_( + reindex["slot_mapping_owned"], non_blocking=True + ) + slot_mapping_owned = self._pcp_slot_mapping_owned[:n_owned] + # The indexer writes owned queries' topk into the shared buffer; slice to + # the owned extent so sparse MLA reads exactly the owned rows. + paged_kv_indices = self.paged_kv_indices[: n_owned * self.topk_tokens] + + get_mla_metadata_v1( + qo_indptr, + paged_kv_indptr, + paged_kv_last_page_len, + self.padded_num_heads, + 1, + True, + self._mla_work_meta_data, + self._mla_work_info_set, + self._mla_work_indptr, + self._mla_reduce_indptr, + self._mla_reduce_final_map, + self._mla_reduce_partial_map, + page_size=1, + kv_granularity=16, + max_seqlen_qo=1, + uni_seqlen_qo=1, + fast_mode=True, + dtype_q=get_aiter_kv_cache_dtype(self.vllm_config), + dtype_kv=get_aiter_kv_cache_dtype(self.vllm_config), + ) + # Owned qo/kv layout differs from the full-batch decode fast-path key, so + # invalidate it (the fast-path is only hit on the non-CP full path). + self._prev_metadata_key = None + + return AiterMlaSparseMetadataForVllm( + num_reqs=common_attn_metadata.num_reqs, + max_query_len=common_attn_metadata.max_query_len, + max_seq_len=common_attn_metadata.max_seq_len, + seq_lens=common_attn_metadata.seq_lens, + num_actual_tokens=n_owned, + query_start_loc=common_attn_metadata.query_start_loc, + slot_mapping=common_attn_metadata.slot_mapping, + block_table=common_attn_metadata.block_table_tensor, + req_id_per_token=req_id_per_token, + block_size=self.kv_cache_spec.block_size, + attn_out_dtype=self.model_dtype, + topk_tokens=self.topk_tokens, + qo_indptr=qo_indptr, + paged_kv_last_page_len=paged_kv_last_page_len, + paged_kv_indices=paged_kv_indices, + paged_kv_indptr=paged_kv_indptr, + work_meta_data=self._mla_work_meta_data, + work_indptr=self._mla_work_indptr, + work_info_set=self._mla_work_info_set, + reduce_indptr=self._mla_reduce_indptr, + reduce_final_map=self._mla_reduce_final_map, + reduce_partial_map=self._mla_reduce_partial_map, + pcp_split=True, + slot_mapping_owned=slot_mapping_owned, + ) + def build(self, common_prefix_len, common_attn_metadata, fast_build=False): num_tokens = common_attn_metadata.num_actual_tokens starts = common_attn_metadata.query_start_loc_cpu.to(torch.int32) @@ -1861,6 +2060,16 @@ def build(self, common_prefix_len, common_attn_metadata, fast_build=False): num_tokens, common_attn_metadata.max_query_len, ) + + # Plugin reuse-TP-as-CP: reduce the query-indexed fields to this rank's + # round-robin 1/cp owned queries (KV columns stay full). Kept in + # lock-step with deepseek_v2._plugin_pcp_active(), which reads the + # pcp_split stamp set here. + if self._pcp_should_split(common_attn_metadata, num_tokens): + return self._build_pcp_split( + common_attn_metadata, num_tokens, sparse_seqlen + ) + torch.cumsum( sparse_seqlen, dim=0, @@ -1974,6 +2183,10 @@ def build(self, common_prefix_len, common_attn_metadata, fast_build=False): class AiterMlaSparseIndexerMetadataBuilder(AttentionMetadataBuilder): _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH reorder_batch_threshold = 1 + # One-shot diagnostic: set True after the first decode-split build so we log + # that the PCP reuse-TP-as-CP decode split actually engaged (used to confirm + # FULL-cudagraph parity/accuracy runs exercise the split path). + _pcp_decode_split_logged = False def __init__( self, @@ -2059,12 +2272,31 @@ def __init__( dtype=torch.int32, device=self.device, ) + # Persistent owned-query buffers for the PCP decode-split path (see + # _build_indexer). pcp_reindex_decode returns fresh per-step tensors; + # under FULL cudagraph the captured indexer decode kernel bakes their + # addresses, so the owned seq_lens / block_table are copied into these + # stable buffers that build() refreshes in place on every replay. + self._pcp_seq_lens_owned = torch.empty( + (max_num_batched_tokens,), dtype=torch.int32, device=self.device + ) + self._pcp_block_table_owned = torch.empty( + (max_num_batched_tokens, max_num_blocks_per_req), + dtype=torch.int32, + device=self.device, + ) self.scheduler_metadata_buffer = torch.empty( (self.num_sms + 1, 2), dtype=torch.int32, device=self.device ) def _build_indexer_one_prefill_chunk( - self, reqs_start, reqs_end, query_start_loc_cpu, seq_lens_cpu, block_table + self, + reqs_start, + reqs_end, + query_start_loc_cpu, + seq_lens_cpu, + block_table, + owned_q=None, ): prefill_query_start_loc = ( query_start_loc_cpu[reqs_start : reqs_end + 1] @@ -2091,6 +2323,20 @@ def _build_indexer_one_prefill_chunk( .to(torch.int32) .to(self.device) ) + if owned_q is not None: + # Plugin reuse-TP-as-CP: keep only this rank's round-robin owned QUERY + # tokens of the chunk. cu_seqlen_ks/ke are per-local-query KV windows + # into the FULL KV -> select owned; token_start/token_end move into the + # compacted owned-q array space (q_fp8/weights/topk are [n_owned]). + # KV-gather fields (block_table, cu_seq_lens, token_to_seq, + # total_seq_lens) stay FULL so every rank still gathers the full KV. + os = int((owned_q < token_start).sum().item()) + oe = int((owned_q < token_end).sum().item()) + local = (owned_q[os:oe] - token_start).to(cu_seqlen_ks.device) + cu_seqlen_ks = cu_seqlen_ks[local].contiguous() + cu_seqlen_ke = cu_seqlen_ke[local].contiguous() + token_start = os + token_end = oe return AiterMlaSparseIndexerPrefillChunkMetadataForVllm( cu_seqlen_ks=cu_seqlen_ks, cu_seqlen_ke=cu_seqlen_ke, @@ -2132,6 +2378,19 @@ def _build_indexer( assert num_decodes + num_prefills == num_reqs assert num_decode_tokens + num_prefill_tokens == num_tokens + # Plugin reuse-TP-as-CP: same batch-global gate as the main sparse builder + # (prefill-only), so the pcp_split stamp and the query reindex stay in + # lock-step with deepseek_v2._plugin_pcp_active() and the model split. + topk_tokens = self.model_config.hf_config.index_topk + pcp_split = _plugin_cp_should_split( + common_attn_metadata, topk_tokens, self.reorder_batch_threshold + ) + owned_q = None + if pcp_split: + owned_q = pcp_round_robin_query_indices( + pcp_pad_len(num_tokens, get_pcp_world_size()), get_pcp_world_size() + ) + prefill_metadata = None if num_prefills > 0: chunk_seq_ids = split_prefill_chunks( @@ -2146,6 +2405,7 @@ def _build_indexer( query_start_loc_cpu, common_attn_metadata.seq_lens_cpu, common_attn_metadata.block_table_tensor, + owned_q=owned_q, ) for reqs_start, reqs_end in chunk_seq_ids ] @@ -2153,25 +2413,106 @@ def _build_indexer( chunks=chunks, ) + # Plugin reuse-TP-as-CP: decide this rank's owned decode queries UP FRONT + # so the decode-metadata block below is skipped entirely when this rank + # owns none. A MIXED batch restricts the shared whole-batch round-robin + # partition (owned_q) to the decode region, which can leave a rank with + # zero owned decode queries; pure decode pads to cp so every rank owns the + # same count. pcp_dec carries the owned (seq_lens, block_table) tensors. + pcp_dec = None + if pcp_split and num_decodes > 0: + seq_lens_full = common_attn_metadata.seq_lens[:num_decodes] + block_table_full = common_attn_metadata.block_table_tensor[ + :num_decodes, ... + ] + # Padded CUDA graph requests have block_table entries of -1. Clamp to 0 + # to prevent OOB access in the DeepGEMM kernel (padded rows have + # seq_len=0, so they produce no meaningful output). + block_table_full.clamp_(min=0) + if num_prefills > 0: + # MIXED: the owned decode queries MUST come from the SAME + # whole-batch round-robin partition (owned_q) that the prefill + # chunks and the sparse-MLA _build_pcp_split reindex use, restricted + # to the decode region [0, num_decode_tokens). The decode region is + # NOT padded separately (unlike pure decode) -- padding it would + # shift the prefill token_start offsets (os/oe are computed over the + # same owned_q in _build_indexer_one_prefill_chunk) and desync the + # topk rows from the owned paged_kv_indptr the sparse layer reads. + # Mixed batches run piecewise/eager, so per-rank-variable owned + # decode counts (differing by at most 1) are fine. + owned_dec_q = owned_q[owned_q < num_decode_tokens].to( + seq_lens_full.device + ) + n_owned_dec = int(owned_dec_q.shape[0]) + seq_lens_owned = seq_lens_full.index_select(0, owned_dec_q).contiguous() + block_table_owned = block_table_full.index_select( + 0, owned_dec_q + ).contiguous() + else: + # PURE decode: pad the decode region to a multiple of cp so every + # rank runs the same owned count (FULL cudagraph uniform shape); + # dummy queries attend nothing and are dropped at the model exit + # all-gather. Same round-robin start/stride as the combined owned_q, + # so the owned decode set matches owned_q < num_decode_tokens. + seq_lens_owned, block_table_owned, _owned_dec_q, n_owned_dec = ( + pcp_reindex_decode( + seq_lens_full, block_table_full, num_decode_tokens + ) + ) + if not AiterMlaSparseIndexerMetadataBuilder._pcp_decode_split_logged: + AiterMlaSparseIndexerMetadataBuilder._pcp_decode_split_logged = True + logger.info( + "[PCP] indexer DECODE-split ENGAGED (%s): pre_split_decode=%d " + "n_owned=%d num_prefills=%d max_seq_len=%d topk=%d", + "mixed" if num_prefills > 0 else "pure-decode", + int(num_decode_tokens), + int(n_owned_dec), + int(num_prefills), + int(common_attn_metadata.max_seq_len), + int(topk_tokens), + ) + num_decodes = n_owned_dec + num_decode_tokens = n_owned_dec + pcp_dec = (seq_lens_owned, block_table_owned) + decode_metadata = None if num_decodes > 0: - torch.diff( - common_attn_metadata.query_start_loc[: num_decodes + 1], - out=self.decode_lens_buffer[:num_decodes], - ) - decode_lens = self.decode_lens_buffer[:num_decodes] - decode_lens_cpu = torch.diff( - common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] - ) + if pcp_dec is not None: + # FULL cudagraph safety: route the owned tensors through persistent + # buffers so a captured indexer decode kernel reads stable addresses + # that are refreshed in place on every replay (see __init__). The + # sparse-MLA slot_mapping stays FULL for the replicated KV write. + seq_lens_owned, block_table_owned = pcp_dec + self._pcp_seq_lens_owned[:num_decodes].copy_( + seq_lens_owned, non_blocking=True + ) + seq_lens = self._pcp_seq_lens_owned[:num_decodes] + bw = block_table_owned.shape[1] + self._pcp_block_table_owned[:num_decodes, :bw].copy_( + block_table_owned, non_blocking=True + ) + block_table = self._pcp_block_table_owned[:num_decodes, :bw] + self.decode_lens_buffer[:num_decodes].fill_(1) + decode_lens = self.decode_lens_buffer[:num_decodes] + decode_lens_cpu = torch.ones(num_decodes, dtype=torch.int32) + else: + torch.diff( + common_attn_metadata.query_start_loc[: num_decodes + 1], + out=self.decode_lens_buffer[:num_decodes], + ) + decode_lens = self.decode_lens_buffer[:num_decodes] + decode_lens_cpu = torch.diff( + common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] + ) - seq_lens = common_attn_metadata.seq_lens[:num_decodes] - block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] + seq_lens = common_attn_metadata.seq_lens[:num_decodes] + block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] - # Padded CUDA graph requests have block_table entries of -1. - # Clamp to 0 to prevent OOB access in the DeepGEMM kernel. - # This is safe because padded requests have seq_lens=0, so the - # kernel produces no meaningful output for those rows. - block_table.clamp_(min=0) + # Padded CUDA graph requests have block_table entries of -1. + # Clamp to 0 to prevent OOB access in the DeepGEMM kernel. + # This is safe because padded requests have seq_lens=0, so the + # kernel produces no meaningful output for those rows. + block_table.clamp_(min=0) max_decode_len = int(decode_lens_cpu.max().item()) if max_decode_len > 1: @@ -2279,6 +2620,7 @@ def _build_indexer( num_prefill_tokens=num_prefill_tokens, prefill=prefill_metadata, decode=decode_metadata, + pcp_split=pcp_split, ) return indexer_metadata diff --git a/atom/plugin/vllm/tp_group_reuse.py b/atom/plugin/vllm/tp_group_reuse.py index c91e6e6893..7c4c874adf 100644 --- a/atom/plugin/vllm/tp_group_reuse.py +++ b/atom/plugin/vllm/tp_group_reuse.py @@ -21,6 +21,76 @@ logger = logging.getLogger("atom") +def _create_aiter_device_comm(vllm_group: Any, comm_unique_name: str) -> Any: + """Create an aiter CudaCommunicator (with its own ca_comm) bound to vLLM's + existing ProcessGroups. Returns None if custom all-reduce is unavailable. + + Each CudaCommunicator owns an INDEPENDENT ``CustomAllreduce`` (ca_comm): its + own IPC pool, rank-data buffer and graph-buffer slot allocator. Creating a + second one over the same ranks is how CP gets a collective communicator that + is isolated from TP's. + """ + from aiter.dist.device_communicators.communicator_cuda import CudaCommunicator + + device_communicator = CudaCommunicator( + cpu_group=vllm_group.cpu_group, + device=vllm_group.device, + device_group=vllm_group.device_group, + unique_name=comm_unique_name, + ) + if device_communicator.ca_comm is None or device_communicator.ca_comm.disabled: + return None + return device_communicator + + +def _build_aiter_adapter( + vllm_group: Any, device_comm: Any, adapter_unique_name: str +) -> Any: + """Wrap vLLM's ProcessGroups + an aiter device_communicator in an aiter + GroupCoordinator, WITHOUT creating new ProcessGroups. + + Inherits all reduce/gather/broadcast methods from GroupCoordinator; only + __init__ is overridden to reuse existing groups. Registers itself under + ``adapter_unique_name`` so aiter's ``outplace_all_gather`` / + ``outplace_reduce_scatter`` custom ops (which resolve ``_groups[group_name]``) + dispatch to THIS group's ca_comm. + """ + from aiter.dist.parallel_state import GroupCoordinator as AiterGroupCoordinator + from aiter.dist.parallel_state import _register_group + + class AiterGroupAdapter(AiterGroupCoordinator): + def __init__(self) -> None: + # Skip GroupCoordinator.__init__ (it creates new ProcessGroups). + self.unique_name = adapter_unique_name + _register_group(self) + self.rank = vllm_group.rank + self.local_rank = vllm_group.local_rank + self.ranks = vllm_group.ranks + self.world_size = vllm_group.world_size + self.rank_in_group = vllm_group.rank_in_group + self.cpu_group = vllm_group.cpu_group + self.device_group = vllm_group.device_group + self.device = vllm_group.device + self.use_device_communicator = True + self.device_communicator = device_comm + self.mq_broadcaster = None + + def destroy(self) -> None: + # cpu_group / device_group are vLLM's -- owned and torn down by vLLM. + # Only release the aiter device_communicator (ca_comm/IPC) we created; + # NEVER destroy the borrowed ProcessGroups (a second reuse adapter + # over the same groups would otherwise double-free them at shutdown). + dc = getattr(self, "device_communicator", None) + if dc is not None: + try: + dc.destroy() + except Exception: # noqa: BLE001 + pass + self.device_communicator = None + + return AiterGroupAdapter() + + def _create_aiter_tp_adapter_from_vllm() -> Any: """Create aiter-compatible TP adapter using vLLM's TP groups and aiter's ca_comm.""" import vllm.distributed.parallel_state as vllm_ps @@ -29,19 +99,8 @@ def _create_aiter_tp_adapter_from_vllm() -> Any: if vllm_tp.world_size == 1: return None - # Import aiter components - must use aiter's CudaCommunicator with aiter's ca_comm - from aiter.dist.device_communicators.communicator_cuda import CudaCommunicator - from aiter.dist.parallel_state import _register_group - - # Create aiter CudaCommunicator with vLLM's ProcessGroups (no new groups created) - device_communicator = CudaCommunicator( - cpu_group=vllm_tp.cpu_group, - device=vllm_tp.device, - device_group=vllm_tp.device_group, - unique_name="tp", - ) - - if device_communicator.ca_comm is None or device_communicator.ca_comm.disabled: + device_communicator = _create_aiter_device_comm(vllm_tp, "tp") + if device_communicator is None: logger.warning( "ATOM tp_group_reuse: aiter ca_comm not available on vLLM's TP group, " "caller will fall back to standard aiter distributed initialization " @@ -49,32 +108,38 @@ def _create_aiter_tp_adapter_from_vllm() -> Any: ) return None - # Inherit from aiter GroupCoordinator - only override __init__ to use existing groups. - # All reduce/gather/broadcast methods are inherited, no need to reimplement. - from aiter.dist.parallel_state import GroupCoordinator as AiterGroupCoordinator + return _build_aiter_adapter(vllm_tp, device_communicator, "tp:0") - class AiterTPAdapter(AiterGroupCoordinator): - """Reuse vLLM's TP groups + aiter's ca_comm. Inherits all methods from GroupCoordinator.""" - def __init__(self, vllm_tp: Any, device_comm: Any): - # Skip GroupCoordinator.__init__ (it creates new ProcessGroups). - # Set attributes directly to match what parent expects. - self.unique_name = "tp:0" - _register_group(self) - self.rank = vllm_tp.rank - self.local_rank = vllm_tp.local_rank - self.ranks = vllm_tp.ranks - self.world_size = vllm_tp.world_size - self.rank_in_group = vllm_tp.rank_in_group - self.cpu_group = vllm_tp.cpu_group - self.device_group = vllm_tp.device_group - self.device = vllm_tp.device - self.use_device_communicator = True - self.device_communicator = device_comm - self.mq_broadcaster = None +def _create_aiter_cp_adapter_from_vllm(tensor_model_parallel_size: int) -> Any: + """Create a DEDICATED Context-Parallel group over the SAME ranks as TP, with + its OWN ca_comm. - adapter = AiterTPAdapter(vllm_tp, device_communicator) - return adapter + Rationale: the reuse-TP-as-CP path previously aliased ``_PCP`` to the TP + adapter, so CP's all-gather / reduce-scatter (token embed, indexer, MoE) + shared TP's single ca_comm -- one IPC pool, one rank-data buffer, one + graph-buffer slot allocator. Under a FULL CUDA graph over a split decode + batch that piles thousands of CP all-gather graph-buffer registrations, + interleaved with TP all-reduce, onto that one slot allocator; the registered + all_gather path then reads wrong peer slots and the graph replays garbage. + A dedicated ca_comm gives CP an isolated slot space (see the FULL-graph + all_gather bug in atom/distributed/pcp_utils.py:pcp_dim0_all_gather). + + Returns None (caller aliases TP) if custom all-reduce is unavailable. + """ + import vllm.distributed.parallel_state as vllm_ps + + vllm_tp = vllm_ps.get_tp_group() + if vllm_tp.world_size == 1: + return None + + device_communicator = _create_aiter_device_comm(vllm_tp, "pcp") + if device_communicator is None: + return None + + cp_adapter = _build_aiter_adapter(vllm_tp, device_communicator, "pcp:0") + _setup_ca_comm_signal(cp_adapter, tensor_model_parallel_size) + return cp_adapter def _setup_ca_comm_signal(adapter: Any, tensor_model_parallel_size: int) -> None: @@ -116,6 +181,52 @@ def init_aiter_dist_from_vllm(tensor_model_parallel_size: int) -> bool: ) # EP may not exist in all vLLM configs _setup_ca_comm_signal(adapter, tensor_model_parallel_size) + # RFC ROCm/ATOM#196: run the Context-Parallel dimension over the TP ranks. + # Setting aiter's `_PCP` makes get_pcp_group() / + # get_prefill_context_model_parallel_{world_size,rank}() and + # pcp_allgather_rerange() operate over the TP ranks, so the native + # round-robin split/gather/indexer paths (all keyed off pcp_is_enabled) + # light up. The reuse-specific model rebuild (full-head attention, MoE + # all-gather/reduce-scatter, no fused-AR norms) is gated separately on + # plugin_attn_cp_enabled() at construction time. + # + # `_PCP` gets a DEDICATED ca_comm (own IPC pool / rank-data / graph-buffer + # slot allocator) instead of aliasing TP's, so CP's all-gather + + # reduce-scatter don't share TP all-reduce's slot allocator -- required + # for the split-decode collectives to be correct inside a FULL CUDA graph + # (see _create_aiter_cp_adapter_from_vllm). Falls back to aliasing TP if + # a second ca_comm can't be created. + from atom.utils import envs + + if envs.ATOM_VLLM_ATTN_CP and adapter.world_size > 1: + cp_adapter = None + try: + cp_adapter = _create_aiter_cp_adapter_from_vllm( + tensor_model_parallel_size + ) + except Exception as e: # noqa: BLE001 + logger.warning( + "ATOM plugin: dedicated CP group init failed (%s); aliasing " + "TP as the CP group instead", + e, + ) + cp_adapter = None + + if cp_adapter is not None: + aiter_ps._PCP = cp_adapter # type: ignore[attr-defined] + logger.info( + "ATOM plugin: created DEDICATED Context-Parallel (_PCP) group " + "with its own ca_comm over %d ranks (ATOM_VLLM_ATTN_CP=1)", + cp_adapter.world_size, + ) + else: + aiter_ps._PCP = adapter # type: ignore[attr-defined] + logger.info( + "ATOM plugin: reused vLLM TP group as the Context-Parallel " + "(_PCP) group over %d ranks (ATOM_VLLM_ATTN_CP=1)", + adapter.world_size, + ) + from aiter.dist.parallel_state import set_custom_all_reduce set_custom_all_reduce(True) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index c3725c7bb4..3f0cd09cae 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -150,6 +150,29 @@ "ATOM_DISABLE_VLLM_PLUGIN": lambda: ( os.getenv("ATOM_DISABLE_VLLM_PLUGIN", "0").lower() == "1" ), + # vLLM plugin Prefill/Context Parallel by *reusing the TP group* as the CP + # group (RFC ROCm/ATOM#196). Unlike native ATOM's separate `-pcp` dimension + # (world = tp x pcp, KV replicated per PCP rank), this runs true sequence + # parallelism over tokens within the SAME TP group: residual/hidden are + # sharded [T/cp], attention runs full-head *replicated* weights on 1/cp of + # the queries against the full (gathered) KV, and MoE does all-gather + + # reduce-scatter around the TP-sharded experts. DSA (index_topk) models only + # (GLM-5.2 / DeepSeek-V3.2). `0` (default) is a bit-exact no-op. + "ATOM_VLLM_ATTN_CP": lambda: (os.getenv("ATOM_VLLM_ATTN_CP", "0").lower() == "1"), + # Extend reuse-TP-as-CP (ATOM_VLLM_ATTN_CP) to also round-robin split the + # DECODE token batch across the CP ranks (default off). With CP on but this + # off, decode runs full-head *replicated* attention on ALL tokens per rank + # (redundant attention + indexer compute, no token split). With this on, a + # PURE-decode batch is split 1/cp the same way prefill is: each rank runs + # q-proj / attention / indexer / MoE on its 1/cp owned decode queries against + # the full (all-gathered) KV, then the hidden states are all-gathered at the + # model exit. This trades an exit all-gather (+ per-layer KV all-gather) for + # cp x less attention/indexer/q-side compute, so it pays off at concurrency + # >> cp and can regress at low concurrency. Plain decode only (max_query_len + # == 1); mixed and spec/MTP batches stay replicated. `0` is a no-op. + "ATOM_VLLM_ATTN_CP_DECODE_SPLIT": lambda: ( + os.getenv("ATOM_VLLM_ATTN_CP_DECODE_SPLIT", "0").lower() == "1" + ), "ATOM_USE_CUSTOM_ALL_GATHER": lambda: ( os.getenv("ATOM_USE_CUSTOM_ALL_GATHER", "1").lower() == "1" ), From b657ec3ad1832f6f8a0493d42ea846ab447bd04f Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Wed, 15 Jul 2026 08:30:15 +0000 Subject: [PATCH 6/7] join all gather Signed-off-by: whx-sjtu --- atom/plugin/vllm/attention/layer_mla.py | 41 +++++++++++++++---------- 1 file changed, 24 insertions(+), 17 deletions(-) diff --git a/atom/plugin/vllm/attention/layer_mla.py b/atom/plugin/vllm/attention/layer_mla.py index 44a10f2668..3640d0b868 100644 --- a/atom/plugin/vllm/attention/layer_mla.py +++ b/atom/plugin/vllm/attention/layer_mla.py @@ -1274,9 +1274,6 @@ def forward_impl_sparse( "positions" ] - positions_full = None - k_c_normed_full = None - k_pe_full = None if pcp: from atom.distributed.pcp_utils import ( get_pcp_world_size, @@ -1290,15 +1287,12 @@ def forward_impl_sparse( n_real = sparse_meta.slot_mapping.shape[0] positions_full = positions[:n_real] n_pad = pcp_pad_len(n_real, pcp_ws) - n_real - # Owned (round-robin) positions rope the owned q and the throwaway - # owned-slot k write; positions_full ropes the full completion write. + # Owned (round-robin) positions rope the owned q and the owned k_pe; + # the full replicated k-cache is completed with a single all-gather + # after the fused kernel (see the `if pcp` write block below). positions = pcp_round_robin_split( pcp_pad_dense(positions_full, n_pad), pcp_ws )[:num_actual_toks] - # Gather raw (un-roped) owned k to full BEFORE the fused kernel ropes - # k_pe in place (k_pe is [n_owned, 1, rope]; k_c_normed [n_owned, lora]). - k_c_normed_full = pcp_allgather_rerange(k_c_normed, pcp_ws)[:n_real] - k_pe_full = pcp_allgather_rerange(k_pe, pcp_ws)[:n_real] write_slot_mapping = sparse_meta.slot_mapping_owned else: positions = positions[:num_actual_toks] @@ -1371,15 +1365,28 @@ def forward_impl_sparse( is_nope_first=True, ) if pcp: - # Complete the FULL replicated k-cache: rope the gathered full k - # in place with full positions, then overwrite every real slot - # (the fused kernel's owned-slot write above is throwaway). The - # rope kernel is 2-component and needs a non-None partner, so pass - # a throwaway of matching length. - self.rotary_emb(positions_full, k_pe_full, torch.empty_like(k_pe_full)) + # Complete the FULL replicated k-cache with a SINGLE all-gather. + # The fused kernel above read the RAW owned k as const inputs (it + # writes roped K to the owned cache slots + Q into q_out, never back + # to k_pe/kv_c), so k_pe is still un-roped here. Rope THIS rank's + # 1/cp owned k_pe in place (owned positions), pack it with the + # rope-free owned kv_c into one [n_owned, kv_lora + pe] tensor, + # all-gather+rerange to the full sequence, then write every real + # slot (overwriting the fused kernel's throwaway owned-slot write). + # Roping the owned shard before the gather (instead of the whole + # gathered sequence on every rank) keeps the rope at 1/cp, and + # merging kv_c + k_pe halves the collective count (one 576-wide + # all-gather vs two). concat_and_cache_mla consumes the gathered + # nope/pe as plain column slices: it honours stride(0) and the + # feature dim stays unit-stride, so no split/contiguous copy is + # needed on the write side. + self.rotary_emb(positions, k_pe, torch.empty_like(k_pe)) + kv_full = pcp_allgather_rerange( + torch.cat([k_c_normed, k_pe.squeeze(1)], dim=-1), pcp_ws + )[:n_real] concat_and_cache_mla( - k_c_normed_full, - k_pe_full.squeeze(1), + kv_full[:, : self.kv_lora_rank], + kv_full[:, self.kv_lora_rank :], kv_cache, sparse_meta.slot_mapping.flatten(), kv_cache_dtype=self.kv_cache_dtype, From 1b3a124c4a514966ad5dfd5169e2e033dfc23464 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Thu, 16 Jul 2026 16:34:37 +0000 Subject: [PATCH 7/7] gather weights Signed-off-by: whx-sjtu --- atom/config.py | 1 + atom/distributed/pcp_utils.py | 50 ---- atom/model_ops/attention_mla.py | 32 ++- atom/model_ops/linear.py | 137 ++++++++++- atom/models/deepseek_v2.py | 114 +++++---- atom/plugin/vllm/attention/cp_gather.py | 296 ++++++++++++++++++++++++ atom/plugin/vllm/attention/layer_mla.py | 162 +++++++++++++ atom/plugin/vllm/attention/metadata.py | 146 +++++------- atom/plugin/vllm/attention/ops.py | 46 ++++ atom/utils/envs.py | 28 +-- recipes/GLM-5.md | 2 - recipes/atom_vllm/GLM-5.md | 28 +++ 12 files changed, 826 insertions(+), 216 deletions(-) create mode 100644 atom/plugin/vllm/attention/cp_gather.py diff --git a/atom/config.py b/atom/config.py index 6aeeb13f1d..7791ef8c85 100644 --- a/atom/config.py +++ b/atom/config.py @@ -254,6 +254,7 @@ def set_splitting_ops_for_v1(self): "aiter.mla_attention", "aiter.atom_vllm_mha_attention", "aiter.atom_vllm_mla_attention", + "aiter.atom_vllm_mla_attention_cp", ] diff --git a/atom/distributed/pcp_utils.py b/atom/distributed/pcp_utils.py index 4589fa339b..f5b1f3767c 100644 --- a/atom/distributed/pcp_utils.py +++ b/atom/distributed/pcp_utils.py @@ -44,19 +44,6 @@ def plugin_attn_cp_enabled() -> bool: return bool(envs.ATOM_VLLM_ATTN_CP) -def plugin_attn_cp_decode_split_enabled() -> bool: - """True when reuse-TP-as-CP additionally round-robin splits the DECODE batch. - - Requires :func:`plugin_attn_cp_enabled` (this flag only extends that path). - When off (default) decode stays full/replicated (every rank runs all tokens - with full heads); when on, a pure-decode batch is split 1/cp exactly like - prefill. See ``ATOM_VLLM_ATTN_CP_DECODE_SPLIT`` in ``envs.py``. - """ - from atom.utils import envs - - return bool(envs.ATOM_VLLM_ATTN_CP_DECODE_SPLIT) and plugin_attn_cp_enabled() - - def get_pcp_world_size() -> int: return get_prefill_context_model_parallel_world_size() @@ -394,43 +381,6 @@ def pcp_sparse_prefill_reindex( } -def pcp_reindex_decode( - seq_lens: torch.Tensor, # [num_decode_tokens] per-query cached KV length - block_table: torch.Tensor, # [num_decode_tokens, max_blocks] per-query blocks - num_decode_tokens: int, - pcp_size: Optional[int] = None, - pcp_rank: Optional[int] = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: - """Reduce plain-decode per-query metadata to this rank's round-robin queries. - - Plain decode has exactly one query per request, so ``num_decode_tokens == - num_decodes`` and owning a round-robin subset of decode QUERIES is the same - as owning that subset of requests. The global decode-token count is padded to - a multiple of ``pcp_size`` with dummy (``seq_len == 0``, ``block_table == 0``) - queries so every rank runs the same ``n_owned == padded / pcp_size`` count; - the dummies attend nothing (their sparse-MLA KV segment is zero-length) and - their hidden output is dropped after the model's exit all-gather. - - Only these per-owned-query rows shrink here; the sparse-MLA ``slot_mapping`` - (KV write) stays FULL, mirroring the prefill reindex, so every rank still - writes the full replicated KV for the new decode tokens. - - Returns ``(seq_lens_owned, block_table_owned, owned_q, n_owned)``. - """ - device = seq_lens.device - padded_total = pcp_pad_len(num_decode_tokens, pcp_size) - n_pad = padded_total - num_decode_tokens - owned_q = pcp_round_robin_query_indices(padded_total, pcp_size, pcp_rank).to(device) - n_owned = int(owned_q.shape[0]) - seq_lens_owned = pcp_pad_dense(seq_lens[:num_decode_tokens], n_pad)[ - owned_q - ].contiguous() - block_table_owned = pcp_pad_dense(block_table[:num_decode_tokens], n_pad)[ - owned_q - ].contiguous() - return seq_lens_owned, block_table_owned, owned_q, n_owned - - def pcp_reindex_ragged( kv_indptr: torch.Tensor, # [T_global + 1] int32 — global per-query prefix sum kv_indices: torch.Tensor, # [kv_indptr[-1]] — ragged packed values diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index de4f9c809d..55b85e02bc 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -184,11 +184,25 @@ class MLAModules: def dynamic_per_batched_tensor_quant( - x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn + x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn, cross_rank_group=None ): + """Per-(whole-)tensor dynamic fp8 quant of an absorbed BMM weight (W_K / W_V). + + ``cross_rank_group`` (reuse-TP-as-CP): when set, the amax is reduced with MAX + across the group so EVERY rank quantizes its head shard with the SAME scalar + scale. That makes the per-rank fp8 shards tile bit-exactly into the full-head + weight after an all-gather (a per-rank local amax would differ across ranks, so + the gathered bytes could not share one scalar). The stored scale is that shared + global scalar, so the runtime gather never has to move the scale. + """ DTYPE_MAX = torch.finfo(dtype).max min_val, max_val = x.aminmax() amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10) + if cross_rank_group is not None: + import torch.distributed as dist + + amax = amax.clone() + dist.all_reduce(amax, op=dist.ReduceOp.MAX, group=cross_rank_group) scale = DTYPE_MAX / amax x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX) return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal() @@ -347,11 +361,23 @@ def process_weights_after_loading(self): ) W_K = W_UK.transpose(0, 1) # 16 512 128 W_V = W_UV.permute(1, 2, 0) # 16 128 512 + # Reuse-TP-as-CP: CP layers gather these per-head-sharded fp8 weights to + # full heads on demand (prefill/mixed), so they must be quantized with a + # SHARED cross-rank scalar scale (see dynamic_per_batched_tensor_quant). + # head dim is dim 0 for both W_K [H, L, P] and W_V [H, P, L]. + cross_rank_group = None + if getattr(self, "_cp_attn", False): + from vllm.distributed.parallel_state import get_tp_group + + tp = get_tp_group() + if tp.world_size > 1: + cross_rank_group = tp.device_group + self._cp_wk_wv_head_dim = 0 self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant( - W_K, dtype=dtypes.fp8 + W_K, dtype=dtypes.fp8, cross_rank_group=cross_rank_group ) self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant( - W_V, dtype=dtypes.fp8 + W_V, dtype=dtypes.fp8, cross_rank_group=cross_rank_group ) @mark_trace(prefix="v_up_proj_and_o_proj", torch_compile=False) diff --git a/atom/model_ops/linear.py b/atom/model_ops/linear.py index 62c920137c..9dc8c047f9 100644 --- a/atom/model_ops/linear.py +++ b/atom/model_ops/linear.py @@ -540,17 +540,98 @@ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): def _gather_full_weight(self, weight): """Gather sharded weight from all TP ranks to reconstruct the full unpartitioned weight.""" + return self._gather_full_weight_via(weight, get_tp_group()) + + def _gather_full_weight_via(self, weight, group): + """Like :meth:`_gather_full_weight` but over an explicit group coordinator. + + ``group`` must expose ``all_gather(tensor, dim=...)`` (a vLLM + GroupCoordinator). Reuse-TP-as-CP runtime weight gather passes a DEDICATED + gather group here (distinct RCCL communicator) so a background gather stream + can run concurrently with the MoE's TP/CP collectives without sharing a + communicator (concurrent collectives on one communicator can deadlock). + """ if self.tp_size <= 1 or self.tp_dim is None: return weight # NCCL cannot all_gather E8M0 scales (MXFP8 source); gather the raw # bytes as uint8 and reinterpret afterwards. The gather only moves # bytes, so this is bit-exact. if weight.dtype == dtypes.fp8_e8m0: - gathered = get_tp_group().all_gather( - weight.view(torch.uint8), dim=self.tp_dim - ) + gathered = group.all_gather(weight.view(torch.uint8), dim=self.tp_dim) return gathered.view(dtypes.fp8_e8m0) - return get_tp_group().all_gather(weight, dim=self.tp_dim) + return group.all_gather(weight, dim=self.tp_dim) + + def gather_full_weight_scale(self, group=None): + """Reconstruct the full unpartitioned ``(weight, weight_scale)`` from this + rank's TP shard, for reuse-TP-as-CP runtime weight gather. + + Bit-exact with the offline full weight iff the shard-local quantized layout + (16x16 preshuffle + per-block/per-tensor scales) tiles cleanly along + ``tp_dim`` -- i.e. concatenating the preshuffled shards along ``tp_dim`` + equals preshuffling the full weight. This holds for ColumnParallel (tp_dim=0, + head-block boundary) and is the top hardware-validation item for RowParallel + (tp_dim=1, o_proj) where the shard boundary is on the reduction dim. + """ + g = group if group is not None else get_tp_group() + w = self._gather_full_weight_via(self.weight.data, g) + ws = None + if getattr(self, "weight_scale", None) is not None: + ws = self.weight_scale.data + # Only gather the scale along tp_dim if it is actually sharded there + # (mirror of _shard_quantized_weight). A per-output-channel per_Token + # scale on a RowParallel layer (o_proj) is computed over the FULL row at + # (online-)quant time and REPLICATED across ranks -- gathering it would + # wrongly tile it and corrupt the GEMM. Block/tensor scales that ARE + # sharded along tp_dim (per_1x128/per_1x32) are still gathered. + if ( + self.tp_dim is not None + and ws.dim() > self.tp_dim + and ws.shape[self.tp_dim] > 1 + ): + ws = self._gather_full_weight_via(ws, g) + return w, ws + + def forward_full( + self, + x, + x_scale, + full_weight, + full_weight_scale, + full_output_size, + otype=dtypes.bf16, + ): + """Run this TP-sharded linear as if it held the full unpartitioned weight, + with NO cross-rank reduce. Used by the reuse-TP-as-CP attention op on the + gather path (q_b projects the full head set; o_proj is a local GEMM on the + already-full attention output). The full weight/scale come from + :meth:`gather_full_weight_scale` (possibly prefetched on a gather stream). + + The full tensors are bound transiently (restored in ``finally``); this runs + inside the Dynamo-opaque CP attention op on the compute stream and only for + prefill/mixed batches (never cudagraph-captured), so the transient rebind is + invisible to torch.compile and safe. + """ + saved_w = self.weight.data + saved_ws = ( + self.weight_scale.data + if getattr(self, "weight_scale", None) is not None + else None + ) + saved_out = self.output_size + saved_reduce = self.reduce_results + try: + self.weight.data = full_weight + if saved_ws is not None and full_weight_scale is not None: + self.weight_scale.data = full_weight_scale + self.output_size = full_output_size + self.reduce_results = False + return self.forward(x, x_scale, otype=otype) + finally: + self.weight.data = saved_w + if saved_ws is not None: + self.weight_scale.data = saved_ws + self.output_size = saved_out + self.reduce_results = saved_reduce def _shard_quantized_weight(self, q_weight, weight_scale): """Split the quantized full weight and scale back to this rank's shard.""" @@ -694,6 +775,21 @@ def online_quantize_weight(self): "quant_dtype": str(online_quant_dtype), } + def _weight_preshuffle_enabled(self) -> bool: + """Whether this layer's fp8 blockscale weight uses aiter's 16x16 preshuffle. + + Defaults to the global env, but reuse-TP-as-CP force-disables it on layers + whose TP shards are gathered at runtime along the REDUCTION dim (RowParallel + ``o_proj``): preshuffled shards do NOT tile there (only ColumnParallel + head-dim shards do -- validated: q_proj gather is bit-exact, o_proj is not), + so the naive gathered full weight is corrupt. Unshuffled fp8 tiles + bit-exactly along both dims and keeps fp8 (the matching non-preshuffle + blockscale GEMM is selected in :meth:`forward`). + """ + if getattr(self, "cp_disable_preshuffle", False): + return False + return bool(envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE) + def process_weights_after_loading(self): # Re-quantize before process_weights if online quantization is enabled if self.quant_config is not None and self.quant_config.online_quant: @@ -764,20 +860,28 @@ def process_weights_after_loading(self): # The triton a8w8 per_Token GEMM consumes the unshuffled (N, K) # weight; only the AITER bpreshuffle fallback needs the shuffle. and not (use_triton_gemm() and gemm_a8w8_triton is not None) + # Reuse-TP-as-CP o_proj must stay UNSHUFFLED so its shards tile + # bit-exactly when gathered along the reduction dim at runtime; + # forward() runs the matching non-preshuffle a8w8 GEMM. + and not getattr(self, "cp_disable_preshuffle", False) ) or ( self.quant_type == QuantType.per_1x32 and (not is_fp4_blockscale or not use_fp4_non_shuffle_triton_gemm()) ) # per_1x128 only needs shuffle when using the preshuffle GEMM path if not need_shuffle and self.quant_type == QuantType.per_1x128: - need_shuffle = envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + need_shuffle = self._weight_preshuffle_enabled() # Modules whose fused forward calls a *preshuffle* blockscale GEMM # directly (e.g. DeepSeek fused qkv_a_proj) need the 16x16-shuffled # weight even under the non-preshuffle path # (ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE=0). Shuffle once here at - # load time instead of per-forward. - if not envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE and getattr( - self, "needs_preshuffled_weight", False + # load time instead of per-forward. Reuse-TP-as-CP layers + # (cp_disable_preshuffle) never take this override -- they MUST stay + # unshuffled so runtime gather along the reduction dim is bit-exact. + if ( + not self._weight_preshuffle_enabled() + and getattr(self, "needs_preshuffled_weight", False) + and not getattr(self, "cp_disable_preshuffle", False) ): need_shuffle = True if need_shuffle: @@ -809,7 +913,7 @@ def forward( # non-preshuffle GEMM expects row-major x_scale quant_func = functools_partial( self.quant_func, - transpose_scale=envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE, + transpose_scale=self._weight_preshuffle_enabled(), ) if self.quant_type.value != QuantType.per_1x32.value: x, x_scale = quant_func( @@ -856,6 +960,19 @@ def forward( bias=self.bias, dtype=otype, ) + elif getattr(self, "cp_disable_preshuffle", False): + # Reuse-TP-as-CP: this o_proj stores UNSHUFFLED per-token fp8 (the + # per-output-channel scale is over the full row and replicated), so + # its shards tile bit-exactly when gathered along the reduction dim. + # Run the matching non-preshuffle a8w8 GEMM (still fully fp8). + y = gemm_a8w8( + x, + self.weight, + x_scale, + self.weight_scale, + self.bias, + dtype=otype, + ) else: y = gemm_a8w8_bpreshuffle( x, @@ -867,7 +984,7 @@ def forward( if self.bias is not None: y += self.bias elif self.quant_type.value == QuantType.per_1x128.value: - if envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE: + if self._weight_preshuffle_enabled(): y = gemm_a8w8_blockscale_preshuffle_impl( x, self.weight, diff --git a/atom/models/deepseek_v2.py b/atom/models/deepseek_v2.py index 51d9cec0e8..e2855e7860 100644 --- a/atom/models/deepseek_v2.py +++ b/atom/models/deepseek_v2.py @@ -2061,6 +2061,7 @@ def __init__( prefix: str = "", layer_num: int = 0, use_indexer_wk_weights_proj_fusion: Optional[bool] = None, + is_mtp_block: bool = False, ) -> None: super().__init__() self.hidden_size = hidden_size @@ -2076,15 +2077,20 @@ def __init__( self.num_heads = num_heads tp_size = get_tensor_model_parallel_world_size() assert num_heads % tp_size == 0 - # RFC ROCm/ATOM#196: when the vLLM plugin reuses the TP group as the CP - # group, attention is token-parallel with *full-head replicated* weights - # (each rank runs ALL heads on its 1/cp query shard against the full, - # gathered KV), so heads are NOT sharded across the group. The q/kv/o - # projections then load the full tensor and o_proj skips its all-reduce. - self._attn_cp_full_head = plugin_attn_cp_enabled() and tp_size > 1 - self.num_local_heads = ( - num_heads if self._attn_cp_full_head else num_heads // tp_size + # RFC ROCm/ATOM#196 (weight-gather revision): the vLLM plugin reuses the TP + # group as the CP group. Attention weights STAY TP-sharded in memory (heads + # sharded 1/tp across the group, exactly like plain TP) -- we do NOT store a + # full-head replicated copy per layer. Instead, prefill / mixed batches + # gather the full q_b/kv_b/o_proj weights on demand per-layer at runtime (see + # AttentionForVllmMLA.forward_impl_cp + cp_gather), run full-head token- + # parallel attention on the 1/cp query shard, then release the transients; + # plain decode keeps running the sharded weights on the classic TP path. + # MTP blocks are excluded (plain TP only, no CP / no gather). + self._attn_cp = ( + plugin_attn_cp_enabled() and tp_size > 1 and not is_mtp_block ) + self._attn_cp_tp_size = tp_size + self.num_local_heads = num_heads // tp_size self.scaling = self.qk_head_dim**-0.5 self.max_position_embeddings = max_position_embeddings @@ -2130,11 +2136,10 @@ def __init__( else: base_quant_config = quant_config - # Head-parallel projections (q_b/q/kv_b) shard the head dim across TP by - # default; under reuse-TP-as-CP they become replicated (full heads). - HeadColParallel = ( - ReplicatedLinear if self._attn_cp_full_head else ColumnParallelLinear - ) + # Head-parallel projections (q_b/q/kv_b) shard the head dim across TP. Under + # reuse-TP-as-CP they stay sharded here and are gathered to full heads on + # demand at runtime (prefill/mixed); nothing is stored replicated. + HeadColParallel = ColumnParallelLinear if self.q_lora_rank is not None: # self.q_a_proj = ReplicatedLinear(self.hidden_size, # self.q_lora_rank, @@ -2194,28 +2199,29 @@ def __init__( source_quant_dtype if is_rocm_aiter_fp4bmm_enabled() else None ), ) - if self._attn_cp_full_head: - # Full-head replicated output projection: each rank already holds the - # complete [num_heads * v_head_dim] attention output for its query - # shard, so o_proj is a local GEMM with no cross-rank all-reduce. - self.o_proj = ReplicatedLinear( - self.num_heads * self.v_head_dim, - self.hidden_size, - bias=False, - quant_config=base_quant_config, - prefix=f"{prefix}.o_proj", - source_quant_dtype=None, - ) - else: - self.o_proj = RowParallelLinear( - self.num_heads * self.v_head_dim, - self.hidden_size, - bias=False, - quant_config=base_quant_config, - reduce_results=not ENABLE_ALLREDUCE_RMSNORM_FUSION, - prefix=f"{prefix}.o_proj", - source_quant_dtype=None, - ) + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + quant_config=base_quant_config, + # Under reuse-TP-as-CP the o_proj reduce is performed *inside* the CP + # attention op (local GEMM on the split/gathered path, manual all-reduce + # on the unsplit decode path), so the module itself never reduces. + reduce_results=( + False if self._attn_cp else (not ENABLE_ALLREDUCE_RMSNORM_FUSION) + ), + prefix=f"{prefix}.o_proj", + source_quant_dtype=None, + ) + if self._attn_cp: + # Reuse-TP-as-CP: o_proj is RowParallel (sharded on the reduction dim). + # Its fp8 shards are gathered to full heads at runtime for the split + # (prefill/mixed) path; the aiter 16x16 preshuffle does NOT tile along + # the reduction dim, so a naive byte-gather of preshuffled shards is + # corrupt (validated: q_proj/ColumnParallel gather is bit-exact, o_proj + # is not). Keep o_proj UNSHUFFLED fp8 (tiles bit-exactly along K) and run + # the matching non-preshuffle blockscale GEMM; still fully fp8-quantized. + self.o_proj.cp_disable_preshuffle = True rope_params = config.rope_parameters rope_theta = rope_params.get("rope_theta") or 10000 @@ -2342,6 +2348,17 @@ def __init__( mla_modules=mla_modules, prefix=prefix, ) + # Reuse-TP-as-CP runtime weight gather: tell the attention op whether this + # layer participates in CP (sharded weights gathered to full heads on + # demand for prefill/mixed) and give it the full (pre-shard) head count and + # TP size it needs to reconstruct the full-head shapes. MTP blocks pass + # _attn_cp=False so they always run the plain sharded TP path. + self.mla_attn.setup_cp_weight_gather( + attn_cp=self._attn_cp, + num_heads_full=self.num_heads, + tp_size=self._attn_cp_tp_size, + hidden_size=self.hidden_size, + ) # Enable q/k RMSNorm + q quant fusion for FP8 and FP4. The larger # qkv_a_proj + reduce + RMSNorm + quant fusion remains gated by @@ -2601,6 +2618,7 @@ def __init__( prefix=f"{prefix}.self_attn", layer_num=layer_num, use_indexer_wk_weights_proj_fusion=use_indexer_wk_weights_proj_fusion, + is_mtp_block=is_mtp_block, ) # Keep input RMSNorm quant fusion narrow: non-Triton FP8 activation quant is only supported for per-token layouts. @@ -2650,14 +2668,17 @@ def uses_quantized_attn_input(layer_quant_config): self.input_norm_quant_type = attn_input_quant_config.quant_type.value self.fuse_input_norm_quant = False self.fuse_ar_input_norm = ENABLE_ALLREDUCE_RMSNORM_FUSION - # RFC ROCm/ATOM#196 (reuse-TP-as-CP): the residual stream is sequence - # parallel [T/cp]. Attention is full-head replicated (o_proj has no - # cross-rank partial) and the FFN is wrapped in all-gather/reduce-scatter - # at this layer (see forward), so NOTHING produces a plain TP partial to - # fold into a fused all-reduce norm. Disable AR fusion so the norms run - # as ordinary (local) RMSNorms on the [T/cp] shard. + # RFC ROCm/ATOM#196 (reuse-TP-as-CP): for prefill/mixed the residual stream + # is sequence parallel [T/cp]. The attention op emits already-reduced output + # (o_proj local on the gathered full heads, or a manual all-reduce on the + # unsplit decode path) and the FFN is wrapped in all-gather/reduce-scatter at + # this layer (see forward), so NOTHING produces a plain TP partial to fold + # into a fused all-reduce norm. Disable AR fusion so the norms run as + # ordinary (local) RMSNorms. MTP blocks stay on plain TP (no CP wrapping). self._attn_cp = ( - plugin_attn_cp_enabled() and get_tensor_model_parallel_world_size() > 1 + plugin_attn_cp_enabled() + and get_tensor_model_parallel_world_size() > 1 + and not is_mtp_block ) self._cp_ffn_layer_name = prefix if self._attn_cp: @@ -2783,6 +2804,15 @@ def _cp_ffn_impl(self, hidden_states: torch.Tensor) -> torch.Tensor: # legs share the CP slot allocator (isolated from TP all-reduce), # which the registered custom all-gather needs to replay correctly. hidden_states = pcp_tp_all_gather_dim0(hidden_states) + # Kick the NEXT CP layer's attention weight all-gather only AFTER the + # pre-MoE activation all-gather above has been issued, so the two + # collectives don't fight for NIC/CU bandwidth. The weight gather runs on + # its dedicated stream/communicator and now overlaps the MoE compute + # below (the gather stream waits on the current stream, i.e. on the + # activation all-gather, before it starts). + from atom.plugin.vllm.attention.cp_gather import prefetch_next_layer + + prefetch_next_layer(self.self_attn.mla_attn) hidden_states = self.mlp(hidden_states) hidden_states = pcp_reduce_scatter_dim0(hidden_states) else: diff --git a/atom/plugin/vllm/attention/cp_gather.py b/atom/plugin/vllm/attention/cp_gather.py new file mode 100644 index 0000000000..9c8db0b6fd --- /dev/null +++ b/atom/plugin/vllm/attention/cp_gather.py @@ -0,0 +1,296 @@ +"""Reuse-TP-as-CP runtime attention weight gather (RFC ROCm/ATOM#196, weight-gather +revision). + +Under reuse-TP-as-CP the attention projection weights (``q_b``/``kv_b``/``o_proj`` +and the absorbed BMM weights ``W_K``/``W_V``) stay TP-sharded in memory -- we do NOT +keep a full-head replicated copy per layer. For prefill / mixed batches (which run +full-head token-parallel attention on this rank's 1/cp query shard) the FULL weights +for a layer are gathered on demand right before that layer's attention, used, then +released. This module owns: + + * the per-layer weight gather itself (``gather_attn_weights``), used both to + prefetch into a slot and as the synchronous fallback; + * the OVERLAP pipeline (intrinsic to CP -- no separate flag): a dedicated RCCL + communicator + a background CUDA stream + an ordered CP-attn layer registry + a + 2-slot double buffer, so layer L+1's gather runs on the side stream concurrently + with layer L's MoE compute. A dedicated communicator is REQUIRED: the MoE's + ``cp_ffn`` all-gather / reduce-scatter runs on the TP/CP communicator, and two + concurrent collectives on one communicator can deadlock. If the dedicated + communicator cannot be built (e.g. single rank) the code falls back to a plain + synchronous gather on the TP coordinator. + +Gather composability (why this is bit-exact): + * ``q_b``/``o_proj`` use per-block / per-channel weight scales (per_1x128 / + per_1x32 / per_token) that are shard-local along ``tp_dim`` -- gathering the + quantized bytes together with those scales reconstructs the full weight + (:meth:`LinearBase.gather_full_weight_scale`). + * fp8 ``W_K``/``W_V`` use a per-(whole-)tensor SCALAR scale; CP layers therefore + quantize them at load with a SHARED cross-rank scalar (see + ``dynamic_per_batched_tensor_quant(..., cross_rank_group=...)``) so the per-rank + fp8 shards tile bit-exactly after an all-gather and the scalar needs no gather. + * fp4 (mxfp4) ``W_K``/``W_V`` use per-1x32 block scales that ARE shard-local, so + both bytes and block scale are gathered along the head dim. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional + +import torch + +__all__ = [ + "GatheredAttnWeights", + "register_cp_attn_layer", + "reset_cp_gather_pipeline", + "get_gathered_weights", + "prefetch_next_layer", + "release_gathered_weights", +] + + +@dataclass +class GatheredAttnWeights: + """Transient full-head attention weights for ONE layer, valid for one forward.""" + + qb_weight: torch.Tensor + qb_scale: Optional[torch.Tensor] + qb_out_size: int + wk: torch.Tensor + wk_scale: torch.Tensor + wv: torch.Tensor + wv_scale: torch.Tensor + o_weight: torch.Tensor + o_scale: Optional[torch.Tensor] + o_out_size: int + num_heads_full: int + # Overlap bookkeeping (unused on the synchronous path). + ready_event: Optional[torch.cuda.Event] = None + + +# --------------------------------------------------------------------------- # +# Ordered CP-attn layer registry (execution order == construction order). +# --------------------------------------------------------------------------- # +_CP_ATTN_LAYERS: List[object] = [] +_CP_ATTN_INDEX: dict = {} + + +def register_cp_attn_layer(op) -> int: + """Register a CP attention op in execution order; returns its index. MTP / non-CP + ops never call this (they run the plain sharded path).""" + if op in _CP_ATTN_INDEX: + return _CP_ATTN_INDEX[op] + idx = len(_CP_ATTN_LAYERS) + _CP_ATTN_LAYERS.append(op) + _CP_ATTN_INDEX[op] = idx + return idx + + +def _next_cp_layer(op): + idx = _CP_ATTN_INDEX.get(op) + if idx is None or idx + 1 >= len(_CP_ATTN_LAYERS): + return None + return _CP_ATTN_LAYERS[idx + 1] + + +# --------------------------------------------------------------------------- # +# Dedicated gather communicator + stream (overlap only). +# --------------------------------------------------------------------------- # +class _RawGroup: + """Minimal ``all_gather(tensor, dim=...)`` shim over a raw process group so the + dedicated gather communicator can be consumed by + :meth:`LinearBase.gather_full_weight_scale` exactly like a vLLM GroupCoordinator. + """ + + def __init__(self, pg, world_size: int): + self._pg = pg + self.world_size = world_size + + def all_gather(self, t: torch.Tensor, dim: int = 0) -> torch.Tensor: + t = t.contiguous() + chunks = [torch.empty_like(t) for _ in range(self.world_size)] + torch.distributed.all_gather(chunks, t, group=self._pg) + return torch.cat(chunks, dim=dim) + + +_gather_group: Optional[_RawGroup] = None +_gather_stream: Optional[torch.cuda.Stream] = None +_gather_group_inited = False + + +def _get_gather_group() -> Optional[_RawGroup]: + """Dedicated communicator over the TP ranks (separate from the TP/CP comm used by + the MoE collectives, so concurrent collectives cannot deadlock).""" + global _gather_group, _gather_group_inited + if _gather_group_inited: + return _gather_group + _gather_group_inited = True + from vllm.distributed.parallel_state import get_tp_group + + tp = get_tp_group() + if tp.world_size <= 1: + _gather_group = None + return None + # new_group is collective across the WHOLE world; every rank must build the same + # subgroup for its TP ranks. Fall back to the TP comm if construction fails. + try: + pg = torch.distributed.new_group(ranks=tp.ranks, backend="nccl") + _gather_group = _RawGroup(pg, tp.world_size) + except Exception: # pragma: no cover - depends on runtime topology + _gather_group = None + return _gather_group + + +def _get_gather_stream() -> Optional[torch.cuda.Stream]: + global _gather_stream + if _gather_stream is None: + _gather_stream = torch.cuda.Stream() + return _gather_stream + + +# --------------------------------------------------------------------------- # +# The gather itself. +# --------------------------------------------------------------------------- # +def _gather_wk_wv(op, group): + """Gather the absorbed BMM weights to full heads along the head dim (dim 0). + + fp8 path: only the fp8 bytes are gathered (scale is a shared global scalar). + fp4 path: bytes AND per-1x32 block scale are gathered (block scale is + shard-local). + """ + hd = getattr(op, "_cp_wk_wv_head_dim", 0) + + def _ag_bytes(t): + # RCCL can't all_gather sub-byte / fp8 dtypes directly; move raw bytes. + return group.all_gather(t.view(torch.uint8), dim=hd).view(t.dtype) + + if getattr(op, "is_aiter_triton_fp4_bmm_enabled", False): + wk = _ag_bytes(op.W_K) + wv = _ag_bytes(op.W_V) + wk_s = group.all_gather(op.W_K_scale.view(torch.uint8), dim=hd).view( + op.W_K_scale.dtype + ) + wv_s = group.all_gather(op.W_V_scale.view(torch.uint8), dim=hd).view( + op.W_V_scale.dtype + ) + else: + wk = _ag_bytes(op.W_K) + wv = _ag_bytes(op.W_V) + # fp8: shared global scalar (see dynamic_per_batched_tensor_quant), no gather. + wk_s = op.W_K_scale + wv_s = op.W_V_scale + return wk, wk_s, wv, wv_s + + +def gather_attn_weights(op, group) -> GatheredAttnWeights: + """Reconstruct this layer's FULL-head attention weights from the TP shards.""" + num_heads_full = op._cp_num_heads_full + qb_weight, qb_scale = op.q_proj.gather_full_weight_scale(group=group) + o_weight, o_scale = op.o_proj.gather_full_weight_scale(group=group) + wk, wk_s, wv, wv_s = _gather_wk_wv(op, group) + return GatheredAttnWeights( + qb_weight=qb_weight, + qb_scale=qb_scale, + qb_out_size=num_heads_full * op.qk_head_dim, + wk=wk, + wk_scale=wk_s, + wv=wv, + wv_scale=wv_s, + o_weight=o_weight, + o_scale=o_scale, + o_out_size=op._cp_hidden_size, + num_heads_full=num_heads_full, + ) + + +# --------------------------------------------------------------------------- # +# Double-buffered overlap pipeline (2 slots, prefetch depth 1). +# --------------------------------------------------------------------------- # +_slots: List[Optional[GatheredAttnWeights]] = [None, None] +_slot_layer: List[Optional[object]] = [None, None] +_slot_ready_ev: List[Optional[torch.cuda.Event]] = [None, None] +_slot_consumed_ev: List[Optional[torch.cuda.Event]] = [None, None] + + +def reset_cp_gather_pipeline() -> None: + """Called at model entry (per forward) so a stale slot from the previous forward + never satisfies this forward's first layer.""" + _slot_layer[0] = _slot_layer[1] = None + + +def _slot_of(op) -> int: + return _CP_ATTN_INDEX[op] % 2 + + +def _gather_into_slot(op, slot: int) -> None: + """Run this layer's gather on the background stream into ``slot`` and record its + ready event. Waits until the slot's previous consumer has finished so the + transient tensors are not overwritten while still in use.""" + group = _get_gather_group() + stream = _get_gather_stream() + if group is None or stream is None: + return + stream.wait_stream(torch.cuda.current_stream()) + if _slot_consumed_ev[slot] is not None: + stream.wait_event(_slot_consumed_ev[slot]) + with torch.cuda.stream(stream): + gw = gather_attn_weights(op, group) + ev = torch.cuda.Event() + ev.record(stream) + gw.ready_event = ev + _slots[slot] = gw + _slot_layer[slot] = op + _slot_ready_ev[slot] = ev + + +def prefetch_next_layer(op) -> None: + """Kick the NEXT CP layer's gather on the background stream so it overlaps this + layer's MoE compute. No-op for the last layer or when no dedicated gather + communicator is available (that layer then gathers synchronously).""" + if _get_gather_group() is None: + return + nxt = _next_cp_layer(op) + if nxt is None: + return + _gather_into_slot(nxt, _slot_of(nxt)) + + +def get_gathered_weights(op) -> GatheredAttnWeights: + """Return this layer's full-head weights. The overlap pipeline prefetches the + NEXT layer's gather during this layer's MoE compute; the first CP layer (and any + pipeline miss) is gathered synchronously into its slot here. Falls back to a + plain synchronous gather on the TP coordinator when no dedicated gather + communicator is available.""" + if _get_gather_group() is None: + # No dedicated communicator (e.g. single rank): plain synchronous gather on + # the aiter TP coordinator (the same one the projections were sharded with) + # so the all_gather concatenation exactly reverses the shard layout. + from aiter.dist.parallel_state import get_tp_group + + return gather_attn_weights(op, get_tp_group()) + + slot = _slot_of(op) + if _slot_layer[slot] is not op: + # Not prefetched (first layer, or pipeline miss): gather synchronously into + # the slot on the background stream and wait. + _gather_into_slot(op, slot) + gw = _slots[slot] + ev = _slot_ready_ev[slot] + if ev is not None: + torch.cuda.current_stream().wait_event(ev) + # Keep the transients alive until the compute stream is done with them. + for t in (gw.qb_weight, gw.wk, gw.wv, gw.o_weight): + if t is not None: + t.record_stream(torch.cuda.current_stream()) + return gw + + +def release_gathered_weights(op) -> None: + """Mark this layer's slot consumed so the background stream may reuse it for a + later layer's gather (double-buffer hand-off).""" + if _get_gather_group() is None: + return + slot = _slot_of(op) + ev = torch.cuda.Event() + ev.record(torch.cuda.current_stream()) + _slot_consumed_ev[slot] = ev diff --git a/atom/plugin/vllm/attention/layer_mla.py b/atom/plugin/vllm/attention/layer_mla.py index 3640d0b868..784bb508e9 100644 --- a/atom/plugin/vllm/attention/layer_mla.py +++ b/atom/plugin/vllm/attention/layer_mla.py @@ -347,6 +347,13 @@ def __init__( self.head_repeat_factor = self.padded_num_heads // self.num_heads self._is_sparse_mla = True self.q_pad_num_heads = getattr(self, "q_pad_num_heads", None) + # Reuse-TP-as-CP runtime weight gather (set by setup_cp_weight_gather from + # the model layer). Default off: plain sharded TP for every batch. + self._cp_attn = False + self._cp_num_heads_full = self.num_heads + self._cp_tp_size = 1 + self._cp_hidden_size = self.num_heads * self.v_head_dim + self._cp_wk_wv_head_dim = 0 _register_vllm_static_forward_context(self) atom_static_context = atom_config.compilation_config.static_forward_context @@ -357,6 +364,146 @@ def __init__( max_num_tokens, dtype=torch.int64, device="cuda" ) + def setup_cp_weight_gather( + self, attn_cp, num_heads_full, tp_size, hidden_size + ) -> None: + """Wire reuse-TP-as-CP runtime weight gather for this attention op. + + Called once at construction from the model layer. When ``attn_cp`` and + ``tp_size > 1`` this op participates in CP: prefill/mixed batches gather the + full-head q_b/W_K/W_V/o_proj on demand (see forward_impl_cp), plain decode + runs the sharded weights. MTP / non-CP layers pass ``attn_cp=False``. + """ + self._cp_attn = bool(attn_cp) and int(tp_size) > 1 + self._cp_num_heads_full = int(num_heads_full) + self._cp_tp_size = int(tp_size) + self._cp_hidden_size = int(hidden_size) + if self._cp_attn: + from atom.plugin.vllm.attention.cp_gather import register_cp_attn_layer + + register_cp_attn_layer(self) + + def _cp_bind_full_head_weights(self, gw): + """Context manager: transiently rebind this op's per-head attributes and + absorbed BMM weights to the gathered FULL-head set so the shared + forward_impl_sparse runs full-head attention, then restore. Runs inside the + Dynamo-opaque CP op on the compute stream for prefill/mixed only (never + cudagraph-captured), so the rebind is invisible to torch.compile. + """ + import contextlib + + from atom.model_ops.attention_mla import _MLA_MIN_HEADS + + @contextlib.contextmanager + def _ctx(): + saved = ( + self.num_heads, + self.W_K, + self.W_K_scale, + self.W_V, + self.W_V_scale, + getattr(self, "padded_num_heads", self.num_heads), + getattr(self, "head_repeat_factor", 1), + self.q_pad_num_heads, + ) + try: + self.num_heads = gw.num_heads_full + self.W_K, self.W_K_scale = gw.wk, gw.wk_scale + self.W_V, self.W_V_scale = gw.wv, gw.wv_scale + self.padded_num_heads = max(gw.num_heads_full, _MLA_MIN_HEADS) + self.head_repeat_factor = self.padded_num_heads // gw.num_heads_full + # Full heads need no q-head padding (the sharded decode kernel may). + self.q_pad_num_heads = None + yield + finally: + ( + self.num_heads, + self.W_K, + self.W_K_scale, + self.W_V, + self.W_V_scale, + self.padded_num_heads, + self.head_repeat_factor, + self.q_pad_num_heads, + ) = saved + + return _ctx() + + def forward_impl_cp(self, query, kv_c_normed, k_pe, kv_cache, attn_metadata, q_scale): + """Reuse-TP-as-CP attention for CP layers. Does q_proj + attention + o_proj + INSIDE this Dynamo-opaque op so the split path can project/absorb with the + gathered FULL-head weights. + + * pcp_split (prefill/mixed): gather full q_b/W_K/W_V/o_proj, run full-head + token-parallel attention on this rank's 1/cp query shard (shared + forward_impl_sparse pcp branch), then a LOCAL o_proj (no all-reduce). + * else (plain decode): sharded q_proj -> sharded attention -> sharded + o_proj + a manual TP all-reduce, i.e. the classic TP path. + """ + from atom.plugin.vllm.attention.cp_gather import ( + get_gathered_weights, + release_gathered_weights, + ) + + pcp = bool(getattr(attn_metadata, "pcp_split", False)) + n_tokens = query.shape[0] + + if pcp: + gw = get_gathered_weights(self) + q = self.q_proj.forward_full( + query, q_scale, gw.qb_weight, gw.qb_scale, gw.qb_out_size + ) + q = q.view(-1, gw.num_heads_full, self.qk_head_dim) + if self.calculate_kv_scales: + self.calc_kv_scales(q, kv_c_normed, k_pe) + # attn_out dtype MUST match the activation dtype (bf16) so the downstream + # o_proj fp8 quant kernel gets a supported input -- NOT get_default_dtype() + # (fp32), which faults the kernel. + attn_out = torch.empty( + (n_tokens, gw.num_heads_full * self.v_head_dim), + dtype=q.dtype, + device=query.device, + ) + with self._cp_bind_full_head_weights(gw): + self.forward_impl_sparse( + q, kv_c_normed, k_pe, kv_cache, attn_metadata, attn_out + ) + out = self.o_proj.forward_full( + attn_out, None, gw.o_weight, gw.o_scale, gw.o_out_size + ) + # Release this layer's gather slot. The NEXT layer's weight gather is now + # kicked from cp_ffn, AFTER the pre-MoE activation all-gather (see + # _cp_ffn_impl), so the weight and activation collectives don't contend -- + # the weight gather overlaps the MoE compute instead of racing the + # activation all-gather. + release_gathered_weights(self) + return out + + # Plain decode: classic sharded TP. o_proj is built with reduce_results=False + # under CP, so reduce manually here with the SAME call o_proj would use + # internally (see LinearBase.forward -> aiter get_tp_group().all_reduce with + # ca_fp8_quant), keeping decode bit-identical to plain TP. + from aiter.dist.parallel_state import get_tp_group + + q = self.q_proj(query, q_scale) + q = q.view(-1, self.num_heads, self.qk_head_dim) + if self.calculate_kv_scales: + self.calc_kv_scales(q, kv_c_normed, k_pe) + # dtype=q.dtype (bf16), see note above; forward_impl carries the profile-run + # (attn_metadata is None) guard and dispatches to the sparse impl. + attn_out = torch.empty( + (n_tokens, self.num_heads * self.v_head_dim), + dtype=q.dtype, + device=query.device, + ) + self.forward_impl( + q, kv_c_normed, k_pe, kv_cache, attn_metadata=attn_metadata, output=attn_out + ) + out = self.o_proj(attn_out) + if self._cp_tp_size > 1: + out = get_tp_group().all_reduce(out, ca_fp8_quant=False) + return out + @property def impl(self): return self @@ -1434,6 +1581,21 @@ def forward( ): kv_c_normed = key k_pe = value + if self._cp_attn: + # Reuse-TP-as-CP layers route q_proj + attention + o_proj through a + # single Dynamo-opaque op so the split (prefill/mixed) path can project + # and absorb with the gathered FULL-head weights and the unsplit (decode) + # path stays on sharded TP. Branching on self._cp_attn (a construction + # constant) is torch.compile-safe; the runtime pcp_split branch lives + # inside the opaque op. + return torch.ops.aiter.atom_vllm_mla_attention_cp( + query, + kv_c_normed, + k_pe, + self.layer_name, + self._cp_hidden_size, + q_scale, + ) q = self.q_proj(query, q_scale) q = q.view(-1, self.num_heads, self.qk_head_dim) if self.calculate_kv_scales: diff --git a/atom/plugin/vllm/attention/metadata.py b/atom/plugin/vllm/attention/metadata.py index 8e74c854f0..d2d848183e 100644 --- a/atom/plugin/vllm/attention/metadata.py +++ b/atom/plugin/vllm/attention/metadata.py @@ -15,7 +15,6 @@ pcp_is_enabled, pcp_pad_dense, pcp_pad_len, - pcp_reindex_decode, pcp_round_robin_query_indices, pcp_sparse_prefill_reindex, plugin_attn_cp_enabled, @@ -1656,29 +1655,33 @@ def _plugin_cp_should_split(common_attn_metadata, topk_tokens, decode_threshold= """Batch-global split gate shared by the plugin sparse builders (main sparse-MLA and sparse indexer), mirrored by deepseek_v2._plugin_pcp_active(). - UNIVERSAL query-dim split: when reuse-TP-as-CP is enabled (``ATOM_VLLM_ATTN_CP``) - and the aiter _PCP group is aliased to the TP group (world > 1), split EVERY - batch type -- pure prefill, pure plain-decode, and MIXED -- over the query dim - using ONE shared whole-batch round-robin partition. Each rank runs the whole - model forward on its ``T/cp`` query shard with full (replicated) KV; the single - exit all-gather restores the full token order before lm_head. This removes the - old replicated-attention fallback (full 64-head attention on every rank). + HAS-PREFILL query-dim split (RFC ROCm/ATOM#196 weight-gather revision): when + reuse-TP-as-CP is enabled (``ATOM_VLLM_ATTN_CP``) and the aiter _PCP group is + aliased to the TP group (world > 1), split ONLY prefill-containing batches -- + pure prefill and MIXED (plain-decode + prefill) -- over the query dim using ONE + shared whole-batch round-robin partition. Each such rank runs the whole model + forward on its ``T/cp`` query shard with full (replicated) KV and full-head + attention whose q_b/kv_b/o_proj weights are gathered on demand; the single exit + all-gather restores the full token order before lm_head. + + This ``pcp_split`` stamp is ALSO the weight-gather signal: split == gather-full + (full-head, ``o_proj`` local) and not-split == the original TP-sharded impl + (1/tp heads, ``o_proj`` all-reduce). Pure plain-decode and spec/MTP verify are + therefore NOT split -- they run classic TP so decode TPOT stays on the sharded + path with no per-step weight gather. Correctness constraints: - * Any decode region must be single-token plain decode built by a non-spec - builder (``decode_threshold == 1``) so token index == request index for the - owned-decode slice. Spec/MTP verify batches (``decode_threshold > 1``) keep - decode replicated. + * Spec/MTP verify batches (``decode_threshold > 1`` with a decode region) are + never split; they run the original TP-sharded impl. * PREFILL has NO seq_len gate. The plugin indexer has no dense-MHA fallback (it always runs the sparse top-k, selecting all keys for short rows), and every plugin CP call site keys off the pcp_split stamp with a length- agnostic K all-gather + owned reindex, so short-context prefill splits correctly. (The native path DOES have a dense fallback and keeps the gate; this plugin gate no longer does.) - * PLAIN-DECODE has NO seq_len gate. The round-robin split is over the QUERY - dim with full replicated KV (mathematically correct at any seq_len) and the - decode path always runs the sparse indexer with a ``min(seq_len, topk)`` - clamp (no dense fallback to desync). Matches native deepseek_v4._pcp_active. + * MIXED splits its decode region via the same whole-batch round-robin + partition (owned_q restricted to the decode tokens); those decode tokens + ride the full-head gathered path along with the prefill tokens. """ if not (plugin_attn_cp_enabled() and pcp_is_enabled()): return False @@ -1697,22 +1700,17 @@ def _plugin_cp_should_split(common_attn_metadata, topk_tokens, decode_threshold= # replicated (return False for pure- and mixed-decode spec batches). if int(num_decode_tokens) > 0 and int(decode_threshold) != 1: return False - if int(num_prefill_tokens) > 0: - # PREFILL-containing (pure prefill or MIXED): ALWAYS split over the query - # dim, at any seq_len. The plugin indexer (sparse_attn_indexer_plugin_mode) - # has NO dense-MHA fallback -- unlike the native indexer it never early- - # returns on max_seqlen_k <= index_topk; it always runs the sparse top-k, - # which selects all keys when the row is shorter than topk (dense- - # equivalent). Every plugin CP call site (entry split, indexer K all-gather - # + separate q/k rope, exit gather, sparse-MLA) keys off the pcp_split - # stamp, not seq_len, and the K all-gather + owned per-query reindex fire - # for any length, so short-context prefill splits correctly (topk_tokens is - # therefore no longer consulted here). - return True - # PURE plain-decode: always split under CP, no seq_len gate. The round-robin - # split is over the query dim with full replicated KV (correct at any seq_len) - # and the decode indexer clamps min(seq_len, topk) with no dense fallback. - return int(common_attn_metadata.max_query_len) == 1 + # HAS-PREFILL gate (weight-gather revision): split == gather-full (full-head, + # o_proj local) iff the batch contains prefill tokens (pure prefill or MIXED). + # The plugin indexer (sparse_attn_indexer_plugin_mode) has NO dense-MHA fallback + # (it always runs the sparse top-k, selecting all keys for short rows), and every + # plugin CP call site keys off the pcp_split stamp with a length-agnostic K + # all-gather + owned reindex, so short-context prefill splits correctly. MIXED + # splits its decode region via the same whole-batch round-robin partition + # (owned_q restricted to the decode tokens). Pure plain-decode returns False and + # runs the original TP-sharded impl (no split, no weight gather, o_proj + # all-reduce), keeping decode TPOT on the classic TP path. + return int(num_prefill_tokens) > 0 class AiterMlaSparseMetadataBuilder(AttentionMetadataBuilder): @@ -2272,19 +2270,6 @@ def __init__( dtype=torch.int32, device=self.device, ) - # Persistent owned-query buffers for the PCP decode-split path (see - # _build_indexer). pcp_reindex_decode returns fresh per-step tensors; - # under FULL cudagraph the captured indexer decode kernel bakes their - # addresses, so the owned seq_lens / block_table are copied into these - # stable buffers that build() refreshes in place on every replay. - self._pcp_seq_lens_owned = torch.empty( - (max_num_batched_tokens,), dtype=torch.int32, device=self.device - ) - self._pcp_block_table_owned = torch.empty( - (max_num_batched_tokens, max_num_blocks_per_req), - dtype=torch.int32, - device=self.device, - ) self.scheduler_metadata_buffer = torch.empty( (self.num_sms + 1, 2), dtype=torch.int32, device=self.device ) @@ -2429,42 +2414,30 @@ def _build_indexer( # to prevent OOB access in the DeepGEMM kernel (padded rows have # seq_len=0, so they produce no meaningful output). block_table_full.clamp_(min=0) - if num_prefills > 0: - # MIXED: the owned decode queries MUST come from the SAME - # whole-batch round-robin partition (owned_q) that the prefill - # chunks and the sparse-MLA _build_pcp_split reindex use, restricted - # to the decode region [0, num_decode_tokens). The decode region is - # NOT padded separately (unlike pure decode) -- padding it would - # shift the prefill token_start offsets (os/oe are computed over the - # same owned_q in _build_indexer_one_prefill_chunk) and desync the - # topk rows from the owned paged_kv_indptr the sparse layer reads. - # Mixed batches run piecewise/eager, so per-rank-variable owned - # decode counts (differing by at most 1) are fine. - owned_dec_q = owned_q[owned_q < num_decode_tokens].to( - seq_lens_full.device - ) - n_owned_dec = int(owned_dec_q.shape[0]) - seq_lens_owned = seq_lens_full.index_select(0, owned_dec_q).contiguous() - block_table_owned = block_table_full.index_select( - 0, owned_dec_q - ).contiguous() - else: - # PURE decode: pad the decode region to a multiple of cp so every - # rank runs the same owned count (FULL cudagraph uniform shape); - # dummy queries attend nothing and are dropped at the model exit - # all-gather. Same round-robin start/stride as the combined owned_q, - # so the owned decode set matches owned_q < num_decode_tokens. - seq_lens_owned, block_table_owned, _owned_dec_q, n_owned_dec = ( - pcp_reindex_decode( - seq_lens_full, block_table_full, num_decode_tokens - ) - ) + # MIXED only: reuse-TP-as-CP now splits the token batch exclusively for + # prefill/mixed batches (pcp_split is gated on num_prefill_tokens > 0), + # so num_prefills > 0 here by construction -- pure decode stays on the + # replicated TP path and never reaches this branch. The owned decode + # queries MUST come from the SAME whole-batch round-robin partition + # (owned_q) that the prefill chunks and the sparse-MLA _build_pcp_split + # reindex use, restricted to the decode region [0, num_decode_tokens). + # The decode region is NOT padded separately -- padding it would shift + # the prefill token_start offsets (os/oe are computed over the same + # owned_q in _build_indexer_one_prefill_chunk) and desync the topk rows + # from the owned paged_kv_indptr the sparse layer reads. Mixed batches + # run piecewise/eager, so per-rank-variable owned decode counts + # (differing by at most 1) are fine. + owned_dec_q = owned_q[owned_q < num_decode_tokens].to(seq_lens_full.device) + n_owned_dec = int(owned_dec_q.shape[0]) + seq_lens_owned = seq_lens_full.index_select(0, owned_dec_q).contiguous() + block_table_owned = block_table_full.index_select( + 0, owned_dec_q + ).contiguous() if not AiterMlaSparseIndexerMetadataBuilder._pcp_decode_split_logged: AiterMlaSparseIndexerMetadataBuilder._pcp_decode_split_logged = True logger.info( - "[PCP] indexer DECODE-split ENGAGED (%s): pre_split_decode=%d " + "[PCP] indexer DECODE-split ENGAGED (mixed): pre_split_decode=%d " "n_owned=%d num_prefills=%d max_seq_len=%d topk=%d", - "mixed" if num_prefills > 0 else "pure-decode", int(num_decode_tokens), int(n_owned_dec), int(num_prefills), @@ -2478,20 +2451,11 @@ def _build_indexer( decode_metadata = None if num_decodes > 0: if pcp_dec is not None: - # FULL cudagraph safety: route the owned tensors through persistent - # buffers so a captured indexer decode kernel reads stable addresses - # that are refreshed in place on every replay (see __init__). The - # sparse-MLA slot_mapping stays FULL for the replicated KV write. - seq_lens_owned, block_table_owned = pcp_dec - self._pcp_seq_lens_owned[:num_decodes].copy_( - seq_lens_owned, non_blocking=True - ) - seq_lens = self._pcp_seq_lens_owned[:num_decodes] - bw = block_table_owned.shape[1] - self._pcp_block_table_owned[:num_decodes, :bw].copy_( - block_table_owned, non_blocking=True - ) - block_table = self._pcp_block_table_owned[:num_decodes, :bw] + # Mixed-only decode split: the owning batch has prefill tokens, so it + # runs eager/piecewise (never FULL cudagraph). Fresh per-step owned + # tensors are therefore safe -- no persistent capture buffers needed. + # The sparse-MLA slot_mapping stays FULL for the replicated KV write. + seq_lens, block_table = pcp_dec self.decode_lens_buffer[:num_decodes].fill_(1) decode_lens = self.decode_lens_buffer[:num_decodes] decode_lens_cpu = torch.ones(num_decodes, dtype=torch.int32) diff --git a/atom/plugin/vllm/attention/ops.py b/atom/plugin/vllm/attention/ops.py index 264a70901e..65e3ba7c04 100644 --- a/atom/plugin/vllm/attention/ops.py +++ b/atom/plugin/vllm/attention/ops.py @@ -94,3 +94,49 @@ def atom_vllm_mla_attention( output=output, ) return output + + +def atom_vllm_mla_attention_cp_fake( + query: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + layer_name: str, + output_hidden_size: int, + q_scale: Optional[torch.Tensor] = None, +) -> torch.Tensor: + # The real op returns the post-o_proj activation, whose dtype is the model + # activation dtype (bf16/fp16) -- NOT torch.get_default_dtype(), which is + # float32 at torch.compile trace time (vLLM only sets the model dtype as the + # default *during* weight load, then reverts). `query` may be fp8 (fused + # q-norm+quant), so key the dtype off `kv_c_normed`, which is always the + # unquantized normed KV in the model activation dtype. + return kv_c_normed.new_empty((query.shape[0], output_hidden_size)) + + +@mark_spliting_op( + is_custom=True, + gen_fake=atom_vllm_mla_attention_cp_fake, + mutates_args=[], +) +def atom_vllm_mla_attention_cp( + query: torch.Tensor, + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + layer_name: str, + output_hidden_size: int, + q_scale: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Reuse-TP-as-CP attention entry: q_proj + attention + o_proj live INSIDE this + Dynamo-opaque op so the split (prefill/mixed) path can project/absorb with + on-demand gathered FULL-head weights, and the runtime pcp_split branch never + bakes under torch.compile. Returns the post-o_proj [tokens, hidden] output. + """ + layer, attn_metadata, kv_cache = _get_layer_context(layer_name) + return layer.forward_impl_cp( + query, + kv_c_normed, + k_pe, + kv_cache, + attn_metadata, + q_scale, + ) diff --git a/atom/utils/envs.py b/atom/utils/envs.py index 3f0cd09cae..1d10488d15 100644 --- a/atom/utils/envs.py +++ b/atom/utils/envs.py @@ -154,25 +154,17 @@ # group (RFC ROCm/ATOM#196). Unlike native ATOM's separate `-pcp` dimension # (world = tp x pcp, KV replicated per PCP rank), this runs true sequence # parallelism over tokens within the SAME TP group: residual/hidden are - # sharded [T/cp], attention runs full-head *replicated* weights on 1/cp of - # the queries against the full (gathered) KV, and MoE does all-gather + - # reduce-scatter around the TP-sharded experts. DSA (index_topk) models only - # (GLM-5.2 / DeepSeek-V3.2). `0` (default) is a bit-exact no-op. + # sharded [T/cp]. Attention weights stay TP-sharded in memory; prefill/mixed + # batches gather the full q_b/kv_b/o_proj (+ absorbed W_K/W_V) on demand + # per-layer at runtime and run full-head token-parallel attention on 1/cp of + # the queries against the full (gathered) KV, while plain decode keeps running + # the sharded weights on the classic TP path. MoE does all-gather + + # reduce-scatter around the TP-sharded experts. The per-layer weight gather runs + # on a DEDICATED communicator + background stream (2-slot double buffer) so it + # overlaps the previous layer's MoE compute -- this is intrinsic to CP, no extra + # flag. DSA (index_topk) models only (GLM-5.2 / DeepSeek-V3.2). `0` (default) is + # a bit-exact no-op. "ATOM_VLLM_ATTN_CP": lambda: (os.getenv("ATOM_VLLM_ATTN_CP", "0").lower() == "1"), - # Extend reuse-TP-as-CP (ATOM_VLLM_ATTN_CP) to also round-robin split the - # DECODE token batch across the CP ranks (default off). With CP on but this - # off, decode runs full-head *replicated* attention on ALL tokens per rank - # (redundant attention + indexer compute, no token split). With this on, a - # PURE-decode batch is split 1/cp the same way prefill is: each rank runs - # q-proj / attention / indexer / MoE on its 1/cp owned decode queries against - # the full (all-gathered) KV, then the hidden states are all-gathered at the - # model exit. This trades an exit all-gather (+ per-layer KV all-gather) for - # cp x less attention/indexer/q-side compute, so it pays off at concurrency - # >> cp and can regress at low concurrency. Plain decode only (max_query_len - # == 1); mixed and spec/MTP batches stay replicated. `0` is a no-op. - "ATOM_VLLM_ATTN_CP_DECODE_SPLIT": lambda: ( - os.getenv("ATOM_VLLM_ATTN_CP_DECODE_SPLIT", "0").lower() == "1" - ), "ATOM_USE_CUSTOM_ALL_GATHER": lambda: ( os.getenv("ATOM_USE_CUSTOM_ALL_GATHER", "1").lower() == "1" ), diff --git a/recipes/GLM-5.md b/recipes/GLM-5.md index deca34c70f..5f1820fdf8 100644 --- a/recipes/GLM-5.md +++ b/recipes/GLM-5.md @@ -232,8 +232,6 @@ 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 & ``` diff --git a/recipes/atom_vllm/GLM-5.md b/recipes/atom_vllm/GLM-5.md index dce2a88b6f..64873706e7 100644 --- a/recipes/atom_vllm/GLM-5.md +++ b/recipes/atom_vllm/GLM-5.md @@ -57,6 +57,34 @@ vllm serve amd/GLM-5.2-MXFP4 \ --no-enable-prefix-caching ``` +### Attention CP: Prefill Context Parallelism (long-context throughput) + +For long-context serving you can reuse Attention CP to optimize performance. Tokens are sharded `T/cp` with sequence parallelism: prefill / mixed batches run full-head, token-parallel attention on each rank's `1/cp` query shard (the `q_b` / `kv_b` / `o_proj` weights stay TP-sharded in memory and are gathered to full heads on demand per layer), the MoE is wrapped in all-gather / reduce-scatter, and plain decode stays on the classic TP path. This cuts prefill attention / indexer compute by `cp` and improves TTFT and throughput for long inputs at high concurrency. It applies only to sparse-MLA (DSA) models (GLM-5.2 / DeepSeek-V3.2), requires `--tensor-parallel-size > 1`, and is a bit-exact no-op when disabled. + +Enable it with a single environment variable: + +- `ATOM_VLLM_ATTN_CP=1` — enable reuse-TP-as-CP. Each layer's attention weight all-gather automatically overlaps the previous layer's MoE compute (dedicated communicator + background CUDA stream + 2-slot double buffer), so no extra flag is needed. + +```bash +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export AITER_USE_FLYDSL_MOE_SORTING=1 +export ATOM_VLLM_ATTN_CP=1 + +vllm serve amd/GLM-5.2-MXFP4 \ + --host localhost \ + --port 8000 \ + --async-scheduling \ + --load-format fastsafetensors \ + --trust-remote-code \ + --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \ + --kv-cache-dtype fp8 \ + --tensor-parallel-size 4 \ + --gpu-memory-utilization 0.9 \ + --max-num-batched-tokens 32768 \ + --no-enable-prefix-caching \ + --additional-config '{"online_quant_config": {"global_quant_config": "ptpc_fp8", "exclude_layer": ["lm_head", "model.embed_tokens", "*.mlp.gate", "model.layers.[0-9].mlp.*expert*", "model.layers.[1-6][0-9].mlp.*expert*", "model.layers.7[0-7].mlp.*expert*"]}}' +``` + ## Step 3: Performance Benchmark Users can use the default vllm bench commands for performance benchmarking. ```bash