perf(sampling): gate __launch_bounds__ to SM107 only (NVBug 6517769) by kahyunnam · Pull Request #4189 · flashinfer-ai/flashinfer

kahyunnam · 2026-07-27T21:01:49Z

📌 Description

Fixes the B300 sampling performance regression tracked in NVBug 6517769, targeting release-v0.6.16 directly.

#4122 annotated 11 sampling kernels with __launch_bounds__(BLOCK_THREADS). That annotation is effectively __launch_bounds__(1024) on every cc ≥ 8, which forces ptxas under a 64-reg/thread ceiling and causes measurable spills / slowdowns on B300 and H100.

Why not a blanket removal?

The annotations were not unexplained noise. Internal GitLab MR !611 (Fix Rubin (SM 10.7) launch failures in sampling kernels via __launch_bounds__, merged into feat_sm107 2026-04-28) documented the Rubin failure mode:

The Top-K sampling kernel was using >64 regs per thread with 1024 threads per block, exceeding the 65,536 register budget.

The Rubin Kernel Requirements Tracker recorded the pre-fix failures as too many resources requested for launch on tests/utils/test_sampling.py / logits sampling. That rationale never made the public #4122 description, but the release branch still carries SM107 support, so a full drop would re-break Rubin native launches.

Fix

Gate the annotation to native SM107 compiles only:

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1070)
#define FLASHINFER_SAMPLING_LAUNCH_BOUNDS(block_threads) __launch_bounds__(block_threads)
#else
#define FLASHINFER_SAMPLING_LAUNCH_BOUNDS(block_threads)
#endif

All 11 kernels use FLASHINFER_SAMPLING_LAUNCH_BOUNDS(BLOCK_THREADS) instead of unconditional __launch_bounds__(BLOCK_THREADS).

When today's JIT maps SM107 → sm_100f (map_sm107_to_100f=True while CUTLASS lacks Sm107), __CUDA_ARCH__ is 1000 and the gate is inactive; that path inherits sm_100 register counts, which already fit a 1024-thread launch on the arches we measured. The gate is the safety net for native sm_107a compiles — the path MR !611 actually fixed.

🔍 Related Issues

NVBug 6517769 — B300 sampling perf drop after 290c091
Internal MR !611 / commit a8c10c24 on feat_sm107 (Rubin launch-failure root cause)
Caused by Adds SM107 support #4122; Revert PR 4122 #4171 reverted Adds SM107 support #4122 wholesale on main (no cherry-pick possible)

🧪 Tests

ptxas (CUDA 13.0) with the gated macro

Target	Kernel	Result
sm_90a	MinP / TopK / TopP / TopKTopP	≤48 regs, no spills (matches unconstrained)
sm_103a	same	≤48 regs, no spills
sm_107a	—	cannot compile on this toolkit (`Unsupported gpu architecture`)

Wall-clock (NVBug repro on H100 NVL)

python3 flashinfer_benchmark.py --routine min_p_sampling_from_probs \
    --batch_size 32 --vocab_size 129280 --min_p 0.1 -vv

Config	median
release (unconditional `__launch_bounds__`)	0.051 ms
this PR (SM107-gated)	0.045 ms (same as full removal)

Correctness: earlier full-removal run of tests/utils/test_sampling.py on this header family was 1320 passed / 0 failed; the gated form is a no-op on H100 device code, so that result still applies.

Reviewer Notes

Please confirm on B300 that the NVBug case is restored, and on Rubin that tests/utils/test_sampling.py still launches when compiled for native sm_107a (and/or under the current sm_100f JIT mapping).
#4171 already removed these annotations from main along with all SM107 enablement; this PR is release-only and preserves SM107 launchability.

