Add fused FP8 KV cache write kernel for TRTLLM MHA backend

[harvenstar]

Summary

This PR introduces a fused FP8 KV cache write kernel for the TRTLLM MHA backend.
It combines FP8 quantization and paged KV cache writes into a single Triton kernel (plus a naive fallback) to reduce kernel launch overhead and avoid intermediate FP8 tensors.

---

Motivation

The current FP8 KV cache path launches multiple small elementwise kernels per step:

• FP8 quantization of K
• FP8 quantization of V
• Writing K to the paged KV cache
• Writing V to the paged KV cache

These kernels all operate on the same data and are memory bound.
The extra launches and intermediate tensors add overhead without bringing additional parallelism.

This PR fuses FP8 quantization and KV cache writes into one kernel for the TRTLLM backend.
RoPE and attention logic remain unchanged, but the extra FP8 elementwise kernels in the KV write path are removed.

---

Changes

New fused FP8 KV kernel

File added

• python/sglang/srt/layers/attention/trtllm_fp8_kv_kernel.py

Triton kernel: _fused_fp8_set_kv_buffer_kernel

• Takes post RoPE K and V in BF16 or FP16.
• Applies FP8 quantization using provided scalar scales.
• Writes directly into the paged FP8 KV cache layout.
• Respects the existing uint8 backed FP8 storage convention.

Python wrapper: fused_fp8_set_kv_buffer

• Supports both:
  • 3D cache layout: [total_slots, num_kv_heads, head_dim]
  • 4D cache layout: [num_pages, page_size, num_kv_heads, head_dim]
• Handles both 2D and 3D K, V inputs via runtime shape checks.
• Provides a naive fallback that exactly reproduces MHATokenToKVPool.set_kv_buffer semantics.
• Naive path is used as the golden implementation in unit tests.

Integration into TRTLLM attention backend

• Wires fused_fp8_set_kv_buffer into:
  • TRTLLMHAAttnBackend.forward_decode
  • TRTLLMHAAttnBackend.forward_extend
    when kv_dtype is FP8.
• Adds a feature flag style switch so non FP8 configurations are unaffected.
• Keeps the original code path as a correctness reference.
• Removes or downgrades all development-only debug logging so runtime logs remain minimal.

---

Accuracy and Testing

The fused path is designed to be numerically equivalent to the existing FP8 KV cache behavior.

Unit tests

1. Naive fused implementation vs MHATokenToKVPool.set_kv_buffer

   • Bit exact match at the written KV cache locations for both K and V.

2. Triton implementation vs naive implementation

   • Matches the naive path when using the same scalar FP8 scales.
   • Uses the same quantization formula as the reference path.

End to end evaluation (GSM8K)

• Setup:
  • Qwen3 MoE, FP8 KV cache, TRTLLM backend.
  • bench_sglang.py, 8 shot, 1319 questions.
• Results:
  • Baseline accuracy: ~0.956
  • Fused FP8 KV cache accuracy: ~0.958
  • Invalid rate: 0.0

No regression on GSM8K accuracy was observed.

---

Benchmarking and Profiling

• Nsight profiling shows that the four small FP8 KV elementwise kernels in the TRTLLM path are replaced by a single fused kernel in forward_decode and forward_extend when FP8 KV cache is enabled.
• KV write latency is reduced and the number of kernel launches per decode step is lower.
• Overall GSM8K benchmark latency remains in the same range, with KV cache writes no longer dominating the FP8 KV path.

Exact speedups depend on GPU, model, batch size, and sequence length.
The main benefit of this PR is to remove unnecessary elementwise launches in the TRTLLM FP8 KV write path while keeping behavior identical.

[gemini-code-assist]

Summary of Changes

Hello @harvenstar, 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 significantly optimizes the FP8 KV cache write process within the TRTLLM MHA backend. By fusing the FP8 quantization and paged KV cache write operations into a single Triton kernel, it addresses performance bottlenecks caused by multiple small kernel launches and the creation of intermediate tensors. This change aims to enhance the efficiency of the KV cache write path, particularly for FP8 configurations, while maintaining numerical accuracy and overall system performance.

Highlights

• Fused FP8 KV Cache Write: Introduces a new Triton kernel that combines FP8 quantization and paged KV cache writes for the TRTLLM MHA backend.
• Performance Optimization: Reduces kernel launch overhead and eliminates intermediate FP8 tensors, replacing multiple small elementwise kernels with a single fused operation.
• New Kernel Implementation: Adds trtllm_fp8_kv_kernel.py containing the Triton kernel _fused_fp8_set_kv_buffer_kernel and its Python wrapper fused_fp8_set_kv_buffer.
• Backend Integration: Integrates the fused kernel into TRTLLMHAAttnBackend.forward_decode and forward_extend for FP8 KV cache configurations, controlled by a feature flag.
• Correctness and Accuracy: Includes a naive fallback implementation for unit testing and demonstrates no regression in accuracy on GSM8K benchmarks.
• Benchmarking: Nsight profiling confirms reduced KV write latency and fewer kernel launches, improving efficiency in the FP8 KV write path.

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. ↩

[gemini-code-assist]

Code Review

This pull request introduces a fused FP8 KV cache write kernel for the TRTLLM MHA backend, which is a great performance optimization. The implementation is well-structured, including a Triton kernel for performance and a naive Python fallback for correctness testing. My review focuses on improving code maintainability and efficiency in a few areas. I've suggested refactoring duplicated code in the Triton kernel and the Python backend, vectorizing a loop in the naive implementation, and improving a type check to be more robust. Overall, this is a solid contribution.

