perf(gemm): update mm_fp4 b12x SM120 NVFP4 dense GEMM kernel - #3560
Merged
Merged
Conversation
…plan Update the vendored b12x SM120 block-scaled dense GEMM (used by mm_fp4(backend="b12x") and the SM12x MoE kernels) to the latest upstream b12x dense.py, and route the dense FP4 default path through upstream's regime-aware plan selection. Kernel (dense_blockscaled_gemm_sm120_b12x.py): - Replace the vendored kernel with upstream b12x dense.py (verbatim), repoint the utils import to flashinfer.cute_dsl.utils, and re-add the FlashInfer integration shims (@cute.jit `wrapper`, `Sm120B12xBlockScaledDenseGemmKernel` alias) that gemm_base / the MoE kernels import. - Backport 6 SF-helper methods (_thrfrg_SFA/B, _get_layoutSFA/B_TV, _partition_fragment_SFA/B) to their inline implementations: upstream delegates to newer cutlass.utils.blackwell_helpers APIs not present in the pinned nvidia-cutlass-dsl 4.5.2. These are exactly the methods the SM12x MoE kernels reuse as a mixin, so MoE keeps working unchanged. Runner (gemm_base.py): - Fix the ctor/can_implement calls for upstream's diverged signatures (ctor args by keyword since mma_k/tile_k were inserted before use_prefetch; can_implement is now M-independent). - Default (non-autotune) path now selects the tile via upstream's _select_default_dense_gemm_plan with expected_m=m, giving the regime-optimal tile (incl. the 16x64 M=1 decode tile). swap_ab is requested by the plan only for narrow-N tiles; since the transposed-store epilogue is not yet wired in this wrapper, the runner falls back to a safe non-swap tile in that case. utils.py: - Relax sm120_make_smem_layout_sfb's N divisibility 64 -> 16 (upstream commit 0daa6ab) so narrow-N tiles are representable once swap_ab is supported. Results on RTX 5080 (SM120), mm_fp4 nvfp4: correctness unchanged (pytest -k b12x 396/0, SM12x MoE passing); decode (m<=8) default path up to 1.15x, overall geomean ~1.00, no static regressions. Not enabled (follow-up): swap_ab narrow-N tiles + cp.async load_path require porting upstream's transposed-store epilogue / load-path handling into the FlashInfer wrapper (enabling swap_ab as-is produces incorrect results). That is the remaining FP4 upside. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The kernel was copied verbatim from upstream b12x, which dragged in a whole launch layer that FlashInfer never uses (it drives the kernel via _compile_block_scaled_gemm + the `wrapper` cute entry, not via dense_gemm). That dead code also (a) registered `b12x::dense_gemm_launch` torch custom ops, polluting the global `b12x::` operator namespace and colliding/shadowing with the real b12x package when both are imported in one process, and (b) kept `from b12x.* import ...`, making FlashInfer depend on the external b12x package being installed. Remove the unused launch layer (~1000 lines): the `b12x::dense_gemm_launch` and `..._functional` custom-op + fake registrations, `_DenseGemmLaunch`, `dense_gemm`, `_get_compiled_dense_gemm`, `_dense_gemm_launch_flat`, the `_DenseGemmPolicy`/`_max_active_clusters_for`/`_dense_gemm_policy_for` helpers, `_cached_alpha_one`, `_empty_dense_gemm_output`, and the Triton split-K finalize kernels. None are referenced by the FlashInfer runner, the SM12x MoE kernels, or the kept `_select_default_dense_gemm_plan` plan path. Drop the now-unused imports (triton, b12x.cute.compiler, b12x.cute.runtime_control) and vendor the 3 tiny `b12x.cute.fp4` helpers (get_ptr_as_int64, scatter_add_bf16, scatter_add_bf16x2 — used only by the opt-in split-K atomic path) so the file has NO `from b12x` imports. Result: `import flashinfer.gemm` no longer pulls in the b12x package (`b12x not in sys.modules`) and no longer registers `b12x::dense_gemm_launch`, so there is no global-namespace collision with the b12x package. mm_fp4 b12x correctness unchanged (pytest -k b12x: 396 passed, 0 failed). File shrinks 3655 -> 2648 lines. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
FlashInfer only drives the vendored b12x SM120 dense GEMM kernel for NVFP4 (mm_fp4 backend="b12x") and the SM12x MoE, never for MXFP8 dense GEMM. Remove the dead MXFP8 warp-MMA path so the file is FP4-only: - _setup_attributes: collapse the FP8/FP4 mma_op branch to MmaMXF4NVF4Op and drop the cute.nvgpu.warp.MmaMXF8Op reference (which does not exist in the public cutlass-dsl 4.5.2 build; it was only reachable via a const_expr branch that is always False for FP4). - can_implement: reject non-Float4E2M1FN up front; remove the FP8 tile and tile_k branches. - _select_default_mma_tiler_mn / _select_default_dense_gemm_plan: drop the is_mxfp8 parameter and both MXFP8 tile regimes. - gemm_base.py: drop the is_mxfp8=False argument at the plan call site. The surviving Float8E4M3FN references are the NVF4 scale-factor dtype, which is unrelated to the removed MXFP8 AB-dtype path. The SM12x MoE kernels borrow only the layout/SF helper methods, none of which touch the FP8 path, so they are unaffected. Verified on RTX 5080 (SM120): tests/gemm/test_mm_fp4.py -k b12x 396 passed / 0 failed; tests/moe/test_b12x_fused_moe.py 95 passed. AI-assisted with Claude Code. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Three changes to the SM120 NVFP4 b12x dense GEMM integration: 1. Enable the device-internal swap_ab path (narrow-N decode tiles, e.g. 64x32). The transposed-store epilogue was already vendored from upstream b12x; the prior garbage output was a C-layout wiring bug, not a missing epilogue. Fix: the wrapper reconstructs C row-major (m,n) unconditionally; the runner honors plan.swap_ab and builds the kernel with the real swap_ab (carried in cache_key), but passes swap_ab=False to the SHARED compile harness -- b12x's swap is device-internal, not the SM100 operand-swap / FFI-transpose convention -- and passes plain `out`. This leaves the shared harness and the SM100/SM103 swap_ab path untouched. Perf is ~neutral on RTX 5080 (m=1 decode is launch/latency- bound); the win upstream reports needs a larger SM120 part. 2. Guard unsupported K in _b12x_gemm_fp4_requirement: the NVFP4 kernel requires K % 128 == 0 (tile_k = sf_vec_size*8), but the default launch path does not call can_implement. Enforcing K at the backend gate makes b12x fall back to cutlass/cudnn under backend="auto" for unsupported K, and raise a clear error under explicit backend="b12x", instead of silently launching a kernel that drops the K-tail. 3. Cleanup (from /simplify): delete the now-dead _select_default_sm120_mma_tiler (the runner switched to _select_default_dense_gemm_plan) and drop the no-op kernel_m/kernel_n aliases. Verified on RTX 5080 (SM120): tests/gemm/test_mm_fp4.py -k b12x 396/0; swap_ab narrow-N probe cos_sim > 0.97; K-guard rejects explicit b12x and falls back under auto (cos_sim 0.99). Port audited vs upstream b12x and correctness-reviewed (only finding was the K-gap, fixed here). AI-assisted with Claude Code. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
get_valid_tactics previously offered only 4 generic swap_ab-free tiles, so enabling autotune could regress BELOW the untuned default: the deterministic plan's regime tiles -- notably the m=1 narrow-N (64,32)+swap_ab decode tile -- were absent from the candidate set, so for those shapes the tuner cached a generic tile and the swap tile became unreachable under tuning. Append the plan's pick to get_valid_tactics so the candidate set always contains the exact default-path tile: autotuning is now monotonic vs the static default (never worse) and the tuner picks up swap_ab on SKUs where it wins. Factor the plan computation into a shared _default_dense_plan() helper used by both forward()'s default path and get_valid_tactics(), making the "candidate set contains the default tile" invariant structural rather than a convention two call sites must keep in sync. Also add a TODO at _heuristic_func_mm_fp4: b12x is supported on SM121 (@supported_compute_capability([120, 121])) but `auto` only routes it at SM120; widen to major == 12 once the DGX Spark bench confirms b12x wins there. Verified on RTX 5080 (SM120): autotune path cos_sim > 0.97 incl. narrow-N swap; tests/gemm/test_mm_fp4.py -k b12x 396/0. Reviewed via /code-review (no correctness findings; the m=1-bucket-only property makes the swap tile safe). AI-assisted with Claude Code. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Remove 12 unused imports, rename unused k_tile loop vars to _k_tile, and add # type: ignore[call-overload] on the CuTe-DSL range(..., unroll=) calls (matching the sibling blackwell_sm12x/moe_*_kernel.py convention). ruff-format normalization of the vendored file. No behavior change (b12x refcheck passes). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
yichengj0
requested review from
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jimmyzho,
kahyunnam,
nv-yunzheq,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
June 10, 2026 04:25
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughRelax SM120 SMEM N-dimension divisibility and refactor the SM120 b12x FP4 GEMM path: add deterministic dense-plan selection, validate packed-K (real_k % 128), revise tactic enumeration/defaulting, and centralize swap_ab handling in the kernel constructor. ChangesSM120 b12x FP4 GEMM Refactoring
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Possibly related PRs
Suggested labels
Suggested reviewers
🚥 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 |
yichengj0
changed the title
![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 refactors the SM120 block-scaled FP4 GEMM implementation to support narrow-N FP4 tiles and swap_ab. It relaxes layout divisibility constraints, removes the old tiler selection function, and introduces a deterministic, regime-aware planning helper to handle tile selection and swap_ab. It also enforces a contraction dimension constraint (K must be a multiple of 128) in the requirements check. A review comment points out a potential NameError in the newly added requirements check because sf_vec_size is referenced in the error message but is not defined in that scope.
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.
![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: 2
🤖 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 5473-5483: The check computing real_k = a.shape[1] * 2 currently
raises for any real_k % 128 != 0 and thus blocks mm_fp4(..., backend="auto")
from falling back to CUTLASS/cuDNN; change the guard so it returns False for
non-explicit b12x selections and only raises when the caller explicitly
requested backend == "b12x". Concretely, after computing real_k, if real_k % 128
!= 0 then: if backend != "b12x" return False (so the fallback path can try other
backends), otherwise raise the existing ValueError with the same message; keep
references to real_k and a as-is to locate the code.
- Around line 5949-5958: Replace the zero-arg lambda assigned to make_kernel
with a proper local function definition to satisfy Ruff E731: define a def
make_kernel(): that returns Sm120B12xBlockScaledDenseGemmKernel(...) using the
same positional args sf_vec_size, mma_tiler_mn, cluster_shape_mn and the keyword
args use_prefetch=use_prefetch, enable_pdl=enable_pdl, swap_ab=swap_ab so the
binding remains identical to the original lambda.
🪄 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: c26c7ceb-7e8f-47f5-8ba3-53e113fe0c78
📒 Files selected for processing (3)
flashinfer/cute_dsl/utils.pyflashinfer/gemm/gemm_base.pyflashinfer/gemm/kernels/dense_blockscaled_gemm_sm120_b12x.py
Comment on lines
5949
to
5958
| # NOTE: upstream b12x ctor inserts mma_k/tile_k/single_work_tile_per_cta | ||
| # before use_prefetch, so pass these by keyword to avoid mis-binding. | ||
| make_kernel = lambda: Sm120B12xBlockScaledDenseGemmKernel( | ||
| sf_vec_size, | ||
| mma_tiler_mn, | ||
| cluster_shape_mn, | ||
| use_prefetch, | ||
| enable_pdl, | ||
| use_prefetch=use_prefetch, | ||
| enable_pdl=enable_pdl, | ||
| swap_ab=swap_ab, | ||
| ) |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# 1) Inspect the exact code around the reported lines
sed -n '5935,5975p' flashinfer/gemm/gemm_base.py | nl -ba | sed -n '1,120p'
# 2) Locate Ruff configuration and check whether E731 is enabled/selected
ls -la | head
rg -n --hidden --no-ignore-vcs "ruff\.toml|tool\.ruff|ruff:|\\bE731\\b|E731" -S . || true
rg -n --hidden --no-ignore-vcs "\\[tool\\.ruff\\]" -S pyproject.toml ruff.toml setup.cfg tox.ini .ruff.toml 2>/dev/null || true
# Print any relevant ruff config sections for quick inspection
for f in pyproject.toml ruff.toml .ruff.toml setup.cfg; do
if [ -f "$f" ]; then
echo "==== $f ===="
sed -n '1,200p' "$f"
sed -n '200,400p' "$f"
fi
doneRepository: flashinfer-ai/flashinfer
Length of output: 111
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# 1) Inspect the exact code around the reported lines (use cat -n instead of nl)
echo "---- flashinfer/gemm/gemm_base.py lines 5945-5965 ----"
sed -n '5945,5965p' flashinfer/gemm/gemm_base.py | cat -n
# 2) Locate Ruff configuration and check whether E731 is enabled/selected
echo "---- Ruff config candidate files ----"
ls -la | rg -n "ruff|pyproject|config" || true
echo "---- Search for E731 and ruff config ----"
rg -n --hidden --no-ignore-vcs "E731" -S . || true
rg -n --hidden --no-ignore-vcs "tool\.ruff|\\[tool\\.ruff\\]|\\bselect\\b|\\bignore\\b|E7[0-9]{2}" -S pyproject.toml ruff.toml .ruff.toml setup.cfg tox.ini 2>/dev/null || true
rg -n --hidden --no-ignore-vcs "\\[tool\\.ruff\\]" -S . || trueRepository: flashinfer-ai/flashinfer
Length of output: 2174
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "---- pyproject.toml [tool.ruff.lint] ----"
sed -n '80,190p' pyproject.toml | cat -n
echo "---- pyproject.toml per-file ignores ----"
rg -n "\\[tool\\.ruff\\.lint\\.per-file-ignores\\]" -n pyproject.toml || true
sed -n '190,260p' pyproject.toml | cat -n || true
echo "---- ruff.toml ----"
sed -n '1,120p' ruff.toml | cat -nRepository: flashinfer-ai/flashinfer
Length of output: 3249
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "---- Search ruff invocation in pre-commit ----"
rg -n "ruff" .pre-commit-config.yaml **/*.sh **/*ci* **/*github* 2>/dev/null || true
echo "---- Show .pre-commit ruff section if present ----"
python3 - <<'PY'
import re, pathlib
p = pathlib.Path(".pre-commit-config.yaml")
if not p.exists():
print("no .pre-commit-config.yaml")
raise SystemExit
txt = p.read_text()
for m in re.finditer(r"ruff[^\n]*", txt):
line = txt.rfind("\n", 0, m.start())+1
end = txt.find("\n", m.end())
print(txt[line:end])
PY
echo "---- ruff config used by any scripts (best-effort) ----"
rg -n --hidden --no-ignore-vcs "ruff(\.toml)?|--config|tool\.ruff" -S . || trueRepository: flashinfer-ai/flashinfer
Length of output: 1551
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "---- .pre-commit-config.yaml (ruff section) ----"
sed -n '35,90p' .pre-commit-config.yaml | cat -n
echo "---- Search GitHub workflows for ruff ----"
rg -n "ruff(-check|-format)?|ruff-pre-commit|config-file" .github -S 2>/dev/null || true
echo "---- Search for other ruff invocations ----"
rg -n "ruff-check|ruff .*--config-file|ruff .*--config|ruff-pre-commit" . -S --hidden --no-ignore-vcs || trueRepository: flashinfer-ai/flashinfer
Length of output: 1361
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "---- ruff version ----"
ruff --version || true
echo "---- ruff check (default config) with --select E731 ----"
ruff check --select E731 flashinfer/gemm/gemm_base.py || true
echo "---- ruff check with explicit --config=pyproject.toml ----"
ruff check --config pyproject.toml --select E731 flashinfer/gemm/gemm_base.py || true
echo "---- ruff check with explicit --config=ruff.toml ----"
ruff check --config ruff.toml --select E731 flashinfer/gemm/gemm_base.py || trueRepository: flashinfer-ai/flashinfer
Length of output: 10111
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '5940,6010p' flashinfer/gemm/gemm_base.py | cat -nRepository: flashinfer-ai/flashinfer
Length of output: 3081
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ruff check flashinfer/gemm/gemm_base.py || trueRepository: flashinfer-ai/flashinfer
Length of output: 16536
Fix Ruff E731 for make_kernel by replacing the zero-arg lambda with a local def (flashinfer/gemm/gemm_base.py:5951).
♻️ Proposed fix
- make_kernel = lambda: Sm120B12xBlockScaledDenseGemmKernel(
- sf_vec_size,
- mma_tiler_mn,
- cluster_shape_mn,
- use_prefetch=use_prefetch,
- enable_pdl=enable_pdl,
- swap_ab=swap_ab,
- )
+ def make_kernel():
+ return Sm120B12xBlockScaledDenseGemmKernel(
+ sf_vec_size,
+ mma_tiler_mn,
+ cluster_shape_mn,
+ use_prefetch=use_prefetch,
+ enable_pdl=enable_pdl,
+ swap_ab=swap_ab,
+ )📝 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.
Suggested change
| # NOTE: upstream b12x ctor inserts mma_k/tile_k/single_work_tile_per_cta | |
| # before use_prefetch, so pass these by keyword to avoid mis-binding. | |
| make_kernel = lambda: Sm120B12xBlockScaledDenseGemmKernel( | |
| sf_vec_size, | |
| mma_tiler_mn, | |
| cluster_shape_mn, | |
| use_prefetch, | |
| enable_pdl, | |
| use_prefetch=use_prefetch, | |
| enable_pdl=enable_pdl, | |
| swap_ab=swap_ab, | |
| ) | |
| # NOTE: upstream b12x ctor inserts mma_k/tile_k/single_work_tile_per_cta | |
| # before use_prefetch, so pass these by keyword to avoid mis-binding. | |
| def make_kernel(): | |
| return Sm120B12xBlockScaledDenseGemmKernel( | |
| sf_vec_size, | |
| mma_tiler_mn, | |
| cluster_shape_mn, | |
| use_prefetch=use_prefetch, | |
| enable_pdl=enable_pdl, | |
| swap_ab=swap_ab, | |
| ) |
🧰 Tools
🪛 Ruff (0.15.15)
[error] 5951-5958: Do not assign a lambda expression, use a def
Rewrite make_kernel as a def
(E731)
🤖 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/gemm/gemm_base.py` around lines 5949 - 5958, Replace the zero-arg
lambda assigned to make_kernel with a proper local function definition to
satisfy Ruff E731: define a def make_kernel(): that returns
Sm120B12xBlockScaledDenseGemmKernel(...) using the same positional args
sf_vec_size, mma_tiler_mn, cluster_shape_mn and the keyword args
use_prefetch=use_prefetch, enable_pdl=enable_pdl, swap_ab=swap_ab so the binding
remains identical to the original lambda.
Source: Linters/SAST tools
bkryu
reviewed
Jun 10, 2026
Comment on lines
+104
to
+113
| _B12X_TIMING = ( | ||
| os.getenv("B12X_TIMING", "0") == "1" or os.getenv("VLLM_B12X_TIMING", "0") == "1" | ||
| ) | ||
| _B12X_TIMING_THRESHOLD_MS = float( | ||
| os.getenv( | ||
| "B12X_TIMING_THRESHOLD_MS", | ||
| os.getenv("VLLM_B12X_TIMING_THRESHOLD_MS", "0"), | ||
| ) | ||
| ) | ||
| _B12X_DENSE_SPLITK_TURBO = os.getenv("B12X_DENSE_SPLITK_TURBO", "0") == "1" |
Collaborator
There was a problem hiding this comment.
Artifact from b12x, can be removed
Contributor
Author
There was a problem hiding this comment.
Removed — dropped the unused _B12X_TIMING / _B12X_TIMING_THRESHOLD_MS / _B12X_DENSE_SPLITK_TURBO globals and the now-unused os import.
Comment on lines
+588
to
+589
| # Relaxed 64 -> 16 (upstream b12x commit 0daa6ab) to allow narrow-N FP4 tiles | ||
| # (64x32 / 64x16) used with swap_ab. |
Collaborator
There was a problem hiding this comment.
Not really a necessary comment
Comment on lines
+5814
to
+5817
| # Single source of truth for the deterministic default tile / swap_ab pick. | ||
| # Used by BOTH forward()'s default path AND get_valid_tactics(), so the | ||
| # autotuner candidate set always contains the exact default-path tile -- | ||
| # i.e. tuning is monotonic vs the static default (it can't pick worse). |
Comment on lines
+5871
to
+5882
| # Generic swap_ab-free candidate set the autotuner profiles. Kept minimal: | ||
| # a larger tile grid made the per-bucket pick noisier (overfits the bucket | ||
| # representative shape), so we only add a few balanced tiles here. | ||
| for mma_tiler_mn in [(64, 64), (64, 128), (128, 64), (128, 128)]: | ||
| _add(mma_tiler_mn, swap_ab=False) | ||
|
|
||
| # Also expose the deterministic plan's regime-optimal pick for this | ||
| # (bucketed) shape -- the same tile forward() uses on the default path, | ||
| # which may be a narrow-N swap_ab tile (e.g. (64,32) at m=1) absent from | ||
| # the generic set above. Including it keeps autotuning monotonic: the | ||
| # tuner can never select something worse than the static default, and it | ||
| # picks up swap_ab on SKUs where it actually wins. |
Collaborator
There was a problem hiding this comment.
Maybe make the comments in these lines shorter for cleanliness
Comment on lines
+5905
to
+5911
| # Deterministic, regime-aware plan from upstream b12x: picks the | ||
| # decode/prefill-optimal tile (via expected_m) and swap_ab for | ||
| # narrow-N, instead of the M-independent default tile. For a single | ||
| # call the actual m IS the representative regime, so expected_m=m. | ||
| plan = _default_dense_plan(m, n, real_k, a.device) | ||
| # Honor the plan's swap_ab (narrow-N decode tiles, e.g. 64x32). | ||
| # b12x swap_ab is device-internal; see the launch + wrapper notes. |
Comment on lines
+5960
to
+5964
| # b12x swap_ab is device-internal (applied via the kernel ctor); the | ||
| # public C stays row-major (m, n). The shared harness's `swap_ab` instead | ||
| # selects the SM100 operand-swap FFI convention (C declared (n, m)), which | ||
| # b12x must not use -- so pass swap_ab=False (keeps c_fake (m, n) + plain | ||
| # `out`). cache_key still carries the real swap_ab (separate caching). |
Contributor
Author
There was a problem hiding this comment.
Shortened — kept the one-line note on why swap_ab=False is passed to the shared harness, since it guards against an incorrect 'fix'. Happy to delete it entirely if you'd prefer.
Comment on lines
+6038
to
+6045
| # TODO(sm121/DGX Spark): b12x is *supported* on SM121 (requirement is | ||
| # @supported_compute_capability([120, 121])) but `auto` only routes it at | ||
| # SM120 here, so on a Spark `backend="auto"` falls back to cutlass/cudnn and | ||
| # never picks b12x. To enable it, widen this to `major == 12` (covers 120+121). | ||
| # Gate the change on the Spark bench: only prefer b12x if it actually beats | ||
| # cutlass/cudnn on GB10. Also verify the kernel's hardcoded | ||
| # get_smem_capacity_in_bytes("sm_120") assumption holds on SM121 (run the | ||
| # e2e handoff's correctness section on the real Spark first). |
Collaborator
There was a problem hiding this comment.
b12x dense GEMM should be runnable on Spark. Also make b12x the preferred backend on Spark. Then remove the comment
Contributor
Author
There was a problem hiding this comment.
Per our discussion, keeping cutlass/cudnn as the SM121 default (cross-SKU benchmarks show b12x isn't faster on GB10 in most cases). Replaced the TODO with a concise note; b12x stays available as an explicit backend on SM121.
| @@ -6015,6 +6035,14 @@ def _heuristic_func_mm_fp4( | |||
| is_sm120 = major == 12 and minor == 0 | |||
Collaborator
There was a problem hiding this comment.
Maybe make this is_sm12x and then just check for major == 12 to address the Spark heuristic comment below
Contributor
Author
There was a problem hiding this comment.
Left is_sm120 as-is per the SM121 decision above — not widening auto to prefer b12x on SM121.
…env-var globals - _b12x_gemm_fp4_requirement: under backend='auto', return False on K%128!=0 (fall back to cutlass/cudnn) instead of raising; explicit backend='b12x' still raises. (per @bkryu / coderabbit) - Trim/condense verbose comments in gemm_base.py and drop 'upstream b12x' references; remove the unnecessary N-divisibility comment in cute_dsl/utils.py. - Remove the dead _B12X_TIMING / _B12X_TIMING_THRESHOLD_MS / _B12X_DENSE_SPLITK_TURBO env-var globals (defined-but-unused artifacts of the stripped b12x timing/launch layer) and the now-unused os import. (per @bkryu) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
yichengj0
force-pushed
the
b12x-new-fp4-gemm
branch
3 times, most recently
from
June 10, 2026 21:15
a8aaef8 to
b2698b2
Compare
Per reviewer: do not make b12x the default backend on SM121 (GB10) -- cutlass/ cudnn are faster there (cross-SKU benchmarks). Replace the SM121 TODO with a concise note (SM121 intentionally excluded; b12x stays an explicit backend) and drop the unnecessary 'warp-level MMA, underfill tile selection' justification on the SM120 line. `auto` behavior unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
yichengj0
force-pushed
the
b12x-new-fp4-gemm
branch
from
June 10, 2026 21:19
b2698b2 to
091b926
Compare
Contributor
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)
flashinfer/gemm/gemm_base.py (1)
5464-5472:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPreserve
backend="auto"fallback for the other b12x-only preconditions.These guards still raise on CUDA
< 13,use_8x4_sf_layout=True, oruse_nvfp4=Falseeven whenbackend="auto". That short-circuits auto-routing before CUTLASS/cuDNN can be considered, unlike thereal_k % 128check just below. ReturnFalsefor non-explicit b12x selection and only raise whenbackend == "b12x".🐛 Proposed fix
# b12x backend requires CUDA 13+, 128x4 scale factor layout, and NVFP4 only. if get_cuda_version().major < 13: + if backend != "b12x": + return False raise ValueError( "b12x FP4 GEMM requires CUDA 13 or later. " f"Current CUDA version: {get_cuda_version()}." ) if use_8x4_sf_layout: + if backend != "b12x": + return False raise ValueError("b12x FP4 GEMM only supports 128x4 scale factor layout.") if not use_nvfp4: + if backend != "b12x": + return False raise ValueError("b12x FP4 GEMM only supports NVFP4 (sf_vec_size=16).")🤖 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/gemm/gemm_base.py` around lines 5464 - 5472, The current checks using get_cuda_version().major, use_8x4_sf_layout, and use_nvfp4 unconditionally raise ValueError and therefore block backend="auto" from falling back; modify those guards inside the b12x selection logic so they only raise when backend == "b12x" and otherwise return False (mirror the behavior of the real_k % 128 check). Concretely, update the branches that reference get_cuda_version(), use_8x4_sf_layout, and use_nvfp4 so they first test backend == "b12x": if backend == "b12x" then raise the ValueError with the same message, else return False. Ensure you keep the exact messages and only change control flow for functions/variables get_cuda_version, use_8x4_sf_layout, use_nvfp4, and backend.
🤖 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 `@flashinfer/gemm/gemm_base.py`:
- Around line 5464-5472: The current checks using get_cuda_version().major,
use_8x4_sf_layout, and use_nvfp4 unconditionally raise ValueError and therefore
block backend="auto" from falling back; modify those guards inside the b12x
selection logic so they only raise when backend == "b12x" and otherwise return
False (mirror the behavior of the real_k % 128 check). Concretely, update the
branches that reference get_cuda_version(), use_8x4_sf_layout, and use_nvfp4 so
they first test backend == "b12x": if backend == "b12x" then raise the
ValueError with the same message, else return False. Ensure you keep the exact
messages and only change control flow for functions/variables get_cuda_version,
use_8x4_sf_layout, use_nvfp4, and backend.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a854caeb-f76f-4f44-b401-51b72a214a7a
📥 Commits
Reviewing files that changed from the base of the PR and between b2698b223de6311a2a92517b1b61e10335ae3a83 and 091b926.
📒 Files selected for processing (1)
flashinfer/gemm/gemm_base.py
Collaborator
|
/bot run |
Collaborator
bkryu
reviewed
Jun 10, 2026
bkryu
left a comment
bkryu
left a comment
Collaborator
There was a problem hiding this comment.
Changes LGTM. Will approve once the CI comes back clean
bkryu
approved these changes
Jun 11, 2026
yichengj0
added a commit
to yichengj0/flashinfer
that referenced
this pull request
Jun 15, 2026
flashinfer-ai#3560 added a b12x precheck rejecting K not a multiple of tile_k=128, which broke explicit backend="b12x" for shapes like K=192 (packed_k=96). This is the vLLM kernels_spark_fp4_flashinfer_nvfp4 regression on DGX Spark / SM121. The kernel's true floor is K % 32 == 0: A is K-major packed FP4 (K/2 bytes) loaded via TMA with assumed_align=16. K need not divide tile_k=128 because the mainloop predicates the partial K-tile and the swizzled SF layout zero-pads scale-groups, so ragged K computes correctly (verified on GB10: cos_sim 0.98-0.99 for K in {64,96,160,192,224,320}; K%32==16 yields garbage). Relax both the gemm_base precheck and the kernel can_implement guard to % 32 so the non-autotune (default-plan) and autotune (candidate-enumeration) paths accept the same shapes. Add ragged-K regression tests and a misaligned-K rejection test. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
bkryu
pushed a commit
that referenced
this pull request
Jun 16, 2026
…#3646) ## Problem PR #3560 added a precheck in `_b12x_gemm_fp4_requirement` that rejects `K` not a multiple of `tile_k = sf_vec_size * 8 = 128`. For explicit `backend="b12x"` this `raise`s, breaking shapes like **K=192** (`packed_k=96`): ``` ValueError: b12x FP4 GEMM requires the contraction dim K to be a multiple of 128 (tile_k = sf_vec_size * 8). Got K=192. ``` These shapes ran before #3560. ## Root cause `K % 128` is the *full MMA K-tile* condition, not the kernel's real correctness floor. The actual floor is **`K % 32 == 0`**: - `A` is K-major packed FP4, `(M, K/2)` bytes, loaded via TMA with `assumed_align=16`, so `K/2 % 16 == 0` → `K % 32 == 0`. - `K` need **not** divide `tile_k=128`: the mainloop predicates the partial final K-tile and the swizzled SF layout zero-pads scale-groups, so ragged K (96, 160, 192, 224, …) computes correctly. Verified on GB10 (SM121): `cos_sim` 0.98–0.99 for K ∈ {64, 96, 160, 192, 224, 320} across M ∈ {1, 8, 256}, narrow-N (swap_ab), and both autotune paths; `K % 32 == 16` (e.g. 112, 176) yields garbage and is correctly rejected. ## Fix Relax both K guards to `% 32`: - `_b12x_gemm_fp4_requirement` precheck (gates the non-autotune default-plan path). - `DenseGemmKernel.can_implement` (gates the autotune candidate enumeration). Both sites move together so the autotune and non-autotune paths accept the same shapes. ## Tests `tests/gemm/test_mm_fp4.py`: - `test_mm_fp4_b12x_ragged_k` — K ∈ {96, 192} × `auto_tuning` ∈ {False, True} (both `auto_tuning` values exercise the distinct guard paths). - `test_mm_fp4_b12x_misaligned_k_raises` — K=112 must reject cleanly. All new tests pass; the 396 existing b12x NVFP4 cases (K=128/256/512) remain green on SM121. ## Note for reviewers #3560's precheck mirrored the b12x donor's `can_implement` contract (which hard-`raise`s on `k % tile_k`). This PR intentionally **extends past** that contract to `% 32`, justified by the correctness data above — the kernel's mainloop already handles ragged K via partial-tile predication. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **New Features** * Relaxed the b12x FP4 GEMM K-dimension alignment requirement from multiples of 128 to multiples of 32 to improve support for more matrix shapes. * **Tests** * Added regression coverage for ragged K values across auto-tuning settings. * Added a test ensuring misaligned K values are rejected with a clear “multiple of 32” error (under supported GPU/CUDA conditions). <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
16 tasks
YangXu1990uiuc
added a commit
that referenced
this pull request
Jul 15, 2026
…out (breadth -> unified fuzzers) GEMM (fuzzer default-on since #3539): bmm_fp8 3456->6 cases (~71 min/leg), mm_fp4 23760->12 + 4 auto (~59 min on Blackwell legs), mm_bf16 7560->13, mm_mxfp8 2688->6 (+large-dim 400->4, stats 18->3), bmm_mxfp8 576->5, bmm_bf16 240->8. Every kept case verified to actually run somewhere (none self-skip into no-ops). Kept in full: error-path / cache-behavior / invariant tests, the #3560 ragged-K anchor, and non-fuzzed paths (8x4 SF layout, bias/pdl epilogues, trtllm weight shuffle). MoE (rides the #3958 fuzzer default-on flip, assumes #3892 routing axes): shape fan-out compressed, the quant x routing x weight-layout kernel-selection matrix kept in FULL -- non-NVFP4 quant numerics are not fuzzed yet, so only shapes were cut, never modes: - renormalize trio: shared axis constants -> boundary tokens [8,3072] + intermediate [1024,384] (trio 4608 -> ~1500 collected) - sigmoid 576->192; deepseekv3 27216->4032 (intermediate axis = hitting set of every routing config's compatible_intermediate_size, each config still runs) - routed-parity 3456->144 (routing x quant x packed/unpacked in full); per-token 4over6 108->16; fp4 tactic sweep 72->24 (odd-token anchor kept); cute-dsl accuracy 144->16 and 24->8 Untouched: error-path / OA-param / routing-replay / kernel-tier / LoRA / CUDA-graph / EP / tactic-pairing tests; #3595 / #3067 anchors. Every touched file carries a header directing future coverage to the fuzzers. AI-assisted (grid analysis + curation by Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Merged
5 tasks
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
3 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.
📌 Description
Summary
This PR refreshes the
b12xdense FP4 GEMM kernel to match the latest upstream for better performance.No public API change. The
autorouting heuristic (_heuristic_func_mm_fp4) is unchanged — b12x is preferred only at SM120, not SM121.Addresses #3517 [Perf][SM12x] Update mm_fp4(backend='b12x') kernels.
Motivation
@lukealonso has made performance improvements to b12x dense GEMMs; we'd like to port over the improvements to FlashInfer. Gains are focused on small-M.
Benchmarks: new b12x vs old b12x
Setup per SKU:
benchmarks/flashinfer_benchmark.py --testlist <104 nvfp4 shapes> --use_nvfp4 --use_128x4_sf_layout --refcheck, 30 iters / 5 warmup, CUPTI timing (harness default). 104 shapes = 13 token counts m∈{1..4096} × 8 (n,k) layer shapes taken from real DeepSeek-R1 and Llama-3 model layers, all K%128==0.The refresh is a consistent small win at small
m(whereswap_abandexpected_mhelp most) and break-even-to-positive overall — no SKU regresses in aggregate.Geomean speedup (old_time / new_time; >1 = new faster), grouped by
m:The win is largest at small
mon PRO 6000 / Spark and break-even on the 5080. A fewn=4096 k=4096/5376shapes regress ≤5% (bolded in the per-shape tables); the PRO 6000 has a sharper8192×2560regression at mediumm(down to 0.873×), which is SKU-specific and flagged as a follow-up.Full per-shape new-vs-old tables (all 104 shapes, sorted by
n, k, m) per SKU:RTX 5080 (SM120)
RTX PRO 6000 (SM120)
DGX Spark / GB10 (SM121)
🔍 Related Issues
Closes #3517.
🚀 Pull Request Checklist
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
tests/gemm/test_mm_fp4.py -k b12x)pytest -k b12x396/0; 104-shape--refcheck104/104 on 5080 / PRO 6000 / GB10.Reviewer Notes
cutlass-dsl 4.5.2(pure upstream b12x needs an internal/newer cutlass-dsl), and the MXFP8 path is pruned (unsupported there).cute_dsl/utils.pyis relaxed 64→16 (upstream commit0daa6ab) to admit the narrow-Nswap_abtiles. This is safe: it only loosens a precondition (every previously-valid N%64==0 tile still passes), the SFB smem layout already rounds N up to a full 128-wide block so the layout is unchanged for existing tiles, and the newly-admitted narrow tiles pass--refcheck.swap_abhere is b12x's device-internal transpose (public C stays row-major) — not the SM100 operand-swap/FFI-transpose convention. The shared_compile_block_scaled_gemmharness used by Sm100/Sm103 is left untouched.Summary by CodeRabbit