Conversation
…emory On SM121 (GB10 / DGX Spark GeForce Blackwell) the per-SM shared-memory opt-in limit is ~99 KB (cudaDevAttrMaxSharedMemoryPerBlockOptin ≈ 101376 bytes). The 256×128 and 128×256 CTA tiles in the SM120 CUTLASS MXFP8 module are compiled with StageCount<2>, which requires ~99 KB of shared memory — right at or beyond SM121's limit. CUTLASS initialize() returns kErrorInternal for those tactics on SM121, causing the autotuner to log four "[MXFP8 SM120 gemm Runner] Failed to initialize … Error: Error Internal" messages on every model load. Root cause: get_valid_tactics() unconditionally returned all 10 tactics (5 CTA shapes × 2 swap-AB variants) without checking whether they can actually initialize on the running device. Fix: add _probe_mxfp8_gemm_tactics() which runs a cheap M=N=K=128 probe for each tactic (initialization failure is shape-independent — it depends only on CTA tile size vs device smem), caches the result once per module instance, and returns only the supported subset. On SM121 this yields tactics [0, 1, 2, 3, 4, 5] (128×32, 128×64, 128×128 × 2 variants); the autotuner never attempts the failing large-tile tactics and the noisy error messages disappear. Also improve the C++ error message in mxfp8_gemm_template_sm120.h: when initialize() fails with kErrorInternal, query the device's smem opt-in limit and include it in the exception text so future diagnostics are actionable rather than generic. Verified on NVIDIA GB10 (SM121, 48 SM, 101376 bytes smem optin): Before: tactics 0-5 OK, tactics 6-9 fail with "Error Internal" After: _probe_mxfp8_gemm_tactics returns [0,1,2,3,4,5] silently Closes: flashinfer-ai#3558 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
tgmerritt
requested review from
aleozlx,
bkryu,
dhiraj113 and
yzh119
as code owners
June 10, 2026 16:23
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdds device-targeted probing for CUTLASS MXFP8 GEMM tactics (tries initializing each tactic on the current CUDA device and caches valid ones), enhances SM120 GEMM initialization errors with a shared-memory diagnostic, and adds tests validating probing behavior across SM120/SM121. ChangesMXFP8 Tactic Probing and Device Validation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a device-capability probe (_probe_mxfp8_gemm_tactics) to pre-filter CUTLASS MXFP8 GEMM tactics, preventing initialization failures on SM121 devices (which have lower shared memory limits than SM120). It also adds detailed error hints for initialization failures in the C++ template and includes corresponding unit tests. The review feedback highlights two critical issues regarding multi-GPU environments: first, hardcoding device index 0 when querying device properties and allocating tensors can lead to incorrect diagnostics and allocations on the wrong GPU; second, caching valid tactics in a single list at the closure scope can cause failures in heterogeneous multi-GPU setups (e.g., mixing SM120 and SM121). It is recommended to use the active device context and key the cache by device_id using a dictionary.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+4008
to
+4060
| def _probe_mxfp8_gemm_tactics(module) -> List[int]: | ||
| """Probe which CUTLASS MXFP8 GEMM tactics can initialize on the current device. | ||
|
|
||
| SM120 (B100/B200 data-centre, ~256 KB shared memory/SM) and SM121 (GB10 / | ||
| DGX Spark GeForce Blackwell, ~99 KB shared memory/SM opt-in) share the same | ||
| compiled module. Large CTA tiles — 256×128 and 128×256 — are compiled with | ||
| ``StageCount<2>``, which requires ~99 KB of shared memory. On SM121 this is | ||
| at or beyond the per-block opt-in limit (``cudaDevAttrMaxSharedMemoryPerBlockOptin | ||
| ≈ 101376 bytes``), so the CUTLASS ``initialize()`` call returns | ||
| ``kErrorInternal`` for those tactics. | ||
|
|
||
| Rather than letting the autotuner discover this noisily (and log confusing | ||
| "[MXFP8 SM120 gemm Runner] Failed to initialize …" messages), we pre-filter | ||
| here with a device-capability probe so the autotuner only sees tactics that | ||
| will actually succeed on this GPU. | ||
|
|
||
| The probe uses square M = N = K = 128 tensors; initialization failure is | ||
| shape-independent (it depends only on CTA tile size and device smem capacity), | ||
| so any valid shapes reach the same CUTLASS ``initialize()`` decision. | ||
| """ | ||
| dev = "cuda" | ||
|
|
||
| def _pad_up(x: int, m: int) -> int: | ||
| return ((x + m - 1) // m) * m | ||
|
|
||
| M = N = K = 128 | ||
| sf_vec = 32 | ||
| k_scales = (K + sf_vec - 1) // sf_vec | ||
|
|
||
| a = torch.zeros(M, K, dtype=torch.float8_e4m3fn, device=dev) | ||
| # b as [N, K] contiguous – the C++ binding expects mat2 with shape [N, K]. | ||
| b = torch.zeros(N, K, dtype=torch.float8_e4m3fn, device=dev) | ||
| sfa = torch.zeros(_pad_up(M, 128) * _pad_up(k_scales, 4), dtype=torch.uint8, device=dev) | ||
| sfb = torch.zeros(_pad_up(N, 128) * _pad_up(k_scales, 4), dtype=torch.uint8, device=dev) | ||
| out = torch.zeros(M, N, dtype=torch.bfloat16, device=dev) | ||
| ws = torch.zeros(8 * 1024 * 1024, dtype=torch.uint8, device=dev) | ||
|
|
||
| valid: List[int] = [] | ||
| for t in range(module.mxfp8_gemm_tactic_num()): | ||
| try: | ||
| module.mxfp8_gemm(a, b, sfa, sfb, out, ws, t) | ||
| valid.append(t) | ||
| except RuntimeError: | ||
| device_name = torch.cuda.get_device_properties(0).name | ||
| logger.debug( | ||
| "mxfp8_gemm tactic %d cannot initialize on %s; skipping. " | ||
| "(Tip: SM121 GB10/DGX Spark has ~99 KB shared memory opt-in per SM; " | ||
| "256×128 and 128×256 CTA tiles require ~99 KB with StageCount<2> " | ||
| "and may fail on that device.)", | ||
| t, | ||
| device_name, | ||
| ) | ||
| return valid |
Contributor
There was a problem hiding this comment.
In multi-GPU environments (especially heterogeneous setups with different GPU models), querying device properties using a hardcoded index 0 can lead to incorrect logging and diagnostics if the active device is not 0. Additionally, the probe should run under the context of the target device to ensure allocations and operations are performed on the correct GPU.\n\nWe can make _probe_mxfp8_gemm_tactics accept an optional device_id (defaulting to the current device) and execute the probe within a torch.cuda.device context manager.
def _probe_mxfp8_gemm_tactics(module, device_id: Optional[int] = None) -> List[int]:\n \"\"\"Probe which CUTLASS MXFP8 GEMM tactics can initialize on the current device.\n\n SM120 (B100/B200 data-centre, ~256 KB shared memory/SM) and SM121 (GB10 /\n DGX Spark GeForce Blackwell, ~99 KB shared memory/SM opt-in) share the same\n compiled module. Large CTA tiles — 256×128 and 128×256 — are compiled with\n ``StageCount<2>``, which requires ~99 KB of shared memory. On SM121 this is\n at or beyond the per-block opt-in limit (``cudaDevAttrMaxSharedMemoryPerBlockOptin\n ≈ 101376 bytes``), so the CUTLASS ``initialize()`` call returns\n ``kErrorInternal`` for those tactics.\n\n Rather than letting the autotuner discover this noisily (and log confusing\n \"[MXFP8 SM120 gemm Runner] Failed to initialize …\" messages), we pre-filter\n here with a device-capability probe so the autotuner only sees tactics that\n will actually succeed on this GPU.\n\n The probe uses square M = N = K = 128 tensors; initialization failure is\n shape-independent (it depends only on CTA tile size and device smem capacity),\n so any valid shapes reach the same CUTLASS ``initialize()`` decision.\n \"\"\"\n if device_id is None:\n device_id = torch.cuda.current_device()\n dev = torch.device(\"cuda\\", device_id)\n\n def _pad_up(x: int, m: int) -> int:\n return ((x + m - 1) // m) * m\n\n M = N = K = 128\n sf_vec = 32\n k_scales = (K + sf_vec - 1) // sf_vec\n\n with torch.cuda.device(device_id):\n a = torch.zeros(M, K, dtype=torch.float8_e4m3fn, device=dev)\n # b as [N, K] contiguous – the C++ binding expects mat2 with shape [N, K].\n b = torch.zeros(N, K, dtype=torch.float8_e4m3fn, device=dev)\n sfa = torch.zeros(_pad_up(M, 128) * _pad_up(k_scales, 4), dtype=torch.uint8, device=dev)\n sfb = torch.zeros(_pad_up(N, 128) * _pad_up(k_scales, 4), dtype=torch.uint8, device=dev)\n out = torch.zeros(M, N, dtype=torch.bfloat16, device=dev)\n ws = torch.zeros(8 * 1024 * 1024, dtype=torch.uint8, device=dev)\n\n valid: List[int] = []\n for t in range(module.mxfp8_gemm_tactic_num()):\n try:\n module.mxfp8_gemm(a, b, sfa, sfb, out, ws, t)\n valid.append(t)\n except RuntimeError:\n device_name = torch.cuda.get_device_properties(device_id).name\n logger.debug(\n \"mxfp8_gemm tactic %d cannot initialize on %s; skipping. \"\n \"(Tip: SM121 GB10/DGX Spark has ~99 KB shared memory opt-in per SM; \"\n \"256×128 and 128×256 CTA tiles require ~99 KB with StageCount<2> \"\n \"and may fail on that device.)\",\n t,\n device_name,\n )\n return valid
Comment on lines
+4066
to
+4079
| # One-time probe: cache valid tactics at closure scope so all runner | ||
| # instances for this module share the result without re-probing. | ||
| _valid_tactics_cache: List[int] = [] | ||
|
|
||
| def cutlass_mxfp8_gemm_runner(): | ||
| class CutlassMxfp8GemmRunner(TunableRunner): | ||
| def get_valid_tactics( | ||
| self, | ||
| inputs: List[torch.Tensor], | ||
| profile: OptimizationProfile, | ||
| ) -> List[int]: | ||
| return list(range(module.mxfp8_gemm_tactic_num())) | ||
| if not _valid_tactics_cache: | ||
| _valid_tactics_cache.extend(_probe_mxfp8_gemm_tactics(module)) | ||
| return list(_valid_tactics_cache) |
Contributor
There was a problem hiding this comment.
Caching the probed tactics in a single list _valid_tactics_cache at the closure scope will cause issues in multi-GPU environments with heterogeneous GPUs (e.g., a system with both SM120 and SM121 GPUs). If one thread probes on an SM120 GPU and populates the cache, a subsequent thread running on an SM121 GPU will reuse the cached SM120 tactics and crash when attempting to run unsupported large CTA tiles.\n\nWe should cache the probed tactics per device_id using a dictionary.
Suggested change
| # One-time probe: cache valid tactics at closure scope so all runner | |
| # instances for this module share the result without re-probing. | |
| _valid_tactics_cache: List[int] = [] | |
| def cutlass_mxfp8_gemm_runner(): | |
| class CutlassMxfp8GemmRunner(TunableRunner): | |
| def get_valid_tactics( | |
| self, | |
| inputs: List[torch.Tensor], | |
| profile: OptimizationProfile, | |
| ) -> List[int]: | |
| return list(range(module.mxfp8_gemm_tactic_num())) | |
| if not _valid_tactics_cache: | |
| _valid_tactics_cache.extend(_probe_mxfp8_gemm_tactics(module)) | |
| return list(_valid_tactics_cache) | |
| # One-time probe: cache valid tactics at closure scope per device ID so all runner\n # instances for this module share the result without re-probing, while supporting multi-GPU.\n _valid_tactics_cache: dict = {}\n\n def cutlass_mxfp8_gemm_runner():\n class CutlassMxfp8GemmRunner(TunableRunner):\n def get_valid_tactics(\n self,\n inputs: List[torch.Tensor],\n profile: OptimizationProfile,\n ) -> List[int]:\n device_id = torch.cuda.current_device()\n if device_id not in _valid_tactics_cache:\n _valid_tactics_cache[device_id] = _probe_mxfp8_gemm_tactics(module, device_id)\n return list(_valid_tactics_cache[device_id]) |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
tests/gemm/test_mm_mxfp8_sm120.py (3)
142-142: 💤 Low valueOptional: Replace
×with*orxto address RUF003.The comment uses
×(MULTIPLICATION SIGN) to distinguish multiplication from tile dimensions (which use lowercasex). While this is readable, Ruff RUF003 flags it as potentially ambiguous. Consider using* 2 variantsorx 2 variantsinstead to avoid the warning. The current form is acceptable if you prefer the visual distinction.🤖 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_mm_mxfp8_sm120.py` at line 142, Replace the MULTIPLICATION SIGN in the comment "At least the three small-tile configs (128x32, 128x64, 128x128) × 2 variants" with a plain ASCII symbol to satisfy RUF003; for example change "× 2 variants" to "* 2 variants" or "x 2 variants" in the comment within tests/gemm/test_mm_mxfp8_sm120.py so the tile-dimension notation remains unambiguous.Source: Linters/SAST tools
154-157: 💤 Low valueClarify tactic descriptions in assertion messages.
The assertion messages identify tactics 6-9 as "(256x128 bf16)", "(256x128 half)", etc., but based on the PR description ("5 CTA shapes × 2 swap-AB variants"), tactics distinguish tile shapes and whether A/B are swapped, not output data types. Since
out_dtypeis a runtime parameter (line 44), consider clarifying these messages to describe the actual tactic properties, e.g., "Tactic 6 (256x128 non-swapped)" or just "Tactic 6 (256x128)".🤖 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_mm_mxfp8_sm120.py` around lines 154 - 157, The assertion messages for tactics 6–9 in tests/gemm/test_mm_mxfp8_sm120.py incorrectly mention data types; update the messages to reflect the actual tactic properties (tile shape and swap-AB) rather than out_dtype (which is a runtime parameter), e.g., change "Tactic 6 (256x128 bf16)" to "Tactic 6 (256x128 non-swapped)" or simply "Tactic 6 (256x128)"; apply the same pattern for tactics 7–9 (use "256x128 swapped" or "128x256 ..." as appropriate) so the messages match the tactic semantics referenced by the valid variable and tactic indices.
137-137: Standardize the import path forgen_gemm_sm120_module_cutlass_mxfp8
flashinfer/jit/gemm/__init__.pyre-exportsgen_gemm_sm120_module_cutlass_mxfp8from.core(and includes it in__all__), so importing fromflashinfer.jit.gemmvsflashinfer.jit.gemm.coreshould refer to the same function. Standardize onfrom flashinfer.jit.gemm import ...for consistency.🤖 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_mm_mxfp8_sm120.py` at line 137, The test currently imports gen_gemm_sm120_module_cutlass_mxfp8 from flashinfer.jit.gemm.core; change that import to use the package re-export by importing from flashinfer.jit.gemm instead (i.e., replace the from flashinfer.jit.gemm.core import gen_gemm_sm120_module_cutlass_mxfp8 line with from flashinfer.jit.gemm import gen_gemm_sm120_module_cutlass_mxfp8) so the test uses the standardized, re-exported symbol.
🤖 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 `@flashinfer/gemm/gemm_base.py`:
- Around line 4008-4059: The probe in _probe_mxfp8_gemm_tactics is pinned to the
default CUDA device and its memoized "valid tactics" list is shared across GPUs;
change it to probe and cache per-GPU (use torch.cuda.current_device() or device
index/name from torch.cuda.get_device_properties) so each device has its own
valid-tactics entry. Concretely, read the active device into a local variable
(instead of hardcoding dev="cuda"), create or use a global cache keyed by device
id/name, and store/lookup the valid List[int] in that per-device entry before
running the probe loop (refer to _probe_mxfp8_gemm_tactics, the dev variable,
and module.mxfp8_gemm / module.mxfp8_gemm_tactic_num to locate where to change
memoization). Ensure tensors (a, b, sfa, sfb, out, ws) are allocated on that
same device and that the exception-logging still uses the correct device
properties.
---
Nitpick comments:
In `@tests/gemm/test_mm_mxfp8_sm120.py`:
- Line 142: Replace the MULTIPLICATION SIGN in the comment "At least the three
small-tile configs (128x32, 128x64, 128x128) × 2 variants" with a plain ASCII
symbol to satisfy RUF003; for example change "× 2 variants" to "* 2 variants" or
"x 2 variants" in the comment within tests/gemm/test_mm_mxfp8_sm120.py so the
tile-dimension notation remains unambiguous.
- Around line 154-157: The assertion messages for tactics 6–9 in
tests/gemm/test_mm_mxfp8_sm120.py incorrectly mention data types; update the
messages to reflect the actual tactic properties (tile shape and swap-AB) rather
than out_dtype (which is a runtime parameter), e.g., change "Tactic 6 (256x128
bf16)" to "Tactic 6 (256x128 non-swapped)" or simply "Tactic 6 (256x128)"; apply
the same pattern for tactics 7–9 (use "256x128 swapped" or "128x256 ..." as
appropriate) so the messages match the tactic semantics referenced by the valid
variable and tactic indices.
- Line 137: The test currently imports gen_gemm_sm120_module_cutlass_mxfp8 from
flashinfer.jit.gemm.core; change that import to use the package re-export by
importing from flashinfer.jit.gemm instead (i.e., replace the from
flashinfer.jit.gemm.core import gen_gemm_sm120_module_cutlass_mxfp8 line with
from flashinfer.jit.gemm import gen_gemm_sm120_module_cutlass_mxfp8) so the test
uses the standardized, re-exported symbol.
🪄 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
Run ID: 06e2a28d-c051-4b2a-b8c8-40dd44d72889
📒 Files selected for processing (3)
flashinfer/gemm/gemm_base.pyinclude/flashinfer/gemm/mxfp8_gemm_template_sm120.htests/gemm/test_mm_mxfp8_sm120.py
…leanups - _probe_mxfp8_gemm_tactics now takes an explicit device_id parameter and probes tensors on that specific device rather than the default CUDA device - _valid_tactics_cache changed from List[int] to Dict[int, List[int]] so heterogeneous multi-GPU systems (e.g. SM120 + SM121) cache tactic sets independently per device - test: import gen_gemm_sm120_module_cutlass_mxfp8 from flashinfer.jit.gemm (package re-export) instead of the internal .core submodule - test: pass torch.cuda.current_device() to _probe_mxfp8_gemm_tactics - test: replace Unicode × with ASCII x in comment (RUF003) - test: fix assertion messages to say "non-swapped"/"swapped" instead of "bf16"/"half" — the swap-AB flag is what distinguishes tactic pairs, not the output dtype Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Member
|
/bot run tests/gemm |
Collaborator
Collaborator
|
[FAILED] Pipeline #54590688: 4/20 passed |
No functional change — line-wrapping only, to satisfy the ruff v0.12.8 and clang-format v19.1.1 pre-commit hooks. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013exd1CEhdnMQoLjhGAykRr
kahyunnam
pushed a commit
that referenced
this pull request
Jul 9, 2026
…n MXFP8xMXFP8) (#3614) ## Description Fixes the original-report half of #3558: the autotuner tuning pass crashes for `cutlass_fused_moe` with MXFP8 activation scaling (CUDA illegal instruction / null TMA SF buffer on the reporter's SM121 build; `!use_block_scaling` kernel assert on current main). **Root cause** (one level deeper than my scoping comment on the issue — `input_sf` is not even a tuning input on this path; the bug is entirely in the C++ gemm profiler): MXFP8×MXFP8 stores activations and weights as plain FP8, so the profiler cannot tell it apart from per-tensor FP8 by dtypes and prepared wrong tuning state in three coordinated places: 1. `GemmProfilerBackend::init` derived `mScalingType` from dtypes alone → `NONE` → `getOffsetActivationSF` sized the activation-SF workspace to **zero** → the TMA warp-specialized MXFP8 kernels were launched with a **null SF pointer** ("gmem address 0 / null TMA buffer" in the report). 2. `getProfilerWorkspaces` reserved only per-tensor float scalars for FP8 weights; MXFP8 needs per-expert weight block SFs + global scales. 3. `prepareQuantParams` fell into per-tensor `QuantParams::FP8`; MXFP8 needs `QuantParams::MXFP8MXFP8`. **Plus tactic hygiene at the source**: `CutlassMoeFCRunner::getTactics` now filters out non-TMA (SM80/SM89-style) fallback configs for MXFP8 instantiations — those paths hard-assert `!use_block_scaling`, so offering them as tuning tactics guarantees a crash on any arch where they are enumerated. This composes with future SM12x enablement (#3463-style): once `isValidSM120MOESpecialisation` admits fp8×fp8, the SM120 TMA configs flow through this filter untouched. The binding passes `mUseMxfp8ActScaling` into `GemmProfilerBackend::init` (new defaulted parameter — other call sites unaffected). ## Validation This box is SM120 (RTX PRO 6000), where MXFP8×MXFP8 is not an enabled arch on main (`test_moe_mxfp8_mxfp8` is SM100-gated), so the honest validation matrix is: - **Compile**: all changes JIT-compile (CUDA 13.0, sm_120). - **No-regression (pure FP8, the dtype-twin path)**: `cutlass_fused_moe` FP8 per-tensor with `autotune(True)` runs the full tuning pass and produces **bit-identical outputs** across {stock main, this PR} × {autotune on, off} (cos identical to 4+ digits on a 128×1024, 8-expert config). `test_moe_fp8` passes. - **Fail-fast improvement on SM120**: MXFP8 + autotune now fails at init with `No valid tactics available for fused moe op...` instead of the deep `Assertion failed: !use_block_scaling (cutlass_fused_moe_kernels.cuh:3153)` mid-tuning. - **E2E regression coverage on SM100**: added `use_autotune` parametrization to the SM100-gated `test_moe_mxfp8_mxfp8` — on a pre-fix tree the `True` cases crash the tuning pass; with this PR they should pass on SM100 CI. @tgmerritt — if you can ride this onto your #3463-based branch on the GB10, your original repro (vLLM `--moe-backend flashinfer_cutlass`, autotune enabled) is the definitive e2e check for the IMA half. The `kErrorInternal` tactic-init failures from your follow-up remain #3568's territory. ## Tests - [x] Tests have been added or updated as needed (`use_autotune` parametrization on `test_moe_mxfp8_mxfp8`) - [x] All tests are passing locally (pure-FP8 autotune regression + `test_moe_fp8` on SM120; MXFP8 cases are SM100-gated) ## Disclosure AI-assisted (Claude); reviewed, compiled, and validated by the submitter. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added MXFP8×MXFP8 quantization support for profiler workspace and parameter initialization. * Added an option to enable MXFP8 activation scaling that influences GEMM tactic selection. * Ensured deterministic activation block-scale values during profiler runs. * **Tests** * Extended MXFP8×MXFP8 tests to run both with and without autotuning to cover tuning and non-tuning paths. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
Collaborator
|
@yichengj0, can you verify the changes in this PR? |
Contributor
|
The GB10 failure is real and reproduces exactly as described. But we verified the PR on SM120 hardware as well (RTX 5080 and RTX PRO 6000 Blackwell), and the results change the picture:
Given that, the simplest fix might be to remove those four configs from the SM120 list entirely, rather than probe-filtering at runtime. If a runtime check is still preferred, filtering configs against the device limit on the C++ side (each config's shared-memory need is a compile-time constant) would avoid the probe's dummy tensors and kernel launches. Either way the test would expect [0..5] on both 12.0 and 12.1. A couple of smaller concerns with the probe as written, in case it stays:
|
5 tasks
bkryu
added a commit
that referenced
this pull request
Sep 3, 2026
…ed-memory limit (#4013) ## 📌 Description The issue: - Four of the ten SM120 MXFP8 CUTLASS tactics fail on every autotuner pass with `[MXFP8 SM120 gemm Runner] Failed to initialize cutlass MXFP8 gemm on sm120. Error: Error Internal`. - The failing tactics (the 256x128 and 128x256 tiles) need more shared memory per block than any SM12x GPU provides, so they can never run anywhere this SM120-only module runs. They are dead code. - Confirmed by measurement on RTX 5080, RTX PRO 6000 Blackwell (CC 12.0) and GB10 / DGX Spark (CC 12.1): the kernels request 111616 bytes per block; all three devices allow at most 101376. The fixes: - Remove the two oversized tile configs and their kernel instantiations. The module now advertises 6 tactics, all usable, and compiles four fewer kernels. - When kernel initialization fails because of shared memory, the error now reports the kernel's need next to the device limit. - Replace the hardcoded tactic-count test with one that runs every advertised tactic (bf16 and fp16 outputs) against a reference, so an unusable config fails loudly and a new valid tactic is covered without test edits. ## 🔍 Related Issues Addresses the MXFP8 item of #3170 (kept open; other items remain). Supersedes #3568; thanks @tgmerritt for the original diagnosis on GB10. Measuring the actual shared-memory requests showed the configs are unusable on all SM12x devices, so this PR removes them at the source instead of filtering them out at runtime. ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Verified on GB10: the full `tests/gemm/test_mm_mxfp8_sm120.py` suite passes and the module reports 6 tactics. The surviving tactics also pass the numeric suite on RTX 5080 and RTX PRO 6000 boards. ## Reviewer Notes If a future SM12x part ships with more shared memory and the large tiles become viable, the place to bring them back is `getConfigs()`, ideally with a static check of each tile's shared-memory need against the device limit. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Performance** * Reduced the SM120/SM121 MXFP8 GEMM CTA tile configurations to the supported set. * Updated SM120 GEMM execution to use a consistent shared-memory staging strategy. * **Bug Fixes** * Workspace estimation failures during SM120 MXFP8 GEMM probing now surface instead of being suppressed. * Added clearer diagnostics when kernel initialization exceeds the device’s shared-memory limit. * **Tests** * Added parameterized validation covering all advertised SM120 MXFP8 tactics for BF16 and FP16 outputs. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Brian K. Ryu <bryu@nvidia.com>
Collaborator
|
Thank you @tgmerritt, #4013 that supersedes this PR has been merged, so I am closing this one. |
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
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.
Problem
On SM121 (GB10 / DGX Spark GeForce Blackwell), 4 of the 10 SM120 MXFP8 CUTLASS tactics fail silently every model load:
Root cause: SM121 has ~99 KB shared-memory opt-in per SM (
cudaDevAttrMaxSharedMemoryPerBlockOptin ≈ 101376 bytes). The 256×128 and 128×256 CTA tiles in the SM120 module are compiled withStageCount<2>, which requires ~99 KB of shared memory — right at SM121's device limit. CUTLASSinitialize()callscudaFuncSetAttribute(..., cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)which fails, returningkErrorInternal.Investigation (NVIDIA GB10 / SM121)
Device:
shared_memory_per_block_optin = 101376 bytes. The failure is shape-independent — it happens at kernel initialization time regardless of M/N/K.Reported in #3558 (comment) and #3170.
Fix
flashinfer/gemm/gemm_base.py: Add_probe_mxfp8_gemm_tactics(module)that runs a cheap M=N=K=128 probe for each tactic (initialization failure is shape-independent). Result cached in closure scope.get_valid_tactics()calls this once and returns only supported tactics. On SM121:[0,1,2,3,4,5](silent).include/flashinfer/gemm/mxfp8_gemm_template_sm120.h: Wheninitialize()fails withkErrorInternal, query device smem opt-in limit and include it in the exception message so the error is actionable.tests/gemm/test_mm_mxfp8_sm120.py: Addtest_probe_mxfp8_gemm_tactics_sm12xverifying ≥6 tactics on SM12x, exactly[0..5]on SM121, all 10 on SM120.Test result on GB10 (SM121)
Before: 4 tactics fail with "Error Internal" on every autotuner run
After:
_probe_mxfp8_gemm_tacticsreturns[0,1,2,3,4,5]silentlyCloses #3170 (partial)
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests