Support NVFP4 KV cache decode on SM120

[Tom-Zheng]

📌 Description

Supports NVFP4 KV cache.

🔍 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

Summary by CodeRabbit

• New Features

  • xqa and decode APIs now accept optional NVFP4 (4-bit) KV-cache scaling tensors (k_sf_cache/v_sf_cache / kv_cache_sf); torch.uint8 KV-cache dtype supported with runtime SM 12.x guard; docstrings updated.

• Tests

  • Added/extended tests exercising NVFP4 KV-cache quantization, scaling factors, uint8 KV-cache paths, and decode integration.

[gemini-code-assist]

Summary of Changes

Hello @Tom-Zheng, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request integrates NVIDIA FP4 (NVFP4) support for Key-Value (KV) cache decoding on SM120 GPUs. This feature aims to significantly reduce the memory footprint of KV caches, enabling larger models or longer context windows with improved performance. The changes involve extensive modifications to CUDA kernels for efficient 4-bit data handling and corresponding updates to the Python API to expose this new capability to users.

Highlights

• NVFP4 KV Cache Support: Introduced support for NVIDIA FP4 (NVFP4) Key-Value (KV) cache decoding, specifically targeting SM120 (Hopper) GPUs for enhanced memory efficiency.
• CUDA Kernel Modifications: Updated core CUDA kernels (mha.cu, mhaUtils.cuh, utils.cuh) to handle NVFP4 data types, including new data structures for scaling factors, specialized matrix load/unpack operations, and adjusted memory access patterns.
• Python API Integration: Modified Python bindings (flashinfer/decode.py, flashinfer/xqa.py, flashinfer/jit/xqa.py) to expose NVFP4 KV cache functionality, allowing users to pass scaling factors and specify torch.uint8 as the KV cache data type.
• Scaling Factor Handling: Implemented mechanisms to load, permute, and apply scaling factors (kSfCacheVLLM, vSfCacheVLLM) alongside the 4-bit KV cache data, crucial for maintaining numerical precision during computation.
• New Test Case: Added a dedicated test case (test_xqa_batch_decode_nvfp4_kv) to validate the correctness and functionality of the NVFP4 KV cache implementation.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• csrc/flashinfer_xqa_binding.cu
  • Updated xqa_wrapper function signature to accept optional KV cache scaling factor tensors (kSfCacheVLLM, vSfCacheVLLM).
• csrc/xqa/defines.h
  • Added ENABLE_4BIT_KV_CACHE macro to conditionally enable 4-bit cache features.
  • Introduced ElemTypeConverter struct with a specialization for __nv_fp4_e2m1 to define NVFP4 data types and properties.
• csrc/xqa/mha.cu
  • Defined grainBytesGmemCache and grainBytesSf for specific 4-bit cache memory handling.
  • Modified SharedMem struct to include KSfSmemBuffer and VSfSmemBuffer for storing KV cache scaling factors in shared memory.
  • Added storeReorderedXTile function for specific data reordering in NVFP4 context.
  • Updated loadInstInMat template to include an is4BitElem parameter, enabling conditional NVFP4 data loading and unpacking.
  • Modified smemQKPartGemm and smemXVPartGemm functions to incorporate loading, permuting, and applying scaling factors for NVFP4 data.
  • Adjusted mergeAndSaveOutTile call based on ENABLE_4BIT_KV_CACHE.
  • Updated launchMHA and launchMHAFlashInfer to pass scaling factor pointers and adjust stride calculations for container elements.
• csrc/xqa/mha.h
  • Introduced CacheElemConverter and updated CacheElem definition to use the converter.
  • Defined GMemCacheHead and GMemCacheHeadSf types to handle packed NVFP4 data and its scaling factors in global memory.
  • Defined PaddedCacheHead and PaddedCacheHeadSf for padded NVFP4 data and scaling factors in shared memory.
  • Updated launchMHA and launchMHAFlashInfer function signatures to include scaling factor pointers.
• csrc/xqa/mhaUtils.cuh
  • Modified copyPartialHeadsAsync and copyHeadsAsync template parameters to include grainBytesSmem and grainBytesGmem for flexible memory access based on data type.
  • Updated KVCacheList struct to include kSfCacheVLLM and vSfCacheVLLM pointers for scaling factors.
• csrc/xqa/utils.cuh
  • Added ldmatrix_16x16_trans_unpack_4b, ldmatrix_8x16_4x_unpack_4b, and ldmatrix_4x_unpack_4b for NVFP4 matrix loading and unpacking from shared memory.
  • Added convertKCacheWordToF16 specializations for __nv_fp4_e2m1 and applyF16ScalingFactors for applying scaling to FP16/BF16 data.
• csrc/xqa/xqa_wrapper.cu
  • Updated xqa_wrapper function signature to accept optional kSfCacheVLLM and vSfCacheVLLM tensors.
  • Passed the new scaling factor pointers to the underlying mha_func.
• flashinfer/decode.py
  • Added kv_cache_sf parameter to trtllm_batch_decode_with_kv_cache and xqa_batch_decode_with_kv_cache functions.
  • Implemented logic to unpack kv_cache_sf into k_cache_sf and v_cache_sf.
  • Passed k_cache_sf and v_cache_sf to the underlying C++ XQA function.
• flashinfer/fp4_quantization.py
  • Modified nvfp4_quantize to conditionally set is_sf_swizzled_layout based on SfLayout.
• flashinfer/jit/xqa.py
  • Added torch.uint8 as a recognized kv_cache_dtype for CACHE_ELEM_ENUM=3, mapping it to NVFP4.
• flashinfer/xqa.py
  • Updated xqa function signature to accept k_sf_cache and v_sf_cache.
  • Added logic to transpose k_sf_cache and v_sf_cache if kv_layout is HND.
  • Included assertions to ensure NVFP4 KV cache is used only on SM120 GPUs and that scaling factors are provided when k_cache.dtype == torch.uint8.
  • Passed k_sf_cache and v_sf_cache to the underlying C++ XQA function.
• tests/attention/test_xqa_batch_decode.py
  • Imported NVFP4 quantization utilities from flashinfer.fp4_quantization.
  • Added to_nvfp4 and nvfp4_to_float helper functions for quantization and dequantization in tests.
  • Modified create_kv_cache to handle 'nvfp4' dtype and return global scales for testing.
  • Introduced a new test test_xqa_batch_decode_nvfp4_kv to verify NVFP4 KV cache functionality on SM120 GPUs.

Activity

• The pull request was authored by Tom-Zheng.
• Pre-commit checks were successfully run, indicating adherence to code style and quality standards.
• Tests have been added or updated and are passing, ensuring the new functionality works as expected and existing functionality is not regressed.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds NVFP4 (4-bit) KV-cache support end-to-end: type-system converters, kernel unpack/scale paths, dual-grain GMEM/SMEM handling, new SF cache pointers threaded through C++ wrappers and Python APIs, and tests exercising NVFP4 KV paths with runtime SM120 checks.

Changes

Cohort / File(s)	Summary
Type system & public types csrc/xqa/defines.h, csrc/xqa/mha.h	Introduce ElemTypeConverter/CacheElemConverter specializations including NVFP4, add ENABLE_4BIT_KV_CACHE flag; change CacheElem/container types and add GF/SM scaling-factor head types (GMemCacheHeadSf, PaddedCacheHeadSf).
Kernel implementation & launchers csrc/xqa/mha.cu, csrc/xqa/mha.h, csrc/xqa/mhaUtils.cuh, csrc/xqa/utils.cuh	Add 4-bit SF GMEM/SMEM buffers, unpack/load/store/reorder paths, FP16/BF16 conversion + scaling, grain-size-aware copy APIs, is4BitElem template paths, and extend launch signatures to accept SF cache pointers (guarded by flag).
Bindings / wrapper layer csrc/flashinfer_xqa_binding.cu, csrc/xqa/xqa_wrapper.cu	Update xqa_wrapper signature to accept optional kSfCacheVLLM/vSfCacheVLLM, derive runtime pointers, and pass SF caches into kernel launches when enabled; reorder kvCachePageList/sequence-related params.
Python API, JIT & decode path flashinfer/xqa.py, flashinfer/decode.py, flashinfer/jit/xqa.py	Expose optional k_sf_cache/v_sf_cache and kv_cache_sf through public APIs and internal calls; add runtime SM120 guard for uint8 KV paths; JIT emits compile flag for torch.uint8 KV dtype.
Quantization behavior flashinfer/fp4_quantization.py	Adjust is_sf_swizzled_layout to depend on chosen SfLayout (linear vs non-linear) rather than always being true.
Tests tests/attention/test_xqa_batch_decode.py, tests/attention/test_xqa.py	Add NVFP4 test utilities (to_nvfp4, nvfp4_to_float), return/propagate KV scaling factors from create_kv_cache, add NVFP4-specific tests, and update existing xqa tests to pass new optional SF args.

Sequence Diagram(s)

sequenceDiagram
    participant Host as Host (Python)
    participant Wrapper as XQA Wrapper (C++)
    participant Kernel as MHA Kernel (CUDA)
    participant GMEM as Global Memory
    participant SMEM as Shared Memory

    Host->>Host: Quantize KV to NVFP4 and compute k_sf_cache / v_sf_cache
    Host->>Wrapper: xqa_wrapper(k_cache, v_cache, k_sf_cache, v_sf_cache, ...)
    Wrapper->>Wrapper: Optional cast to GMemCacheHeadSf*
    Wrapper->>Kernel: Launch kernel with KV + SF cache pointers
    Kernel->>GMEM: Load packed 4-bit K/V and SF arrays
    Kernel->>SMEM: Copy K/V and SF into SMEM buffers (grain-aware)
    Kernel->>Kernel: ldmatrix_4x_unpack_4b() -> convertKCacheWordToF16() -> applyF16ScalingFactor()
    Kernel->>Kernel: Compute attention using unpacked FP16/BF16
    Kernel->>GMEM: Write outputs

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• enable xqa fp8 output #2081 — modifies xqa_wrapper signature in csrc/flashinfer_xqa_binding.cu (similar parameter additions/reordering).
• use scalar for kv_scale in xqa #2033 — updates xqa_wrapper family signatures and kv-cache handling (kvCacheScale-related changes touching same binding paths).
• feat: add xqa backend and completes NHD/HND coverage for trtllm-gen/xqa backend #2001 — changes xqa wrapper/launch interfaces and parameter ordering related to KV cache wiring.

