ci: build jit-cache wheels for CUDA 13.4 with 10.7a (Rubin/SM107)

[Vinnie6167]

Note

Prerequisite now met. This needs pytorch/manylinux2_28-builder:cuda13.4 and
pytorch/manylinuxaarch64-builder:cuda13.4. Both are published on Docker Hub as of
2026-08-14 (6.59 GiB and 6.89 GiB), built by pytorch/pytorch's
build-manywheel-images.yml on main (sha 66c3e509) and rebuilt since — tags through
-v2.14.0-rc4 dated 2026-08-18. The x86 image config confirms BASE_CUDA_VERSION=13.4
with /usr/local/cuda -> /usr/local/cuda-13.4. CUDA 13.4 support landed upstream in
pytorch/pytorch#192256.

Two caveats worth a maintainer's view:

• The toolkit is CUDA 13.4rc1, installed from NVIDIA's preview network repo
  (.ci/docker/common/install_cuda.sh::install_134 notes 13.4 ships no runfile-local
  installer yet). If building official release wheels against a preview CTK is a
  concern, better to say so before this merges.
• Kept as a draft: none of the validation below has been run on real hardware, and
  the aarch64 wheel-size question further down is unanswered.

For the record, a public nvidia/cuda:13.4 would not have substituted here — glibc 2.39
against the manylinux_2_28 tag these wheels declare with no auditwheel check
(flashinfer-jit-cache/build_backend.py::get_tag).

📌 Description

The published flashinfer-jit-cache wheels contain no SM107 (Rubin) code, so Rubin users
hit no kernel image is available for execution on the device.

The cu13.0 wheel ships sm_75 sm_80 sm_89 sm_90a sm_100a sm_103a sm_120f (plus sm_110a sm_121a on aarch64). sm_100a and sm_103a are architecture-specific and will not
load on SM107, and there is no sm_100f family target to fall back to — so nothing in
the wheel runs on Rubin.

This adds a fourth CUDA to the jit-cache matrix rather than altering the existing
ones. CUDA 13.4 is required, not merely preferred: nvcc gains compute_107 only in
13.4 (verified with --list-gpu-arch: 13.2 no, 13.3 no, 13.4 yes). It cannot be
backported to the cu13.0 tier, because modules whose call sites do not pass
map_sm107_to_100f emit -gencode=arch=compute_107a verbatim and would fail the build
there.

Setting 10.7a does two distinct things, and both matter:

1. Modules built through the ~13 call sites passing map_sm107_to_100f=True get
   -gencode=arch=compute_100f,code=sm_100f, because the bundled CUTLASS has no
   struct Sm107 yet (cutlass_supports_sm107()). On a comparable internal build that
   is 777 of 963 modules.
2. Everything else compiles native sm_107a (8 of 963 on that same build), and
   has_sm107 in aot.py flips true — which is what pulls in the four Rubin-only
   modules: fp4_quantization_sm107 and the enable_rubin=True variants of
   trtllm_gen_gemm, trtllm_low_latency_gemm and trtllm_gen_fused_moe.

Point 2 is why simply adding an explicit 10.0f entry to the existing cu13.0 builder is
not an adequate substitute — that would produce family cubins but leave those four
modules out of the cache entirely.

Changes

• jit-cache matrix gains "13.4" in release.yml and nightly-release.yml.
• New >= 13.4 tier in FLASHINFER_CUDA_ARCH_LIST:
  • x86_64: 7.5 8.0 8.9 9.0a 10.0a 10.3a 10.7a 12.0f
  • aarch64: 7.5 8.0 8.9 9.0a 10.0a 10.3a 10.7a 11.0a 12.0f 12.1a
• Upload loop extended: for cuda in 128 129 130 → 128 129 130 134. A matrix entry
  missing from that hardcoded list builds successfully and is then silently never
  attached to the release.

DOCKER_IMAGE needed no change — it already interpolates matrix.cuda.

test-nightly-build is deliberately left at cuda 12.9/13.0. That job is the acceptance
test for these wheels: it installs the freshly built python/cubin/jit-cache artifacts,
sets FLASHINFER_DISABLE_JIT=1 so anything missing from the cache fails rather than
silently source-compiling, and runs the suite in 5 shards. It cannot cover SM107 today,
for two structural reasons that this PR cannot fix:

• it runs on Ampere runners — runs-on: [self-hosted, linux, x64, gpu, sm86, on-demand]
  — and there is no SM107 runner in public CI;
• it downloads only the x86_64 jit-cache artifact, so the aarch64 wheel is not exercised
  by it at all today.

Adding a cu13.4 cell there would therefore cost build time without testing anything new.

The cu12.8 / cu12.9 / cu13.0 arch lists are byte-identical after this change; I resolved
the expression for all 8 matrix cells to confirm only the new tier differs.

❓ Open question for maintainers

Does the cu13.4 aarch64 wheel fit the 2 GiB release-asset ceiling? I have deliberately
not pre-empted this — the list above keeps 12.1a, matching cu13.0.

Measured v0.6.17 release assets:

wheel	size
cu130 x86_64	1.41 GiB
cu130 aarch64	1.58 GiB

On an otherwise identical arch list, adding 10.7a costs roughly 0.34 GiB. That puts
cu13.4 x86_64 near 1.75 GiB (comfortable) but cu13.4 aarch64 near 2.0 GiB — right at
the ceiling this file already documents having hit (#3684 / #4519, which is why 12.1a
is currently cu13.0-aarch64-only).

Worth weighing explicitly: #4363 is reported on aarch64 SM107 hardware, so aarch64 is
simultaneously the arch that most needs this wheel and the one at risk of not fitting.
Whatever gets trimmed to make room, it should not be 10.7a.

