From 729834b2ac8d17422c0e02e3d7e2d58ceebfe674 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sat, 15 Aug 2026 09:30:32 +0000 Subject: [PATCH 1/5] [BI][DSv4] Make mHC tilelang ops batch invariant compute_num_split derives the K-split count from n_sms // cdiv(num_tokens, 64), so the reduction tree changes with the batch (same defect class as the RMSNorm fix in #48391); use_small_fma switches to a second implementation at num_tokens <= 16. Under VLLM_BATCH_INVARIANT=1 the split count is pinned (min(cap, n_sms//4) rounded down to a power of two = 32 on GB200; divides the 256 K-blocks, single wave up to M=256) and the small-fma fork is disabled. Cross-split merge is already a T.serial ordered loop, so pinning the count pins the reduction tree. CUDA-graph cost after pinning: parity with baseline for n<=128 and 193-256, faster at 129-192; only n=1 keeps +2us (fused kernel split into post+GEMM). Tests (10): bitwise stability across batch boundaries for mhc_pre / mhc_fused_post_pre / broadcast / fused-RMSNorm variants, mhc_post and hc_head regressions, negative controls with checkpoint-realistic magnitudes (synthetic small weights wash out real defects), correctness vs the torch reference. Co-Authored-By: Claude Fable 5 Consolidation pass folded in (codex r10/r10b reviewed): shared test helpers, pure Triton key fn, tl.constexpr-instantiated constants (plain global ints fail to compile under Triton 3.7), repo-pinned ruff format. Container suite 39/39 green; topk equivalence probe 240/240 bitwise. --- .../determinism/test_mhc_batch_invariance.py | 373 ++++++++++++++++++ tests/v1/determinism/utils.py | 18 + vllm/model_executor/kernels/mhc/tilelang.py | 10 +- .../kernels/mhc/tilelang_kernels.py | 15 + 4 files changed, 414 insertions(+), 2 deletions(-) create mode 100644 tests/v1/determinism/test_mhc_batch_invariance.py diff --git a/tests/v1/determinism/test_mhc_batch_invariance.py b/tests/v1/determinism/test_mhc_batch_invariance.py new file mode 100644 index 000000000000..f22e6a6f0a71 --- /dev/null +++ b/tests/v1/determinism/test_mhc_batch_invariance.py @@ -0,0 +1,373 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Batch invariance of the mHC tilelang ops. + +A row's outputs must be bitwise identical no matter how many other rows share +the batch. The variant mechanisms are (a) ``compute_num_split`` deriving the +K-split count from the token-tile count and (b) the small-token FMA kernel in +``mhc_fused_post_pre`` switching implementations at ``num_tokens == 16``; both +are disabled under ``VLLM_BATCH_INVARIANT``. +""" + +import pytest +import torch + +import vllm.envs as envs +import vllm.model_executor.kernels.mhc # noqa: F401 +from tests.kernels.test_mhc_kernels import mhc_pre_ref +from tests.v1.determinism.utils import batch_with_victim, skip_if_not_cuda +from vllm.model_executor.kernels.mhc.tilelang_kernels import compute_num_split +from vllm.utils.torch_utils import set_random_seed + +HC_MULT = 4 +HIDDEN_SIZE = 4096 +HC_MULT3 = 2 * HC_MULT + HC_MULT * HC_MULT +RMS_EPS = HC_PRE_EPS = HC_SINKHORN_EPS = 1e-6 +SINKHORN_REPEAT = 20 +HC_POST_ALPHA = 1.0 + +# Dispatch flips only at specific token counts: the small-FMA branch at 8 and +# 16, and compute_num_split whenever n_sms // cdiv(num_tokens, 64) drops. +# Uniform sweeps miss these; enumerate the flip points instead. +BOUNDARIES = [1, 7, 8, 15, 16, 17, 63, 64, 65, 128, 129, 192, 193, 256] + + +@pytest.fixture(autouse=True) +def _fresh_split_cache(): + """compute_num_split caches across the BI toggle; every test (the BI=True + conftest default and the BI=False negative controls alike) must start + with a cleared cache or the toggle silently does nothing.""" + compute_num_split.cache_clear() + + +def _mhc_weights(device): + """Weights at checkpoint-realistic magnitudes. + + The correctness suite's ``fn * 1e-4`` keeps mixes so small that sigmoid + compresses a ~10-ULP GEMM reassociation diff below one fp32 ULP of the + output — the negative control then can't fail. The real checkpoint has + ``hc_attn_scale ~= [2.08, 0.019, 0.245]`` and O(1) mixes, where the same + diff survives into the fp32 outputs. + """ + set_random_seed(0) + fn = ( + torch.randn( + (HC_MULT3, HC_MULT, HIDDEN_SIZE), dtype=torch.float32, device=device + ) + * (HC_MULT * HIDDEN_SIZE) ** -0.5 + ).flatten(1, 2) + hc_scale = torch.tensor([2.0, 0.02, 0.25], dtype=torch.float32, device=device) + hc_base = torch.randn((HC_MULT3,), dtype=torch.float32, device=device) * 0.1 + return fn, hc_scale, hc_base + + +def _run_pre(residual, fn, hc_scale, hc_base): + return torch.ops.vllm.mhc_pre_tilelang( + residual, + fn, + hc_scale, + hc_base, + RMS_EPS, + HC_PRE_EPS, + HC_SINKHORN_EPS, + HC_POST_ALPHA, + SINKHORN_REPEAT, + ) + + +def _pre_row0(n, victim, filler_seed, fn, hc_scale, hc_base): + """Row 0 of every mhc_pre output for a batch of n rows led by victim.""" + (batch,) = batch_with_victim((victim,), n, filler_seed) + outs = _run_pre(batch, fn, hc_scale, hc_base) + return [o[0].clone() for o in outs] + + +@skip_if_not_cuda +def test_mhc_pre_batch_invariance(): + device = "cuda" + fn, hc_scale, hc_base = _mhc_weights(device) + victim = torch.randn((1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device) + + base = _pre_row0(BOUNDARIES[0], victim, 1, fn, hc_scale, hc_base) + for i, n in enumerate(BOUNDARIES[1:], start=2): + outs = _pre_row0(n, victim, i, fn, hc_scale, hc_base) + for name, a, b in zip(("post_mix", "comb_mix", "layer_input"), base, outs): + assert torch.equal(a, b), ( + f"mhc_pre {name} row 0 changed at batch size {n}: " + f"max diff {(a.float() - b.float()).abs().max().item():.3e}" + ) + + +@skip_if_not_cuda +def test_mhc_pre_negative_control(monkeypatch): + """With BI off, the same sweep must show a bitwise difference. + + This proves the harness can detect the defect; if the default path were + already invariant, the BI test above would be vacuous. + """ + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False) + device = "cuda" + fn, hc_scale, hc_base = _mhc_weights(device) + victim = torch.randn((1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device) + + base = _pre_row0(BOUNDARIES[0], victim, 1, fn, hc_scale, hc_base) + diffs = [] + for i, n in enumerate(BOUNDARIES[1:], start=2): + outs = _pre_row0(n, victim, i, fn, hc_scale, hc_base) + if any(not torch.equal(a, b) for a, b in zip(base, outs)): + diffs.append(n) + assert diffs, ( + "default mhc_pre was bitwise invariant across all boundaries; " + "the BI test cannot distinguish fixed from broken" + ) + + +@skip_if_not_cuda +def test_mhc_pre_correctness(): + """The pinned-split path must still match the reference implementation.""" + device = "cuda" + fn, hc_scale, hc_base = _mhc_weights(device) + for n in (1, 16, 129): + set_random_seed(n) + residual = torch.randn( + (n, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device + ) + ref = mhc_pre_ref( + residual, + fn, + hc_scale, + hc_base, + RMS_EPS, + HC_PRE_EPS, + HC_SINKHORN_EPS, + HC_POST_ALPHA, + SINKHORN_REPEAT, + ) + out = _run_pre(residual, fn, hc_scale, hc_base) + # Same tolerance as tests/kernels/test_mhc_kernels.py: the tf32 GEMM + # plus fused sinkhorn diverge from the fp32 reference well below this. + for actual, expected in zip(out, ref, strict=True): + torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2) + + +def _fused_row0(n, victims, filler_seed, fn, hc_scale, hc_base): + args = batch_with_victim(victims, n, filler_seed) + outs = torch.ops.vllm.mhc_fused_post_pre_tilelang( + *args, + fn, + hc_scale, + hc_base, + RMS_EPS, + HC_PRE_EPS, + HC_SINKHORN_EPS, + HC_POST_ALPHA, + SINKHORN_REPEAT, + ) + return [o[0].clone() for o in outs] + + +@skip_if_not_cuda +def test_mhc_fused_post_pre_batch_invariance(): + device = "cuda" + fn, hc_scale, hc_base = _mhc_weights(device) + victims = ( + torch.randn((1, HIDDEN_SIZE), dtype=torch.bfloat16, device=device), + torch.randn((1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device), + torch.randn((1, HC_MULT, 1), dtype=torch.float32, device=device), + torch.randn((1, HC_MULT, HC_MULT), dtype=torch.float32, device=device), + ) + names = ("residual", "post_mix", "comb_mix", "layer_input") + + base = _fused_row0(BOUNDARIES[0], victims, 1, fn, hc_scale, hc_base) + for i, n in enumerate(BOUNDARIES[1:], start=2): + outs = _fused_row0(n, victims, i, fn, hc_scale, hc_base) + for name, a, b in zip(names, base, outs): + assert torch.equal(a, b), ( + f"mhc_fused_post_pre {name} row 0 changed at batch size {n}: " + f"max diff {(a.float() - b.float()).abs().max().item():.3e}" + ) + + +@skip_if_not_cuda +def test_mhc_fused_post_pre_negative_control(monkeypatch): + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False) + device = "cuda" + fn, hc_scale, hc_base = _mhc_weights(device) + victims = ( + torch.randn((1, HIDDEN_SIZE), dtype=torch.bfloat16, device=device), + torch.randn((1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device), + torch.randn((1, HC_MULT, 1), dtype=torch.float32, device=device), + torch.randn((1, HC_MULT, HC_MULT), dtype=torch.float32, device=device), + ) + base = _fused_row0(BOUNDARIES[0], victims, 1, fn, hc_scale, hc_base) + diffs = [] + for i, n in enumerate(BOUNDARIES[1:], start=2): + outs = _fused_row0(n, victims, i, fn, hc_scale, hc_base) + if any(not torch.equal(a, b) for a, b in zip(base, outs)): + diffs.append(n) + assert diffs, ( + "default mhc_fused_post_pre was bitwise invariant across all " + "boundaries; the BI test cannot distinguish fixed from broken" + ) + + +@skip_if_not_cuda +def test_mhc_pre_with_norm_batch_invariance(): + """Production always fuses RMSNorm (norm_weight path); cover that + big_fuse variant too.""" + device = "cuda" + fn, hc_scale, hc_base = _mhc_weights(device) + set_random_seed(3) + norm_w = torch.randn((HIDDEN_SIZE,), dtype=torch.bfloat16, device=device) + victim = torch.randn((1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device) + + def row0(n, seed): + (batch,) = batch_with_victim((victim,), n, seed) + outs = torch.ops.vllm.mhc_pre_tilelang( + batch, + fn, + hc_scale, + hc_base, + RMS_EPS, + HC_PRE_EPS, + HC_SINKHORN_EPS, + HC_POST_ALPHA, + SINKHORN_REPEAT, + norm_weight=norm_w, + ) + return [o[0].clone() for o in outs] + + base = row0(BOUNDARIES[0], 1) + for i, n in enumerate(BOUNDARIES[1:], start=2): + outs = row0(n, i) + for name, a, b in zip(("post_mix", "comb_mix", "layer_input"), base, outs): + assert torch.equal(a, b), ( + f"mhc_pre(norm) {name} row 0 changed at batch size {n}" + ) + + +# The broadcast entry computes its splits from k=hidden_size (cap 16), so the +# count first drops at cdiv(n, 64) = 10, i.e. n = 577. +BROADCAST_BOUNDARIES = [1, 7, 8, 15, 16, 17, 64, 65, 512, 576, 577, 640] + + +def _broadcast_row0(n, victim, seed, fn, fn_b, hc_scale, hc_base, norm_w): + from vllm.model_executor.kernels.mhc.tilelang import mhc_pre_broadcast_tilelang + + (batch,) = batch_with_victim((victim,), n, seed) + outs = mhc_pre_broadcast_tilelang( + batch, + fn, + hc_scale, + hc_base, + RMS_EPS, + HC_PRE_EPS, + HC_SINKHORN_EPS, + HC_POST_ALPHA, + SINKHORN_REPEAT, + norm_weight=norm_w, + fn_broadcast=fn_b, + ) + return [o[0].clone() for o in outs] + + +def _broadcast_setup(device): + fn, hc_scale, hc_base = _mhc_weights(device) + set_random_seed(4) + fn_b = ( + torch.randn((HC_MULT3, HIDDEN_SIZE), dtype=torch.float32, device=device) + * HIDDEN_SIZE**-0.5 + ) + norm_w = torch.randn((HIDDEN_SIZE,), dtype=torch.bfloat16, device=device) + victim = torch.randn((1, HIDDEN_SIZE), dtype=torch.bfloat16, device=device) + return fn, fn_b, hc_scale, hc_base, norm_w, victim + + +@skip_if_not_cuda +def test_mhc_pre_broadcast_batch_invariance(): + """Layer 0's production entry (residual broadcast from (T, H)).""" + fn, fn_b, hc_scale, hc_base, norm_w, victim = _broadcast_setup("cuda") + names = ("residual", "post_mix", "comb_mix", "layer_input") + base = _broadcast_row0(1, victim, 1, fn, fn_b, hc_scale, hc_base, norm_w) + for i, n in enumerate(BROADCAST_BOUNDARIES[1:], start=2): + outs = _broadcast_row0(n, victim, i, fn, fn_b, hc_scale, hc_base, norm_w) + for name, a, b in zip(names, base, outs): + assert torch.equal(a, b), ( + f"mhc_pre_broadcast {name} row 0 changed at batch size {n}" + ) + + +@skip_if_not_cuda +def test_mhc_pre_broadcast_negative_control(monkeypatch): + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False) + fn, fn_b, hc_scale, hc_base, norm_w, victim = _broadcast_setup("cuda") + base = _broadcast_row0(1, victim, 1, fn, fn_b, hc_scale, hc_base, norm_w) + diffs = [] + for i, n in enumerate(BROADCAST_BOUNDARIES[1:], start=2): + outs = _broadcast_row0(n, victim, i, fn, fn_b, hc_scale, hc_base, norm_w) + if any(not torch.equal(a, b) for a, b in zip(base, outs)): + diffs.append(n) + assert diffs, ( + "default mhc_pre_broadcast was bitwise invariant across all " + "boundaries; the BI test cannot distinguish fixed from broken" + ) + + +@skip_if_not_cuda +def test_hc_head_batch_invariance(): + """hc_head has no K-split and each token maps to its own block, so it + should be invariant even without the flag; assert that holds under BI.""" + device = "cuda" + set_random_seed(0) + # Checkpoint-realistic magnitudes (see _mhc_weights): the old ``* 1e-4`` + # weights compress reassociation diffs below one output ULP, making the + # assertion vacuously easy. + fn = ( + torch.randn( + (HC_MULT, HC_MULT * HIDDEN_SIZE), dtype=torch.float32, device=device + ) + * (HC_MULT * HIDDEN_SIZE) ** -0.5 + ) + hc_scale = torch.tensor([2.0], dtype=torch.float32, device=device) + hc_base = torch.randn((HC_MULT,), dtype=torch.float32, device=device) * 0.1 + victim = torch.randn((1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device) + + def row0(n, seed): + (batch,) = batch_with_victim((victim,), n, seed) + out = torch.ops.vllm.hc_head_fused_kernel_tilelang( + batch, fn, hc_scale, hc_base, RMS_EPS, HC_PRE_EPS + ) + return out[0].clone() + + base = row0(BOUNDARIES[0], 1) + for i, n in enumerate(BOUNDARIES[1:], start=2): + out = row0(n, i) + assert torch.equal(base, out), f"hc_head row 0 changed at batch size {n}" + + +@skip_if_not_cuda +def test_mhc_post_batch_invariance(): + """mhc_post (the end-of-loop post-mix catch-up) never goes through + compute_num_split — per-token CTA, reductions only over hc_mult with a + vectorized fixed order — so it must be invariant even without the flag; + pin that fact under BI so a future K-split refactor trips this test.""" + device = "cuda" + set_random_seed(0) + victim_x = torch.randn((1, HIDDEN_SIZE), dtype=torch.bfloat16, device=device) + victim_res = torch.randn( + (1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device + ) + victim_post = torch.randn((1, HC_MULT, 1), dtype=torch.float32, device=device) + victim_comb = torch.randn((1, HC_MULT, HC_MULT), dtype=torch.float32, device=device) + + def row0(n, seed): + x, res, post, comb = batch_with_victim( + (victim_x, victim_res, victim_post, victim_comb), n, seed + ) + out = torch.ops.vllm.mhc_post_tilelang(x, res, post, comb) + return out[0].clone() + + base = row0(BOUNDARIES[0], 1) + for i, n in enumerate(BOUNDARIES[1:], start=2): + out = row0(n, i) + assert torch.equal(base, out), f"mhc_post row 0 changed at batch size {n}" diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index f03ea05b4331..9e1ca88bcead 100644 --- a/tests/v1/determinism/utils.py +++ b/tests/v1/determinism/utils.py @@ -13,6 +13,7 @@ ModelArchConfigConvertorBase, ) from vllm.triton_utils import HAS_TRITON +from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla @@ -69,6 +70,23 @@ class DeviceConfig(NamedTuple): ) +def batch_with_victim( + victims: tuple[torch.Tensor, ...], n: int, seed: int +) -> tuple[torch.Tensor, ...]: + """Batch of ``n`` rows led by each victim row, padded with seeded random + filler rows behind it: batch invariance means the victim's output bits + must not depend on the filler.""" + set_random_seed(seed) + return tuple( + torch.cat( + [v, torch.randn((n - 1, *v.shape[1:]), dtype=v.dtype, device=v.device)] + ) + if n > 1 + else v + for v in victims + ) + + def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str: # Generate more realistic prompts that will actually produce varied tokens # Use a mix of common English text patterns diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index c3c773d07f2f..223f2b20ae35 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch +from vllm import envs from vllm.utils.torch_utils import direct_register_custom_op @@ -39,7 +40,9 @@ def _tilelang_hc_prenorm_gemm( assert x.shape[1] == hc_mult * hidden_size assert x.shape[1] % n_splits == 0 assert (x.shape[1] // n_splits) % n_thr == 0 - use_default_config = tile_n == 12 and n_thr == 512 + # Both fast paths below switch kernels on x.shape[0], which changes a + # row's result when the batch crosses the threshold. + use_default_config = tile_n == 12 and n_thr == 512 and not envs.VLLM_BATCH_INVARIANT if n_splits == 1 and use_default_config and x.shape[0] >= 1024: hc_prenorm_gemm_block_m_tilelang( x, @@ -515,7 +518,10 @@ def mhc_fused_post_pre_tilelang( from vllm.utils.deep_gemm import is_deep_gemm_supported use_deep_gemm = is_deep_gemm_supported() - use_small_fma = num_tokens <= 16 + # The small-token FMA kernel is a second implementation of the same math + # with its own n_splits schedule, so crossing num_tokens == 16 changes a + # row's result. Batch invariance requires a single implementation. + use_small_fma = num_tokens <= 16 and not envs.VLLM_BATCH_INVARIANT if use_small_fma: # TODO(gnovack): investigate autotuning these heuristics tile_n = 2 if num_tokens < 8 else 3 diff --git a/vllm/model_executor/kernels/mhc/tilelang_kernels.py b/vllm/model_executor/kernels/mhc/tilelang_kernels.py index 925f4631a515..67889a421515 100644 --- a/vllm/model_executor/kernels/mhc/tilelang_kernels.py +++ b/vllm/model_executor/kernels/mhc/tilelang_kernels.py @@ -8,6 +8,7 @@ import torch +from vllm import envs from vllm.platforms import current_platform from vllm.tilelang_utils import T, tilelang, tilelang_jit from vllm.utils.math_utils import cdiv @@ -19,6 +20,20 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: device_props = torch.cuda.get_device_properties(0) n_sms = device_props.multi_processor_count + if envs.VLLM_BATCH_INVARIANT: + # grid_size tracks the number of token tiles, so deriving the K-split + # count from it changes the cross-split reduction tree whenever the + # batch grows. Pin to a value that depends only on k and the GPU: + # the largest power of two <= min(k cap, n_sms // 4). Powers of two + # divide the K-blocks evenly, and the n_sms // 4 budget keeps + # grid <= n_sms for token tiles up to 4, which covers decode; on + # GB200 (148 SMs) this lands on 32, measured flat-optimal for + # M <= 256 and within 30% of the best fixed value at prefill sizes. + split_k = n_sms // 4 + if k is not None: + split_k = min(split_k, cdiv(k, block_k) // 4) + split_k = max(split_k, 1) + return 1 << (split_k.bit_length() - 1) split_k = n_sms // grid_size if k is not None: # avoid split_k for small k From a980149582809d4d3ee1646e217de7098f329cf3 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 16 Aug 2026 11:25:31 +0000 Subject: [PATCH 2/5] [BI][DSv4] Run the mHC regressions in CI and scope what they prove The new tests were not listed in the B200 Batch Invariance job, which names its files individually, so they were never executed. Add the file. Also narrow what the change claims. Disabling the two token-count fast paths is a correctness change and applies wherever TileLang runs, ROCm included -- restricting it to CUDA would leave ROCm batch-variant under the flag, which is worse. What was measured is the numerics and the cost of doing so, on SM100 only. Say that at both sites rather than letting a GB200 measurement read as a claim about every platform the code touches. Co-Authored-By: Claude Opus 5 --- .buildkite/test_areas/misc.yaml | 1 + vllm/model_executor/kernels/mhc/tilelang.py | 8 ++++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 8f168cbd367f..5374e3d542a2 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -424,6 +424,7 @@ steps: - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k FLASH_ATTN - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py - pytest -v -s v1/determinism/test_matmul_batch_invariant.py + - pytest -v -s v1/determinism/test_mhc_batch_invariance.py - pytest -v -s v1/determinism/test_cutlass_batch_invariance.py - pytest -v -s v1/determinism/test_online_batch_invariance.py diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 223f2b20ae35..a824528ed88d 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -41,7 +41,10 @@ def _tilelang_hc_prenorm_gemm( assert x.shape[1] % n_splits == 0 assert (x.shape[1] // n_splits) % n_thr == 0 # Both fast paths below switch kernels on x.shape[0], which changes a - # row's result when the batch crosses the threshold. + # row's result when the batch crosses the threshold. Disabling them is a + # correctness change and applies on every platform TileLang runs on; only + # the numerics and the cost of doing so were measured, and only on SM100. + # No performance claim is made for other hardware. use_default_config = tile_n == 12 and n_thr == 512 and not envs.VLLM_BATCH_INVARIANT if n_splits == 1 and use_default_config and x.shape[0] >= 1024: hc_prenorm_gemm_block_m_tilelang( @@ -520,7 +523,8 @@ def mhc_fused_post_pre_tilelang( use_deep_gemm = is_deep_gemm_supported() # The small-token FMA kernel is a second implementation of the same math # with its own n_splits schedule, so crossing num_tokens == 16 changes a - # row's result. Batch invariance requires a single implementation. + # row's result. Batch invariance requires a single implementation. Same + # scope note as above: correctness everywhere, measured only on SM100. use_small_fma = num_tokens <= 16 and not envs.VLLM_BATCH_INVARIANT if use_small_fma: # TODO(gnovack): investigate autotuning these heuristics From 598df3ed6eb39da453132b474dd99a4abac0fb4b Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sun, 16 Aug 2026 11:46:08 +0000 Subject: [PATCH 3/5] [BI][DSv4] Cover the mHC paths production actually calls, and gate them in CI The suite exercised mhc_pre with a fused norm but mhc_fused_post_pre without one, and the fused-norm entry is the shape DSv4 runs between attention and FFN. Its epilogue is not shared with the norm-free one -- it takes a second pass reducing over hidden_size to form the RMS -- so a split-count change reaches the output through a path nothing covered. Add the invariance sweep, a negative control, and a reference composed from the kernel suite's own post and pre references plus an fp32 RMSNorm. The non-DeepGEMM prenorm GEMM has two dispatch flips of its own, keyed on the row count rather than the split count. B200 always takes DeepGEMM, so nothing in the file reached them; call the helper directly instead. Only the < 128 branch turns out to move a row's bits: the >= 1024 block-M variant keeps the same n_thr and tile_n, reassociates K in the same order, and comes out bitwise equal at DSv4's shape. It stays pinned because that equality is a property of the current tile config rather than a guarantee, and the negative control says so rather than asserting a difference that does not exist. Also add vllm/model_executor/kernels/mhc/ to the B200 job's source_file_dependencies: this PR triggers the job only because it adds a test file, and a later change confined to the kernels would not have run any of it. Co-Authored-By: Claude Opus 5 --- .buildkite/test_areas/misc.yaml | 1 + .../determinism/test_mhc_batch_invariance.py | 210 +++++++++++++++++- 2 files changed, 208 insertions(+), 3 deletions(-) diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 5374e3d542a2..713a2dab77f0 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -414,6 +414,7 @@ steps: source_file_dependencies: - vllm/v1/attention - vllm/model_executor/layers + - vllm/model_executor/kernels/mhc/ - tests/v1/determinism/ commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn diff --git a/tests/v1/determinism/test_mhc_batch_invariance.py b/tests/v1/determinism/test_mhc_batch_invariance.py index f22e6a6f0a71..185e2031f6da 100644 --- a/tests/v1/determinism/test_mhc_batch_invariance.py +++ b/tests/v1/determinism/test_mhc_batch_invariance.py @@ -14,7 +14,7 @@ import vllm.envs as envs import vllm.model_executor.kernels.mhc # noqa: F401 -from tests.kernels.test_mhc_kernels import mhc_pre_ref +from tests.kernels.test_mhc_kernels import mhc_post_ref, mhc_pre_ref from tests.v1.determinism.utils import batch_with_victim, skip_if_not_cuda from vllm.model_executor.kernels.mhc.tilelang_kernels import compute_num_split from vllm.utils.torch_utils import set_random_seed @@ -22,14 +22,20 @@ HC_MULT = 4 HIDDEN_SIZE = 4096 HC_MULT3 = 2 * HC_MULT + HC_MULT * HC_MULT -RMS_EPS = HC_PRE_EPS = HC_SINKHORN_EPS = 1e-6 +RMS_EPS = HC_PRE_EPS = HC_SINKHORN_EPS = NORM_EPS = 1e-6 SINKHORN_REPEAT = 20 HC_POST_ALPHA = 1.0 # Dispatch flips only at specific token counts: the small-FMA branch at 8 and # 16, and compute_num_split whenever n_sms // cdiv(num_tokens, 64) drops. # Uniform sweeps miss these; enumerate the flip points instead. -BOUNDARIES = [1, 7, 8, 15, 16, 17, 63, 64, 65, 128, 129, 192, 193, 256] +BOUNDARIES = [1, 7, 8, 15, 16, 17, 63, 64, 65, 127, 128, 129, 192, 193, 256] + +# The non-DeepGEMM prenorm GEMM has two more dispatch flips of its own, on +# x.shape[0] rather than on the split count: a block-M kernel at >= 1024 and a +# wider-tile config below 128. B200 takes the DeepGEMM path, so nothing above +# reaches them; the helper is called directly instead. +PRENORM_BOUNDARIES = [1, 64, 127, 128, 129, 1023, 1024, 1025] @pytest.fixture(autouse=True) @@ -371,3 +377,201 @@ def row0(n, seed): for i, n in enumerate(BOUNDARIES[1:], start=2): out = row0(n, i) assert torch.equal(base, out), f"mhc_post row 0 changed at batch size {n}" + + +def _fused_norm_outs(n, victims, filler_seed, fn, hc_scale, hc_base, norm_w): + args = batch_with_victim(victims, n, filler_seed) + return torch.ops.vllm.mhc_fused_post_pre_tilelang( + *args, + fn, + hc_scale, + hc_base, + RMS_EPS, + HC_PRE_EPS, + HC_SINKHORN_EPS, + HC_POST_ALPHA, + SINKHORN_REPEAT, + norm_weight=norm_w, + norm_eps=NORM_EPS, + ) + + +def _fused_norm_setup(device): + fn, hc_scale, hc_base = _mhc_weights(device) + set_random_seed(5) + norm_w = torch.randn((HIDDEN_SIZE,), dtype=torch.bfloat16, device=device) + victims = ( + torch.randn((1, HIDDEN_SIZE), dtype=torch.bfloat16, device=device), + torch.randn((1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device), + torch.randn((1, HC_MULT, 1), dtype=torch.float32, device=device), + torch.randn((1, HC_MULT, HC_MULT), dtype=torch.float32, device=device), + ) + return fn, hc_scale, hc_base, norm_w, victims + + +@skip_if_not_cuda +def test_mhc_fused_post_pre_with_norm_batch_invariance(): + """The shape production actually calls between attention and FFN. + + The unfused entry above shares the dispatch but not the epilogue: with + norm_weight the kernel takes a second pass that reduces over hidden_size + to form the RMS, so a split-count change reaches the output through a + path the norm-free test does not cover. + """ + fn, hc_scale, hc_base, norm_w, victims = _fused_norm_setup("cuda") + names = ("residual", "post_mix", "comb_mix", "layer_input") + + first = _fused_norm_outs(BOUNDARIES[0], victims, 1, fn, hc_scale, hc_base, norm_w) + base = [o[0].clone() for o in first] + for i, n in enumerate(BOUNDARIES[1:], start=2): + outs = _fused_norm_outs(n, victims, i, fn, hc_scale, hc_base, norm_w) + for name, a, b in zip(names, base, outs): + b = b[0] + assert torch.equal(a, b), ( + f"mhc_fused_post_pre(norm) {name} row 0 changed at batch size {n}: " + f"max diff {(a.float() - b.float()).abs().max().item():.3e}" + ) + + +@skip_if_not_cuda +def test_mhc_fused_post_pre_with_norm_negative_control(monkeypatch): + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False) + fn, hc_scale, hc_base, norm_w, victims = _fused_norm_setup("cuda") + first = _fused_norm_outs(BOUNDARIES[0], victims, 1, fn, hc_scale, hc_base, norm_w) + base = [o[0].clone() for o in first] + diffs = [] + for i, n in enumerate(BOUNDARIES[1:], start=2): + outs = _fused_norm_outs(n, victims, i, fn, hc_scale, hc_base, norm_w) + if any(not torch.equal(a, b[0]) for a, b in zip(base, outs)): + diffs.append(n) + assert diffs, ( + "default mhc_fused_post_pre(norm) was bitwise invariant across all " + "boundaries; the BI test cannot distinguish fixed from broken" + ) + + +def _rmsnorm_ref(x, weight, eps): + xf = x.float() + return ( + xf * torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps) * weight.float() + ).bfloat16() + + +@skip_if_not_cuda +def test_mhc_fused_post_pre_with_norm_correctness(): + """Pinning the dispatch must not change what the fused-norm path computes. + + The reference is composed from the kernel suite's own post and pre + references plus an fp32 RMSNorm, so it shares no code with the fused + implementation. + """ + device = "cuda" + fn, hc_scale, hc_base, norm_w, _ = _fused_norm_setup(device) + for n in (1, 16, 129): + set_random_seed(n) + x = torch.randn((n, HIDDEN_SIZE), dtype=torch.bfloat16, device=device) + residual = torch.randn( + (n, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device + ) + post_layer_mix = torch.randn( + (n, HC_MULT, 1), dtype=torch.float32, device=device + ) + comb_res_mix = torch.randn( + (n, HC_MULT, HC_MULT), dtype=torch.float32, device=device + ) + + residual_ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix) + post_ref, comb_ref, layer_input_ref = mhc_pre_ref( + residual_ref, + fn, + hc_scale, + hc_base, + RMS_EPS, + HC_PRE_EPS, + HC_SINKHORN_EPS, + HC_POST_ALPHA, + SINKHORN_REPEAT, + ) + normed_ref = _rmsnorm_ref(layer_input_ref, norm_w, NORM_EPS) + + out = torch.ops.vllm.mhc_fused_post_pre_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + RMS_EPS, + HC_PRE_EPS, + HC_SINKHORN_EPS, + HC_POST_ALPHA, + SINKHORN_REPEAT, + norm_weight=norm_w, + norm_eps=NORM_EPS, + ) + for actual, expected in zip( + out, (residual_ref, post_ref, comb_ref, normed_ref), strict=True + ): + torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2) + + +def _prenorm_row0(n, victim, seed, fn): + """Row 0 of the non-DeepGEMM prenorm GEMM for a batch of n rows.""" + from vllm.model_executor.kernels.mhc.tilelang import _tilelang_hc_prenorm_gemm + + (batch,) = batch_with_victim((victim,), n, seed) + x = batch.view(n, HC_MULT * HIDDEN_SIZE) + out = torch.empty(1, n, HC_MULT3, dtype=torch.float32, device=x.device) + sqrsum = torch.empty(1, n, dtype=torch.float32, device=x.device) + _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, HIDDEN_SIZE, HC_MULT) + return out[0, 0].clone(), sqrsum[0, 0].clone() + + +@skip_if_not_cuda +def test_prenorm_gemm_no_deep_gemm_batch_invariance(): + """The fallback GEMM's own dispatch, which DeepGEMM hardware never reaches. + + ``mhc_pre_tilelang`` only calls this helper when DeepGEMM is unavailable, + so on B200 every test above goes through tf32_hc_prenorm_gemm instead and + the >= 1024 and < 128 branches stay dark. Call the helper directly. + """ + device = "cuda" + fn, _, _ = _mhc_weights(device) + victim = torch.randn((1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device) + + base = _prenorm_row0(PRENORM_BOUNDARIES[0], victim, 1, fn) + for i, n in enumerate(PRENORM_BOUNDARIES[1:], start=2): + outs = _prenorm_row0(n, victim, i, fn) + for name, a, b in zip(("mul", "sqrsum"), base, outs): + assert torch.equal(a, b), ( + f"hc_prenorm_gemm {name} row 0 changed at batch size {n}: " + f"max diff {(a.float() - b.float()).abs().max().item():.3e}" + ) + + +@skip_if_not_cuda +def test_prenorm_gemm_no_deep_gemm_negative_control(monkeypatch): + """Show which fallback branch actually moves a row's bits. + + Only the < 128 one does. The >= 1024 block-M variant keeps the same + n_thr and tile_n, so it reassociates K in the same order and comes out + bitwise equal to the generic kernel at DSv4's shape -- it is pinned under + the flag because that equality is a property of the current tile config, + not a guarantee, and asserting it here would pin the wrong thing. + """ + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False) + device = "cuda" + fn, _, _ = _mhc_weights(device) + victim = torch.randn((1, HC_MULT, HIDDEN_SIZE), dtype=torch.bfloat16, device=device) + + base = _prenorm_row0(128, victim, 1, fn) + diffs = set() + for i, n in enumerate(PRENORM_BOUNDARIES, start=2): + outs = _prenorm_row0(n, victim, i, fn) + if any(not torch.equal(a, b) for a, b in zip(base, outs)): + diffs.add(n) + assert diffs & {1, 64, 127}, ( + f"the < 128 wide-tile branch left row 0 bitwise unchanged (diffs at " + f"{sorted(diffs)}); the BI test above cannot prove it is pinned" + ) From ce5a8ea25218031a6823bba95ea3bc229f114858 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 19 Aug 2026 00:17:29 +0000 Subject: [PATCH 4/5] [BI][DSv4] Refuse a deep_gemm that cannot pin its configs tf32_hc_prenorm_gemm is DeepGEMM's and the mHC path has no other implementation of it, so a deep_gemm without set_batch_invariant leaves it free to select its config from the batch. The loader's answer to that build -- disable DeepGEMM MoE -- does nothing here, and mHC would keep running under VLLM_BATCH_INVARIANT with a row's result depending on its neighbours. Guard every entry that reaches the kernel, including the two that only consult is_deep_gemm_supported() to choose n_splits. Cached, since this sits on the per-layer path and the answer cannot change once deep_gemm is loaded. Depends on the DeepGEMM branch for deep_gemm_batch_invariant_enabled(). Co-Authored-By: Claude Opus 5 --- .../determinism/test_mhc_batch_invariance.py | 30 ++++++++++++++++ vllm/model_executor/kernels/mhc/tilelang.py | 36 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/tests/v1/determinism/test_mhc_batch_invariance.py b/tests/v1/determinism/test_mhc_batch_invariance.py index 185e2031f6da..57e93342be1a 100644 --- a/tests/v1/determinism/test_mhc_batch_invariance.py +++ b/tests/v1/determinism/test_mhc_batch_invariance.py @@ -575,3 +575,33 @@ def test_prenorm_gemm_no_deep_gemm_negative_control(monkeypatch): f"the < 128 wide-tile branch left row 0 bitwise unchanged (diffs at " f"{sorted(diffs)}); the BI test above cannot prove it is pinned" ) + + +def test_mhc_refuses_a_deep_gemm_that_cannot_be_pinned(monkeypatch): + """Fail closed, do not fall back. + + ``tf32_hc_prenorm_gemm`` is DeepGEMM's and the mHC path has no other + implementation, so a ``deep_gemm`` without ``set_batch_invariant`` would + keep running with its config chosen from the batch. Disabling DeepGEMM MoE, + the loader's answer to such a build, does nothing here. + """ + import vllm.model_executor.kernels.mhc.tilelang as tl + import vllm.utils.deep_gemm as dg + + guard = tl._require_batch_invariant_deep_gemm + monkeypatch.setattr(tl.envs, "VLLM_BATCH_INVARIANT", True) + + monkeypatch.setattr(dg, "deep_gemm_batch_invariant_enabled", lambda: True) + guard.cache_clear() + guard() # pinned: no complaint + + monkeypatch.setattr(dg, "deep_gemm_batch_invariant_enabled", lambda: False) + guard.cache_clear() + with pytest.raises(RuntimeError, match="set_batch_invariant"): + guard() + + # With the flag off the same unpinned build is fine -- nothing was promised. + monkeypatch.setattr(tl.envs, "VLLM_BATCH_INVARIANT", False) + guard.cache_clear() + guard() + guard.cache_clear() diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index a824528ed88d..7518c37266d3 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -1,11 +1,37 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import functools + import torch from vllm import envs from vllm.utils.torch_utils import direct_register_custom_op +@functools.cache +def _require_batch_invariant_deep_gemm() -> None: + """Refuse batch invariance when DeepGEMM cannot deliver it. + + ``tf32_hc_prenorm_gemm`` is DeepGEMM's and the mHC path has no other + implementation of it, so a deep_gemm without ``set_batch_invariant`` leaves + it free to pick its config from the batch. Disabling DeepGEMM MoE, which is + what the loader does with such a build, does nothing here -- the only + fail-closed answer is to refuse. Cached: this sits on the per-layer path and + the answer cannot change once deep_gemm is loaded. + """ + from vllm.utils.deep_gemm import deep_gemm_batch_invariant_enabled + + if not envs.VLLM_BATCH_INVARIANT or deep_gemm_batch_invariant_enabled(): + return + raise RuntimeError( + "VLLM_BATCH_INVARIANT is enabled but the loaded deep_gemm has no " + "set_batch_invariant, so the mHC prenorm GEMM would select its config " + "from the batch. Install a deep_gemm that exposes set_batch_invariant, " + "or disable batch invariance." + ) + + + def _torch_hc_prenorm_gemm( x: torch.Tensor, fn: torch.Tensor, @@ -140,6 +166,8 @@ def mhc_pre_tilelang( from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm from vllm.utils.math_utils import cdiv + _require_batch_invariant_deep_gemm() + assert residual.dtype == torch.bfloat16 assert fn.dtype == torch.float32 assert hc_scale.dtype == torch.float32 @@ -170,6 +198,8 @@ def mhc_pre_tilelang( from vllm.utils.deep_gemm import is_deep_gemm_supported + _require_batch_invariant_deep_gemm() + use_deep_gemm = is_deep_gemm_supported() if use_deep_gemm: # these numbers are from deepgemm kernel impl @@ -377,6 +407,8 @@ def mhc_pre_broadcast_tilelang( from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm + _require_batch_invariant_deep_gemm() + tf32_hc_prenorm_gemm( residual_flat, fn_broadcast, @@ -520,6 +552,8 @@ def mhc_fused_post_pre_tilelang( from vllm.utils.deep_gemm import is_deep_gemm_supported + _require_batch_invariant_deep_gemm() + use_deep_gemm = is_deep_gemm_supported() # The small-token FMA kernel is a second implementation of the same math # with its own n_splits schedule, so crossing num_tokens == 16 changes a @@ -605,6 +639,8 @@ def mhc_fused_post_pre_tilelang( if use_deep_gemm: from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm + _require_batch_invariant_deep_gemm() + tf32_hc_prenorm_gemm( residual_cur_2d, fn, From e084ed206870c9f92e448759c2be7adf47d65cd5 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 19 Aug 2026 00:42:45 +0000 Subject: [PATCH 5/5] [BI][DSv4] Drop the duplicate guard calls mhc_pre_tilelang and mhc_fused_post_pre_tilelang each carried the fail-closed guard twice -- once at the tf32 import and once at the is_deep_gemm_supported site inside the same function. One call per entry point is the whole point of caching it. Co-Authored-By: Claude Opus 5 --- vllm/model_executor/kernels/mhc/tilelang.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 7518c37266d3..7c6316562f73 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -198,8 +198,6 @@ def mhc_pre_tilelang( from vllm.utils.deep_gemm import is_deep_gemm_supported - _require_batch_invariant_deep_gemm() - use_deep_gemm = is_deep_gemm_supported() if use_deep_gemm: # these numbers are from deepgemm kernel impl @@ -639,8 +637,6 @@ def mhc_fused_post_pre_tilelang( if use_deep_gemm: from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm - _require_batch_invariant_deep_gemm() - tf32_hc_prenorm_gemm( residual_cur_2d, fn,