Skip to content

Split trtllm-gen cubin pins per arch - #4191

Merged
aleozlx merged 3 commits into
flashinfer-ai:release-v0.6.16from
Vinnie6167:fix-rubin-features-blackwell-abi
Jul 28, 2026
Merged

aleozlx merged 3 commits into
flashinfer-ai:release-v0.6.16from
Vinnie6167:fix-rubin-features-blackwell-abi

Conversation

@Vinnie6167

@Vinnie6167 Vinnie6167 commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

📌 Description

The BMM & GEMM trtllm-gen cubins in ToT fail on SM 100. This MR introduces a WAR, loading an older set of published cubins when targeting SM 100.

🔍 Related Issues

🚀 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

  • 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.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

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

Reviewer Notes

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

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

@coderabbitai

coderabbitai Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

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: efa8fee7-ca36-45f1-adb3-bebcf81e26a6

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.

@kahyunnam kahyunnam left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In AOT, gen_all_modules() (aot.py) emits both when has_sm100 and has_sm107 (default arch list has 10.0a and 10.7a): each gen_* rewrites the same synlink, then all modules compile afterward. Last writer wins → Blackwell modules (enable_rubin=False, no -DTLLM_RUBIN_FEATURES) can compile against Rubin export headers → same GemmOptions ABI skew this PR is fixing.

Same issue in gen_trtllm_gen_gemm_module, gen_trtllm_low_latency_gemm_module, and gen_trtllm_gen_fused_moe_sm100_module

Image

Comment thread flashinfer/jit/fused_moe.py
Comment thread flashinfer/jit/gemm/core.py Outdated
Comment thread flashinfer/jit/gemm/core.py Outdated
The Blackwell and Rubin variants shared one export-header symlink under
FLASHINFER_GEN_SRC_DIR (independent of enable_rubin). In an AOT build all
modules are generated before anything compiles, so the last gen_* call's
ensure_symlink won, and the other variant could compile against the wrong
(flag-mismatched) export headers -- the same GemmOptions/metadata ABI skew
this PR fixes. Runtime JIT was unaffected (one variant per device).

Give each module its own export-header root (FLASHINFER_GEN_SRC_DIR/
trtllm_export/<module_name>) and put it first on the include path so the
fixed `#include "flashinfer/trtllm/.../trtllmGen_*_export/..."` resolves
per-variant. Also de-conflicts from moe_utils, which shares that path.

Addresses review feedback from @kahyunnam on PR flashinfer-ai#4191.

AI-assisted (Claude Code).
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

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

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

GitLab MR !1058 has been created, and the CI pipeline #59880138 is currently running. I'll report back once the pipeline job completes.

Renaming get_trtllm_gemm_module to the typed _get_trtllm_gemm_module_impl
made mypy check its body, surfacing a pre-existing setup_cubin_loader(Path)
call that expects str. Match the sibling MoE / low-latency-gemm loaders,
which already wrap it in str(). No runtime change.

AI-assisted (Claude Code).
@flashinfer-bot

Copy link
Copy Markdown
Collaborator

[CANCELING] Pipeline #59880138: canceled

@Vinnie6167

Copy link
Copy Markdown
Contributor Author

/bot run

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

GitLab MR !1058 has been updated with latest changes, and the CI pipeline #59881205 is currently running. I'll report back once the pipeline job completes.

@aleozlx
aleozlx merged commit 82b329b into flashinfer-ai:release-v0.6.16 Jul 28, 2026
6 of 7 checks passed
aleozlx added a commit to aleozlx/flashinfer that referenced this pull request Jul 28, 2026
test_get_subdir_file_list was already failing on release-v0.6.16 before this
branch: flashinfer-ai#4191 added TRTLLM_GEN_BMM_RUBIN and TRTLLM_GEN_GEMM_RUBIN to
cubin_dirs, but the test never registered mocks for them. get_checksums()
therefore tried to fetch two unmocked checksums.txt URLs, exhausted the retry
budget (~50s of backoff) and died with FileNotFoundError. Verified failing at
the unmodified base commit c46207c.

Register directory-index and checksums.txt mocks for both Rubin pins, reusing
the non-Rubin index bodies so the mocked pins share kernel filenames exactly as
the published ones do, with every hash differing. Update the meta-info header
count 3 -> 5 (FMHA, GEMM, BMM, GEMM_RUBIN, BMM_RUBIN).

