From 271bd32b8bae3e9fac7f5b021b52cc11b8161b6b Mon Sep 17 00:00:00 2001 From: Xiangyi Zhang Date: Thu, 9 Jul 2026 07:32:10 -0700 Subject: [PATCH] fix(deepseek-v4): route large-token fused mHC around the allinone cliff The vendored allinone kernel dispatches to a tf32 atomic variant above 32 tokens that runs ~10x slower per call (222-240us vs a flat ~23us for the composed post_mapping + prenorm-GEMM + big-fuse path; kernel identity confirmed via profiler, outputs agree within 1.2% relative error over 8 chained layers on all four outputs). Route B>32 through the composed path, keep allinone at B<=32 so the canonical p8 (M=32) path is unchanged, and log each (num_tokens, path) selection once per process. Signed-off-by: Xiangyi Zhang --- .../runtime/layers/deepseek_v4_mhc.py | 41 ++++ .../test_deepseek_v4_mhc_parity_gpu.py | 221 ++++++++++++++++++ test/runtime/test_deepseek_v4_mhc_routing.py | 112 +++++++++ 3 files changed, 374 insertions(+) create mode 100644 test/runtime/test_deepseek_v4_mhc_parity_gpu.py create mode 100644 test/runtime/test_deepseek_v4_mhc_routing.py diff --git a/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py b/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py index 7686fdf01..7236d5532 100644 --- a/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py +++ b/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() diff --git a/test/runtime/test_deepseek_v4_mhc_parity_gpu.py b/test/runtime/test_deepseek_v4_mhc_parity_gpu.py new file mode 100644 index 000000000..474bc09d1 --- /dev/null +++ b/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" diff --git a/test/runtime/test_deepseek_v4_mhc_routing.py b/test/runtime/test_deepseek_v4_mhc_routing.py new file mode 100644 index 000000000..a9c1ef461 --- /dev/null +++ b/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