NVFP4 KV Cache on SM100 #2363
NVFP4 KV Cache on SM100 #2363
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the 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. 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. Comment |
Summary of ChangesHello @samuellees, 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 upgrades the FlashInfer library by integrating native NVFP4 KV Cache support for SM100 GPUs. The changes span across C++ CUDA kernels and Python interfaces, enabling more efficient memory utilization and faster attention computations. Key improvements include a sophisticated kernel selection mechanism that adapts to various generation scenarios, precise handling of FP4 data types and their scaling factors, and refined memory access patterns through updated TMA descriptors. These enhancements aim to boost the performance and scalability of large language models on NVIDIA's latest architecture. Highlights
🧠 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. 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 implements NVFP4 KV Cache for SM100, introducing changes across CUDA kernel launchers, kernel implementation, and Python bindings. The core logic for handling FP4 data types, including scale factors and packed data, appears to be correctly implemented. The kernel selection logic in fmhaKernels.cuh has been significantly refactored for better heuristics and maintainability. However, I've identified a few issues: there are several leftover debugging statements (fflush(stdout)) that should be removed, and more critically, a hardcoded value for mSumOfSeqLensKv in kernelParams.h that seems to be a bug and needs correction.
| // params.mSumOfSeqLensKv = options.mSumOfSeqLensKv; | ||
| params.mSumOfSeqLensKv = 64; |
There was a problem hiding this comment.
The mSumOfSeqLensKv parameter is hardcoded to 64, while the correct assignment from options.mSumOfSeqLensKv is commented out. This will likely cause incorrect behavior for cases where the sum of KV sequence lengths is not 64. Please restore the original logic.
params.mSumOfSeqLensKv = options.mSumOfSeqLensKv;| int64_t workspace_size, Optional<TensorView> attention_sinks, | ||
| Optional<TensorView> cum_seq_lens_q, Optional<TensorView> key_block_scales, | ||
| Optional<TensorView> value_block_scales) { | ||
| fflush(stdout); |
There was a problem hiding this comment.
This fflush(stdout) call appears to be a leftover debugging statement. It should be removed from production code to avoid potential performance impacts and unnecessary console output.
| int numCtasX = numCtasPerSeqQ; | ||
| // printf("[numCtasX trace] Initial: numCtasX = %d (numCtasPerSeqQ = %d)\n", numCtasX, | ||
| // numCtasPerSeqQ); | ||
| fflush(stdout); |
There was a problem hiding this comment.
This fflush(stdout) call seems to be a debugging artifact and should be removed from the final code.
| // printf("[numCtasX trace] After multiCtasKv: numCtasX = %d (was %d, multiplied by | ||
| // numCtasPerSeqKv=%d), maxNumCtasPerSeqKv= %d\n", numCtasX, oldNumCtasX, numCtasPerSeqKv, | ||
| // maxNumCtasPerSeqKv); | ||
| fflush(stdout); |
There was a problem hiding this comment.
Another fflush(stdout) call that looks like a leftover from debugging. Please remove it.
2 participants
📌 Description
This PR implements NVFP4 KV Cache for SM100. Currently it's a draft PR for test.
🔍 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