fix(trtllm): restrict TRT-LLM routed-MoE and GEMM backends to supported architectures - #4177
Smallfu666 wants to merge 1 commit into
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
📝 WalkthroughWalkthroughTRTLLM GEMM runners now filter cubins and tactics by GPU SM compatibility, report architecture-specific errors, and validate runtime configuration indices. Routed MoE backends use narrower architecture support predicates with expanded support and dispatch guard tests. ChangesTRTLLM compatibility
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
e848e38 to
1e3b3a5
Compare
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
csrc/trtllm_gemm_runner.cu (1)
141-150: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftUse the actual execution device for cubin compatibility filtering. Both runners derive SM compatibility from the active CUDA device, while execution can target a separately supplied device. This can select incompatible cubins in multi-GPU processes.
csrc/trtllm_gemm_runner.cu#L141-L150: derive SM froma.device()(or establish a device guard) before constructing and filtering the runner.csrc/trtllm_batched_gemm_runner.cu#L175-L177: pass the target device/SM into construction, or enforce that construction and execution use the same device.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/trtllm_gemm_runner.cu` around lines 141 - 150, The runners currently filter cubins using the active CUDA device instead of the execution target. In csrc/trtllm_gemm_runner.cu:141-150, update getSMVersion usage in the runner construction/filtering path to derive SM from a.device() or establish a device guard; in csrc/trtllm_batched_gemm_runner.cu:175-177, pass the target device/SM into construction or enforce that construction and execution use the same device.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@csrc/trtllm_gemm_runner.cu`:
- Around line 141-150: The runners currently filter cubins using the active CUDA
device instead of the execution target. In csrc/trtllm_gemm_runner.cu:141-150,
update getSMVersion usage in the runner construction/filtering path to derive SM
from a.device() or establish a device guard; in
csrc/trtllm_batched_gemm_runner.cu:175-177, pass the target device/SM into
construction or enforce that construction and execution use the same device.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 2d317abc-3a44-48be-8ecb-6122ef84c6fd
📥 Commits
Reviewing files that changed from the base of the PR and between 4b969c9 and b14024eb421d24dd9cdadeef7fe3b31782cbb708.
📒 Files selected for processing (4)
csrc/trtllm_batched_gemm_runner.cucsrc/trtllm_gemm_runner.cuflashinfer/fused_moe/api.pytests/moe_ep/test_split_fused_moe_kernel_vs_reference.py
|
Arch guard for split path validation on sm120. Looks good to me! + @Anerudhan for vis |
|
Hello @Smallfu666! Thanks for addressing the TensorRT-LLM issue on moe_ep and opening this PR. It looks like there are some merge conflicts to resolve. Once those are taken care of, I'll help get the PR reviewed and checked in! |
|
Thanks @mhoqueanik! Quick check on direction before I resolve the conflicts — they're a symptom of #4171 (Revert PR 4122) rather than a textual clash. That revert rolls the cubin manifest back to the pre-SM107 bundles, drops Worth flagging: only the batched-runner hunk actually conflicts. The Separately, on current Two questions:
Happy to rework as soon as you point me at the preferred direction. |
|
Hello @Smallfu666! Thanks for the detailed analysis! your reading sounds right. The visible conflict is caused by the rollback of #4122, but the underlying architecture-filtering issue still exists independently of SM107 support. Could you rebase the PR around a generic architecture-compatibility check and scope the current mapping to the architectures present on main, {100, 103}? Please remove the Sm107a references for now so the PR does not depend on #4122. If SM107 support re-lands, it should only require extending the compatibility mapping rather than changing the validation mechanism. The proposed validation in both getWorkspaceSizeInBytes() and run() also makes sense (as long as there is no perf drop - assuming it's not in hotpath). That should reject stale cached or autotuner-supplied tactics selected for a different GPU instead of silently dispatching them. It may be best to keep the compatibility rules explicit rather than broadly treating unknown sm10x architectures as compatible. |
…(release port for #4107) (#4230) ## Description Release-specific port of #4177 onto `release-v0.6.16` for #4107. On SM12x (Spark, RTX Pro 6000), TRTLLM routed-MoE backends were incorrectly claiming support (`arch >= 100`) and then dispatching sm100f/sm103a cubins, causing `RuntimeError: Error occurred when running GEMM!` or segfaults in `test_split_fused_moe_kernel_vs_reference`. ## Changes - **`csrc/trtllm_batched_gemm_runner.cu`**: Replace per-SM if-chains with `isArchCompatible()`; reject unknown cubin families; guard `Sm107a` behind `#ifdef TLLM_RUBIN_FEATURES` (only exists in the Rubin cubin pin's headers). - **`csrc/trtllm_gemm_runner.cu`**: Same arch filter for the plain GEMM runner (previously had no arch filtering at all). - **`flashinfer/fused_moe/api.py`**: Tighten `Trtllm*Config.supported()` from `arch >= 100` to explicit allowlists `_TRTLLM_ROUTED_ARCHS = (100, 103, 107)` and `_TRTLLM_ROUTED_FP8_ARCHS = (100, 103)`. - **`tests/moe_ep/test_split_fused_moe_kernel_vs_reference.py`**: Gate GPU tests on `config_cls.supported(arch)`; add CPU contract tests + SM120 regression guard. ## Release-specific notes - Includes **sm107** in the allowlist (unlike the rebased #4177 for `main`, which will drop 107 after the #4122 revert). - `Sm107a` enum case is gated behind `TLLM_RUBIN_FEATURES`, matching the existing release pattern from #4122/#4191. Verified against the actual pinned headers: the default BMM/GEMM pins do not define `Sm107a`; only the Rubin pins do. ## Verification - CPU tests: 36 passed, 2 skipped (`test_split_fused_moe_kernel_vs_reference.py`) - Backend claims: sm120/sm121 now fall back to Cutlass only (no Trtllm* backends) - Compile: `isArchCompatible()` builds cleanly against both default and Rubin BMM export headers; ported batched runner compiles against default pin ## Related - Fixes #4107 - Upstream PR: #4177 (targets `main`, has merge conflicts) - Cherry-pick source commit on `main`: not yet merged ## Pre-existing issue (not in scope) `#4213` on release references `options.mDtypeSfC`, which does not exist in the default BMM cubin pin's headers (only in the Rubin pin). This is a separate release-only compile issue on the non-Rubin module, predating this port.
The trtllm-gen runners selected cubins without checking that the cubin architecture can actually run on the target device: - TrtllmGenBatchedGemmRunner only applied an arch filter when mPatchF2fp was set, so on unsupported GPUs (e.g. SM120/121) sm100f cubins were selected and aborted at dispatch (issue flashinfer-ai#4107). - TrtllmGenGemmRunner had no config.mSm filter at all. - Neither runner validated a cached or user-supplied tactic/configIndex against the filtered config set in getWorkspaceSizeInBytes()/run(), and the batched run() dereferenced configs[configIndex] before any bounds check. - TrtllmMxInt4Config.supported() still claimed arch >= 100, so the routed-MoE dispatcher could pick it on SM120/121 and then hard-fail during runner construction instead of falling back to Cutlass. Fix: - Add an explicit cubin-arch -> SM-version compatibility mapping (isArchCompatible) to both runners, covering the architectures the current manifests ship cubins for: Sm100a/Sm100f on sm100, Sm100f/Sm103a on sm103. Unknown cubin families are rejected rather than assumed compatible, so a future cubin drop fails loudly instead of being silently dispatched. Re-landing SM107 (or adding any new arch) only needs a new enum case here plus the Python allowlist entry; the validation mechanism itself is unchanged. - Validate the tactic/configIndex against the filtered set in both getWorkspaceSizeInBytes() and run() to reject cached or user-supplied incompatible tactics. Both checks run on the host, once per API-level GEMM/MoE invocation. - Re-check the selected cubin against the SM version of the explicit target device argument in run(), so a construction-time/dispatch-time device mismatch fails with a clear error instead of aborting inside the kernel launch. The per-device SM version is cached (static atomic array), so the guard performs no driver query on the dispatch path after first use. - When a GPU has no compatible cubins at all, report that explicitly instead of the generic GEMM-options dump ("No kernel found for the given options"), which pointed users in the wrong direction. Both runners emit the same message format. - Tighten TrtllmMxInt4Config to the same {100, 103} allowlist the other trtllm routed-MoE backends already use; the pinned BMM manifest ships all MxInt4 kernels as sm100f, so the narrowed claim matches actual cubin coverage. - Tests: gate the split fused-MoE GPU tests on the concrete backend contract; CPU tests locking down supported() for all five trtllm routed backends plus a regression guard asserting no trtllm backend claims sm107/110/120/121; a negative-path test asserting the explicit arch error on GPUs without compatible cubins; and a rejection test for a manifest-valid tactic outside the runner's config set (runs on sm100/103, skips elsewhere via an explicit device-arch check). Fixes flashinfer-ai#4107 Signed-off-by: Han-Yin Chang <nick20350@gmail.com>
b14024e to
cccc773
Compare
|
Rebased onto current main (now includes #4180) and implemented as requested:
Also tightened |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
flashinfer/fused_moe/api.py (1)
231-236: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeduplicate the hardcoded
(100, 103)arch tuple.Four
supported()classmethods independently hardcodearch in (100, 103). Given this exact set already churned once (#4122added SM107,#4171reverted it), a single source of truth would prevent a future update from missing one class.♻️ Proposed refactor
+# trtllm-gen routed batched-GEMM cubins currently ship for SM100/SM103 only; +# SM107 support from `#4122` was reverted by `#4171`. See `#4107`. +_TRTLLM_ROUTED_SUPPORTED_ARCHS = (100, 103) + class TrtllmFp4Config: ... `@classmethod` def supported(cls, arch: int) -> bool: - # Current TRTLLM FP4 cubins are supported only on SM100/SM103. - # SM107 support from `#4122` was reverted by `#4171`. - return arch in (100, 103) + return arch in _TRTLLM_ROUTED_SUPPORTED_ARCHSApply the same substitution to
TrtllmFp8BlockConfig,TrtllmBf16Config, andTrtllmMxInt4Config(andTrtllmFp8PerTensorConfigif convenient).Also applies to: 294-299, 395-397, 435-440
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/api.py` around lines 231 - 236, Introduce a single shared constant for the supported TRTLLM architectures `(100, 103)` and update the `supported()` classmethods, including the method shown and those on `TrtllmFp8BlockConfig`, `TrtllmBf16Config`, `TrtllmMxInt4Config`, and `TrtllmFp8PerTensorConfig` if present, to reference it instead of duplicating the tuple. Preserve the existing support set and boolean behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/gemm/test_trtllm_gemm_arch_guard.py`:
- Around line 49-97: Update
test_trtllm_gemm_rejects_in_range_tactic_outside_config_set to use the backend
capability API is_sm100a_supported for its skip condition instead of manually
checking arch values. Skip when that helper reports unsupported hardware or
CUDA/toolkit support, and keep the supported-path assertions and test execution
unchanged.
---
Nitpick comments:
In `@flashinfer/fused_moe/api.py`:
- Around line 231-236: Introduce a single shared constant for the supported
TRTLLM architectures `(100, 103)` and update the `supported()` classmethods,
including the method shown and those on `TrtllmFp8BlockConfig`,
`TrtllmBf16Config`, `TrtllmMxInt4Config`, and `TrtllmFp8PerTensorConfig` if
present, to reference it instead of duplicating the tuple. Preserve the existing
support set and boolean behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: fde2784f-41a1-44a4-9add-b99f996caedd
📥 Commits
Reviewing files that changed from the base of the PR and between b14024eb421d24dd9cdadeef7fe3b31782cbb708 and cccc773.
📒 Files selected for processing (5)
csrc/trtllm_batched_gemm_runner.cucsrc/trtllm_gemm_runner.cuflashinfer/fused_moe/api.pytests/gemm/test_trtllm_gemm_arch_guard.pytests/moe_ep/test_split_fused_moe_kernel_vs_reference.py
🚧 Files skipped from review as they are similar to previous changes (2)
- csrc/trtllm_gemm_runner.cu
- csrc/trtllm_batched_gemm_runner.cu
| def test_trtllm_gemm_rejects_in_range_tactic_outside_config_set(): | ||
| """A manifest index that belongs to a different runner configuration must | ||
| be rejected by the tactic-membership guard, not silently dispatched.""" | ||
| arch = _arch() | ||
| # Explicit device-arch gate: skip only on hardware that is not sm100/103. | ||
| # On sm100/103 this test must run and must not skip. | ||
| if arch not in (100, 103): | ||
| pytest.skip(f"requires an sm100/sm103 GPU, detected sm{arch}") | ||
|
|
||
| op = _load_op() | ||
| fp8_tactics = set( | ||
| op.trtllm_gemm_tactics( | ||
| M, N, K, int(DtypeTrtllmGen.E4m3), int(DtypeTrtllmGen.Bfloat16), False | ||
| ) | ||
| ) | ||
| fp4_tactics = list( | ||
| op.trtllm_gemm_tactics( | ||
| M, N, K, int(DtypeTrtllmGen.E2m1), int(DtypeTrtllmGen.Bfloat16), True | ||
| ) | ||
| ) | ||
| foreign = [t for t in fp4_tactics if t not in fp8_tactics] | ||
| # On supported hardware this test must not skip: the E2m1 and E4m3 cubin | ||
| # families are disjoint in the manifest, so a foreign tactic must exist. | ||
| assert foreign, ( | ||
| "expected at least one E2m1 tactic outside the E4m3 set; " | ||
| f"E4m3={sorted(fp8_tactics)} E2m1={sorted(fp4_tactics)}" | ||
| ) | ||
|
|
||
| a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) | ||
| b = torch.randn(N, K, device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) | ||
| a_scale = torch.ones(1, device="cuda", dtype=torch.float32) | ||
| b_scale = torch.ones(1, device="cuda", dtype=torch.float32) | ||
| out = torch.zeros(M, N, device="cuda", dtype=torch.bfloat16) | ||
| workspace = torch.empty(4 * 1024 * 1024, device="cuda", dtype=torch.int8) | ||
|
|
||
| with pytest.raises(Exception, match="not in this runner's compatible config set"): | ||
| op.trtllm_gemm( | ||
| int(DtypeTrtllmGen.E4m3), | ||
| int(DtypeTrtllmGen.Bfloat16), | ||
| workspace, | ||
| a, | ||
| b, | ||
| a_scale, | ||
| b_scale, | ||
| None, | ||
| out, | ||
| False, | ||
| int(foreign[0]), | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n "not in this runner's compatible config set" csrc include
rg -n "def is_sm100a_supported" flashinfer/utils.py -A 8Repository: flashinfer-ai/flashinfer
Length of output: 812
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== test file excerpt =="
sed -n '1,130p' tests/gemm/test_trtllm_gemm_arch_guard.py
echo
echo "== runner exception excerpt =="
sed -n '170,205p' csrc/trtllm_gemm_runner.cu
echo
echo "== backend/utils related definitions =="
rg -n "is_sm100a_supported|version_at_least|compute_capability|get_compute_capability|torch.version.cuda" flashinfer/utils.py -A 4 -B 4
rg -n "trtllm_gemm|DtypeTrtllmGen|is_sm100a_supported|is_sm100f_supported" csrc include flashinfer tests -g '*.{cu,cuh,h,hpp,cpp,cuh,py}' | head -200
echo
echo "== file list useful =="
fd -a 'test_trtllm_gemm_arch_guard.py|utils.py|trtllm_gemm_runner.cu' .Repository: flashinfer-ai/flashinfer
Length of output: 38553
Use is_sm100a_supported for this GPU skip.
tests/**/*.py should skip unsupported CUDA setups via backend capability APIs. is_sm100a_supported covers SM100 and SM103 and requires CUDA ≥ 12.8, while the manual arch in (100, 103) gate does not require the toolkit version.
♻️ Proposed fix
- arch = _arch()
- # Explicit device-arch gate: skip only on hardware that is not sm100/103.
- # On sm100/103 this test must run and must not skip.
- if arch not in (100, 103):
- pytest.skip(f"requires an sm100/sm103 GPU, detected sm{arch}")
+ from flashinfer.utils import is_sm100a_supported
+
+ arch = _arch()
+ if not is_sm100a_supported(torch.device("cuda")):
+ pytest.skip(f"requires an sm100/sm103 GPU, detected sm{arch}")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def test_trtllm_gemm_rejects_in_range_tactic_outside_config_set(): | |
| """A manifest index that belongs to a different runner configuration must | |
| be rejected by the tactic-membership guard, not silently dispatched.""" | |
| arch = _arch() | |
| # Explicit device-arch gate: skip only on hardware that is not sm100/103. | |
| # On sm100/103 this test must run and must not skip. | |
| if arch not in (100, 103): | |
| pytest.skip(f"requires an sm100/sm103 GPU, detected sm{arch}") | |
| op = _load_op() | |
| fp8_tactics = set( | |
| op.trtllm_gemm_tactics( | |
| M, N, K, int(DtypeTrtllmGen.E4m3), int(DtypeTrtllmGen.Bfloat16), False | |
| ) | |
| ) | |
| fp4_tactics = list( | |
| op.trtllm_gemm_tactics( | |
| M, N, K, int(DtypeTrtllmGen.E2m1), int(DtypeTrtllmGen.Bfloat16), True | |
| ) | |
| ) | |
| foreign = [t for t in fp4_tactics if t not in fp8_tactics] | |
| # On supported hardware this test must not skip: the E2m1 and E4m3 cubin | |
| # families are disjoint in the manifest, so a foreign tactic must exist. | |
| assert foreign, ( | |
| "expected at least one E2m1 tactic outside the E4m3 set; " | |
| f"E4m3={sorted(fp8_tactics)} E2m1={sorted(fp4_tactics)}" | |
| ) | |
| a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) | |
| b = torch.randn(N, K, device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) | |
| a_scale = torch.ones(1, device="cuda", dtype=torch.float32) | |
| b_scale = torch.ones(1, device="cuda", dtype=torch.float32) | |
| out = torch.zeros(M, N, device="cuda", dtype=torch.bfloat16) | |
| workspace = torch.empty(4 * 1024 * 1024, device="cuda", dtype=torch.int8) | |
| with pytest.raises(Exception, match="not in this runner's compatible config set"): | |
| op.trtllm_gemm( | |
| int(DtypeTrtllmGen.E4m3), | |
| int(DtypeTrtllmGen.Bfloat16), | |
| workspace, | |
| a, | |
| b, | |
| a_scale, | |
| b_scale, | |
| None, | |
| out, | |
| False, | |
| int(foreign[0]), | |
| ) | |
| def test_trtllm_gemm_rejects_in_range_tactic_outside_config_set(): | |
| """A manifest index that belongs to a different runner configuration must | |
| be rejected by the tactic-membership guard, not silently dispatched.""" | |
| from flashinfer.utils import is_sm100a_supported | |
| arch = _arch() | |
| if not is_sm100a_supported(torch.device("cuda")): | |
| pytest.skip(f"requires an sm100/sm103 GPU, detected sm{arch}") | |
| op = _load_op() | |
| fp8_tactics = set( | |
| op.trtllm_gemm_tactics( | |
| M, N, K, int(DtypeTrtllmGen.E4m3), int(DtypeTrtllmGen.Bfloat16), False | |
| ) | |
| ) | |
| fp4_tactics = list( | |
| op.trtllm_gemm_tactics( | |
| M, N, K, int(DtypeTrtllmGen.E2m1), int(DtypeTrtllmGen.Bfloat16), True | |
| ) | |
| ) | |
| foreign = [t for t in fp4_tactics if t not in fp8_tactics] | |
| # On supported hardware this test must not skip: the E2m1 and E4m3 cubin | |
| # families are disjoint in the manifest, so a foreign tactic must exist. | |
| assert foreign, ( | |
| "expected at least one E2m1 tactic outside the E4m3 set; " | |
| f"E4m3={sorted(fp8_tactics)} E2m1={sorted(fp4_tactics)}" | |
| ) | |
| a = torch.randn(M, K, device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) | |
| b = torch.randn(N, K, device="cuda", dtype=torch.bfloat16).to(torch.float8_e4m3fn) | |
| a_scale = torch.ones(1, device="cuda", dtype=torch.float32) | |
| b_scale = torch.ones(1, device="cuda", dtype=torch.float32) | |
| out = torch.zeros(M, N, device="cuda", dtype=torch.bfloat16) | |
| workspace = torch.empty(4 * 1024 * 1024, device="cuda", dtype=torch.int8) | |
| with pytest.raises(Exception, match="not in this runner's compatible config set"): | |
| op.trtllm_gemm( | |
| int(DtypeTrtllmGen.E4m3), | |
| int(DtypeTrtllmGen.Bfloat16), | |
| workspace, | |
| a, | |
| b, | |
| a_scale, | |
| b_scale, | |
| None, | |
| out, | |
| False, | |
| int(foreign[0]), | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/gemm/test_trtllm_gemm_arch_guard.py` around lines 49 - 97, Update
test_trtllm_gemm_rejects_in_range_tactic_outside_config_set to use the backend
capability API is_sm100a_supported for its skip condition instead of manually
checking arch values. Skip when that helper reports unsupported hardware or
CUDA/toolkit support, and keep the supported-path assertions and test execution
unchanged.
Source: Path instructions
|
Nice work @Smallfu666 ! I see another PR already addressed the issue in the meantime. Do you think this PR serves additional purpose? Then we can go ahead and process it. |
|
@mhoqueanik Yes—#4230 addresses the core issue on It also adds a target-device architecture check before launch and dedicated plain-GEMM regression coverage. The SM107 handling differs intentionally because Rubin support was reverted on |
3 participants
What this PR does
The trtllm-gen GEMM runners selected cubins without checking that the cubin
architecture can actually run on the current device:
TrtllmGenBatchedGemmRunneronly applied an arch filter whenmPatchF2fpwas set, so on unsupported GPUs (e.g. SM120/SM121) sm100f cubins were
selected and aborted at dispatch ([bug][v0.6.16] test_split_fused_moe_kernel_vs_reference error at /workspace/flashinfer/csrc/trtllm_batched_gemm_runner.cu:305: Error occurred when running GEMM! #4107).
TrtllmGenGemmRunnerhad noconfig.mSmfilter at all.Changes:
isArchCompatible(smVersion, config.mSm)rule set(
Sm100a -> sm100,Sm100f -> sm100/sm103,Sm103a -> sm103,anything unknown -> rejected). Adding a future arch is one enum case, not a
hardcoded SM allowlist edit.
explicit architecture error instead of the GEMM-options dump.
getWorkspaceSizeInBytes()andrun()isvalidated against the filtered config set, so cached or user-supplied
indices for a different configuration fail with a clear error.
run()additionally re-checks the selected cubin against the SM of theexplicit target
deviceargument, so a construction-time/dispatch-timedevice mismatch fails loudly at dispatch rather than inside the kernel
launch. The per-device SM version is cached (static atomic array), so this
guard performs no driver query on the dispatch path after first use.
(The construction-time filter still uses the device current at
construction; deriving the passing set per target device would require
runner-cache/API changes and is left to a follow-up.)
TrtllmMxInt4Config.supported()is tightened fromarch >= 100to(100, 103), matching the other trtllm routed-MoE backends. The pinnedBMM manifest ships MxInt4 kernels exclusively as
sm100f(runs onsm100/sm103), so the narrowed claim matches actual cubin coverage.
supported()contract tests (CPU), a negative-path test forthe explicit arch error on GPUs without compatible cubins, and a direct
rejection test for a manifest-valid tactic outside the runner's config set
(requires SM100/SM103 to construct the runner; on other GPUs it skips via
an explicit device-arch check).
Rebased onto current
main(includes #4180, which touches the samebatched-runner filter loop and updates the BMM artifact pin, and #4237/#4130;
no conflicts).
Fixes #4107
Validation
gen_trtllm_gen_fused_moe_sm100_moduleandgen_trtllm_gen_gemm_modulefrom source at the new BMM pin.(previously: option dump / abort at dispatch).
tests/moe_ep/test_split_fused_moe_kernel_vs_reference.pyandtests/gemm/test_trtllm_gemm_arch_guard.pypass on SM90 (SM100-only casesskip); the SM100/SM103 positive paths require Blackwell hardware I don't
have access to — they are intended to run in this repo's CI when it is
triggered on this PR.
Notes for reviewers
current at construction rather than the tensor's target device is only
partially addressed here (item 4 makes the mismatch fail loudly at
dispatch). A full fix needs the passing set keyed per target device
(public struct + runner-cache changes) and is proposed as a follow-up issue.