From 776367dc868ffd07df7f87631e25c71034040b7c Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 10 Nov 2026 02:25:03 -0800 Subject: [PATCH 1/2] [https://nvbugs/7028060][fix] Stop shimming FlashInfer's W4A16 tile validator once it takes enforce_search_floor FlashInfer 0.7.2 (flashinfer-ai/flashinfer#6292) re-validates the tiles its SM12x W4A16 fused MoE re-pins through force_tile_config with _candidate_tile_fits(..., enforce_search_floor=False). The shim that fused_moe_cute_dsl_b12x.py installs over that validator for FlashInfer 0.6.17 and later looks for the upstream fix as _ultra_wide_fc2_tile_fits, which 0.7.2 does not have, so it stays installed there. Its fixed keyword list then rejects enforce_search_floor and the re-pin fails with "TypeError: ... candidate_tile_fits() got an unexpected keyword argument 'enforce_search_floor'". Leave FlashInfer's validator in place when it takes enforce_search_floor. The shim keeps covering 0.6.17 to 0.7.1, whose validator does not. The existing shim tests now run against a validator that applies the search floor to every tile, as FlashInfer before 0.7.2 does. A new test compiles the W4A16 fused MoE with the ultra-wide FC2 tile forced, after the shim hook, up to the kernel build. Remove the waiver of accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4_w4a16. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../moe/fused_moe/fused_moe_cute_dsl_b12x.py | 14 +++-- tests/integration/test_lists/waives.txt | 1 - .../moe/test_cute_dsl_b12x_moe_backend.py | 63 ++++++++++++++++++- 3 files changed, 69 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index ecd71b4a1170..f2e266db80c4 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import inspect from dataclasses import replace import torch @@ -64,7 +65,7 @@ ActivationType.Swiglu: "silu", } -# FlashInfer's SM12x W4A16 fused MoE (0.6.17 and 0.6.18) selects a TC-decode +# FlashInfer's SM12x W4A16 fused MoE (0.6.17 to 0.7.1) selects a TC-decode # "ultra-wide" FC2 tile of tile_n=512 / tile_k=32 for small m whenever that # collapses FC2 to FC1's wave count (m=3 and m=4 for Qwen3.6-35B-A3B on the # 188-SM RTX PRO 6000). tile_k=32 sits below the tile_k>=64 floor of its @@ -72,10 +73,11 @@ # the ``force_tile_config`` re-pin across the custom-op boundary does not, so # the kernel it just selected is rejected with "force_tile_config fc2 tile # (tile_k=32, tile_n=512) does not fit ..." and CUDA-graph warmup aborts -# executor init (nvbug 6721561). Accepting the ultra tile in the re-validation -# under the same rule that produced it is the upstream fix; this shim mirrors -# it so the pinned wheel behaves like the fixed one. Drop it once the -# flashinfer pin carries ``_ultra_wide_fc2_tile_fits``. +# executor init (nvbug 6721561). FlashInfer 0.7.2 fixes the re-pin: it calls +# ``_candidate_tile_fits(..., enforce_search_floor=False)``. This shim accepts +# the ultra tile in the re-validation of the releases before that, under the +# same rule that produced it, and stands aside on a validator that takes +# ``enforce_search_floor``. _FLASHINFER_W4A16_ULTRA_WIDE_FC2_TILE = (512, 32, 256) # (tile_n, tile_k, cta_threads) _FLASHINFER_W4A16_ULTRA_WIDE_FC2_SINCE = Version("0.6.17") @@ -103,6 +105,8 @@ def _patch_flashinfer_w4a16_ultra_wide_fc2_tile_validation() -> bool: return False if getattr(upstream_fits, "_trtllm_ultra_wide_fc2_shim", False): return True + if "enforce_search_floor" in inspect.signature(upstream_fits).parameters: + return False def candidate_tile_fits( *, diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index fc9a5676b7d7..b06cd3e8e3fc 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -19,7 +19,6 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[throughput_laten accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[tep4_latency_moe_cutlass-torch_compile=True] SKIP (https://nvbugs/6561558) accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.5-fp8kv=True] SKIP (https://nvbugs/6668774) accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4_4gpus_static_eplb[moe_backend=CUTEDSL] SKIP (https://nvbugs/7034340) -accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4_w4a16 SKIP (https://nvbugs/7028060) accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency] SKIP (https://nvbugs/6412098) accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_nvfp4_tp8_mtp3_trtllm SKIP (https://nvbugs/7028063) accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3_8_Flash_Next_VL::test_nvfp4_1gpu_mtp3_trtllm_ple_offload SKIP (https://nvbugs/7028057) diff --git a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index b5d9fecea6c2..94288fa8b08b 100644 --- a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -15,6 +15,7 @@ """B12x MoE backend gating and dispatch tests.""" import dataclasses +import inspect import sys import types from typing import Optional @@ -671,6 +672,22 @@ def _reset_to_stock_validator(kernel, monkeypatch) -> None: ) +def _reset_to_pre_fix_validator(kernel, monkeypatch) -> None: + """Start from a validator that applies the search floor to every tile, as + flashinfer 0.6.17 to 0.7.1 do. Where the installed validator exempts forced + tiles through ``enforce_search_floor``, that validator with the floor + always on, and without the parameter, stands in for them.""" + _reset_to_stock_validator(kernel, monkeypatch) + stock = kernel._candidate_tile_fits + if "enforce_search_floor" not in inspect.signature(stock).parameters: + return + + def pre_fix_fits(**kwargs): + return stock(**kwargs, enforce_search_floor=True) + + monkeypatch.setattr(kernel, "_candidate_tile_fits", pre_fix_fits) + + def _fits(kernel, **overrides) -> bool: kwargs = dict(cta_m_blocks=1, max_shared_mem=_RTX_PRO_6000_SMEM, **_QWEN36_FC2, **_ULTRA_FC2) kwargs.update(overrides) @@ -679,7 +696,7 @@ def _fits(kernel, **overrides) -> bool: def test_w4a16_ultra_wide_fc2_shim_accepts_repinned_ultra_tile(monkeypatch): kernel = _flashinfer_w4a16_kernel() - _reset_to_stock_validator(kernel, monkeypatch) + _reset_to_pre_fix_validator(kernel, monkeypatch) assert not _fits(kernel), ( "flashinfer's stock validator accepts the ultra-wide FC2 tile; the shim in " "fused_moe_cute_dsl_b12x.py is no longer needed" @@ -703,7 +720,7 @@ def test_w4a16_ultra_wide_fc2_shim_accepts_repinned_ultra_tile(monkeypatch): def test_w4a16_ultra_wide_fc2_shim_is_idempotent(monkeypatch): kernel = _flashinfer_w4a16_kernel() - _reset_to_stock_validator(kernel, monkeypatch) + _reset_to_pre_fix_validator(kernel, monkeypatch) stock = kernel._candidate_tile_fits assert _patch_flashinfer_w4a16_ultra_wide_fc2_tile_validation() is True @@ -716,9 +733,49 @@ def test_w4a16_ultra_wide_fc2_shim_is_idempotent(monkeypatch): def test_w4a16_ultra_wide_fc2_shim_yields_to_upstream_fix(monkeypatch): kernel = _flashinfer_w4a16_kernel() - _reset_to_stock_validator(kernel, monkeypatch) + _reset_to_pre_fix_validator(kernel, monkeypatch) stock = kernel._candidate_tile_fits monkeypatch.setattr(kernel, "_ultra_wide_fc2_tile_fits", lambda **_: True, raising=False) assert _patch_flashinfer_w4a16_ultra_wide_fc2_tile_validation() is False assert kernel._candidate_tile_fits is stock + + +def test_w4a16_forced_ultra_wide_fc2_tile_reaches_kernel_build(monkeypatch): + """The custom-op launch recompiles with the tiles the selector chose, + forced through ``force_tile_config``. After the hook every W4A16 B12x leaf + runs at construction, that re-pin of the ultra-wide FC2 tile gets past tile + validation to the kernel build, on the installed flashinfer.""" + kernel = _flashinfer_w4a16_kernel() + _reset_to_stock_validator(kernel, monkeypatch) + _patch_flashinfer_w4a16_ultra_wide_fc2_tile_validation() + + class KernelBuildReached(Exception): + pass + + built = {} + + def build_kernel(**kwargs): + built.update(kwargs) + raise KernelBuildReached + + # The build compiles CuTe DSL code for a GPU; only the tiles it gets matter. + monkeypatch.setattr(kernel, "W4A16FusedMoeKernel", build_kernel) + with pytest.raises(KernelBuildReached): + kernel.compile_w4a16_fused_moe( + size_m=4, + hidden_size=2048, + intermediate_size=512, + num_experts=8, + top_k=2, + activation="silu", + apply_router_weight_on_input=False, + zero_fc2_output=False, + moe_block_size=8, + max_m_blocks=1, + sms=188, + max_shared_mem=101376, + # (fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n), 256 threads each. + force_tile_config=(128, 128, 32, 512), + ) + assert (built["fc2_tile_n"], built["fc2_tile_k"]) == (512, 32) From d1add8510c70c5bc88952e51ec0a27441440e2d9 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 10 Nov 2026 06:44:08 -0800 Subject: [PATCH 2/2] [https://nvbugs/7028060][fix] Install the W4A16 tile shim only over the validator signature it mirrors The shim now installs only when FlashInfer's _candidate_tile_fits has exactly the parameters of 0.6.17 to 0.7.1, the validator it mirrors, as the MLA tuning-config workaround in flashinfer_trtllm_gen.py does for its builder. Any other signature, such as 0.7.2's with enforce_search_floor, leaves FlashInfer's validator in place, so a later signature change cannot re-arm the shim's fixed keyword list. Tests: the pre-0.7.2 stand-in validator takes that exact parameter list; a new test checks that the hook leaves the installed validator in place when it takes enforce_search_floor; the new functions are annotated, and the kernel-build test takes its sizes from the file's named constants. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../moe/fused_moe/fused_moe_cute_dsl_b12x.py | 19 ++++- .../moe/test_cute_dsl_b12x_moe_backend.py | 81 ++++++++++++++----- 2 files changed, 79 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index f2e266db80c4..239b41d58b3a 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py @@ -76,10 +76,23 @@ # executor init (nvbug 6721561). FlashInfer 0.7.2 fixes the re-pin: it calls # ``_candidate_tile_fits(..., enforce_search_floor=False)``. This shim accepts # the ultra tile in the re-validation of the releases before that, under the -# same rule that produced it, and stands aside on a validator that takes -# ``enforce_search_floor``. +# same rule that produced it. It installs only over the validator signature of +# those releases, which it mirrors, and stands aside on any other, such as +# 0.7.2's with ``enforce_search_floor``. _FLASHINFER_W4A16_ULTRA_WIDE_FC2_TILE = (512, 32, 256) # (tile_n, tile_k, cta_threads) _FLASHINFER_W4A16_ULTRA_WIDE_FC2_SINCE = Version("0.6.17") +_FLASHINFER_W4A16_PRE_FIX_FITS_PARAMS = ( + "problem_n", + "problem_k", + "cta_m_blocks", + "tile_n", + "tile_k", + "cta_threads", + "max_shared_mem", + "scale_format", + "weight_layout", + "allow_logical_tail", +) def _patch_flashinfer_w4a16_ultra_wide_fc2_tile_validation() -> bool: @@ -105,7 +118,7 @@ def _patch_flashinfer_w4a16_ultra_wide_fc2_tile_validation() -> bool: return False if getattr(upstream_fits, "_trtllm_ultra_wide_fc2_shim", False): return True - if "enforce_search_floor" in inspect.signature(upstream_fits).parameters: + if tuple(inspect.signature(upstream_fits).parameters) != _FLASHINFER_W4A16_PRE_FIX_FITS_PARAMS: return False def candidate_tile_fits( diff --git a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index 94288fa8b08b..d32f1a22efee 100644 --- a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -18,7 +18,7 @@ import inspect import sys import types -from typing import Optional +from typing import NoReturn, Optional from unittest.mock import patch import pytest @@ -648,11 +648,14 @@ def test_run_moe_joins_the_input_ready_event_before_its_kernel( # flashinfer W4A16 ultra-wide FC2 tile re-validation shim (nvbug 6721561) # --------------------------------------------------------------------------- -# FC2 problem of Qwen3.6-35B-A3B (N=hidden_size, K=moe_intermediate_size) and -# the RTX PRO 6000 opt-in shared memory minus flashinfer's 512-byte reserve: -# the exact geometry the post-merge failure re-pinned. +# FC2 problem of Qwen3.6-35B-A3B (N=hidden_size, K=moe_intermediate_size), the +# SM count and opt-in shared memory of the RTX PRO 6000, and that memory minus +# flashinfer's 512-byte reserve: the exact geometry the post-merge failure +# re-pinned. _QWEN36_FC2 = dict(problem_n=2048, problem_k=512) -_RTX_PRO_6000_SMEM = 101376 - 512 +_RTX_PRO_6000_SMS = 188 +_RTX_PRO_6000_OPTIN_SMEM = 101376 +_RTX_PRO_6000_SMEM = _RTX_PRO_6000_OPTIN_SMEM - 512 _ULTRA_FC2 = dict(tile_n=512, tile_k=32, cta_threads=256) @@ -672,18 +675,42 @@ def _reset_to_stock_validator(kernel, monkeypatch) -> None: ) -def _reset_to_pre_fix_validator(kernel, monkeypatch) -> None: +def _reset_to_pre_fix_validator(kernel: types.ModuleType, monkeypatch: pytest.MonkeyPatch) -> None: """Start from a validator that applies the search floor to every tile, as - flashinfer 0.6.17 to 0.7.1 do. Where the installed validator exempts forced - tiles through ``enforce_search_floor``, that validator with the floor - always on, and without the parameter, stands in for them.""" + flashinfer 0.6.17 to 0.7.1 do, with their signature. Where the installed + validator exempts forced tiles through ``enforce_search_floor``, that + validator with the floor always on stands in for them.""" _reset_to_stock_validator(kernel, monkeypatch) stock = kernel._candidate_tile_fits if "enforce_search_floor" not in inspect.signature(stock).parameters: return - def pre_fix_fits(**kwargs): - return stock(**kwargs, enforce_search_floor=True) + def pre_fix_fits( + *, + problem_n: int, + problem_k: int, + cta_m_blocks: int, + tile_n: int, + tile_k: int, + cta_threads: int, + max_shared_mem: int, + scale_format: str = "e4m3_k16", + weight_layout: str = "packed", + allow_logical_tail: bool = False, + ) -> bool: + return stock( + problem_n=problem_n, + problem_k=problem_k, + cta_m_blocks=cta_m_blocks, + tile_n=tile_n, + tile_k=tile_k, + cta_threads=cta_threads, + max_shared_mem=max_shared_mem, + scale_format=scale_format, + weight_layout=weight_layout, + allow_logical_tail=allow_logical_tail, + enforce_search_floor=True, + ) monkeypatch.setattr(kernel, "_candidate_tile_fits", pre_fix_fits) @@ -741,7 +768,25 @@ def test_w4a16_ultra_wide_fc2_shim_yields_to_upstream_fix(monkeypatch): assert kernel._candidate_tile_fits is stock -def test_w4a16_forced_ultra_wide_fc2_tile_reaches_kernel_build(monkeypatch): +def test_w4a16_ultra_wide_fc2_shim_stands_aside_on_enforce_search_floor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """flashinfer 0.7.2 re-validates forced tiles with + ``_candidate_tile_fits(..., enforce_search_floor=False)``; the hook leaves + that validator in place.""" + kernel = _flashinfer_w4a16_kernel() + _reset_to_stock_validator(kernel, monkeypatch) + stock = kernel._candidate_tile_fits + if "enforce_search_floor" not in inspect.signature(stock).parameters: + pytest.skip("installed flashinfer predates enforce_search_floor") + + assert _patch_flashinfer_w4a16_ultra_wide_fc2_tile_validation() is False + assert kernel._candidate_tile_fits is stock + + +def test_w4a16_forced_ultra_wide_fc2_tile_reaches_kernel_build( + monkeypatch: pytest.MonkeyPatch, +) -> None: """The custom-op launch recompiles with the tiles the selector chose, forced through ``force_tile_config``. After the hook every W4A16 B12x leaf runs at construction, that re-pin of the ultra-wide FC2 tile gets past tile @@ -753,9 +798,9 @@ def test_w4a16_forced_ultra_wide_fc2_tile_reaches_kernel_build(monkeypatch): class KernelBuildReached(Exception): pass - built = {} + built: dict[str, object] = {} - def build_kernel(**kwargs): + def build_kernel(**kwargs: object) -> NoReturn: built.update(kwargs) raise KernelBuildReached @@ -764,8 +809,8 @@ def build_kernel(**kwargs): with pytest.raises(KernelBuildReached): kernel.compile_w4a16_fused_moe( size_m=4, - hidden_size=2048, - intermediate_size=512, + hidden_size=_QWEN36_FC2["problem_n"], + intermediate_size=_QWEN36_FC2["problem_k"], num_experts=8, top_k=2, activation="silu", @@ -773,8 +818,8 @@ def build_kernel(**kwargs): zero_fc2_output=False, moe_block_size=8, max_m_blocks=1, - sms=188, - max_shared_mem=101376, + sms=_RTX_PRO_6000_SMS, + max_shared_mem=_RTX_PRO_6000_OPTIN_SMEM, # (fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n), 256 threads each. force_tile_config=(128, 128, 32, 512), )