feat(mhc): use HIP Mega-MHC by default for supported DeepSeek V4.1

atom/model_ops/deepseek_v41/mhc_pre_delayed.py

@@ -1,23 +1,70 @@
 # SPDX-License-Identifier: MIT
-"""The delayed mHC seam, one AITER kernel pair per sublayer boundary.
+"""The delayed mHC seam, including the sublayer's input RMSNorm.
 
 V4's single `aiter.mhc_pre` cannot serve this model: the collapse applies the
 pre-mix carried from the PREVIOUS sublayer while the one computed here is
-carried to the next. `mhc_fused_post_pre_delayed_rmsnorm` is that seam whole --
+carried to the next. HIP Mega-MHC and the Triton implementation cover that seam --
 the owed post, the gate projection, Sinkhorn, the collapse with the carried
-pre-mix and the sublayer's input RMSNorm -- as a split-K main kernel and a
-per-token reduce; vLLM's ROCm V4.1 runs the same kernel.
+pre-mix and the sublayer's input RMSNorm. HIP uses one compute kernel on gfx950;
+Triton remains available through ATOM_DSV41_MHC_BACKEND=triton.
 
-The norm is the sublayer's own: its output is BF16, and a layer whose first
-GEMM takes FP8 quantizes it there, as the reference model does.
+The norm is the sublayer's own. Attention can request HIP's fused MXFP8
+output alongside BF16 for its other readers; FFN keeps the BF16 interface.
 """
 
+from functools import cache
+
 import torch
 from aiter.jit.utils.torch_guard import torch_compile_guard
 from aiter.ops.triton.fusions.mhc_fused_post_pre_delayed_rmsnorm import (
     mhc_fused_post_pre_delayed_rmsnorm,
 )
 
+from atom.utils import envs
+
+
+@cache
+def _hip_ops():
+    try:
+        from aiter.ops.mega_mhc import mega_mhc, mega_mhc_shuffle_fn
+    except ImportError:
+        return None
+    return mega_mhc, mega_mhc_shuffle_fn
+
+
+def pack_mhc_fn(fn):
+    """Select the backend and pack once after checkpoint loading, like V4."""
+    backend = envs.ATOM_DSV41_MHC_BACKEND
+    if backend not in ("auto", "hip", "triton"):
+        raise ValueError("ATOM_DSV41_MHC_BACKEND must be auto, hip or triton")
+    if backend == "triton":
+        return None
+    supported = (
+        fn.is_cuda
+        and fn.dtype == torch.float32
+        and fn.ndim == 2
+        and fn.shape[0] == 24
+        and fn.shape[1] in (4 * 4096, 4 * 5120, 4 * 6144, 4 * 7168)
+        and fn.is_contiguous()
+        and getattr(
+            torch.cuda.get_device_properties(fn.device), "gcnArchName", ""
+        ).split(":")[0]
+        == "gfx950"
+    )
+    if not supported:
+        if backend == "hip":
+            raise ValueError(
+                "HIP Mega-MHC requires gfx950, contiguous FP32 fn, "
+                "hc_mult=4 and H in {4096, 5120, 6144, 7168}"
+            )
+        return None
+    ops = _hip_ops()
+    if ops is None:
+        if backend == "hip":
+            raise RuntimeError("HIP Mega-MHC requires AITER mega_mhc_shuffle_fn")
+        return None
+    return ops[1](fn)
+
 
 def pre_delayed(
     residual,
@@ -32,6 +79,8 @@ def pre_delayed(
     sinkhorn_iters,
     post_mult,
     norm_eps,
+    w_preshuffle_bf16=False,
+    fuse_quant=False,
     sublayer_output=None,
     post_mix=None,
     combination=None,
@@ -40,6 +89,9 @@ def pre_delayed(
 
     Returns `(residual, normed_input, next_pre_mix, post_mix, combination)`.
     `residual` is the caller's own tensor when no post was requested.
+    `hc_fn` is a HIP BF16 hi/lo pack when `w_preshuffle_bf16=True`,
+    or the original FP32 weight for Triton otherwise.
+    With `fuse_quant`, `normed_input` is `(bf16, fp8, e8m0_scale)`.
     """
     outputs = v41_mhc_pre_delayed(
         residual,
@@ -53,13 +105,20 @@ def pre_delayed(
         sinkhorn_iters=sinkhorn_iters,
         post_mult=post_mult,
         norm_eps=norm_eps,
+        w_preshuffle_bf16=w_preshuffle_bf16,
+        fuse_quant=fuse_quant,
         sublayer_output=sublayer_output,
         post_mix=post_mix,
         combination=combination,
     )
-    # Preserve the no-post input alias outside the custom op. The guarded
-    # implementation returns four new tensors, or five when post updates it.
-    return (residual, *outputs) if sublayer_output is None else tuple(outputs)
+    # Keep the guarded op's outputs flat and preserve the no-post input alias
+    # outside it. Only the attention-facing wrapper groups the norm outputs.
+    if fuse_quant:
+        outputs, quant, scale = outputs[:-2], outputs[-2], outputs[-1]
+    outputs = (residual, *outputs) if sublayer_output is None else tuple(outputs)
+    if fuse_quant:
+        outputs = (outputs[0], (outputs[1], quant, scale), *outputs[2:])
+    return outputs
 
 
 def _pre_delayed_fake(
@@ -71,6 +130,7 @@ def _pre_delayed_fake(
     norm_weight,
     *,
     sublayer_output=None,
+    fuse_quant=False,
     **kwargs,
 ):
     *outer, hc, hidden = residual.shape
@@ -80,6 +140,13 @@ def _pre_delayed_fake(
         residual.new_empty((*outer, hc), dtype=torch.float32),
         residual.new_empty((*outer, hc, hc), dtype=torch.float32),
     ]
+    if fuse_quant:
+        outputs.extend(
+            [
+                residual.new_empty((*outer, hidden), dtype=torch.float8_e4m3fn),
+                residual.new_empty((*outer, hidden // 32), dtype=torch.float8_e8m0fnu),
+            ]
+        )
     return (
         outputs if sublayer_output is None else [torch.empty_like(residual), *outputs]
     )
@@ -99,6 +166,8 @@ def v41_mhc_pre_delayed(
     sinkhorn_iters: int,
     post_mult: float,
     norm_eps: float,
+    w_preshuffle_bf16: bool = False,
+    fuse_quant: bool = False,
     sublayer_output: torch.Tensor | None = None,
     post_mix: torch.Tensor | None = None,
     combination: torch.Tensor | None = None,
@@ -107,7 +176,15 @@ def v41_mhc_pre_delayed(
     *outer, hc_mult, hidden = residual.shape
     rows = residual.numel() // (hc_mult * hidden)
     with_post = sublayer_output is not None
-    new_residual, post, comb, normed, pre = mhc_fused_post_pre_delayed_rmsnorm(
+    kernel = mhc_fused_post_pre_delayed_rmsnorm
+    options = {}
+    if fuse_quant and not w_preshuffle_bf16:
+        raise ValueError("Fused mHC quantization requires the HIP backend")
+    if w_preshuffle_bf16:
+        kernel = _hip_ops()[0]
+        options["w_preshuffle_bf16"] = True
+        options["out_dtype"] = "mxfp8" if fuse_quant else "bf16"
+    new_residual, post, comb, normed, pre = kernel(
         residual.view(rows, hc_mult, hidden),
         hc_fn,
         hc_scale,
@@ -125,11 +202,21 @@ def v41_mhc_pre_delayed(
         ),
         norm_weight=norm_weight,
         norm_eps=norm_eps,
+        **options,
     )
+    if fuse_quant:
+        normed, quant, scale = normed
     outputs = [
         normed.view(*outer, hidden),
         pre.view(*outer, hc_mult),
         post.view(*outer, hc_mult),
         comb.view(*outer, hc_mult, hc_mult),
     ]
+    if fuse_quant:
+        outputs.extend(
+            [
+                quant.view(*outer, hidden),
+                scale.view(torch.float8_e8m0fnu).view(*outer, hidden // 32),
+            ]
+        )
     return outputs if not with_post else [new_residual.view_as(residual), *outputs]

atom/models/deepseek_v41/attention.py

@@ -364,11 +364,15 @@ def _fork_select(self, hidden, qr, qr_scale, cache, step, rope, *, compressed):
             self.indexer.score(*projected, cache, step)
         return joins
 
-    def forward(self, hidden, hidden_scale, cache, step, rope, queries=None):
+    def forward(
+        self, hidden, hidden_scale, cache, step, rope, queries=None, *, hidden_fp8=None
+    ):
         """``queries``: bounded replay's `QueryRows`, the rows of ``step`` this
         layer is asked outputs for; KV is still every row's."""
         compressed = self._fork_compress(hidden, cache, step, rope)
-        q_lora, kv_pre = self.project_qkv(hidden, hidden_scale)
+        q_lora, kv_pre = self.project_qkv(
+            hidden if hidden_fp8 is None else hidden_fp8, hidden_scale
+        )
         qr, qr_scale, kv_normed = self.qk_norm(q_lora, kv_pre)
         selecting = self._fork_select(
             hidden, qr, qr_scale, cache, step, rope, compressed=compressed

atom/models/deepseek_v41/dspark.py

@@ -71,8 +71,10 @@ def project_context(self, hidden, positions, rope, *, packed=False):
         _, kv_pre = self.project_qkv(hidden)
         return self.context_keys(self.kv_norm(kv_pre), positions, rope, packed=packed)
 
-    def forward(self, hidden, hidden_scale, context_kv, step, rope):
-        q_lora, kv_pre = self.project_qkv(hidden, hidden_scale)
+    def forward(self, hidden, hidden_scale, context_kv, step, rope, *, hidden_fp8=None):
+        q_lora, kv_pre = self.project_qkv(
+            hidden if hidden_fp8 is None else hidden_fp8, hidden_scale
+        )
         qr, qr_scale, kv_normed = self.qk_norm(q_lora, kv_pre)
         query = self.wq_b(qr, x_scale=qr_scale).unflatten(
             -1, (self.heads, self.head_dim)

atom/models/deepseek_v41/model.py

@@ -11,7 +11,10 @@
 
 from atom.model_loader.weight_names import WeightsMapper
 from atom.model_ops.deepseek_v41.mhc import SinglePassHCState
-from atom.model_ops.deepseek_v41.mhc_pre_delayed import pre_delayed
+from atom.model_ops.deepseek_v41.mhc_pre_delayed import (
+    pack_mhc_fn,
+    pre_delayed,
+)
 from atom.model_ops.deepseek_v41.rotary import RotaryEmbedding
 from atom.model_ops.embed_head import VocabParallelEmbedding
 from atom.model_ops.engram.device.layer import EngramOp
@@ -84,6 +87,7 @@ def __init__(
                     f"hc_{sublayer}_{suffix}",
                     atom_parameter(torch.empty(shape, dtype=torch.float32)),
                 )
+            self.register_buffer(f"hc_{sublayer}_fn_packed", None, persistent=False)
         self.engram = None
         if spec.layer_id in config.engram_layer_ids:
             width = (
@@ -100,8 +104,23 @@ def __init__(
                 quant_config=native_quant_config(),
             )
 
+    def process_weights_after_loading(self):
+        for sublayer in ("attn", "ffn"):
+            # Mono kernels still consume FP32 fn. Keep one read-only HIP pack
+            # per weight, shared by all token counts and execution streams.
+            setattr(
+                self,
+                f"hc_{sublayer}_fn_packed",
+                pack_mhc_fn(getattr(self, f"hc_{sublayer}_fn")),
+            )
+
     def attention_forward(self, normed, cache, step, rope):
-        """Run attention on its seam's normed BF16 input."""
+        """Keep BF16 for auxiliary readers and reuse the optional QKV quant."""
+        if isinstance(normed, tuple):
+            hidden, hidden_fp8, hidden_scale = normed
+            return self.attn(
+                hidden, hidden_scale, cache, step, rope, hidden_fp8=hidden_fp8
+            )
         return self.attn(normed, None, cache, step, rope)
 
     def engram_forward(self, residual, embeddings, image_mask):
@@ -121,33 +140,38 @@ def prepare_attention(self, state, embeddings, image_mask):
                 state,
                 residual=self.engram_forward(state.residual, embeddings, image_mask),
             )
+        hc_fn = self.hc_attn_fn_packed
         residual, normed, pre, post, comb = pre_delayed(
             state.residual,
             state.pre_mix,
-            self.hc_attn_fn,
+            hc_fn if hc_fn is not None else self.hc_attn_fn,
             self.hc_attn_scale,
             self.hc_attn_base,
             self.attn_norm.weight,
             **self.hc_options,
             norm_eps=self.attn_norm.eps,
+            w_preshuffle_bf16=hc_fn is not None,
+            fuse_quant=hc_fn is not None,
             sublayer_output=state.pending,
             post_mix=state.post_mix,
             combination=state.combination,
         )
         return normed, residual, pre, post, comb
 
     def prepare_ffn(self, output, residual, pre, post, comb):
-        # The attention post folds into this seam: one kernel pair computes the
+        # The attention post folds into this seam: the delayed kernel computes the
         # new residual, the gates and the normed FFN input.
+        hc_fn = self.hc_ffn_fn_packed
         residual, normed, pre, post, comb = pre_delayed(
             residual,
             pre,
-            self.hc_ffn_fn,
+            hc_fn if hc_fn is not None else self.hc_ffn_fn,
             self.hc_ffn_scale,
             self.hc_ffn_base,
             self.ffn_norm.weight,
             **self.hc_options,
             norm_eps=self.ffn_norm.eps,
+            w_preshuffle_bf16=hc_fn is not None,
             sublayer_output=output,
             post_mix=post,
             combination=comb,