-
Notifications
You must be signed in to change notification settings - Fork 1.8k
fix precompiled multi_query_token kernel not having is_fp8_out hash key #6279
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
fix precompiled multi_query_token kernel not having is_fp8_out hash key #6279
Conversation
📝 Walkthrough## Walkthrough
The changes update how the `is_fp8_output` flag is handled in hash key generation for XQA kernels. For precompiled kernels, the flag is now always set to false in the hash key, while for JIT kernels, the original value is used. Comments were added to clarify this logic.
## Changes
| File(s) | Change Summary |
|-------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------|
| cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplCommon.cpp<br>cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQAImplPrecompiled.cpp | Modified hash key generation to conditionally set `is_fp8_output` for precompiled vs JIT kernels; added explanatory comments. |
## Estimated code review effort
🎯 1 (Trivial) | ⏱️ ~3 minutes
## Suggested reviewers
- lucifer1004
- kaiyux
## Poem
> A hash key hops, a flag set anew,
> For JIT, it keeps the value true.
> But precompiled kernels, with logic tight,
> Ignore FP8 and set it right.
> With comments clear and code concise,
> Review is quick—oh, isn’t that nice!
> 🐇✨ Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. 📜 Recent review detailsConfiguration used: .coderabbit.yaml 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
⏰ 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)
✨ Finishing Touches
🧪 Generate unit tests
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
/bot run --disable-fail-fast |
PR_Github #12661 [ run ] triggered by Bot |
PR_Github #12661 [ run ] completed with state |
2abcfd5
to
e1f0e32
Compare
/bot run --disable-fail-fast |
PR_Github #12786 [ run ] triggered by Bot |
PR_Github #12786 [ run ] completed with state |
e1f0e32
to
60a65da
Compare
/bot run --disable-fail-fast |
PR_Github #12885 [ run ] triggered by Bot |
PR_Github #12885 [ run ] completed with state |
Signed-off-by: Jhao-Ting Chen <[email protected]>
60a65da
to
bb94b05
Compare
/bot reuse-pipeline |
PR_Github #13047 [ reuse-pipeline ] triggered by Bot |
PR_Github #13047 [ reuse-pipeline ] completed with state |
…ey (NVIDIA#6279) Signed-off-by: Jhao-Ting Chen <[email protected]> Signed-off-by: Shreyas Misra <[email protected]>
…ey (NVIDIA#6279) Signed-off-by: Jhao-Ting Chen <[email protected]> Signed-off-by: Ransiki Zhang <[email protected]>
…ey (NVIDIA#6279) Signed-off-by: Jhao-Ting Chen <[email protected]> Signed-off-by: Lanyu Liao <[email protected]>
…ey (NVIDIA#6279) Signed-off-by: Jhao-Ting Chen <[email protected]>
…VIDIA#6282 NVIDIA#6279 * [None][infra] Pin the version for triton to 3.3.1 (NVIDIA#6508) Signed-off-by: qqiao <[email protected]> * [None][infra] Pin the version for triton to 3.3.1 (NVIDIA#6508) (NVIDIA#6519) (NVIDIA#6549) Signed-off-by: Yanchao Lu <[email protected]> * [fix]: use safeInitRowMax instead of fp32_lowest to avoid NaN (NVIDIA#7087) Signed-off-by: Yao Yao <[email protected]> * [None][fix] Fix a numerical stability issue for XQA with spec dec Signed-off-by: Yao Yao <[email protected]> * fix typo Signed-off-by: Jhao-Ting Chen <[email protected]> * fix precompiled multi_query_token kernel not having is_fp8_out hash key (NVIDIA#6279) Signed-off-by: Jhao-Ting Chen <[email protected]> * [fix] Fix missing fields in xqa kernel cache key (NVIDIA#6282) Signed-off-by: Yao Yao <[email protected]> --------- Signed-off-by: qqiao <[email protected]> Signed-off-by: Yanchao Lu <[email protected]> Signed-off-by: Yao Yao <[email protected]> Signed-off-by: Jhao-Ting Chen <[email protected]> Co-authored-by: Emma Qiao <[email protected]> Co-authored-by: Yanchao Lu <[email protected]> Co-authored-by: Yao Yao <[email protected]>
Summary by CodeRabbit
Summary by CodeRabbit
Description
Llama 4 Maverick uses 40 attention head and 8 kv head.
This fails to use XQA mha_sm90.cu Hopper style spec-dec kernel due to 64 % (attention head / kv head) != 0.
This falls to precompiled XQA cubin kernel.
However precompiled does not support is_fp8_output flag in hashing key, causing not finding the kernel cubin.
This can be removed when it change to JIT.
Test Coverage
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...
Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]
to print this help message.See details below for each supported subcommand.
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]
Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id
(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test
(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast
(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test
(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"
(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"
(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"
(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test
(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test
(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test
(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge
(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"
(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log
(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug
(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-list
parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.md
and the
scripts/test_to_stage_mapping.py
helper.kill
kill
Kill all running builds associated with pull request.
skip
skip --comment COMMENT
Skip testing for latest commit on pull request.
--comment "Reason for skipping build/test"
is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.