diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py index 5447d6b51afe..a7ff8409fcc6 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 @@ -1150,16 +1151,30 @@ 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, - pre_mix_ext=pre_mix_ext, - ) + 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, + pre_mix_ext=pre_mix_ext, + ) return runner( inputs=[ diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py index 285495103a1f..0a072bcfb4b1 100644 --- a/tests/unittest/_torch/modules/test_mhc.py +++ b/tests/unittest/_torch/modules/test_mhc.py @@ -1198,6 +1198,82 @@ 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]) +@pytest.mark.parametrize("lagged", [False, True]) +def test_mhc_fused_hc_force_deterministic( + monkeypatch: pytest.MonkeyPatch, + num_tokens: int, + hidden_size: int, + fuse_norm: bool, + lagged: bool, +) -> None: + """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") + 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) + + 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 + else None + ) + + pre_mix_ext = torch.rand(num_tokens, 4, device="cuda") if lagged else None + + def run( + values: list[torch.Tensor], external_pre_mix: torch.Tensor | None + ) -> 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, + pre_mix_ext=external_pre_mix, + ) + + expected = tuple(t.clone() for t in run(inputs, pre_mix_ext)) + assert len(expected) == (5 if lagged else 4) + for _ in range(10): + for actual, reference in zip(run(inputs, pre_mix_ext), 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:] + single_pre_mix = pre_mix_ext[row : row + 1].clone() if pre_mix_ext is not None else None + for actual, reference in zip(run(single_inputs, single_pre_mix), 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, pre_mix_ext) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run(inputs, pre_mix_ext) + 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", [