add xqa fp8 mha and fp8 kv cache #1769
Conversation
Summary of ChangesHello @qsang-nv, I'm Gemini Code Assist1! 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 significantly enhances the FlashInfer library by integrating FP8 support for both Multi-Head Attention computations and the Key-Value cache within the XQA framework. These changes are primarily aimed at leveraging the advanced capabilities of NVIDIA Hopper GPUs (SM90+) to achieve substantial performance and memory efficiency gains. The implementation includes new CUDA kernels utilizing GMMA and TMA, along with Python-side modifications to enable configurable FP8 execution paths, ensuring that users can opt into these optimizations while maintaining numerical stability. Highlights
Using Gemini Code AssistThe 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
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 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
|
Summary of ChangesHello @qsang-nv, I'm Gemini Code Assist1! 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 significantly enhances the FlashInfer library by introducing support for FP8 Multi-Head Attention (MHA) and FP8 Key-Value (KV) cache. These additions leverage advanced features of NVIDIA Hopper GPUs, such as GMMA and TMA, to achieve higher performance and memory efficiency for large language model inference. The changes span the CUDA C++ backend, including new kernel implementations and memory management utilities, as well as updates to the Python AOT compilation and testing framework to ensure robust integration and validation of the new FP8 capabilities. Highlights
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for FP8 multi-head attention (MHA) and FP8 KV cache within the XQA kernel, primarily targeting the NVIDIA Hopper architecture. This is a significant feature addition, enabled by new CUDA primitives for Hopper's Tensor Memory Access (TMA) and Grace Hopper MMA (GMMA) instructions. The changes are well-implemented, including new CUDA headers for hardware abstraction, a dispatch mechanism for the new FP8 kernel path, and corresponding updates to the Python build system and tests. The tests have been thoughtfully adjusted with relaxed tolerances for FP8 precision. My review includes one suggestion to refactor a small piece of duplicated code to enhance maintainability.
csrc/xqa/xqa_wrapper.cu
Outdated
| if (run_fp8_mha) { | ||
| launchHopperF8MHAFlashInfer( | ||
| multiProcessorCount, nbKHeads, slidingWinSize, qScale, | ||
| reinterpret_cast<OutputHead*>(output.data_ptr()), | ||
| #if LOW_PREC_OUTPUT | ||
| reinterpret_cast<float const*>(rcpOutScale.data_ptr()), | ||
| reinterpret_cast<float const*>(rcpOutScale.data_ptr()), | ||
| #endif | ||
| reinterpret_cast<InputHead const*>(q.data_ptr()), attentionSinksPtr, | ||
| reinterpret_cast<GMemCacheHead*>(pool.data_ptr()), | ||
| reinterpret_cast<KVCachePageIndex const*>(kvCachePageList.data_ptr()), | ||
| maxSeqLen, reinterpret_cast<uint32_t const*>(seqLen.data_ptr()), batchSize, | ||
| reinterpret_cast<float const*>(kvCacheScale.data_ptr()), | ||
| reinterpret_cast<InputHead const*>(q.data_ptr()), attentionSinksPtr, | ||
| reinterpret_cast<GMemCacheHead*>(pool.data_ptr()), | ||
| reinterpret_cast<KVCachePageIndex const*>(kvCachePageList.data_ptr()), maxSeqLen, | ||
| reinterpret_cast<uint32_t const*>(seqLen.data_ptr()), batchSize, | ||
| reinterpret_cast<float const*>(kvCacheScale.data_ptr()), | ||
| #if SPEC_DEC | ||
| qSeqLen, reinterpret_cast<uint32_t const*>(qCuSeqLens.data_ptr()), | ||
| reinterpret_cast<MaskType const*>(mask.data_ptr()), | ||
| qSeqLen, reinterpret_cast<uint32_t const*>(qCuSeqLens.data_ptr()), | ||
| reinterpret_cast<MaskType const*>(mask.data_ptr()), | ||
| #endif | ||
| reinterpret_cast<uint32_t*>(semaphores.data_ptr()), | ||
| reinterpret_cast<void*>(scratch.data_ptr()), stream); | ||
| reinterpret_cast<uint32_t*>(semaphores.data_ptr()), | ||
| reinterpret_cast<void*>(scratch.data_ptr()), stream); | ||
| } else { | ||
| launchMHAFlashInfer(multiProcessorCount, nbKHeads, slidingWinSize, qScale, | ||
| reinterpret_cast<OutputHead*>(output.data_ptr()), | ||
| #if LOW_PREC_OUTPUT | ||
| reinterpret_cast<float const*>(rcpOutScale.data_ptr()), | ||
| #endif | ||
| reinterpret_cast<InputHead const*>(q.data_ptr()), attentionSinksPtr, | ||
| reinterpret_cast<GMemCacheHead*>(pool.data_ptr()), | ||
| reinterpret_cast<KVCachePageIndex const*>(kvCachePageList.data_ptr()), | ||
| maxSeqLen, reinterpret_cast<uint32_t const*>(seqLen.data_ptr()), batchSize, | ||
| reinterpret_cast<float const*>(kvCacheScale.data_ptr()), | ||
| #if SPEC_DEC | ||
| qSeqLen, reinterpret_cast<uint32_t const*>(qCuSeqLens.data_ptr()), | ||
| reinterpret_cast<MaskType const*>(mask.data_ptr()), | ||
| #endif | ||
| reinterpret_cast<uint32_t*>(semaphores.data_ptr()), | ||
| reinterpret_cast<void*>(scratch.data_ptr()), stream); | ||
| } |
There was a problem hiding this comment.
The if and else blocks contain identical arguments passed to two different functions (launchHopperF8MHAFlashInfer and launchMHAFlashInfer). This code duplication can be reduced to improve maintainability. Since both functions share the same signature, you can use a function pointer to select the appropriate kernel and then call it once with the common set of arguments.
using mha_launcher_t = decltype(&launchMHAFlashInfer);
mha_launcher_t launcher = run_fp8_mha ? &launchHopperF8MHAFlashInfer : &launchMHAFlashInfer;
launcher(multiProcessorCount, nbKHeads, slidingWinSize, qScale,
reinterpret_cast<OutputHead*>(output.data_ptr()),
#if LOW_PREC_OUTPUT
reinterpret_cast<float const*>(rcpOutScale.data_ptr()),
#endif
reinterpret_cast<InputHead const*>(q.data_ptr()), attentionSinksPtr,
reinterpret_cast<GMemCacheHead*>(pool.data_ptr()),
reinterpret_cast<KVCachePageIndex const*>(kvCachePageList.data_ptr()),
maxSeqLen, reinterpret_cast<uint32_t const*>(seqLen.data_ptr()), batchSize,
reinterpret_cast<float const*>(kvCacheScale.data_ptr()),
#if SPEC_DEC
qSeqLen, reinterpret_cast<uint32_t const*>(qCuSeqLens.data_ptr()),
reinterpret_cast<MaskType const*>(mask.data_ptr()),
#endif
reinterpret_cast<uint32_t*>(semaphores.data_ptr()),
reinterpret_cast<void*>(scratch.data_ptr()), stream);
Summary of ChangesHello @qsang-nv, I'm Gemini Code Assist1! 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 significantly enhances the XQA (eXtended Query Attention) kernels by introducing support for FP8 Multi-Head Attention and FP8 Key-Value cache. These changes are designed to optimize performance and memory usage on NVIDIA Hopper (SM90+) GPUs through low-level CUDA programming, including asynchronous tensor memory access and matrix operations. The integration ensures that the system can efficiently handle lower precision data types, with comprehensive testing to maintain accuracy within acceptable bounds. Highlights
Using Gemini Code AssistThe 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
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 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
|
There was a problem hiding this comment.
Code Review
This pull request introduces support for FP8 multi-head attention (MHA) and FP8 KV cache in the XQA kernels, targeting Hopper architecture for performance improvements. The changes include new low-level CUDA files (gmma.cuh, tma.h, tensorMap.cpp) with Hopper-specific WGMMA and TMA instructions, a new FP8 MHA kernel entry point, and updates to the AOT compilation scripts and Python wrappers to handle the new FP8 variants. The tests have also been updated to include FP8 configurations and use a more lenient assertion method to account for precision differences.
My review focuses on code maintainability and clarity. I've suggested refactoring a duplicated code block in the C++ wrapper to improve readability and proposed adding a comment in the Python tests to clarify a magic number used for data scaling. Overall, the changes are well-structured and the addition of FP8 support is a valuable performance enhancement.
csrc/xqa/xqa_wrapper.cu
Outdated
| if (run_fp8_mha) { | ||
| launchHopperF8MHAFlashInfer( | ||
| multiProcessorCount, nbKHeads, slidingWinSize, qScale, | ||
| reinterpret_cast<OutputHead*>(output.data_ptr()), | ||
| #if LOW_PREC_OUTPUT | ||
| reinterpret_cast<float const*>(rcpOutScale.data_ptr()), | ||
| reinterpret_cast<float const*>(rcpOutScale.data_ptr()), | ||
| #endif | ||
| reinterpret_cast<InputHead const*>(q.data_ptr()), attentionSinksPtr, | ||
| reinterpret_cast<GMemCacheHead*>(pool.data_ptr()), | ||
| reinterpret_cast<KVCachePageIndex const*>(kvCachePageList.data_ptr()), | ||
| maxSeqLen, reinterpret_cast<uint32_t const*>(seqLen.data_ptr()), batchSize, | ||
| reinterpret_cast<float const*>(kvCacheScale.data_ptr()), | ||
| reinterpret_cast<InputHead const*>(q.data_ptr()), attentionSinksPtr, | ||
| reinterpret_cast<GMemCacheHead*>(pool.data_ptr()), | ||
| reinterpret_cast<KVCachePageIndex const*>(kvCachePageList.data_ptr()), maxSeqLen, | ||
| reinterpret_cast<uint32_t const*>(seqLen.data_ptr()), batchSize, | ||
| reinterpret_cast<float const*>(kvCacheScale.data_ptr()), | ||
| #if SPEC_DEC | ||
| qSeqLen, reinterpret_cast<uint32_t const*>(qCuSeqLens.data_ptr()), | ||
| reinterpret_cast<MaskType const*>(mask.data_ptr()), | ||
| qSeqLen, reinterpret_cast<uint32_t const*>(qCuSeqLens.data_ptr()), | ||
| reinterpret_cast<MaskType const*>(mask.data_ptr()), | ||
| #endif | ||
| reinterpret_cast<uint32_t*>(semaphores.data_ptr()), | ||
| reinterpret_cast<void*>(scratch.data_ptr()), stream); | ||
| reinterpret_cast<uint32_t*>(semaphores.data_ptr()), | ||
| reinterpret_cast<void*>(scratch.data_ptr()), stream); | ||
| } else { | ||
| launchMHAFlashInfer(multiProcessorCount, nbKHeads, slidingWinSize, qScale, | ||
| reinterpret_cast<OutputHead*>(output.data_ptr()), | ||
| #if LOW_PREC_OUTPUT | ||
| reinterpret_cast<float const*>(rcpOutScale.data_ptr()), | ||
| #endif | ||
| reinterpret_cast<InputHead const*>(q.data_ptr()), attentionSinksPtr, | ||
| reinterpret_cast<GMemCacheHead*>(pool.data_ptr()), | ||
| reinterpret_cast<KVCachePageIndex const*>(kvCachePageList.data_ptr()), | ||
| maxSeqLen, reinterpret_cast<uint32_t const*>(seqLen.data_ptr()), batchSize, | ||
| reinterpret_cast<float const*>(kvCacheScale.data_ptr()), | ||
| #if SPEC_DEC | ||
| qSeqLen, reinterpret_cast<uint32_t const*>(qCuSeqLens.data_ptr()), | ||
| reinterpret_cast<MaskType const*>(mask.data_ptr()), | ||
| #endif | ||
| reinterpret_cast<uint32_t*>(semaphores.data_ptr()), | ||
| reinterpret_cast<void*>(scratch.data_ptr()), stream); | ||
| } |
There was a problem hiding this comment.
There's a large block of duplicated code for launching the MHA kernels. The only difference between the if and else blocks is the function being called (launchHopperF8MHAFlashInfer vs. launchMHAFlashInfer). This could be refactored to improve maintainability and reduce redundancy.
Consider using a function pointer to select the kernel, and then make a single call. This would make the code cleaner and easier to manage if more arguments are added in the future.
For example:
void (*mha_func)(uint32_t, uint32_t, ...); // Using a function pointer type alias
if (run_fp8_mha) {
mha_func = &launchHopperF8MHAFlashInfer;
} else {
mha_func = &launchMHAFlashInfer;
}
mha_func(
multiProcessorCount,
nbKHeads,
slidingWinSize,
// ... other arguments
); using mha_func_t = void (*)(uint32_t, uint32_t, uint32_t, float, OutputHead*,
#if LOW_PREC_OUTPUT
float const*,
#endif
InputHead const*, float const*, GMemCacheHead*,
KVCachePageIndex const*, uint32_t, uint32_t const*, uint32_t,
float const* __restrict__,
#if SPEC_DEC
uint32_t, uint32_t const*, MaskType const*,
#endif
uint32_t*, void*, cudaStream_t);
mha_func_t mha_func = run_fp8_mha ? &launchHopperF8MHAFlashInfer : &launchMHAFlashInfer;
mha_func(multiProcessorCount, nbKHeads, slidingWinSize, qScale,
reinterpret_cast<OutputHead*>(output.data_ptr()),
#if LOW_PREC_OUTPUT
reinterpret_cast<float const*>(rcpOutScale.data_ptr()),
#endif
reinterpret_cast<InputHead const*>(q.data_ptr()), attentionSinksPtr,
reinterpret_cast<GMemCacheHead*>(pool.data_ptr()),
reinterpret_cast<KVCachePageIndex const*>(kvCachePageList.data_ptr()), maxSeqLen,
reinterpret_cast<uint32_t const*>(seqLen.data_ptr()), batchSize,
reinterpret_cast<float const*>(kvCacheScale.data_ptr()),
#if SPEC_DEC
qSeqLen, reinterpret_cast<uint32_t const*>(qCuSeqLens.data_ptr()),
reinterpret_cast<MaskType const*>(mask.data_ptr()),
#endif
reinterpret_cast<uint32_t*>(semaphores.data_ptr()),
reinterpret_cast<void*>(scratch.data_ptr()), stream);
tests/attention/test_xqa.py
Outdated
| if fp8_kv_cache: | ||
| cache_heads /= 4.0 |
There was a problem hiding this comment.
The value 4.0 is used to scale down the cache_heads tensor when fp8_kv_cache is enabled. This appears to be a magic number. To improve code clarity and maintainability, please add a comment explaining the rationale for this specific scaling factor. For example, explaining that it's to prevent overflow and how 4.0 was determined would be very helpful for future readers.
| if fp8_kv_cache: | |
| cache_heads /= 4.0 | |
| if fp8_kv_cache: | |
| # Scale down the cache heads to keep values within the representable range of FP8 | |
| # and prevent overflow during computation. The factor 4.0 is chosen empirically. | |
| cache_heads /= 4.0 |
There was a problem hiding this comment.
Code Review
This pull request introduces support for FP8 multi-head attention (MHA) and FP8 KV cache for Cross-Query Attention (XQA), targeting NVIDIA's Hopper architecture. This is a significant feature addition that leverages low-level hardware capabilities like TMA and GMMA for performance. The changes include new CUDA files for these Hopper-specific features, along with updates to the Python build system and tests to accommodate the new configurations. The review identified a critical bug in the new TMA storeAsync implementation and a high-severity correctness issue related to the handling of masked values in the softmax computation.
| : "memory"); | ||
| } else if constexpr (nbDims == 5) { | ||
| asm volatile( | ||
| "cp.async.bulk.tensor.2d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3, %4, %5}], " | ||
| "[%6];\n" | ||
| : | ||
| : "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), | ||
| "r"(offset[2]), "r"(offset[3]), "r"(offset[4]), "l"(__cvta_generic_to_shared(src)) | ||
| : "memory"); | ||
| } else { | ||
| static_assert(nbDims >= 1 && nbDims <= 5); | ||
| } | ||
| } | ||
|
|
||
| __device__ inline void setTensorMapGlbAddr(CUtensorMap& tensorMap, void* ptr) { | ||
| asm volatile( | ||
| "tensormap.replace.tile.global_address.global.b1024.b64 [%0], %1;\n" ::"l"(&tensorMap), | ||
| "l"(ptr) | ||
| : "memory"); | ||
| } | ||
|
|
||
| __device__ inline void commitGroup() { | ||
| asm volatile("cp.async.bulk.commit_group;\n" : : : "memory"); | ||
| } |
There was a problem hiding this comment.
There appears to be a copy-paste error in the storeAsync template function. For nbDims of 3, 4, and 5, the inline assembly instruction is cp.async.bulk.tensor.2d..., but it should be cp.async.bulk.tensor.3d..., cp.async.bulk.tensor.4d..., and cp.async.bulk.tensor.5d... respectively. This will lead to incorrect memory access patterns and likely data corruption for higher-dimensional tensors.
else if constexpr (nbDims == 3)
{
asm volatile("cp.async.bulk.tensor.3d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3}], [%4];\n"
:
: "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), "r"(offset[2]),
"l"(__cvta_generic_to_shared(src))
: "memory");
}
else if constexpr (nbDims == 4)
{
asm volatile("cp.async.bulk.tensor.4d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3, %4}], [%5];\n"
:
: "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), "r"(offset[2]),
"r"(offset[3]), "l"(__cvta_generic_to_shared(src))
: "memory");
}
else if constexpr (nbDims == 5)
{
asm volatile("cp.async.bulk.tensor.5d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3, %4, %5}], [%6];\n"
:
: "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), "r"(offset[2]),
"r"(offset[3]), "r"(offset[4]), "l"(__cvta_generic_to_shared(src))
: "memory");
}| ? true | ||
| : packedMask & (1u << ((col + actualQSeqLen - nbValidCols) - maskPosStart)); | ||
| acc(m, n)(i, j) = maskFlag && col < nbValidCols ? acc(m, n)(i, j) : -INFINITY; | ||
| acc(m, n)(i, j) = maskFlag && col < nbValidCols ? acc(m, n)(i, j) : safeInitRowMax; |
There was a problem hiding this comment.
Using safeInitRowMax for masked elements can lead to incorrect results. When an entire row/sequence is masked, all attention scores become safeInitRowMax. In the softmax computation, maxVal also becomes safeInitRowMax, and exp(score - maxVal) evaluates to 1 for all masked positions. This results in a uniform attention distribution over masked tokens, and the output becomes the average of values in V, instead of zero.
A correct implementation should ensure that the softmax output for masked tokens is zero. If the entire row is masked, the final output should also be zero. This might require changes in the softmax function to handle safeInitRowMax specially, and in the final normalization step to handle a row sum of zero.
| @@ -16,8 +16,8 @@ | |||
|
|
|||
| #include "pytorch_extension_utils.h" | |||
|
|
|||
| void xqa_wrapper(int64_t multiProcessorCount, int64_t nbKHeads, int64_t slidingWinSize, | |||
| double qScale, at::Tensor output, | |||
| void xqa_wrapper(bool run_fp8_mha, int64_t multiProcessorCount, int64_t nbKHeads, | |||
There was a problem hiding this comment.
Instead of making this a flag, could we pass a dtype?
Same for the other places where we pass:
- the type of the input (only bf16 and fp16 supported I think)
- the type of the kv-cache (fp8 or bf16)
- the type in which we perform arithmetic (the same type as the kv-cache I think?)
|
|
||
| inline constexpr float log2e = 1.4426950408889634; // std::log2(M_E) | ||
| inline constexpr float safeInitRowMax = -1e+30F; | ||
| // we used an optimization where exp(x-rowMax) is computed as: |
There was a problem hiding this comment.
That's interesting: what were the symptoms of the instability? Accuracy loss?
There was a problem hiding this comment.
This is copied from NVIDIA/TensorRT-LLM@c1aa7f3, you may ask the author, I am not sure about this question.
tests/attention/test_xqa.py
Outdated
| @@ -354,4 +364,21 @@ def cache_head_at( | |||
| kernel_output = output[req][b][ | |||
| idx_k_head * head_grp_size : (idx_k_head + 1) * head_grp_size | |||
| ].to(torch.float32) | |||
| assert torch.allclose(ref_output, kernel_output, atol=0.01, rtol=0.01) | |||
| if fp8_kv_cache or run_fp8_mha: | |||
| atol = 0.05 | |||
There was a problem hiding this comment.
How did you tune this tolerance? Can it be smaller?
There was a problem hiding this comment.
From 0.01 to 0.05, add 0.01 every step. And it can't be smaller from my test.
Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com>
Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com>
Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com>
Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com>
Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com>
| else: | ||
| flag_sliding_window = ["-DSLIDING_WINDOW=0"] | ||
|
|
||
| if sm_version == 100: |
There was a problem hiding this comment.
Is it possible to add SM103 support by targeting SM100f instead of SM100a?
And similarly, can we add SM121 support by targeting SM120f instead of SM120a?
There was a problem hiding this comment.
What's the difference in those archs? I mean SM103/SM100f/SM100a, and SM121/SM120f/SM120a.
There was a problem hiding this comment.
The "a" means "arch-specific" and the "f" means "family".
SM100 and SM103 are in the "SM100f family".
SM100a (SM100 arch-specific) will only run on SM100 devices.
But I believe SM103 devices have a superset of the SM100 features, and therefore if you target SM100f instead of SM100a during compilation, your cubin will be able to run on SM103 as well, without any loss of optimization on either device. So I think it's strictly better than targeting SM103a.
SM121 and SM120 have a similar story: it's better to target SM120f as a compilation target, yielding a cubin that will run on both SM120 and SM121 devices without any compromise to performance.
See this documentation for details: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#family-specific-features
@aleozlx can you confirm my understanding?
There was a problem hiding this comment.
yes sm_100f is known as family specific or family conditional. this is important to enhance device compatibility in a sort of out of the box fashion.
i'd like to point out a few things tho based on my experience:
- with strictly jit compilation where the premise is that at runtime fewer devices are in the compatibility question, the arch conditionals may be the safest way to target the instruction supersets. (when the target is available to compile by the toolkit at the time of implementation)
- family conditionals are important for compatibility story (indeed aligning with your understanding conceptually) but it is not without inherent engineering complexity. from an engineer's perspective i naturally experience a slightly more complicated story at the levels beneath. i'll spare the details but leave it as a reliability intuition (in the context of jit).
however, i want to bring this CompilationContext.get_nvcc_flags_list() up for consideration to not hard code any guidance either way but have it abstracted so we can adjust if situation changes. briefly, how this works is for each op/backend if you whitelist your supported targets, this function shall serve as the mapping to provide the recommended flags.
we can put in sm103 support (likely fine for attn) supposing our cicd will catch the issue if not
There was a problem hiding this comment.
Makes sense, let me clarify my guidance then:
- Since SM120 and SM121 are identical architectures, naively I would think that supporting both is only marginally harder than supporting one.
- The story seems even a little harder for SM100 and SM103, because they are not identical, but if you don't care about SM103-specific features, I would think the marginal effort of supporting SM103 shouldn't be massive.
- Therefore the default design for any solution for SM100 and SM120 should at least try to include SM103 and SM121 support, or should at least be designed with SM103 and SM121 in mind, even if some of the details are left as a future TODO.
- What compilation targets you use, and how you architect the code to query for those compilation targets, is an engineering implementation detail, and I probably shouldn't be opinionated about that. :)
The problem here is that the PR doesn't address (3), maybe because @qsang-nv didn't know about (1) and (2). I'm only proposing that we re-think the design here with the above in mind.
Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com>
Signed-off-by: Qidi Sang <200703406+qsang-nv@users.noreply.github.com>
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughThis pull request adds FP8 MHA support to FlashInfer's XQA kernels, introduces conditional paged KV cache layout handling, refactors MHA launcher selection, and adds tensor map utilities for asynchronous memory operations. Support extends to SM100 and SM120 architectures. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant xqa_wrapper as xqa_wrapper<br/>(C++ binding)
participant MHA_Dispatcher as MHA Launcher<br/>Dispatcher
participant FP8_MHA as FP8 MHA<br/>Kernel
participant STD_MHA as Standard MHA<br/>Kernel
User->>xqa_wrapper: call with run_fp8_mha flag
xqa_wrapper->>MHA_Dispatcher: select launcher based on run_fp8_mha
alt run_fp8_mha == true
MHA_Dispatcher->>FP8_MHA: launchHopperF8MHAFlashInfer
Note over FP8_MHA: Execute FP8 matrix ops
FP8_MHA-->>xqa_wrapper: return
else run_fp8_mha == false
MHA_Dispatcher->>STD_MHA: launchMHAFlashInfer
Note over STD_MHA: Execute standard precision ops
STD_MHA-->>xqa_wrapper: return
end
xqa_wrapper-->>User: output attention results
sequenceDiagram
participant Python_API as Python API<br/>(flashinfer.xqa)
participant Module_Gen as Module<br/>Generator
participant JIT_Compiler as JIT Compiler
participant KV_Layout as KV Cache<br/>Layout Handler
Python_API->>Module_Gen: get_xqa_module(fp16_input, fp8_kv_cache, sm_version)
activate Module_Gen
Module_Gen->>JIT_Compiler: gen_xqa_module with flags
activate JIT_Compiler
alt PAGED_KV_CACHE_LAYOUT == 1
JIT_Compiler->>KV_Layout: Use kCacheVLLM, vCacheVLLM<br/>(separate K/V tensors)
KV_Layout->>JIT_Compiler: return paged layout tensors
Note over KV_Layout: VLLM paging scheme
else
JIT_Compiler->>KV_Layout: Use pool<br/>(contiguous layout)
KV_Layout->>JIT_Compiler: return pool tensor
end
JIT_Compiler->>JIT_Compiler: select NVCC flags<br/>for sm_version
JIT_Compiler-->>Module_Gen: compiled module
deactivate JIT_Compiler
Module_Gen-->>Python_API: XQA module handle
deactivate Module_Gen
Python_API->>Python_API: xqa(..., kCacheVLLM,<br/>vCacheVLLM, run_fp8_mha)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Rationale: This change spans multiple heterogeneous concerns—FP8 MHA support, conditional KV cache layout switching (paged vs. contiguous), multi-SM architecture support, new tensor map async operations, and refactored parameter passing across C++ bindings and Python layers. Each cohort (CUDA kernels, Python JIT/AOT generation, test suite) requires separate reasoning due to architecture-specific guards, conditional compilation blocks, and semantic shifts in memory layout. While individual edits are focused, the breadth of files and the interaction of preprocessor conditionals, parameter reordering, and layout changes demand careful cross-file validation. Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/attention/test_xqa.py (1)
28-33: Avoid GPU property access at import time.Accessing
torch.cuda.get_device_properties(0)during import can break test discovery on CPU/multi-device envs. Move it inside the test after skip checks.-props = torch.cuda.get_device_properties(0) -sm_count = props.multi_processor_count +sm_count = None # set inside test to avoid import-time CUDA queries
♻️ Duplicate comments (4)
csrc/xqa/tma.h (1)
208-229: Bug: 3D/4D/5D storeAsync use 2D opcode (will corrupt data).The cp.async store paths for nbDims 3–5 incorrectly use tensor.2d. Must be tensor.3d/4d/5d.
Apply this diff:
@@ } else if constexpr (nbDims == 3) { - asm volatile( - "cp.async.bulk.tensor.2d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3}], [%4];\n" - : - : "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), - "r"(offset[2]), "l"(__cvta_generic_to_shared(src)) - : "memory"); + asm volatile( + "cp.async.bulk.tensor.3d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3}], [%4];\n" + : + : "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), + "r"(offset[2]), "l"(__cvta_generic_to_shared(src)) + : "memory"); } else if constexpr (nbDims == 4) { - asm volatile( - "cp.async.bulk.tensor.2d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3, %4}], [%5];\n" - : - : "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), - "r"(offset[2]), "r"(offset[3]), "l"(__cvta_generic_to_shared(src)) - : "memory"); + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3, %4}], [%5];\n" + : + : "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), + "r"(offset[2]), "r"(offset[3]), "l"(__cvta_generic_to_shared(src)) + : "memory"); } else if constexpr (nbDims == 5) { - asm volatile( - "cp.async.bulk.tensor.2d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3, %4, %5}], " - "[%6];\n" - : - : "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), - "r"(offset[2]), "r"(offset[3]), "r"(offset[4]), "l"(__cvta_generic_to_shared(src)) - : "memory"); + asm volatile( + "cp.async.bulk.tensor.5d.global.shared::cta.bulk_group.tile [%0, {%1, %2, %3, %4, %5}], [%6];\n" + : + : "l"(reinterpret_cast<uint64_t>(&tensorMap)), "r"(offset[0]), "r"(offset[1]), + "r"(offset[2]), "r"(offset[3]), "r"(offset[4]), "l"(__cvta_generic_to_shared(src)) + : "memory"); }csrc/xqa/mha.cu (1)
479-479: Critical: Masked position initialization may cause incorrect attention output.Using
safeInitRowMaxfor masked elements can lead to incorrect results. When an entire row is masked, all scores becomesafeInitRowMax, and in softmax computationexp(score - maxVal)evaluates to1for all positions, producing a uniform distribution over masked tokens instead of zero output.As noted in the previous review, the softmax function should handle
safeInitRowMaxspecially to ensure masked tokens contribute zero to the output, or alternatively masked positions should use a different sentinel value that results in zero after softmax.csrc/flashinfer_xqa_binding.cu (1)
19-21: Prefer a typed precision enum over a new boolean flag.Using
bool run_fp8_mhadoes not scale. Replace with a small enum (e.g.,int32_t precision: {bf16, fp16, fp8}) and, similarly, pass/cache element/compute dtypes as enums instead of separate flags. This reduces combinatorial overload and ABI churn.tests/attention/test_xqa.py (1)
241-246: Good: FP8 cache scaling is documented.Comment explains the 4.0 factor and overflow concerns.
🧹 Nitpick comments (10)
csrc/xqa/tma.h (1)
74-83: Comment/code mismatch for nbDims==1 path.The comment says “nbDims==1 does not need tensormap,” but the code uses the tensor.1d variant taking a tensor map. Either drop the map for 1D linear copies or update the comment.
Also applies to: 129-138
csrc/xqa/gmma.cuh (1)
27-56: Bitfield layout is implementation-defined; prefer explicit packing.Relying on 64‑bit bitfield layout and reinterpret_cast to Raw can be brittle across compilers/ABIs. Recommend encoding/decoding with shifts/masks into a uint64_t to guarantee layout and endianness. Keep sizeof(MatDesc)==8 as a guard.
csrc/xqa/tensorMap.h (1)
3-3: cuda.h include: make header robust to non-CUDA analysis/compiles.Static analysis flagged ‘cuda.h’ not found. If this header is transitively included by non‑CUDA TU(s), guard the include or move these declarations behind a build flag. Example: wrap with a small shim header included only from .cpp, or add a dedicated config that ensures CUDA include paths are present in CI.
csrc/xqa/tensorMap.cpp (1)
43-73: Tensor map for contiguous KV cache looks correct.The function properly constructs a tensor map for contiguous KV cache layout:
- Global dimensions and strides are configured appropriately
- Swizzle selection based on cache line size (128B or 64B)
- Error handling via
checkCuwrapperMinor suggestion: The error message on line 64 "unsupported cache head size" could be more specific about expected values.
- throw std::runtime_error("unsupported cache head size"); + throw std::runtime_error("unsupported partElems: " + std::to_string(partElems) + + ", expected 128 or 64");flashinfer/jit/xqa.py (1)
76-100: SM version selection and build configuration verified; optional refactor still recommended for clarity and error handling.The changes are correct:
- New source files (
mha_sm90.cu,tensorMap.cpp) exist incsrc/xqa/- Build configuration properly references and links them
- Required CUDA Driver API linker flag (
-lcuda) and cache layout flag includedHowever, the SM version selection logic could be improved for maintainability. The current code defaults to
sm90a_nvcc_flagsfor unrecognized versions, which implicitly handlessm_version=90but obscures intent and provides no validation for truly unsupported architectures.Consider making the SM90 case explicit and adding validation:
- if sm_version == 100: + if sm_version == 90: + sm_nvcc_flags = sm90a_nvcc_flags + elif sm_version == 100: sm_nvcc_flags = sm100a_nvcc_flags elif sm_version == 120: sm_nvcc_flags = sm120a_nvcc_flags else: - sm_nvcc_flags = sm90a_nvcc_flags + raise ValueError(f"Unsupported sm_version: {sm_version}")This makes supported architectures explicit and catches invalid SM versions early.
csrc/flashinfer_xqa_binding.cu (1)
25-35: KV cache params behind preprocessor guards: keep Python and C++ signatures locked.Since params differ when
PAGED_KV_CACHE_LAYOUT!=1, ensure JIT always defines it (the JIT path does) and document this contract near the binding to avoid accidental ABI mismatches. Consider adding a static assert/log print on init when it’s not set.Also applies to: 28-29
tests/attention/test_xqa.py (3)
263-275: Remove unusedbeam_widtharg (lint: ARG001).
beam_widthincache_head_atis unused; drop it and update call sites.-def cache_head_at( +def cache_head_at( batch, is_k, idx_kv_head, pos, - cache_k_heads, - cache_v_heads, - page_list, - beam_width, + cache_k_heads, + cache_v_heads, + page_list, nb_k_heads, tokens_per_page, ): @@ - cache_head = cache_head_at( + cache_head = cache_head_at( batch, kv == 0, idx_kv_head, pos, cache_k_heads, cache_v_heads, - page_list_arg, - beam_width, + page_list_arg, nb_k_heads, tokens_per_page, )Also applies to: 291-303
317-319: Make scratch size configurable; 256 MiB may OOM CI.Read from an env var with a sane default to reduce flakiness.
- scratch_size = 256 << 20 + import os + scratch_mb = int(os.environ.get("FLASHINFER_TEST_SCRATCH_MB", "256")) + scratch_size = scratch_mb << 20You can validate with different values in CI matrix.
392-397: Stable epsilon for relative diff.Optional: use dtype-aware epsilon via
torch.finfoto avoid hard-coded 1e-8.- diff_rel = diff_abs / (torch.abs(ref_output) + 1e-8) + eps = torch.finfo(torch.float32).eps + diff_rel = diff_abs / (torch.abs(ref_output) + eps)flashinfer/xqa.py (1)
147-150: Avoid repeated capability queries and shorten the error.Cache CC once and use a shorter exception message (addresses Ruff TRY003).
- if get_compute_capability(torch.device(device="cuda"))[0] not in [9, 10, 12]: - raise RuntimeError("XQA is only supported on SM90, SM100, SM120 GPUs") - sm_version = int(get_compute_capability(torch.device(device="cuda"))[0] * 10) + cc_major, _ = get_compute_capability(torch.device(device="cuda")) + if cc_major not in (9, 10, 12): + raise RuntimeError("Unsupported GPU (require SM90/100/120)") + sm_version = int(cc_major * 10)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
csrc/flashinfer_xqa_binding.cu(1 hunks)csrc/xqa/gmma.cuh(1 hunks)csrc/xqa/mha.cu(5 hunks)csrc/xqa/mha.h(2 hunks)csrc/xqa/tensorMap.cpp(1 hunks)csrc/xqa/tensorMap.h(1 hunks)csrc/xqa/tma.h(1 hunks)csrc/xqa/utils.cuh(2 hunks)csrc/xqa/xqa_wrapper.cu(2 hunks)flashinfer/aot.py(3 hunks)flashinfer/jit/xqa.py(2 hunks)flashinfer/xqa.py(6 hunks)tests/attention/test_xqa.py(10 hunks)
🧰 Additional context used
🧬 Code graph analysis (10)
csrc/xqa/tensorMap.h (1)
csrc/xqa/tensorMap.cpp (6)
getElemBytes(10-41)getElemBytes(10-10)makeTensorMapForContiguousKVCache(43-73)makeTensorMapForContiguousKVCache(43-47)makeTensorMapForPagedKVCache(75-117)makeTensorMapForPagedKVCache(75-78)
csrc/xqa/tensorMap.cpp (1)
csrc/xqa/utils.h (1)
checkCu(39-48)
flashinfer/jit/xqa.py (1)
flashinfer/jit/core.py (2)
JitSpec(181-280)gen_jit_spec(283-347)
flashinfer/xqa.py (3)
flashinfer/jit/xqa.py (1)
gen_xqa_module(38-101)flashinfer/jit/core.py (1)
build_and_load(268-280)flashinfer/utils.py (3)
register_custom_op(266-275)register_custom_op(285-304)get_compute_capability(245-248)
csrc/xqa/xqa_wrapper.cu (2)
csrc/xqa/mha_sm90.cu (4)
launchHopperF8MHAFlashInfer(3168-3275)launchHopperF8MHAFlashInfer(3168-3185)scratch(506-513)scratch(506-506)csrc/xqa/mha.cu (2)
launchMHAFlashInfer(2657-2749)launchMHAFlashInfer(2657-2674)
flashinfer/aot.py (1)
flashinfer/jit/core.py (1)
JitSpec(181-280)
csrc/xqa/mha.h (1)
csrc/xqa/mha_sm90.cu (4)
launchHopperF8MHAFlashInfer(3168-3275)launchHopperF8MHAFlashInfer(3168-3185)scratch(506-513)scratch(506-506)
tests/attention/test_xqa.py (1)
flashinfer/utils.py (1)
get_compute_capability(245-248)
csrc/xqa/tma.h (1)
csrc/xqa/mha_sm90.cu (16)
void(548-577)void(579-584)void(588-598)void(1693-1727)void(1765-1797)void(1799-1816)void(1841-1887)void(1976-1997)void(1999-2017)void(2049-2131)void(2180-2254)void(2256-2275)void(2278-2296)void(2316-2332)void(2336-2359)void(2396-2420)
csrc/flashinfer_xqa_binding.cu (1)
csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_sm100_binding.cu (4)
output(230-396)output(230-238)output(398-566)output(398-409)
🪛 Clang (14.0.6)
csrc/xqa/tensorMap.h
[error] 3-3: 'cuda.h' file not found
(clang-diagnostic-error)
🪛 Ruff (0.14.0)
flashinfer/xqa.py
148-148: Avoid specifying long messages outside the exception class
(TRY003)
tests/attention/test_xqa.py
271-271: Unused function argument: beam_width
(ARG001)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Deploy Docs
🔇 Additional comments (21)
csrc/xqa/utils.cuh (2)
34-41: Numerical-stability note: initialize rowMax safely but validate ranges seen in practice.Lowering safeInitRowMax to -1e5 avoids FMA overflow in x*log2e - bias, but it changes the effective lower bound. Please validate on adversarial logits (very negative rows) to ensure no early saturation and no accuracy regressions. Consider guarding the optimization per-arch or switching to compute (x - rowMax) before scaling to avoid FMA on large magnitudes.
49-51: Code is correct; review comment contains incorrect assumptions.SM100 (Blackwell) opt-in dynamic shared memory per block is 227 KB, which matches the value at line 50. SM120 (Hopper Next) is 99 KB, which is already correctly configured on line 46—not on lines 49-51 as the review suggests.
The conditional structure properly segregates architectures:
- Line 45-46: SM120 (
__CUDA_ARCH__ == 1200) → 99 KB ✓- Line 49-50: SM100 (
__CUDA_ARCH__ == 1000) → 227 KB ✓Lines 49-51 handle only SM90 and SM100; SM120 is on a separate branch.
Likely an incorrect or invalid review comment.
csrc/xqa/gmma.cuh (1)
60-66: Assumption: shared address fits 18 bits (0x3FFFF) — verify on SM100/SM120.The assert/encode mask assumes __cvta_generic_to_shared(data) fits 18 bits (>>4 stored). Please validate against SM100/SM120 SMEM addressing; increase field width or add debug fallback if this trips on larger SMEM configs.
csrc/xqa/mha.h (2)
138-150: API surface change verified — no action required.Declaration (mha.h:133), definition (mha.cu:2657), and the single call site (xqa_wrapper.cu:43–68) all use identical
#if PAGED_KV_CACHE_LAYOUT == 1preprocessor guards. Arguments are correctly conditionalized at the call site to match the conditional signature. Python bindings invoke throughxqa_wrapper, which handles the dispatch internally, isolating the FFI boundary from the conditional signature. ODR/ABI consistency is maintained within the compilation unit.
194-211: All concerns verified successfully.The function declaration (mha.h:194) and definition (mha_sm90.cu:3168) match exactly, including all conditional parameters (LOW_PREC_OUTPUT, PAGED_KV_CACHE_LAYOUT, SPEC_DEC). The function is properly integrated: it's used in xqa_wrapper.cu:43 with correct FP8 conditional logic (
run_fp8_mha ? &launchHopperF8MHAFlashInfer : &launchMHAFlashInfer), defined in the correct TU (mha_sm90.cu), and linked properly without issues.csrc/xqa/mha.cu (2)
95-95: LGTM: SM100 architecture support added correctly.The addition of
__CUDA_ARCH__ == 1000appropriately extends support to SM100 devices with the same configuration as SM800/870/900, using larger kHeadPartBytes (128) and cacheVTileSeqLen (64).
2657-2722: Based on my analysis of the kernel signature and the function call, there is a parameter mismatch in the launchMHAFlashInfer function call when SPEC_DEC is defined.Looking at the kernel_mha signature (mha_sm90.cu:615-652), the kernel expects:
nbKHeads(always)slidingWinSize(if SLIDING_WINDOW)qScale, output, rcpOutScale, q, attentionSinks, cacheListbeamSearchParams(if USE_BEAM_SEARCH)batchSize, kvCacheScale- Tensor maps via grid_constant (not passed as regular parameters)
specDecParams(if SPEC_DEC)semaphores, scratchHowever, at line 2707-2722, when SPEC_DEC is defined, the call passes
qSeqLen, nbKHeads, headGrpSize, qCuSeqLensas four separate parameters, but the kernel expects onlynbKHeadsat that position. Additionally, the call passesmask(line 2722) but the kernel has nomaskparameter—it expectsspecDecParamsinstead.The review comment requires verification of how SpecDecParams and BeamSearchParams should be constructed and passed, since the current call site appears to pass individual fields separately rather than properly constructed structs.
flashinfer/jit/xqa.py (4)
18-24: LGTM: Imports updated appropriately.The added imports for SM-specific NVCC flags enable proper multi-architecture support.
26-35: LGTM: NVCC flags configured correctly.The flags properly enable paged KV cache with layout 1, consistent with the conditional compilation paths in the C++ code.
47-55: LGTM: Flag generation logic is correct.The conditional flag generation properly handles:
- FP16 vs BF16 input (
DTYPEandINPUT_FP16)- FP8 vs FP16/BF16 KV cache (
CACHE_ELEM_ENUM)
38-46: All call sites are already updated with the new signature.Verification confirms that:
- The new
fp16_input,fp8_kv_cache, andsm_versionparameters are consistently used across the codebase- Both call sites (
flashinfer/aot.py:404andflashinfer/xqa.py:40) correctly pass the new parameters- Wrapper functions (
get_xqa_moduleandxqa) use the updated signature- No
use_fp16parameter exists anywhere in the codebaseThe API changes are complete and properly integrated.
csrc/xqa/xqa_wrapper.cu (2)
22-38: LGTM: Function signature updated appropriately.The signature changes are well-designed:
run_fp8_mhaparameter enables runtime selection between FP8 and standard MHAOptional<TensorView>forattentionSinksis more idiomatic than raw pointers- Conditional KV cache parameters based on
PAGED_KV_CACHE_LAYOUTproperly support both layout modes
39-65: Function pointer approach is type-safe; signatures are compatible.Verification confirms that both
launchHopperF8MHAFlashInferandlaunchMHAFlashInferhave identical signatures with matching conditional compilation blocks and parameter lists, making the function pointer assignment safe and correct.flashinfer/aot.py (3)
358-372: LGTM: gen_xqa signature updated for multi-architecture support.The function signature changes are consistent with the JIT module updates:
- Parameter renaming improves clarity
- SM gating ensures generation only when supported architectures are available
- New
fp8_kv_cache_parameter enables FP8 KV cache configurations
373-412: Multi-SM architecture support implemented correctly.The iteration logic properly:
- Constructs
sm_versionslist based on available architectures- Iterates over SM versions along with other configuration parameters
- Validates configurations before generating modules
- Passes all parameters to
gen_xqa_moduleconsistently
527-546: LGTM: gen_all_modules updated consistently.The changes to
gen_all_modulesproperly wire through the new parameters and SM version support to the XQA generator.csrc/xqa/tensorMap.cpp (2)
10-41: LGTM: Data type size lookup implemented correctly.The
getElemBytesfunction provides comprehensive coverage of CUDA tensor map data types with appropriate error handling.
75-117: Paged KV cache tensor map correctly supports two layout modes with consistent stride calculations.The implementation correctly configures tensor map dimensions and strides for two distinct layouts:
- VLLM Layout (PAGED_KV_CACHE_LAYOUT == 1): dimensions
{headElems, nbKHeads, tokensPerPage, pages}with strides accounting for head-first ordering- XQA Layout (PAGED_KV_CACHE_LAYOUT == 0, default): dimensions
{headElems, tokensPerPage, nbKHeads, pages}with strides accounting for token-first orderingThe dimension ordering aligns with memory access patterns throughout the codebase (verified in mha.cu, mhaUtils.cuh, and mha_sm90.cu). Both layouts apply the same swizzle modes and error handling. No issues identified.
csrc/flashinfer_xqa_binding.cu (2)
24-25: Good: Optional attention sinks.Switching to
tvm::ffi::Optional<TensorView>makes the API safer and clearer.
21-23: No changes needed; LOW_PREC_OUTPUT=0 is already set in compilation flags.The codebase already includes
"-DLOW_PREC_OUTPUT=0"in thexqa_nvcc_flagslist withinflashinfer/jit/xqa.py. This flag is passed toextra_cuda_cflagsin thegen_jit_spec()call, ensuring thercpOutScaleparameter is not included in the C++ function signature. There is no ABI drift risk because the conditional parameter is compiled out consistently.flashinfer/xqa.py (1)
50-72: Signature wiring looks consistent with the binding.Param order matches
xqa_wrapper(includingrun_fp8_mha, optionalattentionSinks, and separate K/V caches).If
LOW_PREC_OUTPUTis ever enabled, extend these call sites to passrcpOutScaleor force-DLOW_PREC_OUTPUT=0in JIT.Also applies to: 73-91
| compute_capability = get_compute_capability(torch.device(device="cuda")) | ||
| if compute_capability[0] != 9: | ||
| pytest.skip("XQA only supports on Hopper at this moment") | ||
| if compute_capability[0] != 9 and run_fp8_mha: | ||
| pytest.skip("XQA supports fp8 mha only on Hopper GPUs") | ||
| set_random_seed(42) |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Compute sm_count inside the test.
Set SM count after capability checks to avoid premature CUDA access.
def test_xqa(
@@
- compute_capability = get_compute_capability(torch.device(device="cuda"))
+ compute_capability = get_compute_capability(torch.device(device="cuda"))
if compute_capability[0] != 9 and run_fp8_mha:
pytest.skip("XQA supports fp8 mha only on Hopper GPUs")
set_random_seed(42)
+ props = torch.cuda.get_device_properties(torch.cuda.current_device())
+ sm_count = props.multi_processor_countAlso applies to: 329-330
🤖 Prompt for AI Agents
In tests/attention/test_xqa.py around lines 180-183, the test currently calls
into CUDA to compute sm_count before checking compute capability and may access
CUDA prematurely; move the sm_count computation so it runs after the
compute_capability check and any pytest.skip decision (i.e., compute sm_count
only after verifying compute_capability and run_fp8_mha), and apply the same
change to the other occurrence around lines 329-330; ensure you call the
sm_count helper (or get_sm_count) with the CUDA device only after the skip logic
and after set_random_seed(42) if that ordering is required.
4 participants
📌 Description
Add xqa fp8 mha and fp8 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Tests