feat(prezero): wire split-K GEMM prezero into Kimi MLA/MoE decode via a torch.compile pass

atom/config.py

@@ -237,6 +237,7 @@ def compute_hash(self) -> str:
         factors.append(self.local_cache_dir)
         factors.append(self.cudagraph_capture_sizes)
         factors.append(self.cuda_graph_sizes)
+        factors.append(bool(os.getenv("ATOM_ENABLE_PREZERO")))
 
         return hashlib.sha256(str(factors).encode()).hexdigest()
 

atom/model_ops/attention_mla.py

@@ -31,6 +31,9 @@
     fused_qk_rope_concat_and_cache_mla_seg = None
 from aiter.dist.parallel_state import get_dp_group
 from aiter.mla import mla_decode_fwd, mla_prefill_fwd
+from aiter.tuned_gemm import tgemm
+
+from atom.model_ops.prezero import prezero_active
 from aiter.ops.triton.attention.mla import (
     mla_decode_fwd as triton_shuffle_mla_decode_fwd,
 )
@@ -350,7 +353,7 @@ def process_weights_after_loading(self):
             )
 
     @mark_trace(prefix="v_up_proj_and_o_proj", torch_compile=False)
-    def _v_up_proj_and_o_proj(self, x):
+    def _v_up_proj_and_o_proj(self, x, oproj_prezero=None):
         # Convert from (B, N, L) to (N, B, L)
         x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
         # Multiply (N, B, L) x (N, L, V) -> (N, B, V), Convert from (N, B, V) to (B, N, V)
@@ -382,14 +385,22 @@ def _v_up_proj_and_o_proj(self, x):
             )
             # Convert from (B, N, V) to (B, N * V)
             x = x.reshape(-1, self.num_heads * self.v_head_dim)
+        if oproj_prezero is not None:
+            active = prezero_active(self.q_proj.prezero_n_total)
+            tgemm.mm(x, self.o_proj.weight, None, zero_init=not active, out=oproj_prezero)
+            return oproj_prezero
         return self.o_proj(x)
 
     @mark_trace(prefix="q_proj_and_k_up_proj", torch_compile=False)