[Qiaolin-Yu]

/tag-and-rerun-ci

[Qiaolin-Yu]

Why do we need to check it here?

[harvenstar]

Good catch.

[Qiaolin-Yu]

So it's only enabled for Qwen? Is it possible to make it more general so that other models can use it?

[harvenstar]

The check should be in the backend, not model-specific code. I'll move the fusion detection
logic entirely into TRTLLMHAAttnBackend._should_use_fused_fp8_path() and remove all the qwen3_moe
changes. This way any model using TRTLLM backend with FP8 KV cache gets the optimization automatically. Does the change sound good to you?

[Qiaolin-Yu]

sg

[harvenstar]

Done.✅ Fusion logic now lives in backend, and work for all models.

[ispobock]

1. share profile figure comparison to see the improvement after fusion
2. share end to end throughput improvement on Qwen3-235B

qq: Why trtllm_mla don't need this fusion?

[ispobock]

Why don't parallelize heads? (i.e. add head to grid)

[harvenstar]

Done. Moved head iteration to grid dimension via head_block_id = tl.program_id(1).

[ispobock]

Can we also parallelize k and v?

[harvenstar]

Done. Grid is now (tokens, head_blocks, 2) where kv_idx = tl.program_id(2) selects K or V.

[ispobock]

It seems not needed?

[harvenstar]

Removed in latest commit.

[ispobock]

It seems not needed?

[harvenstar]

Removed in latest commit.

[ispobock]

Could you add unit test for this kernel?

[harvenstar]

Added test/srt/test_trtllm_fp8_kv_kernel.py with serveral test cases covering
different input/cache dimensions, scaling options, and edge cases. All
passing.

[harvenstar]

1. share profile figure comparison to see the improvement after fusion
2. share end to end throughput improvement on Qwen3-235B

qq: Why trtllm_mla don't need this fusion?

Thank you for reviewing. Here is the analysis and comparison of the performance metrics,.

1. Throughput Improvement Summary

The fusion optimization led to a noticeable improvement in overall token throughput (tok/s) and a reduction in inter-token latency (ITL) on the Qwen3-235B model with FP8 KV cache enabled.

Metric	Before Fusion	After Fusion	Improvement (%)
Total token throughput (tok/s)	1635.66	1693.52	+3.54%
Output token throughput (tok/s)	818.01	846.94	+3.54%
Mean TPOT (Time per Output Token) (ms)	19.07	18.42	-3.41% (Faster)
Mean ITL (Inter-Token Latency) (ms)	19.08	18.43	-3.41% (Faster)
Request throughput (req/s)	0.81	0.84	+3.70%
Mean E2E Latency (ms)	19737.67	19064.12	-3.41% (Faster)

---

2. Full Benchmark Results and Profile Figure (For Reference)

To maintain the exact formatting and alignment of the raw data, the full results are enclosed in separate code blocks.

Before Fusion

After Fusion

Before Fusion

============ Serving Benchmark Result ============
Backend:                     sglang-oai
Traffic request rate:        inf
Max request concurrency:     16
Successful requests:         80
Benchmark duration (s):      99.13
Total input tokens:          81050
Total input text tokens:     81050
Total input vision tokens:   0
Total generated tokens:      81085
Total generated tokens (retokenized): 81047
Request throughput (req/s):  0.81
Input token throughput (tok/s): 817.65
Output token throughput (tok/s): 818.01
Peak output token throughput (tok/s): 992.00
Peak concurrent requests:    23
Total token throughput (tok/s): 1635.66
Concurrency:                 15.93
----------------End-to-End Latency----------------
Mean E2E Latency (ms):       19737.67
Median E2E Latency (ms):     19797.55
---------------Time to First Token----------------
Mean TTFT (ms):              430.65
Median TTFT (ms):            303.05
P99 TTFT (ms):               842.41
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):              19.07
Median TPOT (ms):            19.15
P99 TPOT (ms):               20.51
---------------Inter-Token Latency----------------
Mean ITL (ms):               19.08
Median ITL (ms):             17.27
P95 ITL (ms):                17.68
P99 ITL (ms):                34.68
Max ITL (ms):                473.31

After Fusion

============ Serving Benchmark Result ============
Backend:                                 sglang-oai
Traffic request rate:                    inf       
Max request concurrency:                 16        
Successful requests:                     80        
Benchmark duration (s):                  95.74     
Total input tokens:                      81050     
Total input text tokens:                 81050     
Total input vision tokens:               0         
Total generated tokens:                  81085     
Total generated tokens (retokenized):    81047     
Request throughput (req/s):              0.84      
Input token throughput (tok/s):          846.58    
Output token throughput (tok/s):         846.94    
Peak output token throughput (tok/s):    1040.00   
Peak concurrent requests:                23        
Total token throughput (tok/s):          1693.52   
Concurrency:                             15.93     
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   19064.12  
Median E2E Latency (ms):                 19127.56  
---------------Time to First Token----------------
Mean TTFT (ms):                          415.68    
Median TTFT (ms):                        305.58    
P99 TTFT (ms):                           764.02    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          18.42     
Median TPOT (ms):                        18.50     
P99 TPOT (ms):                           19.88     
---------------Inter-Token Latency----------------
Mean ITL (ms):                           18.43     
Median ITL (ms):                         16.62     
P95 ITL (ms):                            16.98     
P99 ITL (ms):                            33.48     
Max ITL (ms):                            459.80    
==================================================

[ispobock]

/tag-and-rerun-ci

[harvenstar]

/tag-and-rerun-ci