Optimize GLM5.2 MTP

atom/config.py

@@ -1288,6 +1288,13 @@ def compute_hash(self) -> str:
                 ),
             )
         )
+        # Vocab-embedding replication (ATOM_REPLICATE_VOCAB_EMBED) changes both the
+        # embed weight shape ([vocab] vs [vocab/tp]) and the embed op (local
+        # F.embedding vs masked-embedding + all-reduce), so it alters the compiled
+        # graph and MUST be part of its key. Without this, toggling the flag — or
+        # deploying it on top of a cache built with the other setting — reuses a
+        # stale artifact and trips assert_size_stride at runtime.
+        factors.append(bool(envs.ATOM_REPLICATE_VOCAB_EMBED))
 
         hash_str = hashlib.md5(
             str(factors).encode(), usedforsecurity=False

atom/model_ops/attentions/aiter_mla.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 import torch
+import triton
 from atom.utils import envs
 from aiter import (
     decode_update_mla_metadata_v1,
@@ -20,6 +21,7 @@
 from atom.utils import CpuGpuBuffer
 from atom.utils.block_convert import (
     kv_indices_generate_triton,
+    mtp_prepare_decode_mla_kernel,
 )
 from atom.utils.forward_context import AttentionMetaData, Context
 
@@ -99,6 +101,11 @@ def get_impl_cls() -> Type["MLAAttention"]:
 
 
 class AiterMLAMetadataBuilder(CommonAttentionBuilder):
+    # EagleProposer folds the per-draft-step position/context bump into
+    # prepare_mtp_decode's fused kernel when this is set (matches the MHA
+    # backend). The fused kernel handles both sparse and dense MLA.
+    fuse_mtp_decode_position_update = True
+
     def __init__(self, model_runner):
         if envs.ATOM_MLA_PAGE_SIZE > 1:
             self.block_size = envs.ATOM_MLA_PAGE_SIZE
@@ -111,6 +118,10 @@ def __init__(self, model_runner):
                 f"got --block-size {model_runner.block_size}"
             )
         CommonAttentionBuilder.__init__(self, model_runner)
+        # Single-program block for the fused MTP-decode metadata kernel. Sized
+        # to the max batch (runtime bs <= max_bs) so one tl.cumsum spans the
+        # whole batch in a single launch.
+        self._mtp_fuse_block = triton.next_power_of_2(self.max_bs + 1)
         config = model_runner.config
         hf_config = config.hf_config
         # `self.num_attention_heads` set by CommonAttentionBuilder.__init__.
@@ -589,21 +600,53 @@ def prepare_mtp_decode(
         positions: torch.Tensor,  # [total_tokens] int32
         only_update: bool = False,
         num_reject_tokens: torch.Tensor = None,
+        *,
+        update_context_lens: bool = False,
+        positions_out: torch.Tensor | None = None,
+        last_token_indices: torch.Tensor | None = None,
     ):
+        """Per-draft-step MLA metadata update, fused into a single kernel.
+
+        One ``_mtp_prepare_decode_mla_kernel`` launch performs, in place:
+          - ``kv_indptr += cu_seqlens_q`` (needed by kv_indices + slot_mapping),
+          - (sparse) per-seq ``min(kv_count, index_topk)`` cumsum ->
+            ``sparse_kv_indptr``,
+          - (fused position update) ``positions += 1`` when ``positions_out`` is
+            given, and ``context_lens += 1`` when ``update_context_lens`` is set.
+
+        ``fuse_mtp_decode_position_update`` makes EagleProposer route the
+        per-step position/context bumps through here instead of launching them
+        as separate kernels. ``last_token_indices`` is accepted for signature
+        parity with the MHA backend but unused (MLA's ``positions`` is already
+        one entry per sequence at this point).
+        """
+        del last_token_indices  # MLA positions are already per-seq (1 per token)
         var = self.model_runner.forward_vars
         kv_indptr = var["kv_indptr"].gpu[: bs + 1]
+        cu_seqlens_q = var["cu_seqlens_q"].gpu[: bs + 1]
         if self.is_sparse:
-            # Update dense kv_indptr (needed for kv_indices generation and slot_mapping)
-            kv_indptr += var["cu_seqlens_q"].gpu[: bs + 1]
-            # Recompute sparse_kv_indptr: per-seq sparse count = min(dense_kv_count, index_topk)
             sparse_kv_indptr = var["sparse_kv_indptr"].gpu[: bs + 1]
-            kv_counts = kv_indptr[1 : bs + 1] - kv_indptr[:bs]
-            sparse_counts = torch.clamp(kv_counts, max=self.index_topk)
-            sparse_kv_indptr[0] = 0
-            sparse_kv_indptr[1 : bs + 1] = torch.cumsum(sparse_counts, dim=0)
         else:
             assert self.block_size == 1
-            kv_indptr += var["cu_seqlens_q"].gpu[: bs + 1]
+            sparse_kv_indptr = None
+
+        update_positions = positions_out is not None
+        context_lens = var["context_lens"].gpu[:bs] if update_context_lens else None
+
+        mtp_prepare_decode_mla_kernel[(1,)](
+            kv_indptr,
+            cu_seqlens_q,
+            sparse_kv_indptr if self.is_sparse else kv_indptr,
+            positions_out if update_positions else kv_indptr,
+            context_lens if update_context_lens else kv_indptr,
+            bs,
+            self.index_topk if self.is_sparse else 0,
+            positions_out.stride(0) if update_positions else 1,
+            IS_SPARSE=self.is_sparse,
+            UPDATE_POSITIONS=update_positions,
+            UPDATE_CONTEXT_LENS=update_context_lens,
+            BLOCK=self._mtp_fuse_block,
+        )
 
         kv_indices_generate_triton(
             var["block_tables"].gpu[:bs],

atom/model_ops/embed_head.py

@@ -103,6 +103,24 @@ def masked_embedding(
     return _masked_embedding_launcher(x, weight, vocab_start_idx, vocab_end_idx)
 
 
+def _replicated_embedding_fake(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    return torch.empty(
+        x.numel(),
+        weight.shape[1],
+        dtype=weight.dtype,
+        device=weight.device,
+    )
+
+
+@torch_compile_guard(gen_fake=_replicated_embedding_fake)
+def replicated_embedding(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    # Keep the lookup opaque to torch.compile: inductor otherwise fuses the
+    # embedding gather into the surrounding graph, which corrupts the MTP draft
+    # rollout (acceptance collapses ~69%->45%) — the same reason
+    # VocabParallelEmbedding routes through the masked_embedding custom op.
+    return F.embedding(x, weight)
+
+
 class VocabParallelEmbedding(nn.Module):
 
     def __init__(
@@ -184,7 +202,7 @@ def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
         param.data.copy_(loaded_weight)
 
     def forward(self, x: torch.Tensor):
-        return F.embedding(x, self.weight)
+        return replicated_embedding(x, self.weight)
 
 
 class ParallelLMHead(VocabParallelEmbedding):

atom/model_ops/layernorm.py

@@ -219,6 +219,7 @@ def __init__(
         x_pad_to_multiple: int = 0,
         fused_allreduce: bool = False,
         fused_quant: bool = False,
+        fused_quant_emit_bf16: bool = False,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
     ) -> None:
@@ -229,6 +230,12 @@ def __init__(
         self.x_pad_to_multiple = x_pad_to_multiple
         self.fused_allreduce = fused_allreduce
         self.use_fused_quant = fused_quant
+        # When the fused AllReduce+RMSNorm+quant path is active AND a downstream
+        # consumer also needs the unquantized (bf16) normed activation (e.g. the
+        # GLM-5.2 DSA indexer, whose wk/weights_proj GEMMs run in BF16), the
+        # kernel additionally emits that bf16 mirror. forward() then returns
+        # ((fp8, scale, bf16), residual) instead of ((fp8, scale), residual).
+        self.fused_quant_emit_bf16 = fused_quant_emit_bf16
         self.tp_size = get_tensor_model_parallel_world_size()
         self.quant_config = quant_config
         self.prefix = prefix
@@ -335,6 +342,44 @@ def forward(
             assert (
                 residual is not None
             ), "fused_allreduce_rmsnorm requires residual input!"
+            if self.use_fused_quant and self.quant_type.value in (
+                _QV_PER_1X128,
+                _QV_PER_TOKEN,
+            ):
+                # Combined AllReduce + RMSNorm + FP8 quant: the downstream GEMM
+                # (e.g. fused_qkv_a_proj) consumes the (fp8, scale) output
+                # directly, dropping a standalone per-token/per-group quant kernel
+                # from the hot path. `_aiter_transpose_scale` (resolved at init)
+                # selects the scale layout for that GEMM. The fused kernel does
+                # not support non-contiguous input.
+                if self.fused_quant_emit_bf16:
+                    # Also emit the pre-quant bf16 normed activation for a
+                    # co-consumer that runs in BF16 (GLM-5.2 DSA indexer).
+                    x, residual, x_scale, x_bf16 = (
+                        tensor_model_parallel_fused_allreduce_rmsnorm_quant(
+                            x.contiguous(),
+                            residual,
+                            self.weight,
+                            self.eps,
+                            quant_type=self.quant_type,
+                            group_size=128,
+                            transpose_scale=self._aiter_transpose_scale,
+                            emit_bf16=True,
+                        )
+                    )
+                    return (x, x_scale, x_bf16), residual
+                x, residual, x_scale = (
+                    tensor_model_parallel_fused_allreduce_rmsnorm_quant(
+                        x.contiguous(),
+                        residual,
+                        self.weight,
+                        self.eps,
+                        quant_type=self.quant_type,
+                        group_size=128,
+                        transpose_scale=self._aiter_transpose_scale,
+                    )
+                )
+                return (x, x_scale), residual
             # tensor_model_parallel_fused_allreduce_rmsnorm does not support non-contiguous input
             x, residual = tensor_model_parallel_fused_allreduce_rmsnorm(
                 x.contiguous(),
@@ -1054,6 +1099,109 @@ def __call__(
         )
 
 
+# ---------------------------------------------------------------------------
+# Fused dual RMSNorm + concat — single Triton launch.
+#
+# Two independent RMSNorms over the SAME hidden dim (different inputs, different
+# weights, shared eps) whose bf16 results are concatenated on the last axis:
+#     out = cat([rmsnorm(xe, we), rmsnorm(xh, wh)], dim=-1)   # [M, 2H]
+# One program per row does both norms and writes each into its half of the
+# [M, 2H] output, so the concat is free — it eliminates the separate cat
+# kernel's read+write of 2*M*H (halving traffic vs two norms + torch.cat) and
+# folds three launches (enorm, hnorm, cat) into one. Used by the DeepSeek/GLM
+# MTP eh_proj input (enorm(embed) ++ hnorm(prev_hidden)).
+# ---------------------------------------------------------------------------
+
+
+@triton.jit
+def _fused_dual_rmsnorm_cat_kernel(
+    xe_ptr,
+    xh_ptr,
+    we_ptr,
+    wh_ptr,
+    out_ptr,
+    H,
+    stride_xe_m,
+    stride_xh_m,
+    stride_out_m,
+    eps,
+    BLOCK_H: tl.constexpr,
+):
+    row = tl.program_id(0)
+    cols = tl.arange(0, BLOCK_H)
+    mask = cols < H
+
+    # enorm half -> out[row, :H]
+    xe = tl.load(xe_ptr + row * stride_xe_m + cols, mask=mask, other=0.0).to(tl.float32)
+    we = tl.load(we_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    rstd_e = tl.rsqrt(tl.sum(xe * xe, axis=-1) / H + eps)
+    out_e = (xe * rstd_e * we).to(out_ptr.dtype.element_ty)
+    tl.store(out_ptr + row * stride_out_m + cols, out_e, mask=mask)
+
+    # hnorm half -> out[row, H:2H]
+    xh = tl.load(xh_ptr + row * stride_xh_m + cols, mask=mask, other=0.0).to(tl.float32)
+    wh = tl.load(wh_ptr + cols, mask=mask, other=0.0).to(tl.float32)
+    rstd_h = tl.rsqrt(tl.sum(xh * xh, axis=-1) / H + eps)
+    out_h = (xh * rstd_h * wh).to(out_ptr.dtype.element_ty)
+    tl.store(out_ptr + row * stride_out_m + H + cols, out_h, mask=mask)
+
+
+def _fused_dual_rmsnorm_cat_fake(
+    xe: torch.Tensor,
+    we: torch.Tensor,
+    xh: torch.Tensor,
+    wh: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    M, H = xe.shape
+    return torch.empty((M, 2 * H), dtype=xe.dtype, device=xe.device)
+
+
+@torch_compile_guard(gen_fake=_fused_dual_rmsnorm_cat_fake)
+def fused_dual_rmsnorm_cat(
+    xe: torch.Tensor,
+    we: torch.Tensor,
+    xh: torch.Tensor,
+    wh: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """cat([rmsnorm(xe, we), rmsnorm(xh, wh)], dim=-1) in one Triton launch.
+
+    Args:
+        xe, xh: (M, H) bf16/fp16 inputs (same shape).
+        we, wh: (H,) RMSNorm weights (one per input).
+        eps: shared RMSNorm epsilon.
+    Returns:
+        (M, 2H) tensor: [:, :H] = rmsnorm(xe, we), [:, H:] = rmsnorm(xh, wh).
+    """
+    assert xe.shape == xh.shape, f"shape mismatch {xe.shape} vs {xh.shape}"
+    assert xe.dim() == 2, f"expected 2-D inputs, got {xe.dim()}-D"
+    xe = xe.contiguous()
+    xh = xh.contiguous()
+    M, H = xe.shape
+    out = torch.empty((M, 2 * H), dtype=xe.dtype, device=xe.device)
+    if M == 0:
+        return out
+    BLOCK_H = triton.next_power_of_2(H)
+    num_warps = 8 if BLOCK_H >= 4096 else 4
+    _fused_dual_rmsnorm_cat_kernel[(M,)](
+        xe,
+        xh,
+        we,
+        wh,
+        out,
+        H,
+        xe.stride(0),
+        xh.stride(0),
+        out.stride(0),
+        eps,
+        BLOCK_H=BLOCK_H,
+        num_warps=num_warps,
+        num_stages=2,
+    )
+    return out
+
+
 @torch_compile_guard()
 def layernorm2d_fwd_(
     x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, eps: float, dim: int