fix(deepseek-v4): route large-token fused mHC around the allinone cliff

python/tokenspeed/runtime/layers/deepseek_v4_mhc.py

@@ -574,6 +574,27 @@ def mhc_fused_hc(
     return residual_cur, layer_input, post_cur, comb_cur
 
 
+# Above this token count the vendored allinone kernel (TRT-LLM snapshot
+# 54efe57e) dispatches to a tf32 atomic variant with a ~10x per-call cliff
+# (222-240us vs a flat ~23us on B200); the composed two-stage path
+# (post_mapping + prenorm GEMM + big fuse) takes over there.
+_MHC_FUSED_ALLINONE_MAX_TOKENS = 32
+_MHC_FUSED_ROUTING_LOGGED: set[tuple[int, str]] = set()
+
+
+def _log_mhc_fused_routing(num_tokens: int, path: str) -> None:
+    key = (num_tokens, path)
+    if key in _MHC_FUSED_ROUTING_LOGGED:
+        return
+    _MHC_FUSED_ROUTING_LOGGED.add(key)
+    logger.info(
+        "V4 mHC fused_hc routing: num_tokens=%d path=%s threshold=%d",
+        num_tokens,
+        path,
+        _MHC_FUSED_ALLINONE_MAX_TOKENS,
+    )
+
+
 def _trtllm_mhc_fused_hc(
     x_prev: torch.Tensor,
     residual_prev: torch.Tensor,
@@ -611,6 +632,26 @@ def _trtllm_mhc_fused_hc(
             ),
         )
 
+    if B > _MHC_FUSED_ALLINONE_MAX_TOKENS:
+        # The vendored allinone kernel dispatches to a tf32 atomic variant
+        # above 32 tokens that runs ~10x slower (216-242us vs a flat ~23us
+        # for the composed two-stage path, microbenched on B200). Route large
+        # token counts through post_mapping + prenorm-GEMM + big-fuse instead;
+        # outputs agree within bf16 tolerance across chained layers.
+        _log_mhc_fused_routing(B, "composed")
+        residual_cur = _trtllm_mhc_post(x_prev, residual_prev, post_prev, comb_prev)
+        layer_input, post_cur, comb_cur = _trtllm_mhc_pre(
+            residual_cur,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_eps,
+            sinkhorn_iters,
+        )
+        return residual_cur, layer_input, post_cur, comb_cur
+    _log_mhc_fused_routing(B, "allinone")
+
     x_flat = x_prev.reshape(B, hidden_size).contiguous()
     res_flat = residual_prev.reshape(B, hc_mult, hidden_size).contiguous()
     post_flat = post_prev.reshape(B, hc_mult).float().contiguous()

test/runtime/test_deepseek_v4_mhc_parity_gpu.py

@@ -0,0 +1,221 @@
+# Copyright (c) 2026 LightSeek Foundation
+
+"""GPU parity for fused-mHC allinone vs composed routing at production dims.
+
+Compares all four outputs (residual, layer_input, post, comb) of the allinone
+kernel against the composed post_mapping + prenorm-GEMM + big-fuse path on
+identical inputs, across the routing boundary shapes and over an 8-layer
+chained iteration, with absolute and relative error bounds.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import pytest
+import torch
+
+# CI registration (AST-parsed, runtime no-op). The TRT-LLM mHC kernels only
+# support NVIDIA sm100, so restrict to B200-class runners.
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from ci_system.ci_register import register_cuda_ci  # noqa: E402
+
+register_cuda_ci(
+    est_time=120,
+    suite="runtime-1gpu",
+    disabled_on_runners=["h100-*", "b300-*", "*mi3*", "linux-mi*"],
+    disabled_on_runners_reason="TRT-LLM mHC kernels are sm100-only",
+)
+
+os.environ.setdefault("TOKENSPEED_V4_MHC_BACKEND", "trtllm")
+
+from tokenspeed.runtime.layers import deepseek_v4_mhc as mhc  # noqa: E402
+
+
+def _is_sm100() -> bool:
+    return (
+        torch.cuda.is_available()
+        and torch.version.hip is None
+        and torch.cuda.get_device_capability(0) == (10, 0)
+    )
+
+
+pytestmark = pytest.mark.skipif(not _is_sm100(), reason="requires NVIDIA sm100")
+
+HC, HID = 4, 7168
+MIX = (2 + HC) * HC
+RMS_EPS, HC_EPS, SINK = 1e-6, 1e-2, 4
+REL_TOL = 0.03
+
+
+def _inputs(m, seed):
+    torch.manual_seed(seed)
+    dev = "cuda:0"
+    return (
+        torch.randn(m, HID, device=dev, dtype=torch.bfloat16),
+        torch.randn(m, HC, HID, device=dev, dtype=torch.bfloat16) * 0.5,
+        torch.rand(m, HC, 1, device=dev),
+        torch.rand(m, HC, HC, device=dev) / HC,
+        torch.randn(MIX, HC * HID, device=dev) * 0.02,
+        torch.ones(3, device=dev),
+        torch.zeros(MIX, device=dev),
+    )
+
+
+def _allinone(x, res, post, comb, fn, scale, base, ws):
+    if res.shape[0] > mhc._MHC_FUSED_ALLINONE_MAX_TOKENS:
+        pytest.skip("allinone reference only defined for B<=threshold")
+    return mhc._trtllm_mhc_fused_hc(
+        x, res, post, comb, fn, scale, base, RMS_EPS, HC_EPS, SINK, ws
+    )
+
+
+def _composed(x, res, post, comb, fn, scale, base):
+    res_cur = mhc._trtllm_mhc_post(x, res, post, comb)
+    layer_in, post_cur, comb_cur = mhc._trtllm_mhc_pre(
+        res_cur, fn, scale, base, RMS_EPS, HC_EPS, SINK
+    )
+    return res_cur, layer_in, post_cur, comb_cur
+
+
+def _check(oa, ob, m, layer):
+    names = ("residual", "layer_input", "post", "comb")
+    for name, ta, tb in zip(names, oa, ob):
+        fa, fb = ta.float(), tb.float()
+        abs_err = (fa - fb).abs().max().item()
+        rel_err = ((fa - fb).norm() / fb.norm().clamp_min(1e-9)).item()
+        if name in ("residual", "layer_input"):
+            # bf16 outputs: bound by bf16 ULP at the reference magnitude
+            # (the two paths quantize to bf16 at different points), floored
+            # for near-zero references and growing with chain depth since
+            # ULP-level drift compounds per chained layer.
+            abs_tol = max(0.02, fb.abs().max().item() * 2**-7 * (2 + layer))
+        else:
+            # fp32 post/comb: O(1) sigmoid-scale values computed through
+            # bf16-quantized intermediates; fixed small bound.
+            abs_tol = 0.02 * (1 + 0.5 * layer)
+        assert (
+            abs_err <= abs_tol
+        ), f"M={m} L{layer} {name} abs={abs_err:.4f} tol={abs_tol:.4f}"
+        assert rel_err < REL_TOL, f"M={m} L{layer} {name} rel={rel_err:.4f}"
+
+
+@pytest.mark.parametrize("m", [16, 32])
+@pytest.mark.parametrize("seed", [0, 1, 2])
+def test_mhc_allinone_vs_composed_chained_parity(m, seed):
+    # At/below the threshold both paths are reachable; the composed path must
+    # agree with allinone through an 8-layer chained iteration.
+    x, res, post, comb, fn, scale, base = _inputs(m, seed)
+    ws = mhc.MhcFusedWorkspace()
+    ws.reset()
+    fa = (x, res, post, comb)
+    fb = (x, res, post, comb)
+    for layer in range(8):
+        oa = tuple(t.clone() for t in _allinone(*fa[:4], fn, scale, base, ws))
+        ob = _composed(*fb[:4], fn, scale, base)
+        _check(oa, ob, m, layer)
+
+        # Normalize the re-injected sublayer output like a real network's
+        # norms would, so synthetic magnitudes stay bounded across the chain.
+        def _renorm(t):
+            return (t.float() / t.float().std().clamp_min(1e-3)).to(torch.bfloat16)
+
+        fa = (
+            _renorm(oa[1]),
+            oa[0],
+            oa[2].view(m, HC, 1),
+            oa[3].view(m, HC, HC),
+        )
+        fb = (
+            _renorm(ob[1]),
+            ob[0],
+            ob[2].view(m, HC, 1),
+            ob[3].view(m, HC, HC),
+        )
+
+
+@pytest.mark.parametrize("m", [36, 40, 64])
+def test_mhc_routed_path_matches_composed_reference(m):
+    # Above the threshold the wrapper must return exactly the composed-path
+    # results (routing correctness at the boundary shapes).
+    x, res, post, comb, fn, scale, base = _inputs(m, 0)
+    ws = mhc.MhcFusedWorkspace()
+    ws.reset()
+    routed = mhc._trtllm_mhc_fused_hc(
+        x, res, post, comb, fn, scale, base, RMS_EPS, HC_EPS, SINK, ws
+    )
+    ref = _composed(x, res, post, comb, fn, scale, base)
+    for ta, tb in zip(routed, ref):
+        assert torch.equal(ta.float(), tb.float())
+
+
+@pytest.mark.parametrize("m", [36, 40, 64])
+@pytest.mark.parametrize("seed", [0, 1])
+def test_mhc_large_m_allinone_vs_composed_chained_parity(m, seed):
+    # The path being REPLACED: force the allinone kernel above the routing
+    # threshold (it is ~10x slower there but numerically defined) and compare
+    # against the composed replacement over an 8-layer chained iteration, so
+    # cross-layer drift on the affected shapes is bounded too.
+    from unittest.mock import patch
+
+    x, res, post, comb, fn, scale, base = _inputs(m, seed)
+    ws = mhc.MhcFusedWorkspace()
+    ws.reset()
+    fa = (x, res, post, comb)
+    fb = (x, res, post, comb)
+
+    def _renorm(t):
+        return (t.float() / t.float().std().clamp_min(1e-3)).to(torch.bfloat16)
+
+    for layer in range(8):
+        with patch.object(mhc, "_MHC_FUSED_ALLINONE_MAX_TOKENS", 1 << 20):
+            oa = tuple(
+                t.clone()
+                for t in mhc._trtllm_mhc_fused_hc(
+                    *fa[:4], fn, scale, base, RMS_EPS, HC_EPS, SINK, ws
+                )
+            )
+        ob = _composed(*fb[:4], fn, scale, base)
+        _check(oa, ob, m, layer)
+        fa = (_renorm(oa[1]), oa[0], oa[2].view(m, HC, 1), oa[3].view(m, HC, HC))
+        fb = (_renorm(ob[1]), ob[0], ob[2].view(m, HC, 1), ob[3].view(m, HC, HC))
+
+
+@pytest.mark.parametrize("m", [36, 64])
+def test_mhc_composed_graph_capture_replay(m):
+    # The composed branch allocates temporaries inside CUDA graph capture;
+    # replay must be stable, respond to static-input mutation, match eager,
+    # and not grow the graph pool across replays.
+    x, res, post, comb, fn, scale, base = _inputs(m, 0)
+    ws = mhc.MhcFusedWorkspace()
+    ws.reset()
+
+    def call():
+        return mhc._trtllm_mhc_fused_hc(
+            x, res, post, comb, fn, scale, base, RMS_EPS, HC_EPS, SINK, ws
+        )
+
+    for _ in range(3):
+        call()
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = call()
+    for _ in range(10):
+        graph.replay()
+    torch.cuda.synchronize()
+    eager = _composed(x, res, post, comb, fn, scale, base)
+    for ta, tb in zip(out, eager):
+        assert torch.equal(ta.float(), tb.float())
+    # Mutate a static input and replay again; the captured graph must track it.
+    x.copy_(torch.randn_like(x))
+    graph.replay()
+    torch.cuda.synchronize()
+    eager2 = _composed(x, res, post, comb, fn, scale, base)
+    assert torch.equal(out[1].float(), eager2[1].float())
+    reserved = torch.cuda.memory_reserved()
+    for _ in range(50):
+        graph.replay()
+    torch.cuda.synchronize()
+    assert torch.cuda.memory_reserved() == reserved, "graph pool grew on replay"

test/runtime/test_deepseek_v4_mhc_routing.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2026 LightSeek Foundation
+
+"""Routing tests for the fused-mHC allinone/composed token threshold."""
+
+from __future__ import annotations
+
+import os
+import sys
+from unittest.mock import patch
+
+# CI registration (AST-parsed, runtime no-op). Routing logic is mocked and
+# GPU-free, but the tokenspeed import graph requires a CUDA environment.
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from ci_system.ci_register import register_cuda_ci  # noqa: E402
+
+register_cuda_ci(est_time=30, suite="runtime-1gpu")
+
+import torch
+
+from tokenspeed.runtime.layers import deepseek_v4_mhc as mhc
+
+HC, HID = 4, 64
+
+
+def _inputs(m):
+    x = torch.randn(m, HID, dtype=torch.bfloat16)
+    res = torch.randn(m, HC, HID, dtype=torch.bfloat16)
+    post = torch.rand(m, HC, 1)
+    comb = torch.rand(m, HC, HC)
+    fn = torch.randn((2 + HC) * HC, HC * HID)
+    return x, res, post, comb, fn, torch.ones(3), torch.zeros((2 + HC) * HC)
+
+
+def _call(m):
+    x, res, post, comb, fn, scale, base = _inputs(m)
+    return mhc._trtllm_mhc_fused_hc(
+        x, res, post, comb, fn, scale, base, 1e-6, 1e-2, 4, mhc.MhcFusedWorkspace()
+    )
+
+
+def test_fused_hc_routes_large_tokens_to_composed():
+    mhc._MHC_FUSED_ROUTING_LOGGED.clear()
+    m = mhc._MHC_FUSED_ALLINONE_MAX_TOKENS + 1
+    sentinel_res = torch.zeros(m, HC, HID, dtype=torch.bfloat16)
+    with (
+        patch.object(mhc, "_trtllm_mhc_post", return_value=sentinel_res) as post_fn,
+        patch.object(
+            mhc,
+            "_trtllm_mhc_pre",
+            return_value=(
+                torch.zeros(m, HID, dtype=torch.bfloat16),
+                torch.zeros(m, HC, 1),
+                torch.zeros(m, HC, HC),
+            ),
+        ) as pre_fn,
+        patch.object(mhc, "trtllm_mhc_fused_hc") as allinone,
+    ):
+        out = _call(m)
+    post_fn.assert_called_once()
+    pre_fn.assert_called_once()
+    allinone.assert_not_called()
+    assert out[0] is sentinel_res
+
+
+def test_fused_hc_keeps_allinone_at_threshold_and_below():
+    mhc._MHC_FUSED_ROUTING_LOGGED.clear()
+    for m in (1, mhc._MHC_FUSED_ALLINONE_MAX_TOKENS):
+        with (
+            patch.object(mhc, "trtllm_mhc_fused_hc") as allinone,
+            patch.object(mhc, "_trtllm_mhc_post") as post_fn,
+            patch.object(mhc.MhcFusedWorkspace, "get") as ws_get,
+        ):
+            buf = lambda *s, dt=torch.bfloat16: torch.zeros(*s, dtype=dt)  # noqa: E731
+            ws_get.return_value.get.return_value = (
+                buf(m, HC, HID),
+                buf(m, HC, 1, dt=torch.float32),
+                buf(m, HC, HC, dt=torch.float32),
+                buf(m, HID),
+                buf(1, m, (2 + HC) * HC, dt=torch.float32),
+                buf(1, m, dt=torch.float32),
+                buf(1, dt=torch.int32),
+            )
+            _call(m)
+        allinone.assert_called_once()
+        post_fn.assert_not_called()
+
+
+def test_fused_hc_empty_batch_short_circuits():
+    out = _call(0)
+    assert out[0].shape == (0, HC, HID)
+
+
+def test_fused_hc_routing_logs_once_per_shape(caplog):
+    mhc._MHC_FUSED_ROUTING_LOGGED.clear()
+    m = mhc._MHC_FUSED_ALLINONE_MAX_TOKENS + 8
+    with (
+        patch.object(mhc, "_trtllm_mhc_post", return_value=torch.zeros(m, HC, HID)),
+        patch.object(
+            mhc,
+            "_trtllm_mhc_pre",
+            return_value=(
+                torch.zeros(m, HID),
+                torch.zeros(m, HC, 1),
+                torch.zeros(m, HC, HC),
+            ),
+        ),
+        caplog.at_level("INFO", logger=mhc.logger.name),
+    ):
+        _call(m)
+        _call(m)
+    routing_logs = [r for r in caplog.records if "fused_hc routing" in r.getMessage()]
+    assert len(routing_logs) == 1