Add regression coverage for the checksum collision fixed in the previous
commit: assert paths are unique and that a filename present in both a plain
and a Rubin pin resolves to its own pin's hash. Against the unfixed
get_checksums() this fails with

  AssertionError: Bmm_..._schedS_bN_clmp_dynBatch_sm100f.cubin resolved to the
  same checksum for both pins (bbbb111122223333) -- the per-pin hashes collided

tests/test_artifacts.py: 1 failed, 3 passed -> 4 passed.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
aleozlx added a commit that referenced this pull request Jul 28, 2026
…ild break (#4200)

<!-- .github/pull_request_template.md -->

## 📌 Description

Fixes both failing jobs in the 0.6.16rc2 release run

([30330417436](https://github.com/flashinfer-ai/flashinfer/actions/runs/30330417436))
on `c46207c2`.
Both are deterministic consequences of introducing a second,
Rubin-specific cubin pin
(`TRTLLM_GEN_BMM_RUBIN` / `TRTLLM_GEN_GEMM_RUBIN`) alongside the
existing one — neither is flaky.

**1. `build-flashinfer-cubin` — `Failed to download cubins: checksum
mismatch`**

`get_checksums()` flattened every pin's `checksums.txt` into one dict
keyed by **bare filename**,
namespacing only `.h` files. The plain and Rubin pins ship the **same**
sm100f/sm103a kernel names
built from different sources, so the pin processed last silently
overwrote the earlier one's hashes.
`get_subdir_file_list()` then looked cubins up by bare name while
headers already used the full path.

Measured against the live published manifests:

| pin pair | shared filenames | same hash | **conflicting** |
|---|---|---|---|
| `TRTLLM_GEN_BMM` vs `_RUBIN` | 2534 | 0 | **2534** |
| `TRTLLM_GEN_GEMM` vs `_RUBIN` | 108 | 3 | **105** |

The Rubin pins come later in `cubin_dirs`, so they win and 2640 of 8055
files are verified against
the wrong pin's hash. Replaying both keying schemes over the real
manifests:

```
OLD (bare filename)    entries= 5412  files=8055  MISMATCHED=2640
NEW (full path)        entries= 8055  files=8055  MISMATCHED=0
```

**Fix:** key every entry by full path, and look cubins up the same way
headers already were.

This is **not** a stale pinned hash — all 12 `checksums.txt` manifests
still match their
`CheckSumHash` SHA256 exactly. It is also **not** download flakiness:
the log shows two transient
`403 Forbidden`s on
`fmhaSm100fKernel_QkvE4m3OE2m1H256...Q{8,32}Kv128StaticSwapsAbForGen.cubin`,
but
both recovered on retry (no `Max retries reached`) and both URLs serve
`200` now. A checksum-retry
loop would have re-downloaded 2640 files, failed identically, and turned
a fast failure into a slow one.

**2. `build-flashinfer-jit-cache` (all 6 matrix entries) — compile
error**

```
csrc/trtllm_batched_gemm_runner.cu(137): error:
enum "batchedGemm::trtllm::gen::CudaArch" has no member "Sm107a"
```

The `Sm107a` comparison added in #4122 is unguarded, but that enumerator
exists only in the Rubin
pin's generated header — the default pin's `flashinferMetaInfo.h` has
`Sm100f`/`Sm103a` and **zero**
occurrences of `Sm107a` (the Rubin pin has 2552). So the non-Rubin
module `fused_moe_trtllm_sm100`
cannot compile. This is why `main`, where #4122 was reverted, is
unaffected.

**Fix:** guard with `TLLM_RUBIN_FEATURES`, which
`gen_trtllm_gen_fused_moe_sm100_module` already
defines (`rubin_flags`) exactly when the Rubin pin is selected — so the
guard is by construction in
sync with the enum's availability. `sm_version == 107` is unreachable in
the non-Rubin module
regardless: `enable_rubin` is chosen from the same device compute
capability that `getSMVersion()` reports.

**3. `tests/test_artifacts.py::test_get_subdir_file_list` — already
broken on the release branch**

Independent of the above, #4191 added the two Rubin pins to `cubin_dirs`
without registering mocks for
them, so the test exhausted its retry budget (~50 s of backoff) and died
with `FileNotFoundError`.
Verified failing at the unmodified base `c46207c2`. Mocks added,
meta-info header count corrected
3 → 5, plus regression coverage for the collision above.

## 🔍 Related Issues

- Release run:
https://github.com/flashinfer-ai/flashinfer/actions/runs/30330417436
- Introduced by the per-arch pin split (#4191) interacting with SM107
support (#4122)

## 🚀 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

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

`tests/test_artifacts.py`: **1 failed, 3 passed → 4 passed**.

The new regression assertion is verified to actually catch the bug —
reverting only `artifacts.py`
to the buggy base while keeping the new test yields:

```
AssertionError: Bmm_..._schedS_bN_clmp_dynBatch_sm100f.cubin resolved to the same
checksum for both pins (bbbb111122223333) -- the per-pin hashes collided
```

## Reviewer Notes

- **Scope of local testing.** `tests/test_artifacts.py` is fully
exercised locally (it is
`responses`-mocked, no GPU needed). The two build jobs themselves are
not reproducible locally and
  need CI to confirm.
- **Why the jit-cache fix is sufficient for the whole matrix.** All 4
failed jit-cache jobs report
**exactly one** compiler error each — the same `Sm107a` one, same target
`csrc_trtllm_batched_gemm_runner.cuda.o`. The matrix builds `7.5 8.0 8.9
9.0a 10.0a 10.3a 12.0f`
(no 10.7), so `has_sm107` is false and the Rubin modules are not
compiled there — this change
  unblocks the job without pulling in any `sm_107a` codegen.
- **Two alternatives considered and rejected** for the `Sm107a` guard:
- A new `FLASHINFER_HAS_SM107A` macro — nothing would define it, so the
`#else` branch would
    always compile, silently disabling Rubin kernel selection.
- Falling back to `Sm103a`/`Sm100f` under `sm_version == 107` — worse
than the build error: on real
Rubin silicon it would select wrong-arch cubins instead of failing
loudly.
- **Follow-up worth considering (not in this PR):** the `checksum
mismatch` error does not name the
offending file. Including the path plus expected/actual hashes would
have made this diagnosable
  from the log alone.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
aleozlx pushed a commit that referenced this pull request Jul 29, 2026
…4225)

## Summary

Fixes two test regressions introduced by #4191 (release-only, not on
`main`):

- **#4208** — `test_make_tuning_config_reuses_topk_ids_initializer`
called `get_trtllm_moe_sm100_module.cache_clear()`, but #4191 moved
`@functools.cache` to `_get_trtllm_moe_sm100_module_impl`.
- **#4166 (re-open)** — #4187 updated the symlink test to assert under
`GEN_SRC_DIR/flashinfer/...`, but #4191 moved the export symlink to
`GEN_SRC_DIR/trtllm_export/fused_moe_trtllm_sm100/flashinfer/...`.

Both failures appeared on all Blackwell jobs in Test Cycle 2 (pipeline
59912149, rc3).

## Test plan

- [x] `pytest
tests/autotuner/test_autotuner_core.py::test_make_tuning_config_reuses_topk_ids_initializer`
— pass
- [x] `pytest tests/utils/test_gen_module_symlink_race_condition.py` —
pass
- [ ] B300/GB200 CI rerun to confirm cycle-2 green on these two files
kahyunnam added a commit that referenced this pull request Jul 29, 2026
…(release port for #4107) (#4230)

## Description

Release-specific port of #4177 onto `release-v0.6.16` for #4107.

On SM12x (Spark, RTX Pro 6000), TRTLLM routed-MoE backends were
incorrectly
claiming support (`arch >= 100`) and then dispatching sm100f/sm103a
cubins,
causing `RuntimeError: Error occurred when running GEMM!` or segfaults
in
`test_split_fused_moe_kernel_vs_reference`.

## Changes

- **`csrc/trtllm_batched_gemm_runner.cu`**: Replace per-SM if-chains
with
`isArchCompatible()`; reject unknown cubin families; guard `Sm107a`
behind
`#ifdef TLLM_RUBIN_FEATURES` (only exists in the Rubin cubin pin's
headers).
- **`csrc/trtllm_gemm_runner.cu`**: Same arch filter for the plain GEMM
runner
  (previously had no arch filtering at all).
- **`flashinfer/fused_moe/api.py`**: Tighten `Trtllm*Config.supported()`
from
`arch >= 100` to explicit allowlists `_TRTLLM_ROUTED_ARCHS = (100, 103,
107)`
  and `_TRTLLM_ROUTED_FP8_ARCHS = (100, 103)`.
- **`tests/moe_ep/test_split_fused_moe_kernel_vs_reference.py`**: Gate
GPU tests
on `config_cls.supported(arch)`; add CPU contract tests + SM120
regression guard.

## Release-specific notes

- Includes **sm107** in the allowlist (unlike the rebased #4177 for
`main`, which
  will drop 107 after the #4122 revert).
- `Sm107a` enum case is gated behind `TLLM_RUBIN_FEATURES`, matching the
existing
release pattern from #4122/#4191. Verified against the actual pinned
headers:
the default BMM/GEMM pins do not define `Sm107a`; only the Rubin pins
do.

## Verification

- CPU tests: 36 passed, 2 skipped
(`test_split_fused_moe_kernel_vs_reference.py`)
- Backend claims: sm120/sm121 now fall back to Cutlass only (no Trtllm*
backends)
- Compile: `isArchCompatible()` builds cleanly against both default and
Rubin BMM
  export headers; ported batched runner compiles against default pin

## Related

- Fixes #4107
- Upstream PR: #4177 (targets `main`, has merge conflicts)
- Cherry-pick source commit on `main`: not yet merged

## Pre-existing issue (not in scope)

`#4213` on release references `options.mDtypeSfC`, which does not exist
in the
default BMM cubin pin's headers (only in the Rubin pin). This is a
separate
release-only compile issue on the non-Rubin module, predating this port.
Vinnie6167 added a commit that referenced this pull request Jul 29, 2026
…ackages carry mDtypeSfC (#4235)

## 📌 Description

Fixes two coupled problems on `release-v0.6.16`:

**1. The branch tip cannot compile the trtllm-gen fused-MoE module — on
any architecture.**
#4213 checks `BatchedGemmOptions::mDtypeSfC`, but that field does not
exist in either pinned cubin package's exported headers (Blackwell
`b368d003`, Rubin `46d3f356` — verified 0 occurrences in both published
`BatchedGemmOptions.h`, and in the internal `81a53cff` package as well).
Any JIT or AOT rebuild of `fused_moe_trtllm_*` fails with:
```
trtllm_batched_gemm_runner.cu(128): error: class "BatchedGemmOptions" has no member "mDtypeSfC"
```
(The original PR #4168 shows the same failure in its CI on main.)

**2. The FP4-MoE illegal-memory-access #4213 targets (#4164 / NVBug
6517914) is still live on Rubin.**
The offending `bmm_E2m1xFp32_*` kernels emit linear FP32 output
scale-factors into the E4M3-sized (1 B/block) buffer the MoE pipeline
allocates → device-side overflow → async IMA, near-deterministic under
autotune. Verified against the published catalogs: **132 such kernels,
all in the sm_107a package (`46d3f356`); zero in the Blackwell package**
— which is why the per-arch pin split (#4191) healed Blackwell while
Rubin still crashes.

## What changed

Replace the un-compilable typed check with a kernel-**name** filter
carrying the same intent: reject configs whose `mFunctionName` contains
the `E2m1xFp32` token. This compiles against the *currently published*
packages and removes the bad candidates before autotune can run them.
#4213's other changes (per-token-scaling restructure, launcher
tactic-enumeration alignment) are inherited unchanged.

**Scope note:** the name filter is deliberately narrower than the typed
check — it rejects only the known-bad family. Against the published
catalogs this is behaviorally equivalent (the `E2m1xFp32` kernels are
the only offenders present). The TODO in the code ties restoring the
typed `mDtypeSfC` check to publishing + pinning packages that carry the
field.

## Validation

- Build break reproduced (sbsa AOT wheel build, CTK 13.4) at the branch
tip; gone with this patch.
- VR200 (Rubin) re-run of the affected MoE test files in progress;
results will be posted here before marking ready for review.

Refs: #4164, #4213, #4168 · NVBug 6517914

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Vinnie6167 <Vinnie6167@users.noreply.github.com>
jimmyzho added a commit to jimmyzho/flashinfer that referenced this pull request Jul 30, 2026
Select the Rubin (sm_107 / cc 10.7) deep-gemm artifact directory at
runtime, mirroring the enable_rubin split added for the trtllm-gen
GEMM/BMM/MoE modules in PR flashinfer-ai#4191.

- Add is_rubin_arch() / get_deepgemm_artifact_path() helpers and use them
  in load_all(), load(), and KernelMap.init_indices() so cubins and
  kernel_map.json are fetched from DEEPGEMM_RUBIN on Rubin.
- Add KernelMap.KERNEL_MAP_HASH_RUBIN for the Rubin kernel_map.json.
- Register DEEPGEMM_RUBIN in get_subdir_file_list()'s download list.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jimmyzho added a commit that referenced this pull request Jul 30, 2026
## Summary

Follow-up to #4258 (which added `ArtifactPath.DEEPGEMM_RUBIN` and its
checksum). That PR merged the path + checksum only; this PR adds the
arch-based split that actually routes deep-gemm cubin loading to the
Rubin directory on sm_107, mirroring the `enable_rubin` pattern from
#4191.

### `flashinfer/deep_gemm.py`
- New `is_rubin_arch()` / `get_deepgemm_artifact_path()` helpers
selecting `DEEPGEMM_RUBIN` on arch `107a` and `DEEPGEMM` otherwise.
- Used in `load_all()`, `load()`, and `KernelMap.init_indices()` so both
the kernel cubins and `kernel_map.json` are fetched from the correct
per-arch directory.
- Added `KernelMap.KERNEL_MAP_HASH_RUBIN = f8bf2b1b…a36e27` for the
Rubin `kernel_map.json` (a separate manifest from the default one),
selected by arch.

### `flashinfer/artifacts.py`
- Added `DEEPGEMM_RUBIN` to `get_subdir_file_list()`'s download list so
`flashinfer artifacts download` fetches it, and updated the DEEPGEMM
comment to reference the Rubin map hash.

Both Rubin hashes were verified by downloading the artifacts directly
from the cubin repository.

Targets the `release-v0.6.16` release branch.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aleozlx added a commit that referenced this pull request Jul 30, 2026
<!-- .github/pull_request_template.md -->

## 📌 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)

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

## 🔍 Related Issues

<!-- Link any related issues here -->

#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

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## 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.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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>
dierksen pushed a commit to dierksen/flashinfer that referenced this pull request Sep 11, 2026
flashinfer-ai#4191 gave gen_trtllm_gen_gemm_module() and
gen_trtllm_low_latency_gemm_module() an enable_rubin switch and used it for
three things -- the cubin pin (TRTLLM_GEN_GEMM_RUBIN), the
-DTLLM_RUBIN_FEATURES define, and the module name -- but left the gencode
flags hardcoded:

    + sm100a_nvcc_flags,

So on Rubin both modules were named *_sm107, pointed at the Rubin pin and
compiled with the Rubin define, yet targeted compute_100a. Switch the arch
flags on the same condition as everything else enable_rubin already selects.

This is metadata, not kernels. csrc/trtllm_gemm_runner.cu and
csrc/trtllm_low_latency_gemm_runner.cu contain no __global__, __device__ or
<<<>>>, and neither do the trtllmGen_gemm_export headers; the GEMMs are
precompiled cubins fetched from the pin and launched through the driver API
(loadCubinData -> cuModuleGetFunction -> launchKernel). Only the device-link
stub carried the wrong target, and nothing loads it, which is why Rubin has
been functionally correct all along and no test caught this. PR flashinfer-ai#5096's
per-provider SASS validation is the first thing to read that label.

Controlled A/B on an SM107 (GR100, compute_cap 10.7), CUDA 13.5, same tree,
only these lines differing, cuobjdump --list-elf on the built .so:

    enable_rubin=True   (aot.py takes this path under has_sm107)
      before: GENCODE=['-gencode=arch=compute_100a,code=sm_100a']
              trtllm_gemm_sm107.1.sm_100a.cubin
              trtllm_low_latency_gemm_sm107.1.sm_100a.cubin
      after:  GENCODE=['-gencode=arch=compute_107a,code=sm_107a']
              trtllm_gemm_sm107.1.sm_107a.cubin
              trtllm_low_latency_gemm_sm107.1.sm_107a.cubin

    enable_rubin=False  (Blackwell, unchanged in both arms)
      trtllm_gemm.1.sm_100a.cubin
      trtllm_low_latency_gemm.1.sm_100a.cubin

Both arms build clean, so CUDA 13.5 accepts compute_107a for these TUs.

sm107a_nvcc_flags already existed in flashinfer/jit/core.py and was simply
unused at these two sites. The Blackwell branch is the previous behaviour
verbatim, so the only path that changes is has_sm107.

Note the parenthesisation: a few lines above, gen_tgv_gemm_sm10x_module()
writes this same idiom without parens, where `if/else` binds looser than `+`
and silently drops the cflags list in the else branch. That is a separate
pre-existing bug and is not touched here.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants