[https://nvbugs/7028060][fix] Stop shimming FlashInfer's W4A16 tile validator once it takes enforce_search_floor

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,20 +65,34 @@
     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
 # generic ``_candidate_tile_fits`` check, which the auto-selection skips but
 # 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. 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:
@@ -103,6 +118,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 tuple(inspect.signature(upstream_fits).parameters) != _FLASHINFER_W4A16_PRE_FIX_FITS_PARAMS:
+        return False
 
     def candidate_tile_fits(
         *,

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)

tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py

@@ -15,9 +15,10 @@
 """B12x MoE backend gating and dispatch tests."""
 
 import dataclasses
+import inspect
 import sys
 import types
-from typing import Optional
+from typing import NoReturn, Optional
 from unittest.mock import patch
 
 import pytest
@@ -647,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)
 
 
@@ -671,6 +675,46 @@ def _reset_to_stock_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, 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(
+        *,
+        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)
+
+
 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 +723,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 +747,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 +760,67 @@ 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_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
+    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: dict[str, object] = {}
+
+    def build_kernel(**kwargs: object) -> NoReturn:
+        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=_QWEN36_FC2["problem_n"],
+            intermediate_size=_QWEN36_FC2["problem_k"],
+            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=_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),
+        )
+    assert (built["fc2_tile_n"], built["fc2_tile_k"]) == (512, 32)