Revert "Add fused FP8 KV cache write kernel for TRTLLM MHA backend"#14550
Revert "Add fused FP8 KV cache write kernel for TRTLLM MHA backend"#14550Qiaolin-Yu wants to merge 1 commit intomainfrom
Conversation
Qiaolin-Yu
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
Ying1123,
ch-wan,
ispobock and
merrymercy
as code owners
Summary of ChangesHello @Qiaolin-Yu, 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 completely rolls back a recent feature that aimed to optimize FP8 KV cache writes using a fused kernel for the TRTLLM MHA backend. The decision to revert was prompted by persistent failures observed in the B200 continuous integration environment. By removing the problematic code and restoring the system to its prior state, this change ensures stability and resolves the CI pipeline issues. 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
|
|
/tag-and-rerun-ci |
![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 reverts the addition of a fused FP8 KV cache write kernel for the TRTLLM MHA backend, which was reportedly causing CI failures on b200 systems. The revert has been executed correctly: the kernel implementation and its corresponding tests have been removed, and the TRTLLM MHA backend has been updated to remove all usage of the fused kernel, falling back to the previous implementation. Additionally, some minor refactoring and cleanup have been applied to python/sglang/srt/models/qwen3_moe.py, which improves code readability. The changes appear correct and complete for the stated purpose of the revert.
|
@Qiaolin-Yu Thanks for reverting. I reproduced the B200 Triton error locally with a minimal script, root cause is k_scale / v_scale being 0D tensors seen as pointer in the FP8 KV kernel. A fix is already in progress. I am preparing a follow up PR with a corrected and |
Fix B200 CI failure that caused sgl-project#14093 to be reverted in sgl-project#14550. Original issue discovered in CI run: https://github.com/sgl-project/sglang/actions/runs/19985523381/job/57335123826?pr=14493 Root cause: - layer.k_scale/v_scale are torch.nn.Parameter (0-D tensors) - Triton interprets tensor arguments as pointers - Expression '1.0 / k_scale' in kernel caused IncompatibleTypeError Solution: - Convert tensor scales to Python float in wrapper before kernel launch - Handles None, Tensor, and scalar cases explicitly - Preserves per-tensor scaling semantics Testing: - Added regression test test_fp8_kv_kernel_accepts_tensor_scales - Verified on local environment with ATTN_BACKEND=trtllm_mha - All 11 kernel tests pass, including new regression test - test_llama31_fp4.py::TestLlama31FP4B200::test_gsm8k passes
Fix B200 CI failure that caused sgl-project#14093 to be reverted in sgl-project#14550. Original issue discovered in CI run: https://github.com/sgl-project/sglang/actions/runs/19985523381/job/57335123826?pr=14493 Root cause: - layer.k_scale/v_scale are torch.nn.Parameter (0-D tensors) - Triton interprets tensor arguments as pointers - Expression '1.0 / k_scale' in kernel caused IncompatibleTypeError Solution: - Convert tensor scales to Python float in wrapper before kernel launch - Handles None, Tensor, and scalar cases explicitly - Preserves per-tensor scaling semantics Testing: - Added regression test test_fp8_kv_kernel_accepts_tensor_scales - Verified on local environment with ATTN_BACKEND=trtllm_mha - All 11 kernel tests pass, including new regression test - test_llama31_fp4.py::TestLlama31FP4B200::test_gsm8k passes
3 participants
Reverts #14093 because it breaks b200 ci