Suggested labels

run-ci

Suggested reviewers

• cyx-6
• jiahanc
• nvmbreughe
• yzh119
• djmmoss
• kahyunnam

Poem

🐰 Tiny bits in a carrot-case,
SF scales tucked in their place,
Kernels nibble, reorder, hop,
SMEM hums and outputs drop,
A rabbit cheers — small bytes embrace!

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is incomplete. While it states 'Supports NVFP4 KV cache' and marks checklist items as complete, it lacks substantive detail about what was changed, why the changes were needed, related issues, and specific implementation notes.	Expand the description section to explain the NVFP4 KV cache implementation, why SM120 support is needed, and provide context on key architectural changes made across multiple files.
Docstring Coverage	⚠️ Warning	Docstring coverage is 26.92% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Support NVFP4 KV cache decode on SM120' accurately and clearly summarizes the main purpose of the PR - adding NVFP4 KV cache support for SM120 architecture.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

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 and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces support for NVFP4 KV cache on the SM120 architecture. The changes are extensive, modifying CUDA C++ kernels, helper utilities, and Python wrappers to handle the new 4-bit data type and its associated scaling factors. The refactoring to use ElemTypeConverter is a good design choice for managing different data types. However, I've identified a critical issue in the CUDA code involving an out-of-scope variable that would likely prevent compilation, and a logic error in a test case that could lead to incorrect test validation. My review includes suggestions to fix these issues.

[coderabbitai]

Actionable comments posted: 4

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

tests/attention/test_xqa_batch_decode.py (1)

107-175: ⚠️ Potential issue | 🟡 Minor

Use v_global_scale when dequantizing V cache.

The NVFP4 reference path currently uses k_global_scale for V, which will skew the reference when K/V scales differ.

🛠️ Suggested fix

-        v_cache_dq = nvfp4_to_float(v_cache, v_scale, k_global_scale)
+        v_cache_dq = nvfp4_to_float(v_cache, v_scale, v_global_scale)

csrc/xqa/mha.cu (1)

2885-2893: ⚠️ Potential issue | 🟡 Minor

launchMHA is dead code and never invoked—stride calculation issue confirmed but latent.

The stride calculation difference is real: launchMHA divides by validElemsPerHead (lines 2890–2893), while launchMHAFlashInfer divides by container_elems_per_head = validElemsPerHead / CacheElemConverter::ElemsPerContainer (lines 2974–2979). For 4-bit KV cache (ElemsPerContainer == 2), this is a 2× difference.

However, launchMHA is never called from anywhere in the codebase—neither from xqa_wrapper.cu (which always dispatches to launchMHAFlashInfer) nor from any other entry point. This is dead code. The bug is real but unreachable.

If launchMHA is intentionally kept for backward compatibility or future use, apply the suggested fix to align stride calculations with launchMHAFlashInfer.

Suggested fix — align launchMHA stride calculation with launchMHAFlashInfer

   // Convert stride from elements to Heads
-  uint32_t const stride_page_in_heads = static_cast<uint32_t>(kv_stride_page / validElemsPerHead);
-  uint32_t const stride_token_in_heads = static_cast<uint32_t>(kv_stride_token / validElemsPerHead);
-  uint32_t const stride_head_in_heads = static_cast<uint32_t>(kv_stride_head / validElemsPerHead);
+  uint32_t const container_elems_per_head =
+      validElemsPerHead / CacheElemConverter::ElemsPerContainer;
+  uint32_t const stride_page_in_heads =
+      static_cast<uint32_t>(kv_stride_page / container_elems_per_head);
+  uint32_t const stride_token_in_heads =
+      static_cast<uint32_t>(kv_stride_token / container_elems_per_head);
+  uint32_t const stride_head_in_heads =
+      static_cast<uint32_t>(kv_stride_head / container_elems_per_head);

🤖 Fix all issues with AI agents

