fix: key CUDA property caches by device index

[abligail]

📌 Description

This PR fixes CUDA device property helpers whose cached results could be incorrectly reused across different GPUs when called with an unindexed torch.device("cuda").

The changes:

• Resolve an unindexed CUDA device to torch.cuda.current_device() before cache lookup.
• Key compute capability, shared-memory capacity, SM count, PDL support, default generator, and memory-bandwidth caches by concrete CUDA device index.
• Resolve NVML devices by GPU/MIG UUID so that CUDA_VISIBLE_DEVICES remapping is handled correctly.
• Use the maximum NVML memory clock when calculating theoretical peak memory bandwidth instead of caching the current clock.
• Add regression tests covering current-device switching, explicit CUDA devices, CPU rejection, UUID normalization, and per-device cache isolation.

🔍 Related Issues

N/A

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

Validation completed on a Linux CUDA environment with two NVIDIA GeForce RTX 5090 GPUs:

• pre-commit run --all-files: all 14 hooks passed.
• pytest -q tests/utils/test_device_utils.py: 10 passed.
• Dual-GPU validation with CUDA_VISIBLE_DEVICES=1,0: passed, including correct logical-to-physical UUID mapping and separate cache entries for both devices.
• Isolated tests/utils run: 19,099 passed and 973 skipped.

The remaining 57 failures and one collection error reproduce identically on origin/main:

• FP4 quantization: 56 failures on both this branch and origin/main.
• cuDNN logging replay: 1 failure on both this branch and origin/main.
• IPC test collection: the same Python 3.11 array.array[int] annotation error occurs on both branches.

A repository-wide collection check discovered 524,276 tests with no branch-specific collection regressions.

Reviewer focus areas:

• Resolution of unindexed CUDA devices before caching.
• NVML UUID-based lookup under CUDA_VISIBLE_DEVICES remapping.
• Cache isolation across multiple CUDA devices.

Summary by CodeRabbit

• Bug Fixes
  • Improved CUDA device selection for unindexed devices and explicit device indices.
  • Ensured GPU properties, bandwidth, shared-memory, and generator queries target the correct active device.
  • Improved handling of GPU identifiers, including string inputs and MIG-style UUIDs.
  • Added clearer validation for CUDA-only operations and CPU device inputs.
  • Corrected device capability and PDL checks to reflect the selected device.

[gemini-code-assist]

Caution

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

[coderabbitai]

📝 Walkthrough

Walkthrough

CUDA utility functions now resolve concrete CUDA indices, use device-specific caching where required, validate CUDA-only inputs, resolve NVML bandwidth by UUID, and select default generators by index. Tests cover device switching, cache behavior, MIG UUIDs, generator selection, and CPU handling.

Changes

CUDA device utilities

Layer / File(s)	Summary
Device index and cached queries flashinfer/utils.py, tests/utils/test_device_utils.py	Device-property and shared-memory queries use concrete CUDA indices. Dynamic SM-count, name, and PDL helpers no longer use stale caches. Tests cover device switching, cache sizes, and CPU behavior.
Per-device bandwidth lookup flashinfer/utils.py, tests/utils/test_device_utils.py	GPU memory bandwidth accepts string or device inputs and resolves NVML handles by GPU or MIG UUID. Tests cover indexed devices and mocked NVML results.
Indexed default generators flashinfer/utils.py, tests/utils/test_device_utils.py	Default generator lookup validates CUDA devices and selects generators by concrete index. Tests cover current-device changes and explicit cuda:0 selection.

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

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant DeviceUtils
  participant TorchCUDA
  participant NVML
  Caller->>DeviceUtils: request device-specific utility
  DeviceUtils->>TorchCUDA: resolve concrete CUDA index
  TorchCUDA-->>DeviceUtils: return device index or properties
  DeviceUtils->>NVML: resolve UUID for bandwidth lookup
  NVML-->>DeviceUtils: return bandwidth data
  DeviceUtils-->>Caller: return indexed utility result

Loading

Possibly related PRs

• flashinfer-ai/flashinfer#4144: Both PRs update CUDA device-query handling and shared runtime GPU helpers.

Suggested reviewers: sricketts, bkryu, guangyunh-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: isolating CUDA property caches by device index.
Description check	✅ Passed	The description follows the template, explains the changes, documents validation, and clearly identifies baseline failures and remaining test status.
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

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

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@flashinfer/utils.py`:
- Around line 1450-1452: Add `@functools.cache` to _get_default_generator so
generator lookups are cached per device_index and the function exposes
cache_info() for the regression test. Ensure functools is imported or otherwise
available in the module.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 366889d3-772b-4d47-9fff-57e69960bd5e

📥 Commits

Reviewing files that changed from the base of the PR and between f2e0440 and 7a43a92.

📒 Files selected for processing (2)

• flashinfer/utils.py
• tests/utils/test_device_utils.py

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

• tests/utils/test_device_utils.py