[DCP][Kimi-k3] Decode on the gathered head width

atom/model_ops/attention_mla.py

@@ -214,13 +214,20 @@ def mla_kernel_num_heads(num_heads: int) -> int:
     return -(-num_heads // _MLA_MIN_HEADS) * _MLA_MIN_HEADS
 
 
-# Gathered widths aiter serves with a dedicated kernel. The other multiples of
-# 16 (48, 80, 96, 112) are folded onto the 16-head kernel instead, and that fold
-# reinterprets head groups as extra sequence rows (total_s *= ori_nhead//16)
-# without touching kv_indptr, which desynchronises the row -> global position
-# mapping the round-robin causal mask runs on. DCP decode must avoid them.
+# Gathered widths aiter serves with a dedicated kernel.
 _MLA_DCP_KERNEL_WIDTHS = (16, 32, 64, 128)
 
+# Widest gathered query aiter has any MLA decode dispatch for; past it
+# mla_decode_fwd asserts rather than falling back to anything.
+#
+# A persistent decode reaches every multiple of 16 up to that: the widths
+# without a dedicated kernel (48, 80, 96, 112) fold onto the 16-head one, which
+# reinterprets head groups as extra sequence rows and rebuilds qo_indptr and the
+# kv indptrs to match, so the round-robin mask survives it. A NON-persistent one
+# does not -- aiter's fold is guarded on persistent mode -- so it stays on the
+# width sets below, which only hold widths that have their own kernel.
+_MLA_DCP_MAX_KERNEL_HEADS = 128
+
 _MLA_DCP_KERNEL_WIDTHS_NON_PERSISTENT = (16, 32, 128)
 _MLA_DCP_KERNEL_WIDTHS_NON_PERSISTENT_FP8 = (16, 128)
 _MLA_DCP_SPARSE_PREFILL_WIDTHS = (16, 128)
@@ -298,43 +305,42 @@ def mla_dcp_kernel_num_heads(
     kv_cache_dtype: str,
     persistent: bool,
 ) -> int:
-    """Width to pad the GATHERED query heads to for a DCP decode.
+    """Width to gather the query heads to for a DCP decode.
 
     DCP decode all-gathers Q on the head dim before calling the kernel, so what
     gets dispatched on is ``num_heads * dcp_world_size``; a single rank's head
-    count is never seen and is the wrong thing to pad. Round that gathered width
-    up to one aiter serves natively for the mode this decode actually runs in --
-    the folded widths are no use here because the fold breaks the round-robin
-    causal mask (see above), and gqa=64 is off the table on any
-    non-persistent decode (see _MLA_DCP_KERNEL_WIDTHS_NON_PERSISTENT).
-
-    ``kv_cache_dtype`` no longer selects whether gqa=64 is excluded -- it is
-    excluded for both dtypes -- but it still selects the non-persistent set,
-    because fp8 lacks a gqa=32 kernel there and bf16 does not.
+    count is never seen and is the wrong thing to round. A persistent decode
+    takes that width as-is once it is a multiple of 16, dedicated kernel or
+    fold; a non-persistent one has no fold to fall back on and must be padded
+    onto a width that has its own kernel.
+
+    ``kv_cache_dtype`` only selects the non-persistent set: gqa=64 is excluded
+    for both dtypes there (fp8 aborts on it, bf16 silently miscomputes it), and
+    fp8 lacks a gqa=32 kernel on top of that while bf16 does not.
     """
-    gathered = max(num_heads * dcp_world_size, min_kernel_heads)
-    # gqa=64 is dropped for BOTH dtypes without persistent mode (fp8 aborts on
-    # it, bf16 silently miscomputes it); fp8 drops 32 on top of that.
-    widths = _MLA_DCP_KERNEL_WIDTHS
-    if not persistent:
+    gathered = mla_kernel_num_heads(max(num_heads * dcp_world_size, min_kernel_heads))
+    if persistent:
+        if gathered <= _MLA_DCP_MAX_KERNEL_HEADS:
+            return gathered
+    else:
         widths = (
             _MLA_DCP_KERNEL_WIDTHS_NON_PERSISTENT_FP8
             if kv_cache_dtype.startswith("fp8")
             else _MLA_DCP_KERNEL_WIDTHS_NON_PERSISTENT
         )
-    for width in widths:
-        if width >= gathered:
-            return width
+        for width in widths:
+            if width >= gathered:
+                return width
     global _dcp_kernel_width_warned
     if not _dcp_kernel_width_warned:
         _dcp_kernel_width_warned = True
         logger.warning(
-            f"DCP decode gathers {gathered} query heads, past the widest natively "
-            f"dispatched MLA kernel ({widths[-1]}); falling back to "
-            "the folded kernel, which is incorrect for MTP (round-robin causal "
-            "mask). Lower decode_context_parallel_size or raise tp."
+            f"DCP decode gathers {gathered} query heads, past the widest MLA "
+            f"kernel aiter dispatches ({_MLA_DCP_MAX_KERNEL_HEADS}); it serves "
+            "neither a kernel nor a fold that wide and will abort in "
+            "mla_decode_fwd. Lower decode_context_parallel_size or raise tp."
         )
-    return mla_kernel_num_heads(gathered)
+    return gathered
 
 
 def mla_dcp_sparse_prefill_num_heads(

atom/models/kimi_k3.py

@@ -27,7 +27,7 @@
 # Dynamo-opaque custom op that dispatches the MoE between single- and dual-stream
 # forwards (shared with deepseek_v2/v4). Imported for the registration only.
 from atom.model_ops import module_dispatch_ops as _module_dispatch_ops  # noqa: F401
-from atom.model_ops.attention_mla import MLAModules
+from atom.model_ops.attention_mla import MLAModules, qrep_tp_override
 from atom.model_ops.attention_residual import AttnRes
 from atom.model_ops.base_attention import Attention
 from atom.model_ops.embed_head import ParallelLMHead, VocabParallelEmbedding
@@ -622,12 +622,14 @@ def __init__(
         self.q_a_layernorm = RMSNorm(
             self.q_lora_rank, eps=1e-6, prefix=f"{prefix}.q_a_layernorm"
         )
+        # DCP Query Replication: {} unless QREP is on -- see qrep_tp_override.
         self.q_b_proj = ColumnParallelLinear(
             self.q_lora_rank,
             self.num_heads * self.q_head_dim,
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.q_b_proj",
+            **qrep_tp_override(self.tp_size),
         )
         self.kv_a_layernorm = RMSNorm(
             self.kv_lora_rank, eps=1e-6, prefix=f"{prefix}.kv_a_layernorm"

atom/models/kimi_k3_dspark.py

@@ -561,14 +561,19 @@ def write_context_kv(self, ctx_hidden, positions) -> None:
         self.self_attn.write_context_kv(ctx_hidden, positions, slot_mapping)
 
     def forward(
-        self, positions: torch.Tensor, hidden_states: torch.Tensor
-    ) -> torch.Tensor:
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-        hidden_states = residual + self.self_attn(positions, hidden_states)
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        return residual + self.mlp(hidden_states)
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+        hidden_states = self.self_attn(positions, hidden_states)
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        return self.mlp(hidden_states), residual
 
 
 class KimiK3DSpark(DSparkDraftModel):
@@ -712,9 +717,10 @@ def forward_spec(
         draft_ids[:, 0] = input_ids
         hidden = self.embed_tokens(draft_ids.view(-1))
 
+        residual = None
         for layer in self.layers:
-            hidden = layer(positions, hidden)
-        hidden = self.final_norm(hidden)
+            hidden, residual = layer(positions, hidden, residual)
+        hidden, _ = self.final_norm(hidden, residual)
 
         base_logits = self.lm_head(hidden).view(bs, T, -1)
         return self._sample_block(base_logits, input_ids), None

atom/plugin/vllm/models/kimi_k3_dspark.py

@@ -28,9 +28,11 @@ def forward(
         hidden_states = (
             inputs_embeds if inputs_embeds is not None else self.embed_tokens(input_ids)
         )
+        residual = None
         for layer in self.layers:
-            hidden_states = layer(positions, hidden_states)
-        return self.final_norm(hidden_states)
+            hidden_states, residual = layer(positions, hidden_states, residual)
+        hidden_states, _ = self.final_norm(hidden_states, residual)
+        return hidden_states
 
     def write_combined_context_kv(
         self,