In `@csrc/xqa/mha.cu`:
- Around line 1169-1185: The vSfPrefetch register buffer is populated but never
used; either remove the prefetch block or make the main consumer read from
vSfPrefetch. To fix: if removing, delete the Array2D<uint32_t,...> vSfPrefetch
declaration and the entire prefetch loop (the `#if` ENABLE_4BIT_KV_CACHE block
that fills vSfPrefetch); otherwise change the consumer loop that currently calls
vSf(rowIdx, colIdx) to index into vSfPrefetch using the same sfRowIdx
calculation (use sfRowIdxInSlice / sfRowIdx mapping produced in the prefetch
loop and nbSfPrefetchBuffers for column indexing) so the registers filled by
vSfPrefetch are consumed instead of reloading vSf. Ensure you keep
ENABLE_4BIT_KV_CACHE gating consistent and remove any now-unused variables if
you choose removal.
- Around line 2113-2121: In the loadPages lambda remove the accidental reference
to idxNextSMemVBuf and the unused dstSf: do not call getSmemVSfTile or declare
dstSf inside loadPages (this symbol is defined only in loadVTilePart), and
ensure the ENABLE_4BIT_KV_CACHE branch only declares/uses dst (the
smem.vCachePages reference) — delete the unused dstSf/idxNextSMemVBuf lines so
the lambda compiles when BEAM_WIDTH > 1; keep loadPages calling
loadPagesForBeamSearchAsync with dst unchanged.

In `@flashinfer/xqa.py`:
- Around line 306-311: The check that calls
get_compute_capability(torch.device(device="cuda")) queries the default CUDA
device and can be wrong for tensors on other GPUs; update the capability checks
in the NVFP4 path to call get_compute_capability with the actual tensor device
(e.g. use q.device or k_cache.device) instead of torch.device(device="cuda") so
the assertions around k_cache.dtype, k_sf_cache, and v_sf_cache validate the
correct GPU for the tensor in functions/blocks referencing
get_compute_capability and the NVFP4 KV path.

In `@tests/attention/test_xqa_batch_decode.py`:
- Around line 582-765: The test function test_xqa_batch_decode_nvfp4_kv binds an
unused variable in_kv_lens from generate_seq_lens_decode; remove the unused
binding by changing the assignment to only capture needed values (e.g., q_lens,
seq_lens = generate_seq_lens_decode(...)) or replace in_kv_lens with an
underscore (_) so the unused return value is ignored; update the call site of
generate_seq_lens_decode in test_xqa_batch_decode_nvfp4_kv and ensure no other
references to in_kv_lens remain.

🧹 Nitpick comments (1)

csrc/xqa/mha.h (1)

84-87: Inconsistent use of plain division vs. exactDiv for PaddedCacheHeadSf.

Line 64 uses exactDiv(validElemsPerHead, CacheElemConverter::QuantVectorSize) for GMemCacheHeadSf, but line 86 uses plain / for PaddedCacheHeadSf. Using exactDiv here would provide a compile-time safety net against future misconfigurations where headElems might not be exactly divisible by QuantVectorSize.

Suggested fix

 `#if` ENABLE_4BIT_KV_CACHE
 using PaddedCacheHeadSf =
-    Vec<CacheElemConverter::ScalingFactorType, headElems / CacheElemConverter::QuantVectorSize>;
+    Vec<CacheElemConverter::ScalingFactorType, exactDiv(headElems, CacheElemConverter::QuantVectorSize)>;
 `#endif`

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

get_compute_capability uses default CUDA device instead of q.device.

Lines 307 and 313 call get_compute_capability(torch.device(device="cuda")) which queries the default CUDA device. If q resides on a non-default GPU in a multi-GPU setup, this could yield an incorrect capability check. The same pattern exists at line 300, so this is pre-existing, but worth noting since it gates the new NVFP4 path.

Suggested fix

