From c2a48da9ab1d96c7cbca8f90604c7e1f05a81c08 Mon Sep 17 00:00:00 2001 From: Shiki Wu Date: Mon, 5 Oct 2026 03:37:54 -0700 Subject: [PATCH 1/2] [None][fix] Honor deterministic mode in fused mHC dispatch Signed-off-by: Shiki Wu --- tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 33 +++++++--- tests/unittest/_torch/modules/test_mhc.py | 68 +++++++++++++++++++++ 2 files changed, 92 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py index 3efeff5f409a..51ef3ed03b3f 100644 --- a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py +++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py @@ -17,6 +17,7 @@ Falls back to FMA when DeepGEMM is unavailable or autotuner cache misses. """ +import os from functools import lru_cache, partial from typing import Any, List @@ -1069,15 +1070,29 @@ def mhc_fused_hc( sinkhorn_repeat=sinkhorn_repeat, ) - tuner = AutoTuner.get() - _, best_tactic = tuner.choose_one( - "trtllm::mhc_fused_hc", - [runner], - MhcFusedHcRunner.tuning_config, - [x_prev, residual_prev, post_mix_prev, comb_mix_prev, w_t_cur, hc_scale_cur, hc_base_cur], - norm_weight=norm_weight, - norm_eps=norm_eps, - ) + if os.getenv("FORCE_DETERMINISTIC", "0") == "1": + # This fixed FMA tactic has no cross-CTA floating-point atomics and + # uses the same reduction order for context and generation shapes. + # Bypass tuning so an existing cache cannot reintroduce atomic split-K. + best_tactic = _FUSED_HC_FALLBACK_TACTIC_FMA + else: + tuner = AutoTuner.get() + _, best_tactic = tuner.choose_one( + "trtllm::mhc_fused_hc", + [runner], + MhcFusedHcRunner.tuning_config, + [ + x_prev, + residual_prev, + post_mix_prev, + comb_mix_prev, + w_t_cur, + hc_scale_cur, + hc_base_cur, + ], + norm_weight=norm_weight, + norm_eps=norm_eps, + ) return runner( inputs=[ diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py index 0ef432e33e50..4fec582bb5fb 100644 --- a/tests/unittest/_torch/modules/test_mhc.py +++ b/tests/unittest/_torch/modules/test_mhc.py @@ -1194,6 +1194,74 @@ def _make_fused_hc_runner_case(n: int, hidden_size: int, hc_mult: int, seed: int return runner, inputs +@pytest.mark.parametrize("num_tokens", [1, 129]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) +@pytest.mark.parametrize("fuse_norm", [False, True]) +def test_mhc_fused_hc_force_deterministic( + monkeypatch: pytest.MonkeyPatch, num_tokens: int, hidden_size: int, fuse_norm: bool +) -> None: + """Deterministic dispatch must survive cached atomic tactics and graph replay.""" + from tensorrt_llm._torch.modules.mhc import mhc_cuda + + monkeypatch.setenv("FORCE_DETERMINISTIC", "1") + runner, inputs = _make_fused_hc_runner_case(num_tokens, hidden_size, 4, seed=73) + inputs[0].mul_(hidden_size) + inputs[1].mul_(hidden_size) + inputs[4].mul_(500) + # A pre-existing autotune cache may contain a split-K tactic. The public + # deterministic path must not execute it, including at a cache hit. + if _mhc_fused_hc_mma_available(): + cached_tactic = ("fused_half_mma", 0, 4, 256, 1) + else: + cached_tactic = ("fused_all_fma", 2, 2, 0, 1) + monkeypatch.setattr( + mhc_cuda.AutoTuner.get(), "choose_one", lambda *args, **kwargs: (runner, cached_tactic) + ) + norm_weight = ( + torch.randn(hidden_size, dtype=torch.bfloat16, device="cuda") * 0.1 + 1 + if fuse_norm + else None + ) + + def run(values: list[torch.Tensor]) -> tuple[torch.Tensor, ...]: + return mhc_cuda.mhc_fused_hc( + *values, + n=4, + hidden_size=hidden_size, + rms_eps=runner.rms_eps, + hc_pre_eps=runner.hc_pre_eps, + hc_sinkhorn_eps=runner.hc_sinkhorn_eps, + hc_post_mult_value=runner.hc_post_mult_value, + sinkhorn_repeat=runner.sinkhorn_repeat, + norm_weight=norm_weight, + norm_eps=1e-6, + ) + + expected = tuple(t.clone() for t in run(inputs)) + for _ in range(10): + for actual, reference in zip(run(inputs), expected, strict=True): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + # Every row is independent: moving a row to a one-token decode shape + # must not change the arithmetic selected by deterministic dispatch. + row = num_tokens // 2 + single_inputs = [v[row : row + 1].clone() for v in inputs[:4]] + inputs[4:] + for actual, reference in zip(run(single_inputs), expected, strict=True): + torch.testing.assert_close(actual, reference[row : row + 1], rtol=0, atol=0) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run(inputs) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run(inputs) + graph.replay() + for actual, reference in zip(captured, expected, strict=True): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + @pytest.mark.parametrize( "backend,num_k_splits,tile_m,expected_shapes", [ From dd313887513c4c46d7c3cbfb3e70a9fc55f76701 Mon Sep 17 00:00:00 2001 From: Shiki Wu Date: Tue, 20 Oct 2026 04:03:41 -0700 Subject: [PATCH 2/2] [None][test] Assert deterministic mHC bypasses autotuning Signed-off-by: Shiki Wu --- tests/unittest/_torch/modules/test_mhc.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py index ef93645cc5c9..0a072bcfb4b1 100644 --- a/tests/unittest/_torch/modules/test_mhc.py +++ b/tests/unittest/_torch/modules/test_mhc.py @@ -1209,7 +1209,7 @@ def test_mhc_fused_hc_force_deterministic( fuse_norm: bool, lagged: bool, ) -> None: - """Deterministic dispatch must survive cached atomic tactics and graph replay.""" + """Deterministic dispatch must bypass autotuning and remain stable across graph replay.""" from tensorrt_llm._torch.modules.mhc import mhc_cuda monkeypatch.setenv("FORCE_DETERMINISTIC", "1") @@ -1217,15 +1217,11 @@ def test_mhc_fused_hc_force_deterministic( inputs[0].mul_(hidden_size) inputs[1].mul_(hidden_size) inputs[4].mul_(500) - # A pre-existing autotune cache may contain a split-K tactic. The public - # deterministic path must not execute it, including at a cache hit. - if _mhc_fused_hc_mma_available(): - cached_tactic = ("fused_half_mma", 0, 4, 256, 1) - else: - cached_tactic = ("fused_all_fma", 2, 2, 0, 1) - monkeypatch.setattr( - mhc_cuda.AutoTuner.get(), "choose_one", lambda *args, **kwargs: (runner, cached_tactic) - ) + + def unexpected_choose_one(*args: object, **kwargs: object) -> None: + raise AssertionError("Deterministic mHC dispatch must bypass autotuning") + + monkeypatch.setattr(mhc_cuda.AutoTuner.get(), "choose_one", unexpected_choose_one) norm_weight = ( torch.randn(hidden_size, dtype=torch.bfloat16, device="cuda") * 0.1 + 1 if fuse_norm