Only a real cu13.4 build settles it. If it overshoots, dropping 12.1a from the cu13.4
aarch64 list is the obvious lever — sm121 users would stay on the cu130 wheel, which this
PR leaves untouched — but that is a maintainer's call about which architectures matter on
a Rubin-era stack, not something I wanted to decide silently.

🔍 Related Issues

Addresses #4363 (P0) — "CUDA 13/aarch64 JIT-cache package lacks SM107/10.7a coverage" —
which reports no kernel image is available for execution on the device for the BF16
paged-MLA specialization used by Kimi K3 (head_dim_ckv=512, head_dim_kpe=64) on
aarch64 SM107 under SGLang.

That issue's control experiment is what makes this a packaging fix rather than a kernel
one: uninstalling flashinfer-jit-cache, setting FLASHINFER_CUDA_ARCH_LIST=10.7a and
letting FlashInfer source-JIT the same operator passed numerical validation and completed
two-node TP8 Kimi K3 serving with full GSM8K accuracy. Source-build SM107 support landed
in #4122; this is its packaging counterpart.

I have deliberately not used a closing keyword, for two reasons:

• [Bug] CUDA 13/aarch64 JIT-cache package lacks SM107/10.7a coverage #4363 asks for a published package. This PR only changes the workflow — nothing ships
  until the cuda13.4 builder images exist and a release runs.
• It also suggests an SM107 regression test proving the published wheel can execute BF16
  paged MLA without uninstalling the jit-cache. The harness for that already exists —
  test-nightly-build does exactly this, via FLASHINFER_DISABLE_JIT=1 — but it runs on
  sm86 runners and consumes only x86_64 artifacts. Satisfying that half of the issue needs
  an SM107 runner plus aarch64 artifact coverage in that job, which is separate work from
  this PR.

Maintainers should feel free to add a closing keyword for that issue if you would rather
it auto-close on the release that first carries the wheel — I have avoided writing one
here so that merging this PR does not close it prematurely.

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

CI-only change; there is no unit test that exercises a release matrix. The validation
that matters needs the 13.4 builder image and is listed below — none of it has been run
yet, since the image does not exist:

1. nvcc --list-gpu-arch | grep compute_107 inside the builder, both arches.
2. Build one cu13.4 wheel per arch and measure the asset size against the 2 GiB limit.
3. cuobjdump --list-elf census over all modules: expect sm_107a in the ~8 modules that
   do not take the mapping, sm_100f broadly present, no other arch lost.
4. Confirm the four has_sm107-gated modules are present by name.
5. End-to-end on SM107 hardware: the BF16 paged-MLA case (CKV=512, KPE=64) that currently
   raises "no kernel image".
6. Blackwell regression: sm_100a / sm_103a paths unchanged.

Reviewer Notes

The FLASHINFER_CUDA_ARCH_LIST ternary is now four tiers deep and close to its
readability limit. I kept the in-idiom extension to minimise review surface, but a matrix
include: refactor that names the arch list per cell would be cleaner — happy to do that
instead if you'd prefer.

Independently of this PR: once the bundled CUTLASS gains struct Sm107, the ~13
map_sm107_to_100f=True call sites switch from sm_100f to native sm_107a
automatically with no code change, per the comment in compilation_context.py.

Summary by CodeRabbit

• New Features
  • Added CUDA 13.4 support to nightly and regular release builds.
  • Published CUDA 13.4 JIT-cache artifacts.
  • Expanded CUDA architecture coverage, including additional aarch64 configurations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: d5c90332-7cf2-4e45-a868-caf604326234

📥 Commits

Reviewing files that changed from the base of the PR and between b2d0236 and 575bfa1.

📒 Files selected for processing (2)

• .github/workflows/nightly-release.yml
• .github/workflows/release.yml

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

Changes

CUDA 13.4 JIT-cache release support

Layer / File(s)	Summary
CUDA 13.4 build matrix and architecture selection .github/workflows/nightly-release.yml, .github/workflows/release.yml	The workflows add CUDA 13.4 JIT-cache builds. Architecture lists now vary by CUDA version and platform, including expanded aarch64 coverage.
CUDA 13.4 artifact publication .github/workflows/nightly-release.yml, .github/workflows/release.yml	Release upload loops include CUDA 13.4 JIT-cache artifacts.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: 🔵 Low · up to 575bf

This PR adds CUDA 13.4 jit-cache builds for Rubin coverage, but raw 13.4 inputs should be verified to select the new cu134 artifact and the aarch64 wheel must be measured against the 2 GiB release-asset limit before publication. The change is mergeable with explicit maintainer follow-up on these bounded release risks.

Possibly related issues

• [Bug] CUDA 13/aarch64 JIT-cache package lacks SM107/10.7a coverage #4363 — Updates CUDA and aarch64 JIT-cache architecture coverage.
• [RFC] Align the CI/build CUDA toolkit matrix with upstream: target 12.9 · 13.0 · 13.4 #4453 — Adds CUDA 13.4 JIT-cache builds and publication.

Possibly related PRs

• flashinfer-ai/flashinfer#4469 — Introduces related CUDA JIT-cache matrix and artifact handling.
• flashinfer-ai/flashinfer#4527 — Modifies CUDA architecture selection in the same release workflows.

Suggested labels: run-ci

Suggested reviewers: dierksen, kahyunnam, aleozlx

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description covers the required sections, explains the changes and risks, links the related issue, and clearly documents pending validation.
Title check	✅ Passed	The title clearly and concisely identifies the CI change, CUDA 13.4 support, JIT-cache wheels, and SM107 coverage.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 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

• X
• Mastodon
• Reddit
• LinkedIn

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