[None][fix] Honor deterministic mode in fused mHC dispatch

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=[

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",
     [