-    def _q_proj_and_k_up_proj(self, x, x_scale=None):
-        q_nope, q_pe = (
-            self.q_proj(x, x_scale)
-            .view(-1, self.num_heads, self.qk_head_dim)
-            .split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+    def _q_proj_and_k_up_proj(self, x, x_scale=None, qb_prezero=None):
+        if qb_prezero is not None:
+            active = prezero_active(self.q_proj.prezero_n_base)
+            tgemm.mm(x, self.q_proj.weight, None, zero_init=not active, out=qb_prezero)
+            q_proj_out = qb_prezero
+        else:
+            q_proj_out = self.q_proj(x, x_scale)
+        q_nope, q_pe = q_proj_out.view(-1, self.num_heads, self.qk_head_dim).split(
+            [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
         )
 
         # Convert from (B, N, P) to (N, B, P)
@@ -956,6 +967,7 @@ def _forward_decode(
         q: torch.Tensor,
         kv_c_and_k_pe_cache: torch.Tensor,
         attn_metadata: AttentionMetaData,
+        oproj_prezero: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         assert kv_c_and_k_pe_cache.numel() > 0
         assert attn_metadata is not None
@@ -1116,7 +1128,7 @@ def _forward_decode(
         if self.head_repeat_factor > 1:
             o = o[:, :: self.head_repeat_factor, :].contiguous()
 
-        return self._v_up_proj_and_o_proj(o)
+        return self._v_up_proj_and_o_proj(o, oproj_prezero=oproj_prezero)
 
     def forward_impl(
         self,
@@ -1125,6 +1137,8 @@ def forward_impl(
         k_rope: torch.Tensor,
         positions: torch.Tensor = None,
         q_scale: Optional[torch.Tensor] = None,
+        qb_prezero: Optional[torch.Tensor] = None,
+        oproj_prezero: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         # kv_cache = self.kv_cache
         forward_context: ForwardContext = get_forward_context()
@@ -1206,7 +1220,9 @@ def forward_impl(
                     prefill_q, k_nope, k_rope, kv_cache, attn_metadata
                 )
         else:
-            q_nope, q_rope = self._q_proj_and_k_up_proj(q, x_scale=q_scale)
+            q_nope, q_rope = self._q_proj_and_k_up_proj(
+                q, x_scale=q_scale, qb_prezero=qb_prezero
+            )
 
             if self.use_seg_mla:
                 # Seg path: allocate q_out with a padded last dim so each head row
@@ -1296,7 +1312,9 @@ def forward_impl(
             if context.is_prefill:
                 output = self._forward_prefill_mla(q_out, kv_cache, attn_metadata)
             else:
-                output = self._forward_decode(q_out, kv_cache, attn_metadata)
+                output = self._forward_decode(
+                    q_out, kv_cache, attn_metadata, oproj_prezero=oproj_prezero
+                )
 
         return output
 
@@ -1310,6 +1328,8 @@ def forward(
         positions: torch.Tensor = None,
         q_scale: Optional[torch.Tensor] = None,
         output: torch.Tensor = None,
+        qb_prezero: Optional[torch.Tensor] = None,
+        oproj_prezero: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
         return self.forward_impl(
@@ -1318,6 +1338,8 @@ def forward(
             k_rope=k_rope,
             positions=positions,
             q_scale=q_scale,
+            qb_prezero=qb_prezero,
+            oproj_prezero=oproj_prezero,
         )
 
 

atom/model_ops/base_attention.py

@@ -327,6 +327,8 @@ def fake_(
     layer_name: str,
     use_mla: bool,
     qkv: torch.Tensor,
+    qb_prezero: Optional[torch.Tensor] = None,
+    oproj_prezero: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     output_shape = list(q.shape)
     # If we fusion rmsnorm and quant, the input dtype is fp8, but actually we use bf16 for output.
@@ -342,7 +344,9 @@ def fake_(
 # Dynamo will not try to inspect any of the internal operations for prefill or decode
 # This way, although attention operation is complicated,
 # we can still capture the model's computation graph as a full-graph
-@mark_spliting_op(is_custom=True, gen_fake=fake_, mutates_args=[])
+@mark_spliting_op(
+    is_custom=True, gen_fake=fake_, mutates_args=["qb_prezero", "oproj_prezero"]
+)
 def unified_attention_with_output_base(
     q: torch.Tensor,
     q_scale: Optional[torch.Tensor],
@@ -352,6 +356,8 @@ def unified_attention_with_output_base(
     layer_name: str,
     use_mla: bool,
     qkv: torch.Tensor,
+    qb_prezero: Optional[torch.Tensor] = None,
+    oproj_prezero: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     atom_config = get_current_atom_config()
     self = atom_config.compilation_config.static_forward_context[layer_name]
@@ -362,6 +368,8 @@ def unified_attention_with_output_base(
             k_rope=v,
             positions=positions,
             q_scale=q_scale,
+            qb_prezero=qb_prezero,
+            oproj_prezero=oproj_prezero,
         )
     else:
         return self.impl.forward(

atom/model_ops/layernorm.py

@@ -318,6 +318,7 @@ def forward(
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
         x_scale: Optional[torch.Tensor] = None,
+        zero_fill: Optional[torch.Tensor] = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         if self.x_pad_to_multiple > 0:
             assert (
@@ -341,6 +342,7 @@ def forward(
                 residual,
                 self.weight,
                 self.eps,
+                zero_fill=zero_fill,
             )
             return x, residual
         else:

atom/model_ops/paged_attention.py

@@ -116,9 +116,20 @@ def forward(
         positions: torch.Tensor = None,
         q_scale: Optional[torch.Tensor] = None,
         qkv: torch.Tensor = None,
+        qb_prezero: Optional[torch.Tensor] = None,
+        oproj_prezero: Optional[torch.Tensor] = None,
         **kwargs,
     ):
         output = torch.ops.aiter.unified_attention_with_output_base(
-            query, q_scale, key, value, positions, self.layer_name, self.use_mla, qkv
+            query,
+            q_scale,
+            key,
+            value,
+            positions,
+            self.layer_name,
+            self.use_mla,
+            qkv,
+            qb_prezero,
+            oproj_prezero,
         )
         return output

atom/model_ops/prezero.py

@@ -0,0 +1,89 @@
+"""Split-K GEMM "prezero" helpers.
+
+A split-K GEMM can skip its in-kernel output zero-init (NOZINIT) and atomic-add
+into a buffer the preceding fused allreduce-rmsnorm zeroed for free, when that
+pays off for the current batch. The decision lives inside opaque custom ops so
+it re-runs at each per-bs cudagraph capture instead of being frozen by the one
+prefill-time compile.
+"""
+
+from typing import Optional
+
+import torch
+from aiter import is_prezero_free
+from aiter.dist.communication_op import tensor_model_parallel_fused_allreduce_rmsnorm
+from aiter.jit.utils.torch_guard import torch_compile_guard
+from aiter.tuned_gemm import tgemm
+
+from atom.config import get_current_atom_config
+from atom.utils.forward_context import get_forward_context
+
+
+def _prezero_hidden() -> int:
+    # AR out_hidden_dim == hidden_size; sets the zero-fill CTA width.
+    return get_current_atom_config().hf_config.hidden_size
+
+
+def prezero_active(n_total: int) -> bool:
+    import os
+
+    if not os.getenv("ATOM_ENABLE_PREZERO"):
+        return False
+    ctx = get_forward_context().context
+    return (not ctx.is_prefill) and is_prezero_free(
+        ctx.graph_bs, n_total, _prezero_hidden()
+    )
+
+
+@torch_compile_guard(
+    mutates_args=["out"],
+    gen_fake=lambda out, a, weight, n_total, bias=None: None,
+)
+def mm_maybe_prezero_(
+    out: torch.Tensor,
+    a: torch.Tensor,
+    weight: torch.Tensor,
+    n_total: int,
+    bias: Optional[torch.Tensor] = None,
+) -> None:
+    active = prezero_active(n_total)
+    tgemm.mm(a, weight, bias, zero_init=not active, out=out)
+
+
+@torch_compile_guard(
+    mutates_args=["zero_buf"],
+    gen_fake=lambda x, residual, weight, eps, zero_buf, n_total, n_base=0: (
+        torch.empty_like(x),
+        torch.empty_like(residual),
+    ),
+)
+def ar_rmsnorm_maybe_prezero_(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    weight: torch.Tensor,
+    eps: float,
+    zero_buf: torch.Tensor,
+    n_total: int,
+    n_base: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # zero the largest free cumulative prefix; drop only the overflowing tail.
+    import os
+
+    zf = None
+    if os.getenv("ATOM_ENABLE_PREZERO"):
+        ctx = get_forward_context().context
+        if not ctx.is_prefill:
+            bs = ctx.graph_bs
+            h = x.shape[-1]
+            if is_prezero_free(bs, n_total, h):
+                zf = zero_buf
+            elif n_base and is_prezero_free(bs, n_base, h):
+                m = x.shape[0]
+                zf = zero_buf.view(-1)[: m * n_base].view(m, n_base)
+    return tensor_model_parallel_fused_allreduce_rmsnorm(
+        x.contiguous(),
+        residual,
+        weight,
+        eps,
+        zero_fill=zf,
+    )

atom/models/deepseek_v2.py

@@ -1941,6 +1941,26 @@ def __init__(
             if layer_quant_dtype in (dtypes.fp8, dtypes.fp4x2):
                 self.fuse_qknorm_quant = True
 
+        # split-K prezero widths, read by the pre-split compile pass. qkv_a + q_b
+        # ride the input_layernorm AR donor (base); o_proj adds an overflowing
+        # tail gated independently (greedy). Only bf16 (a16w16) GEMMs qualify.
+        if self.q_lora_rank is not None:
+            n_qb = (
+                self.q_b_proj.weight.shape[0]
+                if self.q_b_proj.quant_type.value == QuantType.No.value
+                else 0
+            )
+            n_oproj = (
+                self.o_proj.weight.shape[0]
+                if self.o_proj.quant_type.value == QuantType.No.value
+                else 0
+            )
+            self.q_b_proj.prezero_n_qkva = self.fused_qkv_a_proj.weight.shape[0]
+            self.q_b_proj.prezero_n_qb = n_qb
+            self.q_b_proj.prezero_n_oproj = n_oproj
+            self.q_b_proj.prezero_n_base = self.q_b_proj.prezero_n_qkva + n_qb
+            self.q_b_proj.prezero_n_total = self.q_b_proj.prezero_n_base + n_oproj
+
     def forward(
         self,
         positions: torch.Tensor,