From 0a63726f1ae5324a42d0ba60e34b351fd90ff9ae Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Thu, 27 Aug 2026 15:19:55 +0000 Subject: [PATCH 1/2] remove 96 to 128 padding Signed-off-by: whx-sjtu --- atom/model_ops/attention_mla.py | 62 ++++++++++++++++++--------------- atom/models/kimi_k3.py | 4 ++- 2 files changed, 37 insertions(+), 29 deletions(-) diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index e7b0816f20..60c093af50 100644 --- a/atom/model_ops/attention_mla.py +++ b/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( diff --git a/atom/models/kimi_k3.py b/atom/models/kimi_k3.py index 06dc645c04..2bbd55eab5 100644 --- a/atom/models/kimi_k3.py +++ b/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" From 5a8c47f2b5c13449d146681edeb1efc37481ab67 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Thu, 20 Aug 2026 05:34:34 +0000 Subject: [PATCH 2/2] fuse add rmsnorm Signed-off-by: whx-sjtu --- atom/models/kimi_k3_dspark.py | 26 ++++++++++++++--------- atom/plugin/vllm/models/kimi_k3_dspark.py | 6 ++++-- 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/atom/models/kimi_k3_dspark.py b/atom/models/kimi_k3_dspark.py index b08e2649a5..777b9e061a 100644 --- a/atom/models/kimi_k3_dspark.py +++ b/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 diff --git a/atom/plugin/vllm/models/kimi_k3_dspark.py b/atom/plugin/vllm/models/kimi_k3_dspark.py index 64258c3b42..8a355939ec 100644 --- a/atom/plugin/vllm/models/kimi_k3_dspark.py +++ b/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,