feat: CuTe DSL kernels for Rubin (SM107) and batched FP8 GEMM for Blackwell - #4526
Merged
Vinnie6167 merged 32 commits intoAug 19, 2026
Merged
Conversation
Adds CuTe DSL kernel support for NVIDIA Rubin (SM107), together with the Python CuTe DSL batched FP8 GEMM implementation the Rubin kernels build on. Batched FP8 GEMM (CuTe DSL) - new functionality on Blackwell as well as Rubin; FlashInfer previously had only a C++ bmm_fp8 path: - gemm/kernels/bmm_fp8_blackwell.py: persistent dense GEMM for SM100/SM103 - gemm/kernels/bmm_fp8_rubin.py: the SM107 variant - gemm/kernels/bmm_fp8_wrapper.py: dispatch and autotuning across both - gemm/kernels/epilogue_utils.py: shared epilogue helpers (scaled TMA store, scaled store, alpha-scaled store) Rubin fused-MoE kernels under fused_moe/cute_dsl/rubin/: - gather + grouped GEMM + SwiGLU fusion - grouped GEMM + finalize fusion - a cp.async producer / UMMA consumer pipeline, plus helpers Helpers shared by the Blackwell and Rubin MoE paths move to fused_moe/cute_dsl/common/kernel_utils.py. Rubin dense and grouped GEMM: - gemm/kernels/dense_blockscaled_gemm_sm107.py: blockscaled dense GEMM - gemm/kernels/grouped_gemm_masked_rubin.py: masked grouped GEMM - gemm/kernels/grouped_gemm_masked_wrapper.py: dispatch wrapper Integration: - Tactic generation and selection extended to SM107 (fused_moe/cute_dsl/tuner.py, gemm/gemm_base.py). SM107 encodes additional tactic parameters, so the autotuner's tactic tuple is arch-aware. - SM107 added to the CuTe DSL attention compatibility table (TMEM size 576). - Kernel launches updated for the CuTe DSL API, which now requires an explicit smem= argument. The CuTe DSL requirement stays a floor (>=4.7.0) rather than becoming a pin. The SM107 kernels need CuTe DSL >= 4.8 (cutlass.utils.rubin_helpers), so they are imported lazily behind is_rubin_cute_dsl_available(): on an older DSL the rest of FlashInfer works unchanged and only the SM107 CuTe DSL paths are unavailable, and users on CTK 13.4 with CuTe DSL 4.8+ get them automatically. The pyproject cu12/ cu13 extras are relaxed from ==4.7.0 to >=4.7.0 for the same reason - an exact pin there would have prevented installing 4.8 at all. Test coverage is extended in tests/gemm/test_bmm_fp8.py and tests/moe/test_cute_dsl_fused_moe.py.
`pre-commit run` reported failures from ruff-format, ruff-check (21) and mypy (18). All 39 lint/type errors were confined to three CuTe DSL kernel bodies: flashinfer/gemm/kernels/dense_blockscaled_gemm_sm107.py flashinfer/fused_moe/cute_dsl/rubin/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py flashinfer/fused_moe/cute_dsl/rubin/blockscaled_contiguous_grouped_gemm_finalize_fusion.py Those are excluded rather than rewritten, matching how the repo already treats equivalent files (cute_dsl/attention/fmha, moe_ep/kernel_src/.../src). The existing comment in pyproject.toml gives the reason: reflowing or "simplifying" DSL code can change traced semantics - e.g. `x in (a, b)` yields a Python bool rather than the DSL Boolean predicate `x == a or x == b` produces. The reported rules (SIM210 `True if x else False`, B008 `cutlass.Int64(0)` defaults, B007 unused `cutlass.range` induction variables) are exactly that class of change, so "fixing" them in traced kernel bodies would be a behavioural risk for no gain. Exclusions are added in all three places the repo maintains them: the mypy `exclude` and ruff `extend-exclude` lists in pyproject.toml, and the mypy hook's `exclude` regex in .pre-commit-config.yaml (the hook passes filenames explicitly, which overrides the pyproject setting). ruff-format is applied to the seven non-excluded files it touched. The three excluded files are left byte-identical, since formatting them would defeat the exclusion. `pre-commit run --files <changed>` is now clean: 14 hooks, 0 failures.
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: Your plan includes up to 8 reviews per rolling hour; 7 remain after this review. 📝 WalkthroughWalkthroughThis change adds SM107/Rubin support across CuTe DSL attention, fused MoE, FP8 and block-scaled GEMM paths. It adds Rubin kernels, architecture-specific dispatch and tuning, shared utilities, backend registration, tests, and launch configuration updates. ChangesRubin SM107 CuTe DSL support
Estimated code review effort: 5 (Critical) | ~120 minutes Merge Risk: 🟠 High · up to The PR adds new SM107 and batched FP8 execution paths, but unresolved issues can silently select the wrong kernel, fail with unclear runtime errors, or process an extra tile using potentially uninitialized expert metadata; important tests are also bypassed. Merge should wait until these correctness and validation-path issues are fixed or explicitly accepted by the owners. Sequence Diagram(s)sequenceDiagram
participant API
participant ArchitectureDispatch
participant RubinWrapper
participant RubinKernel
API->>ArchitectureDispatch: request SM107 CuTe DSL operation
ArchitectureDispatch->>RubinWrapper: validate tensors and select tactic
RubinWrapper->>RubinKernel: compile or load cached kernel
RubinKernel-->>API: execute GEMM or fused MoE result
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
Contributor
Author
|
@flashinfer-bot run |
bkryu
reviewed
Aug 14, 2026
bkryu
left a comment
bkryu
left a comment
Collaborator
There was a problem hiding this comment.
Thanks @Vinnie6167, left some mechanical comments in a cursory first pass
Resolves conflicts with 35 upstream commits, principally the fused-MoE CuTe DSL rework that adds MXFP8 x MXFP4 support. The substantive upstream change is that the Blackwell kernels' can_implement signatures split `ab_dtype` into separate `a_dtype`/`b_dtype` (A and B may now differ). Our SM107 kernels already used the split form, so the resolution keeps our Rubin/Blackwell dispatch and adopts the new signature on the Blackwell side: - tuner.py: kept the `_is_rubin_tactic` branch; updated both Blackwell can_implement calls to a_dtype/b_dtype. - act_fusion / finalize_fusion: kept the `is_rubin` branches and the arch-specific tactic parameters (3D `mma_tiler` + `mma_inst_shape` for SM107 alongside Blackwell's `mma_tiler_mn`), unioned them with upstream's new dtype parameters, and renamed our `ab_dtype_cutlass` uses to `a_dtype_cutlass` / `b_dtype_cutlass` to match the auto-merged regions. - Cache keys keep the "sm107"/"sm100" discriminator and gain upstream's split dtypes. - __init__ docstring: unioned (upstream's dtype list, our Rubin mention). Also threads `mma_tiler`/`mma_inst_shape` through the `*_nvfp4` convenience wrappers. Upstream restructured those into back-compat shims that accept only `mma_tiler_mn`, so the Rubin tactic parameters had no path from fused_moe.py to the kernels; mypy caught this as an unexpected-keyword error rather than it surfacing at runtime.
Contributor
🚨 POTENTIAL BREAKING PUBLIC API CHANGE DETECTED 🚨Caution THIS PR APPEARS TO BREAK THE PUBLIC API. AUTHORS AND REVIEWERS: DO NOT MISS THIS. This is an advisory warning and does not gate merging. Confirm compatibility and provide a deprecation or migration path, or track the fix in a follow-up PR. 1 public API finding(s):
|
…fp8 fallback
Two review findings on the SM107 paths.
TGV: the arch gate had been widened to [100, 103, 107] without any
corresponding kernel support. tgv_gemm_cute_ext.py still requests shared memory
sized for sm_100 (`get_smem_capacity_in_bytes("sm_100")`), and every tgv case in
tests/gemm/test_mm_bf16.py fails on Rubin with
`CUDA_LAUNCH_INVALID_CONFIG: cudaErrorInvalidValue` - 432 of 432, confirmed on
VR200 hardware. A candidate fix that corrects the smem capacity for the running
arch was tried and did *not* resolve it (same 432 failures, same error), so the
cause is not yet understood. Revert the gate to [100, 103] until there is
working SM107 support to justify widening it: the decorator on
_tgv_gemm_requirement, the runtime _match_sm_version guard, and the docstring.
tests/gemm/test_tgv_gemm.py returns to its upstream contents.
bmm_fp8: the untuned SM107 path hardcoded tactic 0. That config is a 2-CTA
256x256 tile ("best for large problems"); _can_implement_config_sm107 rejects it
for small problems, and running an invalid 2-CTA config is an illegal
instruction (mma_tiler M>=256 is a hardware constraint). This is exactly what
the neighbouring get_valid_tactics comment warns about - "DO NOT return [0] as
fallback". Select a validated config via get_valid_configs instead, mirroring
the sm100 branch, and raise a clear error naming the problem shape when none is
valid rather than launching an invalid kernel.
The cute-dsl bmm_fp8 backend only exists for Rubin: the heuristic appends "cute-dsl_sm107" solely when _match_sm_version(["107"]) holds, and its own comment records that Blackwell "aligns with main, which has no cute-dsl bmm_fp8 backend". Two places advertised it more widely. _cute_dsl_bmm_fp8_requirement declared [100, 103, 107]. On sm100/sm103 an explicit backend="cute-dsl" therefore passed the capability check, reached fp8_gemm_sm100 with an empty runner list and died on `assert runners, "No suitable runners found"` instead of reporting an unsupported backend. Narrow it to [107]. tests/gemm/test_bmm_fp8.py skipped cute-dsl only when the major version was not 10, so every sm100/sm103 run exercised the backend and hit that assert - a bulk failure on Blackwell CI for a Rubin-only feature. Gate on the exact (10, 7) capability instead.
int(None) in the disk-cache kernel name: on SM107 the use_tma_store slot is repurposed to carry (inst_m, inst_n, inst_k, tiler_k, prefetch_dist), and get_valid_tactics enumerates prefetch_dist as (0, 2, None). The name builder mapped int() over the tuple, so every None-prefetch tactic raised TypeError as soon as the name was needed - a third of the SM107 tactic space was uncompilable. Render None as its own symbol; it means "auto" and compiles differently from 0, so it must stay distinct rather than collapse to it. NameError across the SM103/SM107 tactic blocks: batch_size, m_aligned and n_aligned were assigned inside `if sm_version in [103, 107] and Sm103Kernel is not None`, but the SM107 block that reads them is guarded independently on Sm107Kernel. The two kernels are imported under separate try/except ImportError and have different requirements - SM103 needs the internal cutlass-dsl wheel, SM107 needs rubin_helpers - so on Rubin with a public DSL >= 4.8 the SM103 import fails, the SM107 block runs, and mm_fp4 autotuning dies with NameError: n_aligned. Hoist the three above both guards.
The cute-dsl bmm_fp8 wrapper shipped an SM100 dispatch arm that could never have run: _compile_and_run_bmm_sm100 passed `batch` as a 13th positional argument to _get_compiled_bmm_sm100, which takes 12 parameters, so any call raised TypeError before compiling anything. Nothing reached it - the heuristic only ever appends "cute-dsl_sm107", no test or benchmark calls the public entry with arch="sm100", and the requirement decorator is now [107] - so the defect was invisible. The whole arm is new in this PR rather than pre-existing upstream code, so it is removed rather than repaired: shipping a broken, unexercised path invites someone to trust it, and a repair could not be validated here (the reviewer notes the fakes use a symbolic batch, so `batch` would also need to come out of the sm107 cache key). Removed: SM100_AUTOTUNE_CONFIGS, _DEFAULT_SM100_CONFIG, _get_compiled_bmm_sm100, _compile_and_run_bmm_sm100, _can_implement_config_sm100, get_valid_sm100_configs, and the sm100 tactic selector plus its bucket helpers in kernels/utils.py. The public bmm_fp8_cute_dsl entry point and the runner factory are now sm107-only, and auto-detect raises for any other architecture instead of silently choosing the broken path. bmm_fp8_blackwell.py is untouched: the SM107 kernel derives from it, and it remains the Blackwell CuTe DSL kernel this PR adds. Net -473 lines. pre-commit clean across the repo.
…7 import Trace registration: moving grouped_gemm_nt_masked into kernels/grouped_gemm_masked_wrapper.py dropped the trace= argument it carried at its previous home, leaving @flashinfer_api bare. That orphans grouped_gemm_nt_masked_trace in trace/templates/gemm.py, stops fi_trace() recording the API, and fails the registry inventory test, which AST-scans the modules listed in _TRACE_REGISTRATION_MODULES. Restore the decorator argument and repoint the registry entry at the wrapper - grouped_gemm_masked_blackwell no longer defines a decorated API, so the entry moves rather than being duplicated. SM107 import gating: flashinfer/gemm/__init__.py imported the SM107 masked-GEMM kernel in the same try/except as the Blackwell ones, guarded only by is_cute_dsl_available(). grouped_gemm_masked_rubin imports cutlass.utils.rubin_helpers at module scope, which requires CuTe DSL >= 4.8; on 4.7 - which requirements.txt permits - that import raises, the whole block is abandoned and _cute_dsl_kernels stays empty, so grouped_gemm_nt_masked, Sm100BlockScaledPersistentDenseGemmKernel and create_scale_factor_tensor disappear as well, despite importing cleanly. Gate the SM107 import on is_rubin_cute_dsl_available() and build the export list incrementally, matching kernels/__init__.py and the lazy-import contract requirements.txt documents.
…_size _cute_dsl_bmm_fp8_requirement was inserted between @supported_compute_capability([100, 103, 107]) and the function it decorated, so the decorator was transplanted onto the new function and _check_bmm_fp8_problem_size lost its capability gate entirely. Restore it. The cute-dsl requirement keeps the [107] gate it should have had all along, and _check_bmm_fp8_problem_size returns to upstream's [100, 103, 107].
Cleanup pass over the CuTe DSL files, plus one licensing fix. Restored the BSD-3-Clause header and the cutlass provenance line on fused_moe/cute_dsl/blackwell/utils.py. Moving helpers out to common/kernel_utils.py had dropped the second licence block; the file's header is now byte-identical to upstream again. This was the only comment with consequences beyond tidiness. Removed the standalone-runner scaffolding from two kernel files - the argparse CLIs, benchmark-file parsers, tensor factories, reference checks and `if __name__ == "__main__"` blocks carried over from the CuTe DSL examples these were ported from. Nothing imports them: the only names taken from those modules are the kernel classes themselves, which is verified above the cut. That is 1,312 lines, and it also cleared two of the mypy errors the files were excluded for; ruff's B007/B008/SIM210 remain, so the ruff exclusion stays while the mypy exclusion is kept for the rest. Also removed: the self-referential "Location: <this file>" docstring lines, the "UPSTREAM KERNEL CODE / Last synced" sync banners, the CLI "Example usage" blocks, and eight imports orphaned by the deletions above. Trimmed the mm_fp4 fallback comment to the design rationale - why the sm100 selector is the right one on sm107 - and dropped the end-to-end throughput tables and tuning-split measurements from it. requirements.txt is restored to upstream's `nvidia-cutlass-dsl==4.7.0` and the pyproject cu12/cu13 extras to `==4.7.0`, pending the separate discussion about how to constrain the DSL version. The SM107 kernels stay lazily imported behind is_rubin_cute_dsl_available(), so with that pin they are simply unavailable rather than breaking anything.
Vinnie6167
marked this pull request as ready for review
August 14, 2026 22:18
Vinnie6167
requested review from
Aneureka,
IwakuraRein,
aleozlx,
feih-nv,
jiahanc,
nv-yunzheq,
samuellees and
yzh119
as code owners
August 14, 2026 22:18
…-cute-dsl-kernels One conflict, in tests/gemm/test_bmm_fp8.py. main (flashinfer-ai#4511 and follow-ups) replaced the parametrize matrix with a curated _SMOKE_CASES list, moving randomized breadth to tests/gemm/test_unified_gemm_fuzz.py; this branch had added m=256 to that matrix so the SM107 cute-dsl backend had an executable shape. Resolved onto main's structure, re-adding the Rubin coverage it has no reason to carry: - a cute-dsl smoke case at b=1, m=256, n=10304, k=2688. Every entry in SM107_AUTOTUNE_CONFIGS is 2-CTA with mma_tiler M=256, so the CTA tile is 256x128 and the problem needs m >= 256 and n >= 128; main's existing cases (m=48, n=80) all fall below that floor. - the cute-dsl skip guards (SM107 only, 16-alignment, matching input dtypes, and the tile floor), mirroring the other backends' guards. - cute-dsl added to the auto_tuning allowlist.
Both sides merged upstream main independently (the GitHub 'Update branch' button produced 69cd870 while this branch merged 693fed4 locally). No conflicts; the local resolution of tests/gemm/test_bmm_fp8.py onto main's new _SMOKE_CASES structure is preserved, along with the cute-dsl smoke case, its SM107 guards, and the autotune allowlist entry.
Contributor
Author
|
@flashinfer-bot run |
Merged
5 tasks
Drops the base requirement from >=4.7.0a0 to >=4.6.2a0. 4.6.1 is still excluded; 4.6.2, 4.7.x and the 4.8 pre-releases all satisfy it. Safe on a DSL without Rubin support: the SM107 kernels are imported behind is_rubin_cute_dsl_available(), and the Arch.sm_107* references in the MLA decode paths resolve through a getattr fallback to sm_103f, so an older DSL keeps its pre-Rubin behaviour rather than raising AttributeError. Note the cu12/cu13 extras in pyproject.toml still carry >=4.7.0a0, so for anyone installing flashinfer[cu12] or [cu13] the two combine and the effective floor remains 4.7.0a0.
Contributor
Author
|
@flashinfer-bot run |
kahyunnam
added a commit
that referenced
this pull request
Aug 25, 2026
…(>=4.6.2a0) (#4715) ## Summary `requirements.txt` on `release-v0.6.18` requires `nvidia-cutlass-dsl>=4.7.0a0`, while `main` requires `>=4.6.2a0`. This lowers the release branch to match, making the two files byte-identical. The divergence is an artifact of the rc4 cherry-pick: it captured an intermediate head of #4526, which had only relaxed the previous `==4.7.0` hard pin to `>=4.7.0a0`. The head that actually merged to `main` (f092274) lowered the floor to `>=4.6.2a0`. This line is the only remaining difference between the two branches' `requirements.txt`. ## Why it matters `requirements.txt` is the base `install_requires` — `pyproject.toml` sets `dynamic = ["dependencies"]` with `dependencies = {file = ["requirements.txt"]}`. A 4.7-only floor there makes a plain `pip install` unsatisfiable alongside `quack-kernels` 0.6.4, which hard-pins `nvidia-cutlass-dsl==4.6.2` (the same conflict #4555 / #4556 worked around with `--no-deps` in CI). Nothing that currently gets 4.7 loses it: | Install path | DSL requirement | Changed? | | --- | --- | --- | | base (`requirements.txt`) | `>=4.6.2a0` | yes, was `>=4.7.0a0` | | `[cu12]` / `[cu13]` extras (`pyproject.toml`) | `>=4.7.0a0` | no — already matches `main` | | CI (`scripts/test_utils.sh`) | `>=4.7.0a0`, installed explicitly | no | Kernels that genuinely need the newer DSL are gated at runtime, not by this floor: `is_rubin_cute_dsl_available()` plus the arch probes now consulted by #4649 and #4710. On a 4.6.2 environment those paths deselect the backend and fall back instead of failing with `KeyError: 'sm_107a'` from inside the DSL. ## Test plan - [x] `requirements.txt` is byte-identical to `upstream/main` - [x] `pyproject.toml` cu12/cu13 extras unchanged and already equal to `main` - [x] pre-commit clean (including the `fix requirements.txt` hook, so no reordering needed) - [ ] CI on this branch — the DSL version CI resolves is unaffected, since `scripts/test_utils.sh` installs `>=4.7.0a0` explicitly Metadata-only change; no code paths touched.
Vinnie6167
added a commit
to Vinnie6167/flashinfer
that referenced
this pull request
Aug 25, 2026
…workaround Replaces the previous two commits' gating with an explicit guard, and restores moe_sort to its pre-flashinfer-ai#4526 form. The even-tile rounding and the tile-metadata zero-fill are two halves of one workaround: rounding the count up keeps 2-CTA clusters uniform at the barrier, and initializing the buffers is what makes the resulting padding tile safe, because the kernel bounds-checks the metadata against that same rounded count. Gating both on cluster_shape_m > 1 was correct but left a coupled pair of code paths that nothing executes -- and the coupling being invisible is what produced the bug in the first place. No reachable Rubin tactic uses a multi-CTA cluster: tile_size is restricted to 128, forcing mma_tiler_m == 128 and hence cluster_shape_m == 1. Enumerating the tactic list gives 4 entries, all cluster_shape_mn == (1, 1). So raise NotImplementedError naming both requirements. Re-enabling tile_size=256 then fails immediately with a pointer to what must be handled, rather than silently reading uninitialized memory. Note the rounding was NOT previously dead: it ran on every Rubin call, since the old condition was just `if is_rubin:` with no cluster check. On odd tile counts it produced a real padding tile that occupied a scheduler slot and ran an MMA which -- with mn_limit == 0 -- loaded and stored nothing. Removing it drops that work too, not just the two fill launches. moe_sort is now byte-identical to its pre-flashinfer-ai#4526 state. Validated on SM107 (GR100, CuTe DSL 4.8.0a0): tests/moe/test_cute_dsl_fused_moe.py gives 82 failed / 299 passed / 161 skipped both before and after, with identical FAILED node-id sets. The 82 are pre-existing on that image. AI-assisted: analysis and patch drafted with Claude Code.
Vinnie6167
added a commit
to Vinnie6167/flashinfer
that referenced
this pull request
Aug 27, 2026
This was referenced Sep 1, 2026
5 tasks
Vinnie6167
added a commit
to Vinnie6167/flashinfer
that referenced
this pull request
Sep 4, 2026
…ti-CTA tactics Review feedback. 1. Both finalize-fusion schedulers loaded tile_idx_to_expert_idx[tile_idx] before checking tile_idx < num_valid_tiles, while the neighbouring mn_limit load was already inside it. The value is discarded and the index stays inside the allocation, so this is not a wrong-answer bug, but with torch.empty it is an uninitialized read that initcheck would flag. Moved inside the guard at all five sites -- three in the Blackwell kernel, two in the Rubin one (tile_idx_to_group_idx). The gather/activation kernels already loaded inside the check on both architectures. Pre-existing: these buffers were torch.empty before flashinfer-ai#4526 too. The zero-fill that PR added masked it; removing the zero-fill re-exposes it, so it is fixed here. 2. Multi-CTA Rubin tactics are now rejected in get_valid_tactics rather than only at runtime, so the autotuner never selects one. The NotImplementedError in _moe_core_impl stays as a backstop for callers that pass tactic parameters directly. Rejection still happens before moe_sort. 3. Documented the initialization contract on the moe_sort buffers, which is only accurate once (1) lands. AI-assisted: analysis and patch drafted with Claude Code.
3 tasks
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
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.
📌 Description
Adds CuTe DSL kernel support for NVIDIA Rubin (SM107), together with the Python
CuTe DSL batched FP8 GEMM implementation the Rubin kernels build on.
Batched FP8 GEMM (CuTe DSL) — new functionality on Blackwell as well as
Rubin; FlashInfer previously had only a C++
bmm_fp8path:gemm/kernels/bmm_fp8_blackwell.py— persistent dense GEMM for SM100/SM103gemm/kernels/bmm_fp8_rubin.py— the SM107 variantgemm/kernels/bmm_fp8_wrapper.py— dispatch and autotuning across bothgemm/kernels/epilogue_utils.py— shared epilogue helpers (scaled TMA store,scaled store, alpha-scaled store)
Rubin fused-MoE kernels under
fused_moe/cute_dsl/rubin/:Helpers shared by the Blackwell and Rubin MoE paths move to
fused_moe/cute_dsl/common/kernel_utils.py.Rubin dense and grouped GEMM:
gemm/kernels/dense_blockscaled_gemm_sm107.py— blockscaled dense GEMMgemm/kernels/grouped_gemm_masked_rubin.py— masked grouped GEMMgemm/kernels/grouped_gemm_masked_wrapper.py— dispatch wrapperIntegration:
fused_moe/cute_dsl/tuner.py,gemm/gemm_base.py). SM107 encodes additional tactic parameters, so theautotuner's tactic tuple is arch-aware.
smem=argument.Dependency handling
The CuTe DSL requirement stays a floor (
>=4.7.0), not a pin — FlashInfercontinues to work with older CuTe DSL and CUDA toolkit releases exactly as
before.
The SM107 kernels additionally need CuTe DSL >= 4.8 (they use
cutlass.utils.rubin_helpers). They are imported lazily behindis_rubin_cute_dsl_available(), so:CuTe DSL paths are unavailable
The
cu12/cu13extras inpyproject.tomlare relaxed from==4.7.0to>=4.7.0for the same reason — an exact pin would have made 4.8 uninstallable,so the SM107 kernels could never have been reached.
🔍 Related Issues
None.
🚀 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
The batched FP8 GEMM kernels are not Rubin-specific.
bmm_fp8_blackwell.pyand
epilogue_utils.pyare new functionality for Blackwell as well, andbmm_fp8_rubin.pyimports directly from both. They are worth reviewing on theirown merits rather than as Rubin scaffolding.
Backward compatibility with CuTe DSL 4.7 needs confirming. This PR adds the
smem=launch argument to three shared, non-Rubin files(
quantization/kernels/{mxfp4,mxfp8,nvfp4}_quantize.py), which run on both 4.7and 4.8. If 4.7's
launch()does not accept that keyword, those paths wouldbreak for 4.7 users — the lazy-import gating above does not cover them. If 4.8
tolerates omitting
smem=, the cleanest fix is to leave those three filesunchanged.
Lint/type exclusions.
ruffandmypyreported 39 findings, all confined tothree CuTe DSL kernel bodies (
dense_blockscaled_gemm_sm107.pyand the twofused_moe/cute_dsl/rubin/*_fusion.pykernels). Following the convention alreadyused for
cute_dsl/attention/fmha/andmoe_ep/kernel_src/.../src, these areexcluded rather than rewritten — the existing note in
pyproject.tomlwarns thatreflowing or "simplifying" traced DSL code can change semantics (a DSL Boolean
predicate becoming a Python bool). The reported rules (
SIM210,B008oncutlass.Int64(0)defaults,B007oncutlass.rangeinduction variables) areexactly that class of change.
ruff formatis applied to every non-excluded file.Validation status. Static checks only so far: every changed Python file
parses, all intra-package imports resolve, and no module-scope import reaches
the CuTe DSL 4.8-only API (so the package stays importable on older DSL).
A build plus the GEMM and MoE suites on SM100/SM103 and SM107 hardware are
still needed.
Summary by CodeRabbit
New Features
Bug Fixes
Tests