You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[https://nvbugs/7028060][fix] Stop shimming FlashInfer's W4A16 tile validator once it takes enforce_search_floor - #22004
FlashInfer 0.7.2 fixed the SM12x W4A16 fused MoE force_tile_config re-pin
(flashinfer-ai/flashinfer#6292): the re-pin re-validates the tiles the selector
chose with _candidate_tile_fits(..., enforce_search_floor=False), so the
TC-decode ultra-wide FC2 tile (tile_n=512, tile_k=32) passes.
fused_moe_cute_dsl_b12x.py carries a shim for the same problem, which every
W4A16 NVFP4 B12x MoE leaf installs over that validator at construction for
FlashInfer 0.6.17 and later. It recognized an upstream fix only by a _ultra_wide_fc2_tile_fits attribute, which FlashInfer 0.7.2 does not have, so
with the flashinfer-python pin at 0.7.2 (#20911) the shim stays installed and its
fixed keyword list rejects the new argument. Every W4A16 fused MoE launch through run_w4a16_moe goes through that re-pin, so the first one, during executor
warmup, fails:
TypeError: _patch_flashinfer_w4a16_ultra_wide_fc2_tile_validation.<locals>.candidate_tile_fits() got an unexpected keyword argument 'enforce_search_floor'
This change installs the shim only over the validator signature of FlashInfer
0.6.17 to 0.7.1, the one it mirrors, and leaves any other validator in place, such
as 0.7.2's, which adds enforce_search_floor. The hook runs for every W4A16 NVFP4
B12x MoE leaf, on SM120 and SM121 alike. The waiver of accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4_w4a16 is removed.
Test Coverage
tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py::test_w4a16_forced_ultra_wide_fc2_tile_reaches_kernel_build
(new, cpu_only, collected by unittest/_torch/moe in l0_cpu.yml): after the
shim hook runs, compiles the W4A16 fused MoE with the ultra-wide FC2 tile forced
through force_tile_config, with the kernel build stubbed, and checks the forced
tile reaches the kernel build.
test_w4a16_ultra_wide_fc2_shim_stands_aside_on_enforce_search_floor (new, cpu_only): with the installed validator taking enforce_search_floor, the hook
returns False and leaves that validator in place.
The existing test_w4a16_ultra_wide_fc2_shim_* tests now start from a validator
with the 0.6.17 to 0.7.1 signature that applies the search floor to every tile, as
those releases do, so they keep covering the shim with FlashInfer 0.7.2 installed.
accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4_w4a16
(l0_rtx_pro_6000.yml, post-merge), no longer waived.
Verified with flashinfer-python 0.7.2, with main at 0bcbcfe built in the CI
image with the CI build flags (CUDA arch 120 only):
On an SM 12.0 GPU, TestQwen3_6_35B_A3B::test_nvfp4_w4a16 fails without this
change with the TypeError above and passes with it (GSM8K 95.830, threshold
91.797), run in one pytest process after TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutlass-torch_compile=False] and [...-torch_compile=True], which pass as well.
On a GPU-less node set up like the CPU stage: the two new unit tests fail without
this change, the kernel-build test with the same TypeError, and pass with it; the
three existing shim tests pass with and without it. unittest/_torch/moe with -m cpu_only passes with this change (751 passed, 1 skipped by its waiver).
…alidator 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>
No actionable comments were generated in the recent review.
Walkthrough
The FlashInfer tile-validation shim now checks the validator signature before installation. Tests cover pre-fix and updated validator behavior, including compilation with an ultra-wide FC2 tile. The waiver for the corresponding integration test was removed.
Changes
FlashInfer Tile Validator Compatibility
Layer / File(s)
Summary
Gate shim installation on validator signature tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py
The module checks the validator parameter names before installing the shim. Comments describe the affected FlashInfer versions and the validator behavior through 0.7.1.
Test validator compatibility and tile compilation tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py, tests/integration/test_lists/waives.txt
Tests set up pre-fix validator behavior and check shim acceptance, idempotency, updated-validator handling, and compilation with FC2 tile dimensions (512, 32). The waiver entry for TestQwen3_6_35B_A3B::test_nvfp4_w4a16 was removed.
NVIDIA/TensorRT-LLM#21778: This PR retains the B12x W4A16 FlashInfer tile-validation shim; this PR changes when the shim applies.
Merge Risk: Minimal · up to d1add
The compatibility change has no established merge-blocking issue; merge after normal checks.
Pre-merge checks: 5
Check name
Status
Explanation
Title check
The title follows the required ticket-and-type format and clearly states that the shim stops when FlashInfer adds enforce_search_floor.
Docstring Coverage
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…
Linked Issues check
Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check
Check skipped because no linked issues were found for this pull request.
Description check
The description explains the issue, the fix, and the test coverage. It also reports verification results. The PR checklist section is not included, but the description is otherwise mostly complete.
Full details: Docstring Coverage
Explanation
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 3 functions across 2 files.
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Checked against the FlashInfer 0.7.2 and 0.7.3 wheels: both validators take enforce_search_floor and the force_tile_config re-pin passes enforce_search_floor=False, so standing the shim aside lets FlashInfer accept the ultra-wide FC2 tile itself, and the new kernel-build test fails on the parent with the reported TypeError. The CPU stages that collect unittest/_torch/moe do run these shim tests with FlashInfer 0.7.2. Counts: 0 blockers, 0 possibly blocking, 2 required, 4 nice to have.
[Required] PR description: the auto-generated CodeRabbit block ("Dev Engineer Review" / "QA Engineer Review" / "Per-File QA Perspective") contradicts the change. It says the changed unit test file "is not listed in test-db/ or qa/", but tests/integration/test_lists/test-db/l0_cpu.yml collects it through unittest/_torch/moe (CPU-Generic-x86-1/arm-1 run these shim tests on main today). It also says no runtime or CI results are available, while your Test Coverage section reports the passing runs. The PR template asks that a CodeRabbit summary make sense; please delete the block or correct it.
[Nice to have] The un-waived accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4_w4a16 runs only in RTXPro6000D-PyTorch-Post-Merge-1 (the post_merge block of l0_rtx_pro_6000.yml), so pre-merge CI never runs this fix on GPU. Please run /bot run --extra-stage "RTXPro6000D-PyTorch-Post-Merge-1" before merge and link the passing run. The DGX Spark (SM121) QA cases reported in the bug (Nemotron-3.5-Lightning NVFP4, the Qwen3.6 perf case) go through the same hook; mentioning them in the description lets QA close them with this fix.
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.
…he 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>
Thanks for the review. Pushed d1add85 with the inline points, each answered in its thread, and updated the description:
Removed the generated summary block. It said the changed test file is in no test list, while l0_cpu.yml collects it through unittest/_torch/moe, and that no results were available.
The description now lists the new stand-aside test and notes that the hook runs for every W4A16 NVFP4 B12x MoE leaf, on SM120 and SM121 alike; only SM 12.0 was run for this PR.
Repair Bot cannot start CI on this repository, so it cannot run RTXPro6000D-PyTorch-Post-Merge-1 itself.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
sunnyqgg
left a comment
Description
FlashInfer 0.7.2 fixed the SM12x W4A16 fused MoE
force_tile_configre-pin(flashinfer-ai/flashinfer#6292): the re-pin re-validates the tiles the selector
chose with
_candidate_tile_fits(..., enforce_search_floor=False), so theTC-decode ultra-wide FC2 tile (
tile_n=512,tile_k=32) passes.fused_moe_cute_dsl_b12x.pycarries a shim for the same problem, which everyW4A16 NVFP4 B12x MoE leaf installs over that validator at construction for
FlashInfer 0.6.17 and later. It recognized an upstream fix only by a
_ultra_wide_fc2_tile_fitsattribute, which FlashInfer 0.7.2 does not have, sowith the flashinfer-python pin at 0.7.2 (#20911) the shim stays installed and its
fixed keyword list rejects the new argument. Every W4A16 fused MoE launch through
run_w4a16_moegoes through that re-pin, so the first one, during executorwarmup, fails:
This change installs the shim only over the validator signature of FlashInfer
0.6.17 to 0.7.1, the one it mirrors, and leaves any other validator in place, such
as 0.7.2's, which adds
enforce_search_floor. The hook runs for every W4A16 NVFP4B12x MoE leaf, on SM120 and SM121 alike. The waiver of
accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4_w4a16is removed.Test Coverage
tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py::test_w4a16_forced_ultra_wide_fc2_tile_reaches_kernel_build(new,
cpu_only, collected byunittest/_torch/moeinl0_cpu.yml): after theshim hook runs, compiles the W4A16 fused MoE with the ultra-wide FC2 tile forced
through
force_tile_config, with the kernel build stubbed, and checks the forcedtile reaches the kernel build.
test_w4a16_ultra_wide_fc2_shim_stands_aside_on_enforce_search_floor(new,cpu_only): with the installed validator takingenforce_search_floor, the hookreturns
Falseand leaves that validator in place.test_w4a16_ultra_wide_fc2_shim_*tests now start from a validatorwith the 0.6.17 to 0.7.1 signature that applies the search floor to every tile, as
those releases do, so they keep covering the shim with FlashInfer 0.7.2 installed.
accuracy/test_llm_api_pytorch.py::TestQwen3_6_35B_A3B::test_nvfp4_w4a16(
l0_rtx_pro_6000.yml, post-merge), no longer waived.Verified with flashinfer-python 0.7.2, with main at 0bcbcfe built in the CI
image with the CI build flags (CUDA arch 120 only):
TestQwen3_6_35B_A3B::test_nvfp4_w4a16fails without thischange with the
TypeErrorabove and passes with it (GSM8K 95.830, threshold91.797), run in one pytest process after
TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutlass-torch_compile=False]and[...-torch_compile=True], which pass as well.this change, the kernel-build test with the same
TypeError, and pass with it; thethree existing shim tests pass with and without it.
unittest/_torch/moewith-m cpu_onlypasses with this change (751 passed, 1 skipped by its waiver).Fixes https://nvbugs/7028060