fix(gdn): defer CUDA device queries in gdn_decode_bf16_state to lazy helpers

[AmirF194]

📌 Description

gdn_decode_bf16_state.py computed NUM_SMS/_GPU_MAJOR/_USE_PACKED_FMA via eager
torch.cuda.get_device_properties(0)/get_device_capability(0) calls at module scope, so
import flashinfer crashed on any machine without a visible CUDA device.

#3293 fixed this exact pattern for #3262 by moving the calls into lazy, device-parameterized
helpers (get_num_sm/get_sm_version, already the convention in every other gdn_kernels/*.py
file). #3502 added new code to the same file and reintroduced the eager module-scope pattern
without reusing that fix. git log -G on the two lines shows the full lifecycle: introduced
(#2679) -> fixed (#3293) -> reintroduced (#3502).

This PR reapplies #3293's fix to the code #3502 added: _select_tile_v_for_mtp and
_get_bf16_mtp_config now take a device keyword, threaded from q.device at each of the three
call sites (all three functions assert q is not None earlier in the same function body). The
module-scope globals are removed.

🔍 Related Issues

Fixes #4143

🚀 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.).

Reviewer Notes

Verified in a clean python:3.12-slim container (CPU-only torch, CUDA_VISIBLE_DEVICES=, no
GPU hardware): import flashinfer and import flashinfer.gdn_decode both crash on unpatched
main (AssertionError: Torch not compiled with CUDA enabled at the NUM_SMS line) and both
succeed after this change. tests/test_gdn_decode_cpu_import.py pins this as a subprocess test
(fails on main, passes on the branch, confirmed both ways) so a future addition to this file
can't reintroduce the eager pattern a third time.

I have not run this against an actual GPU: the three call sites already had assert q is not None in scope before reaching the changed lines, and get_num_sm/get_sm_version are the
same helpers already exercised by every sibling kernel file, but the packed-FMA / SM-count
dispatch itself is unverified on real hardware in this session.

Summary by CodeRabbit

• Bug Fixes

  • Improved GPU-aware kernel selection across systems with multiple GPU types.
  • Fixed wide-vector and multi-token prediction execution paths to select appropriate hardware optimizations at runtime.
  • Improved reliability when importing the package without visible CUDA devices.
  • Ensured device-specific decisions are based on the active query tensor’s GPU.

• Tests

  • Added regression coverage for importing the package when GPUs are hidden.
  • Added validation for CUDA-dependent behavior in environments without visible devices.

[gemini-code-assist]

Caution

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

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: debd8920-b4c6-4b6d-b61e-99944e8fdb51

📥 Commits

Reviewing files that changed from the base of the PR and between f32f740 and 88c7feb.

📒 Files selected for processing (2)

• flashinfer/gdn_kernels/gdn_decode_bf16_state.py
• tests/gdn/test_decode_pretranspose_bf16_padded_pool.py

🚧 Files skipped from review as they are similar to previous changes (1)

• flashinfer/gdn_kernels/gdn_decode_bf16_state.py

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

---

📝 Walkthrough

Walkthrough

The GDN BF16 decode module now detects SM properties at runtime for the target device. MTP tile selection and packed-FMA dispatch are device-aware. Subprocess tests verify imports succeed without visible CUDA devices.

Changes

GDN device-aware dispatch

Layer / File(s)	Summary
Runtime device configuration flashinfer/gdn_kernels/gdn_decode_bf16_state.py	Replaces import-time SM detection with runtime helpers. MTP tile selection receives the device, and wide-vector paths derive packed-FMA support from q.device.
Import regression coverage tests/gdn/test_decode_pretranspose_bf16_padded_pool.py	Adds subprocess tests for importing flashinfer.gdn_decode and flashinfer with CUDA_VISIBLE_DEVICES cleared.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to 88c7f

CUDA device detection is deferred until decode kernels run, allowing flashinfer and GDN imports to succeed when no CUDA device is visible. The CPU-only import regression coverage passes, with no concrete current-head merge-blocking risk identified.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 12 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: deferring CUDA device queries in the GDN decode state module.
Description check	✅ Passed	The description explains the regression, implementation, related issue, tests, checklist status, and reviewer notes. It also clearly states that real-GPU validation was not performed.
Linked Issues check	✅ Passed	The changes satisfy issue #4143 by removing eager module-scope CUDA queries, threading device-aware lazy helpers through all relevant call sites, and adding CPU-only import regression tests.
Out of Scope Changes check	✅ Passed	The modified implementation and regression tests directly support the linked issue and stated objectives. No unrelated code changes are identified.

• Fix all pre-merge checks with AI

✨ 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.}

[AmirF194]

Checking in after a week with no review yet. Happy to split the fix further or add more test coverage if that would help; no rush.

[bkryu]

Thank you for the fix @AmirF194, I left one comment about the organization of added tests. Can you move the newly added tests into an existing file instead of adding a new file?

[bkryu]

There is a tests/gdn directory and we'd like to organize our test files.

However, what might be better is to simply put the tests in this file inside a suitable existing file because it covers only two cases.

[AmirF194]

Moved the two import-crash tests into tests/gdn/test_decode_pretranspose_bf16_padded_pool.py and deleted the standalone file. Reran them under CPU-only Docker (CUDA_VISIBLE_DEVICES unset): both still pass in the new location.

[bkryu]

/bot run tests/gdn

[flashinfer-bot]

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

[flashinfer-bot]

[FAILED] Pipeline #60922613 — 16/18 executed test jobs passed

Compared with nightly #60831563.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Notes
5090	✅ Pass	✅ Pass	—
B300	✅ Pass	✅ Pass	—
GB200	✅ Pass	✅ Pass	—
GB300	✅ Pass	✅ Pass	—
H100	✅ Pass	✅ Pass	—
RTX Pro 6000 Blackwell	✅ Pass	✅ Pass	—

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Multi-GPU and Multi-Node Tests — 4/6 passed

GPU	CUDA 12.9	CUDA 13.0	Notes
B300 (multi-GPU)	❔ Failed	❔ Failed	—
GB200 (multi-node)	✅ Pass	✅ Pass	—
GB300 (multi-node)	✅ Pass	✅ Pass	—

No individual test or infrastructure failures could be extracted.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[AmirF194]

Rebased past main to resolve a conflict: main added dtype_compat.as_bf16 and a _dtype_key helper in this same file (unrelated to this fix) since the internal CI ran. Kept both changes; they don't overlap functionally with the lazy device-query fix here.

Re-verified in Docker (CPU-only, CUDA_VISIBLE_DEVICES= unset): import flashinfer, import flashinfer.gdn_decode, and import flashinfer.gdn_kernels.gdn_decode_bf16_state all succeed post-rebase, and both regression tests in tests/gdn/test_decode_pretranspose_bf16_padded_pool.py (test_import_gdn_decode_without_cuda_device, test_import_flashinfer_without_cuda_device) pass. ruff check/ruff format --check clean. Both commits SSH-signed and verified.

I don't have GPU hardware here to re-run the actual kernel numerics, so the internal CI pipeline is the source of truth for that (it previously reported 16/18, all unit tests green across every GPU/CUDA combination; the two B300 (multi-GPU) jobs came back with no individual test or infrastructure failure extractable from the report, so I can't tell from here whether that's related to this change or not). Happy to have /bot run re-triggered on the rebased head whenever convenient.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[AmirF194]

Rebased past main again to clear a second conflict in the same file: main's new build_and_load_cute_dsl_kernel/make_kernel_name imports landed next to this fix's get_sm_version/get_num_sm imports, same shape as the 08-12 rebase. Kept both, no functional overlap (checked call sites). CPU-only import scope unchanged.

[flashinfer-bot]

[FAILED] Pipeline #60922613 — 16/18 executed test jobs passed

No usable JUnit artifact was available; individual tests and nightly comparison could not be recovered.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Notes
5090	✅ Pass	✅ Pass	—
B300	✅ Pass	✅ Pass	—
GB200	✅ Pass	✅ Pass	—
GB300	✅ Pass	✅ Pass	—
H100	✅ Pass	✅ Pass	—
RTX Pro 6000 Blackwell	✅ Pass	✅ Pass	—

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Multi-GPU and Multi-Node Tests — 4/6 passed

GPU	CUDA 12.9	CUDA 13.0	Notes
B300 (multi-GPU)	❔ Unknown	❔ Unknown	Unknown: script failed before producing a JUnit report (2 jobs; CUDA 12.9, CUDA 13.0)
GB200 (multi-node)	✅ Pass	✅ Pass	—
GB300 (multi-node)	✅ Pass	✅ Pass	—

Failure details

Timeouts, infrastructure, or incomplete jobs

• Unknown: script failed before producing a JUnit report — B300 / CUDA 12.9 / multi-GPU, B300 / CUDA 13.0 / multi-GPU
  • Jobs: multi_gpu_test_b300: [cu130], multi_gpu_test_b300: [cu129]

[flashinfer-bot]

[FAILED] Pipeline #60922613 — 16/18 executed test jobs passed

No usable JUnit artifact was available; individual tests and nightly comparison could not be recovered.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Notes
5090	✅ Pass	✅ Pass	—
B300	✅ Pass	✅ Pass	—
GB200	✅ Pass	✅ Pass	—
GB300	✅ Pass	✅ Pass	—
H100	✅ Pass	✅ Pass	—
RTX Pro 6000 Blackwell	✅ Pass	✅ Pass	—

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Multi-GPU and Multi-Node Tests — 4/6 passed

GPU	CUDA 12.9	CUDA 13.0	Notes
B300 (multi-GPU)	❔ Unknown	❔ Unknown	Unknown: script failed before producing a JUnit report (2 jobs; CUDA 12.9, CUDA 13.0)
GB200 (multi-node)	✅ Pass	✅ Pass	—
GB300 (multi-node)	✅ Pass	✅ Pass	—

Failure details

Timeouts, infrastructure, or incomplete jobs

• Unknown: script failed before producing a JUnit report — B300 / CUDA 12.9 / multi-GPU, B300 / CUDA 13.0 / multi-GPU
  • Jobs: multi_gpu_test_b300: [cu130], multi_gpu_test_b300: [cu129]

[AmirF194]

Closing this one: #4507 just merged and removes these exact two module-scope calls (NUM_SMS/_GPU_MAJOR) as part of a broader per-device target refactor, so the import crash from #4143 is fixed there instead. Verified the new gdn_decode_bf16_state.py has zero torch.cuda calls left at module scope.