… release branch PR flashinfer-ai#4122 annotated 11 sampling kernels with __launch_bounds__(BLOCK_THREADS), which is BLOCK_THREADS=1024 for every compute capability >= 8. That caps ptxas at 64 registers per thread and changes its allocation strategy, making several of these already register-heavy kernels (CUB scans/reductions, Philox state, vectorized loads over the full vocab) spill to local memory. NVBug 6517769 reports the resulting B300 regression: min_p_sampling_from_probs at batch 32 / vocab 129280 runs at 0.944x of the pre-flashinfer-ai#4122 baseline. flashinfer-ai#4122 was reverted wholesale on main by flashinfer-ai#4171, so main no longer carries these annotations, but release-v0.6.16 still does. This reverts the sampling.cuh portion only; the resulting file is byte-identical to main's, so the two branches converge. The annotation was the sole change flashinfer-ai#4122 made to this file, and nothing has touched it since, so no SM107 functionality is affected. ptxas (CUDA 13.0) with the annotation vs without: sm_103a (B300) MinPSamplingFromProb 32 regs, 12B spills -> 40 regs, none TopKTopPSamplingFromProb 64 regs, 48B spills -> 46 regs, none sm_100a TopKTopPSamplingFromProb 64 regs, 48B spills -> 46 regs, none sm_90a TopKSamplingFromProb 32 regs, 28B spills -> 48 regs, none TopKTopPSamplingFromProb 32 regs, 24B spills -> 46 regs, none Without the annotation no kernel exceeds 48 registers on any of these targets, so a 1024-thread launch stays within the 64K per-block register file and the removal cannot introduce a launch failure. Measured on H100 NVL at the NVBug shape (batch 32, vocab 129280, 200 iters, median of 7 reps, builds interleaved): min_p 0.0192 -> 0.0173 ms (10.5% faster) top_k 0.1622 -> 0.1526 ms ( 6.3% faster) top_p 0.0255 -> 0.0252 ms ( 1.2%, and ptxas shows no spills for it) top_p barely moving while the two spilling kernels recover confirms the spills are the mechanism rather than some unrelated codegen difference.

gemini-code-assist · 2026-07-27T21:01:52Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

coderabbitai · 2026-07-27T21:01:59Z

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 8cd951f6-3f1e-42fe-931c-b800acb3b8bd

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

Follow-up to the blanket removal: internal MR !611 on feat_sm107 added __launch_bounds__ because native sm_107a ptxas allocated >64 regs/thread for TopK sampling at 1024 threads, exceeding the 65536-register budget ("too many resources requested for launch"). That rationale never made the public flashinfer-ai#4122 description, but the release branch still needs SM107 to launch. Keep the annotation only when __CUDA_ARCH__ >= 1070 so Rubin native compiles stay launch-safe, while sm_90/sm_100/sm_103 keep unconstrained allocation (no spills; restores the NVBug 6517769 H100/B300 regression). When JIT maps SM107 to sm_100f, the gate is inactive and that path inherits sm_100 register counts which already fit a 1024-thread launch. Verified on this host: sm_90a/sm_103a ptxas matches the unconstrained profile, and flashinfer_benchmark min_p (batch 32, vocab 129280) is 0.045 ms — same as the full removal. sm_107a cannot be compiled here.