-    if k_cache.dtype == torch.uint8:
-        assert get_compute_capability(torch.device(device="cuda"))[0] in [12], (
+    if k_cache.dtype == torch.uint8:
+        assert get_compute_capability(q.device)[0] in [12], (
             "XQA NVFP4 KV is only supported on SM120 GPUs"
         )
         assert k_sf_cache is not None, "K SF cache is required when NVFP4 KV is used"
         assert v_sf_cache is not None, "V SF cache is required when NVFP4 KV is used"

-    if get_compute_capability(torch.device(device="cuda"))[0] not in [9, 10, 12]:
+    if get_compute_capability(q.device)[0] not in [9, 10, 12]:
         raise RuntimeError("XQA is only supported on SM90, SM100, SM120 GPUs")

🤖 Prompt for AI Agents

In `@flashinfer/xqa.py` around lines 306 - 311, The check that calls
get_compute_capability(torch.device(device="cuda")) queries the default CUDA
device and can be wrong for tensors on other GPUs; update the capability checks
in the NVFP4 path to call get_compute_capability with the actual tensor device
(e.g. use q.device or k_cache.device) instead of torch.device(device="cuda") so
the assertions around k_cache.dtype, k_sf_cache, and v_sf_cache validate the
correct GPU for the tensor in functions/blocks referencing
get_compute_capability and the NVFP4 KV path.

[samuellees]

I have verified the functions through SGLang Qwen model. The PR for SGLang sgl-project/sglang#18314

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

flashinfer/xqa.py (1)

154-160: ⚠️ Potential issue | 🟠 Major

Keep k_sf_cache / v_sf_cache optional in public API to avoid breakage.

Adding them as required positional parameters breaks existing xqa(...) callers that do not use NVFP4 KV.

💡 Proposed fix

 def xqa(
     q: torch.Tensor,
     k_cache: torch.Tensor,
     v_cache: torch.Tensor,
-    k_sf_cache: Optional[torch.Tensor],
-    v_sf_cache: Optional[torch.Tensor],
+    k_sf_cache: Optional[torch.Tensor] = None,
+    v_sf_cache: Optional[torch.Tensor] = None,
     page_table: torch.Tensor,
     seq_lens: torch.Tensor,
     output: torch.Tensor,

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/xqa.py` around lines 154 - 160, The xqa function signature
currently requires k_sf_cache and v_sf_cache as positional parameters which
breaks callers that don't use NVFP4 KV; make these parameters optional by giving
them default value None in the xqa(...) signature (keep their types as
Optional[torch.Tensor]) and update the function body to handle None (e.g., skip
NVFP4-specific logic when k_sf_cache or v_sf_cache is None); ensure the symbols
to change are the xqa(...) function header and any internal branches that
reference k_sf_cache / v_sf_cache so behavior remains unchanged for callers that
pass them and is safe for callers that omit them.

csrc/xqa/mha.cu (1)

2884-2892: ⚠️ Potential issue | 🟠 Major

Use container-based stride conversion in launchMHA for 4-bit KV.

launchMHAFlashInfer already switched to container-aware stride normalization, but launchMHA still divides by validElemsPerHead. For containerized 4-bit cache this can misaddress KV pages/tokens/heads.

💡 Proposed fix

-  uint32_t const stride_page_in_heads = static_cast<uint32_t>(kv_stride_page / validElemsPerHead);
-  uint32_t const stride_token_in_heads = static_cast<uint32_t>(kv_stride_token / validElemsPerHead);
-  uint32_t const stride_head_in_heads = static_cast<uint32_t>(kv_stride_head / validElemsPerHead);
+  uint32_t const container_elems_per_head =
+      validElemsPerHead / CacheElemConverter::ElemsPerContainer;
+  uint32_t const stride_page_in_heads =
+      static_cast<uint32_t>(kv_stride_page / container_elems_per_head);
+  uint32_t const stride_token_in_heads =
+      static_cast<uint32_t>(kv_stride_token / container_elems_per_head);
+  uint32_t const stride_head_in_heads =
+      static_cast<uint32_t>(kv_stride_head / container_elems_per_head);

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@csrc/xqa/mha.cu` around lines 2884 - 2892, In launchMHA the stride conversion
still divides kv_stride_page/token/head by validElemsPerHead (producing
stride_page_in_heads, stride_token_in_heads, stride_head_in_heads) which breaks
addressing for containerized 4-bit KV; change the conversion to use the
KVCacheList container-aware element size (the same approach used in
launchMHAFlashInfer) — read the container's valid-elements-per-head or
element-size helper from KVCacheList (constructed as cacheList) and use that
value when normalizing kv_stride_page/token/head to head units so
page/token/head strides are computed correctly for 4-bit containers.

♻️ Duplicate comments (4)

flashinfer/xqa.py (1)

306-314: ⚠️ Potential issue | 🟠 Major

Compute capability checks should use tensor device, not default CUDA device.

Line 307 and Line 313 query torch.device("cuda"), which can check the wrong GPU in multi-device runs.

💡 Proposed fix

-    if (
-        k_cache.dtype == torch.float8_e4m3fn
-        and get_compute_capability(torch.device(device="cuda"))[0] == 9
-    ):
+    device_cc_major = get_compute_capability(q.device)[0]
+    if k_cache.dtype == torch.float8_e4m3fn and device_cc_major == 9:
         run_sm90_fp8_mha = True
     else:
         run_sm90_fp8_mha = False

     if k_cache.dtype == torch.uint8:
-        assert get_compute_capability(torch.device(device="cuda"))[0] in [12], (
+        assert device_cc_major in [12], (
             "XQA NVFP4 KV is only supported on SM120 GPUs"
         )
         assert k_sf_cache is not None, "K SF cache is required when NVFP4 KV is used"
         assert v_sf_cache is not None, "V SF cache is required when NVFP4 KV is used"

-    if get_compute_capability(torch.device(device="cuda"))[0] not in [9, 10, 12]:
+    if device_cc_major not in [9, 10, 12]:
         raise RuntimeError("XQA is only supported on SM90, SM100, SM120/SM121 GPUs")

#!/bin/bash
set -euo pipefail
rg -n 'get_compute_capability\(torch\.device\(device="cuda"\)\)' flashinfer/xqa.py

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/xqa.py` around lines 306 - 314, The compute-capability checks call
get_compute_capability(torch.device(device="cuda")) which can target the wrong
GPU in multi-device runs; update those calls to use the actual tensor device
(e.g., k_cache.device) instead. Specifically, replace the
get_compute_capability(torch.device(device="cuda")) usages inside the
k_cache.dtype == torch.uint8 branch and the subsequent supported-SM check with
get_compute_capability(k_cache.device) (or the appropriate tensor like
q_cache.device if q is the primary tensor) so the checks use the same CUDA
device as the tensors being validated.

csrc/xqa/mha.cu (2)

2114-2116: ⚠️ Potential issue | 🔴 Critical

Undefined symbol in loadPages causes compilation failure.

idxNextSMemVBuf is not in scope in this lambda, so getSmemVSfTile(idxNextSMemVBuf) cannot compile when BEAM_WIDTH > 1 and 4-bit KV is enabled.

💡 Proposed fix

 `#if` BEAM_WIDTH == 1
   ...
 `#else`
   auto& dst = smem.vCachePages[grpLoadV ? warpGrpIdx : warpIdx.x];
-#if ENABLE_4BIT_KV_CACHE
-  auto& dstSf = getSmemVSfTile(idxNextSMemVBuf);
-#endif

   loadPagesForBeamSearchAsync<grpLoadV ? gemm1WarpsPerGrp : 1U>(
       grpLoadV ? warpIdxInGrp : 0U, dst, cacheList, false, idxReq, idxPageBeg, nbPages);
 `#endif`

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@csrc/xqa/mha.cu` around lines 2114 - 2116, The lambda in loadPages references
idxNextSMemVBuf which is out of scope when BEAM_WIDTH>1 with
ENABLE_4BIT_KV_CACHE enabled; update the lambda to use a visible index or
capture/pass idxNextSMemVBuf into the lambda and then call getSmemVSfTile with
that captured variable (or compute the correct buffer index visible in
loadPages), e.g., ensure the lambda captures idxNextSMemVBuf or replace it with
the in-scope idx (used elsewhere in loadPages) before calling getSmemVSfTile to
initialize dstSf.

---

1168-1184: ⚠️ Potential issue | 🟠 Major

Remove or consume vSfPrefetch; it is currently dead code.

The prefetch buffer is populated but not used in subsequent computation, adding avoidable overhead.

💡 Proposed fix (remove dead prefetch block)

-#if ENABLE_4BIT_KV_CACHE
-  // Prefetch buffer for SF
-  constexpr uint32_t nbSfPrefetchBuffers = exactDiv(warpTile.x, 16 * sizeof(uint32_t));
-  Array2D<uint32_t, exactDiv(cacheVTileSeqLen, 4), nbSfPrefetchBuffers> vSfPrefetch;
-#pragma unroll
-  for (uint32_t i = 0; i < vSfPrefetch.rows; i += 4) {
-#pragma unroll
-    for (uint32_t j = 0; j < nbSfPrefetchBuffers; j++) {
-      uint32_t const sfRowIdxInSlice = (laneId() % 4) * 4;
-      uint32_t const sfRowIdx = (i / 4) * 16 + i % 4 + sfRowIdxInSlice;
-      vSfPrefetch(i, j) = reinterpret_cast<const uint32_t&>(vSf.template at(sfRowIdx, j));
-    }
-  }
-#endif

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@csrc/xqa/mha.cu` around lines 1168 - 1184, Remove the dead prefetch buffer:
delete the nbSfPrefetchBuffers declaration and the Array2D vSfPrefetch
declaration plus the nested loops that populate it (the entire `#if`
ENABLE_4BIT_KV_CACHE prefetch block that references nbSfPrefetchBuffers,
vSfPrefetch, warpTile, cacheVTileSeqLen, vSf, and laneId()). After removal,
ensure there are no remaining references to vSfPrefetch or nbSfPrefetchBuffers
elsewhere; if a prefetch is actually needed, instead replace this block with a
true consumer of the populated data where vSf is used.

tests/attention/test_xqa_batch_decode.py (1)

632-632: ⚠️ Potential issue | 🟡 Minor

Drop the unused in_kv_lens binding.

Line 632 binds in_kv_lens but it is not used in this test.

🧹 Proposed cleanup

-    q_lens, in_kv_lens, seq_lens = generate_seq_lens_decode(
+    q_lens, _in_kv_lens, seq_lens = generate_seq_lens_decode(
         batch_size, q_len_per_req, max_in_kv_len
     )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/attention/test_xqa_batch_decode.py` at line 632, The test binds an
unused variable from generate_seq_lens_decode; remove the unnecessary in_kv_lens
binding and only capture the used values (e.g., q_lens and seq_lens) from the
call to generate_seq_lens_decode so the test no longer declares an unused
variable; update the assignment that currently reads "q_lens, in_kv_lens,
seq_lens = generate_seq_lens_decode(...)" to only bind the needed symbols.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@flashinfer/decode.py`:
- Around line 2536-2548: Ensure NVFP4 scale-factor requirements are enforced and
fix the single-tensor branch error message: if the runtime/backend requires
NVFP4 KV scale-factor tensors, do not accept kv_cache_sf == None and raise a
clear ValueError; when kv_cache_sf is a tuple, validate it has exactly two
elements before unpacking into k_cache_sf and v_cache_sf; when kv_cache_sf is a
single tensor assert kv_cache_sf.shape[1] == 2 (change the message to state
"must be 2" rather than "1 or 2") and then unbind as before to assign
k_cache_sf, v_cache_sf.

In `@tests/attention/test_xqa_batch_decode.py`:
- Around line 164-173: The V cache is being dequantized with the K global scale
causing incorrect reference KV values; in the block handling kv_dtype == "nvfp4"
(variables and helpers: k_cache, v_cache, to_nvfp4, nvfp4_to_float, k_scale,
v_scale, k_global_scale, v_global_scale, ref_kv_cache) change the dequantization
call for v_cache to use v_global_scale instead of k_global_scale so v_cache_dq =
nvfp4_to_float(v_cache, v_scale, v_global_scale) before stacking into
ref_kv_cache.

---

Outside diff comments:
In `@csrc/xqa/mha.cu`:
- Around line 2884-2892: In launchMHA the stride conversion still divides
kv_stride_page/token/head by validElemsPerHead (producing stride_page_in_heads,
stride_token_in_heads, stride_head_in_heads) which breaks addressing for
containerized 4-bit KV; change the conversion to use the KVCacheList
container-aware element size (the same approach used in launchMHAFlashInfer) —
read the container's valid-elements-per-head or element-size helper from
KVCacheList (constructed as cacheList) and use that value when normalizing
kv_stride_page/token/head to head units so page/token/head strides are computed
correctly for 4-bit containers.

In `@flashinfer/xqa.py`:
- Around line 154-160: The xqa function signature currently requires k_sf_cache
and v_sf_cache as positional parameters which breaks callers that don't use
NVFP4 KV; make these parameters optional by giving them default value None in
the xqa(...) signature (keep their types as Optional[torch.Tensor]) and update
the function body to handle None (e.g., skip NVFP4-specific logic when
k_sf_cache or v_sf_cache is None); ensure the symbols to change are the xqa(...)
function header and any internal branches that reference k_sf_cache / v_sf_cache
so behavior remains unchanged for callers that pass them and is safe for callers
that omit them.

---

Duplicate comments:
In `@csrc/xqa/mha.cu`:
- Around line 2114-2116: The lambda in loadPages references idxNextSMemVBuf
which is out of scope when BEAM_WIDTH>1 with ENABLE_4BIT_KV_CACHE enabled;
update the lambda to use a visible index or capture/pass idxNextSMemVBuf into
the lambda and then call getSmemVSfTile with that captured variable (or compute
the correct buffer index visible in loadPages), e.g., ensure the lambda captures
idxNextSMemVBuf or replace it with the in-scope idx (used elsewhere in
loadPages) before calling getSmemVSfTile to initialize dstSf.
- Around line 1168-1184: Remove the dead prefetch buffer: delete the
nbSfPrefetchBuffers declaration and the Array2D vSfPrefetch declaration plus the
nested loops that populate it (the entire `#if` ENABLE_4BIT_KV_CACHE prefetch
block that references nbSfPrefetchBuffers, vSfPrefetch, warpTile,
cacheVTileSeqLen, vSf, and laneId()). After removal, ensure there are no
remaining references to vSfPrefetch or nbSfPrefetchBuffers elsewhere; if a
prefetch is actually needed, instead replace this block with a true consumer of
the populated data where vSf is used.

In `@flashinfer/xqa.py`:
- Around line 306-314: The compute-capability checks call
get_compute_capability(torch.device(device="cuda")) which can target the wrong
GPU in multi-device runs; update those calls to use the actual tensor device
(e.g., k_cache.device) instead. Specifically, replace the
get_compute_capability(torch.device(device="cuda")) usages inside the
k_cache.dtype == torch.uint8 branch and the subsequent supported-SM check with
get_compute_capability(k_cache.device) (or the appropriate tensor like
q_cache.device if q is the primary tensor) so the checks use the same CUDA
device as the tensors being validated.

In `@tests/attention/test_xqa_batch_decode.py`:
- Line 632: The test binds an unused variable from generate_seq_lens_decode;
remove the unnecessary in_kv_lens binding and only capture the used values
(e.g., q_lens and seq_lens) from the call to generate_seq_lens_decode so the
test no longer declares an unused variable; update the assignment that currently
reads "q_lens, in_kv_lens, seq_lens = generate_seq_lens_decode(...)" to only
bind the needed symbols.

---

ℹ️ Review info

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c8d4769 and 8618527.

📒 Files selected for processing (12)

• csrc/flashinfer_xqa_binding.cu
• csrc/xqa/defines.h
• csrc/xqa/mha.cu
• csrc/xqa/mha.h
• csrc/xqa/mhaUtils.cuh
• csrc/xqa/utils.cuh
• csrc/xqa/xqa_wrapper.cu
• flashinfer/decode.py
• flashinfer/fp4_quantization.py
• flashinfer/jit/xqa.py
• flashinfer/xqa.py
• tests/attention/test_xqa_batch_decode.py

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

• csrc/xqa/xqa_wrapper.cu
• csrc/xqa/defines.h

[flashinfer-bot]

/bot run

[flashinfer-bot]

@flashinfer-bot is not authorized to trigger this CI job. cc: @yzh119, @sricketts, @yongwww

[yzh119]

/bot run

[flashinfer-bot]

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

[yzh119]

LGTM overall, left some comments.

[yzh119]

Please also annotate the expected shape of scaling factors.

[Tom-Zheng]

Done, thanks

[yzh119]

ditto

[Tom-Zheng]

Done, thanks

[flashinfer-bot]

[FAILED] Pipeline #45302193: 6/20 passed

[Tom-Zheng]

@yzh119 all comments are resolved and would you pls take a look and trigger CI?

[jiahanc]

/bot run

[flashinfer-bot]

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

[flashinfer-bot]

[FAILED] Pipeline #45502683: 6/20 passed

[coderabbitai]

♻️ Duplicate comments (1)

flashinfer/xqa.py (1)

308-313: ⚠️ Potential issue | 🟡 Minor

get_compute_capability uses default CUDA device instead of tensor device.

The capability checks use torch.device(device="cuda") which queries the default CUDA device. In multi-GPU setups where tensors reside on non-default GPUs, this could yield incorrect results.

Suggested fix

     if k_cache.dtype == torch.uint8:
-        assert get_compute_capability(torch.device(device="cuda"))[0] in [12], (
+        assert get_compute_capability(k_cache.device)[0] in [12], (
             "XQA NVFP4 KV is only supported on SM120 GPUs"
         )

-    if get_compute_capability(torch.device(device="cuda"))[0] not in [9, 10, 12]:
+    if get_compute_capability(q.device)[0] not in [9, 10, 12]:
         raise RuntimeError("XQA is only supported on SM90, SM100, SM120/SM121 GPUs")

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@flashinfer/xqa.py` around lines 308 - 313, The compute-capability check
incorrectly queries the default CUDA device; change the get_compute_capability
call to use the actual tensor/device where k_cache lives (e.g., pass
torch.device(k_cache.device) or the existing local variable device that points
to the tensor's device) instead of torch.device(device="cuda") so the SM check
(in the block guarding k_cache.dtype == torch.uint8) uses the correct GPU for
multi-GPU setups; update the get_compute_capability(...) call accordingly and
keep the surrounding assertions (k_sf_cache, v_sf_cache) intact.

🧹 Nitpick comments (1)

tests/attention/test_xqa_batch_decode.py (1)

630-634: Prefix unused in_kv_lens with underscore.

Static analysis flagged in_kv_lens as unused. Prefix with underscore to indicate intentional discard.

Suggested fix

-    q_lens, in_kv_lens, seq_lens = generate_seq_lens_decode(
+    q_lens, _in_kv_lens, seq_lens = generate_seq_lens_decode(
         batch_size, q_len_per_req, max_in_kv_len
     )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/attention/test_xqa_batch_decode.py` around lines 630 - 634, The
variable in_kv_lens returned from generate_seq_lens_decode is unused and should
be marked as intentionally discarded; change the assignment to unpack it into a
prefixed-underscore name (e.g. _in_kv_lens) while keeping q_lens and seq_lens
as-is so the call to generate_seq_lens_decode and uses of q_lens/seq_lens remain
unchanged.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Duplicate comments:
In `@flashinfer/xqa.py`:
- Around line 308-313: The compute-capability check incorrectly queries the
default CUDA device; change the get_compute_capability call to use the actual
tensor/device where k_cache lives (e.g., pass torch.device(k_cache.device) or
the existing local variable device that points to the tensor's device) instead
of torch.device(device="cuda") so the SM check (in the block guarding
k_cache.dtype == torch.uint8) uses the correct GPU for multi-GPU setups; update
the get_compute_capability(...) call accordingly and keep the surrounding
assertions (k_sf_cache, v_sf_cache) intact.

---

Nitpick comments:
In `@tests/attention/test_xqa_batch_decode.py`:
- Around line 630-634: The variable in_kv_lens returned from
generate_seq_lens_decode is unused and should be marked as intentionally
discarded; change the assignment to unpack it into a prefixed-underscore name
(e.g. _in_kv_lens) while keeping q_lens and seq_lens as-is so the call to
generate_seq_lens_decode and uses of q_lens/seq_lens remain unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 439dcace-8e6a-48ee-8fd5-7f0a747ab41f

📥 Commits

Reviewing files that changed from the base of the PR and between 8618527 and 2104ef7.

📒 Files selected for processing (4)

• csrc/xqa/mha.cu
• flashinfer/xqa.py
• tests/attention/test_xqa.py
• tests/attention/test_xqa_batch_decode.py

[samuellees]

/bot run

[flashinfer-bot]

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

[flashinfer-bot]

[SUCCESS] Pipeline #45680141: 10/20 passed