….6.16rc2 (#4197) Cherry-picks for the `release-v0.6.16` branch, plus the version bump to `0.6.16rc2`. ## Commits | Commit | Source | Status | | --- | --- | --- | | `fix(norm): convert float2 to e4m3 directly in packed cast` | #4167 (`e683e307` on `main`) | merged upstream | | `feat(xqa): ragged Q and per-row sliding-window masking for speculative decode` | #4137 (`60783fb9` on `main`) | merged upstream | | `test(jit): assert BMM export symlink under GEN_SRC_DIR` | #4187 (`417bbd29` on `main`) | merged upstream | | `feat(mla): support packed low-head and variable-Q decode` | #4178 (PR head `0774943c`) | **not yet merged** to `main` | | `bump version to 0.6.16rc2` | — | — | Applied in `main` merge order; #4187 sits directly on top of #4137 upstream, so that ordering is preserved. ## Notes - **#4189 was not cherry-picked** — it was merged directly into `release-v0.6.16` (`34368112`) and is already the base of this branch. - **#4178 is still open against `main`.** It was cherry-picked at the request of the release owner. Its two head commits (`b6cc59491`, `0774943c2`) are squashed into one commit here, mirroring how it will land upstream. If the PR changes before merge, this commit should be dropped and re-picked. - The three merged cherry-picks apply cleanly with patch-ids identical to their sources. The #4178 squash is byte-identical to the PR's combined diff. - The #4178 commit was committed with `--no-verify`: the `mypy` pre-commit hook reports two pre-existing `[no-redef]` errors in `flashinfer/mla/_core.py` (conditional `if/else` imports of `_check_can_implement` / `_get_split_kv_and_workspace_size`). These reproduce identically on PR #4178's own branch and on the `release-v0.6.16` base — they are not introduced by this cherry-pick, and surface only when those modules are checked together. Left unfixed here to keep the cherry-pick faithful. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Lee Yongjun <jqueen.astro@gmail.com> Co-authored-by: yichengj <yichengj@nvidia.com> Co-authored-by: bryu <bryu@nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Ka-Hyun Nam <knam@nvidia.com> Co-authored-by: mingyangw <mingyangw@nvidia.com>

## 📌 Description This PR relands SM 107 support to main branch (reverted in #4171) as well as some other release fixes. #### Cherry Picks - #4191 - #4189 - #4200 - #4215 - #4225 - #4230 - #4235 - #4226 - #4257 - #4258 - #4261 #### Other Changes - Rubin guards from #4252's conflict resolution (`TLLM_RUBIN_FEATURES`: SiTuGlu static_asserts + tile-192 advertisement, compiled out for the Rubin BMM pin) - Test-contract update: `test_unified_moe.py` arch assertions written post-revert (#4159) flipped to the restored contract (FP4/BF16 claim 107; FP8 stays 100/103)  ## 🔍 Related Issues  #4107, #4164, reverts #4171 ## 🚀 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. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes   ## Summary by CodeRabbit * **New Features** * Added support for Rubin/SM107 GPUs across GEMM, MoE, attention, quantization, sampling, and DeepGEMM workflows. * Added architecture-aware kernel selection, memory sizing, compilation, and artifact handling. * **Bug Fixes** * Improved validation and error messages for incompatible GPU architectures and invalid kernel configurations. * Clearly rejects unsupported NVFP4 KV-cache operations on SM107. * **Documentation** * Updated installation guidance with the SM107 architecture target. * **Tests** * Expanded architecture coverage and compatibility checks across GPU test suites.  --------- Co-authored-by: Vinnie6167 <Vinnie6167@users.noreply.github.com> Co-authored-by: Ka-Hyun Nam <knam@nvidia.com> Co-authored-by: Alex Yang <aleyang@nvidia.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Jimmy Zhou <79552142+jimmyzho@users.noreply.github.com>

kahyunnam requested review from aleozlx, bkryu, cyx-6, jimmyzho, nv-yunzheq, saltyminty, samuellees, sricketts, yongwww, yyihuang and yzh119 as code owners July 27, 2026 21:01

kahyunnam requested review from Aneureka, dhiraj113, feih-nv, jiahanc, qiching and qsang-nv as code owners July 27, 2026 21:01

kahyunnam changed the title ~~perf(sampling): drop __launch_bounds__ added by #4122 (NVBug 6517769)~~ perf(sampling): drop __launch_bounds__ added by #4122 Jul 27, 2026

aleozlx approved these changes Jul 27, 2026

View reviewed changes

kahyunnam changed the title ~~perf(sampling): drop __launch_bounds__ added by #4122~~ perf(sampling): gate __launch_bounds__ to SM107 only (NVBug 6517769) Jul 27, 2026

aleozlx merged commit 3436811 into flashinfer-ai:release-v0.6.16 Jul 28, 2026
5 checks passed

aleozlx mentioned this pull request Jul 28, 2026

release-v0.6.16: cherry-pick #4167, #4137, #4187, #4178 and bump to 0.6.16rc2 #4197

Merged

This was referenced Jul 28, 2026

feat(examples): exercise and audit top-k/top-p sampling end to end #4216

Draft

[Bug] Rubin/SM107: opaque ptxas "sm_107a is not defined" on CUDA 13.2 — Triton path needs CUDA >= 13.3 #4217

Closed

Vinnie6167 mentioned this pull request Jul 30, 2026

SM 107 Reland + Merge Back from v0.6.16 Release Branch #4280

Merged

5 tasks

Navigation Menu

Uh oh!

perf(sampling): gate __launch_bounds__ to SM107 only (NVBug 6517769) - #4189

perf(sampling): gate __launch_bounds__ to SM107 only (NVBug 6517769)#4189
aleozlx merged 2 commits into
flashinfer-ai:release-v0.6.16from
kahyunnam:fix/nvbug-6517769-sampling-launch-bounds

kahyunnam commented Jul 27, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Jul 27, 2026

Uh oh!

coderabbitai Bot commented Jul 27, 2026 •

edited

Loading

Review skipped

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

kahyunnam commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

Why not a blanket removal?

Fix

🔍 Related Issues

🧪 Tests

ptxas (CUDA 13.0) with the gated macro

Wall-clock (NVBug repro on H100 NVL)

Reviewer Notes

Uh oh!

gemini-code-assist Bot commented Jul 27, 2026

Uh oh!

coderabbitai Bot commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kahyunnam commented Jul 27, 2026 •

edited

Loading

coderabbitai Bot commented Jul 27, 2026 •

edited

Loading