[Performance][MLA][ROCm] Remove redundant D2D copy in deepseek

[ganyi1996ppo]

Purpose

We found there are two redundant D2D copy in deepseek, which can be removed by in-place write inside the kernel.

Before

After

This PR replace those two index copy to inplace write inside the kernel, and we witness there are about 3% performance gain over deepseek on AMD Mi300 with 3.5k/1k inputs

Test Plan

gsm8k

Test Result

# Before
============ Serving Benchmark Result ============
Successful requests:                     64        
Failed requests:                         0         
Maximum request concurrency:             64        
Benchmark duration (s):                  64.74     
Total input tokens:                      229312    
Total generated tokens:                  65536     
Request throughput (req/s):              0.99      
Output token throughput (tok/s):         1012.31   
Peak output token throughput (tok/s):    1664.00   
Peak concurrent requests:                64.00     
Total Token throughput (tok/s):          4554.42   
---------------Time to First Token----------------
Mean TTFT (ms):                          13440.84  
Median TTFT (ms):                        13691.91  
P99 TTFT (ms):                           23198.25  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          50.01     
Median TPOT (ms):                        49.77     
P99 TPOT (ms):                           60.45     
---------------Inter-token Latency----------------
Mean ITL (ms):                           50.01     
Median ITL (ms):                         40.45     
P99 ITL (ms):                            46.96     
==================================================

# This PR
============ Serving Benchmark Result ============
Successful requests:                     64        
Failed requests:                         0         
Maximum request concurrency:             64        
Benchmark duration (s):                  63.03     
Total input tokens:                      229312    
Total generated tokens:                  65536     
Request throughput (req/s):              1.02      
Output token throughput (tok/s):         1039.81   
Peak output token throughput (tok/s):    1685.00   
Peak concurrent requests:                64.00     
Total Token throughput (tok/s):          4678.13   
---------------Time to First Token----------------
Mean TTFT (ms):                          13330.10  
Median TTFT (ms):                        13578.90  
P99 TTFT (ms):                           23058.80  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          48.45     
Median TPOT (ms):                        48.21     
P99 TPOT (ms):                           58.89     
---------------Inter-token Latency----------------
Mean ITL (ms):                           48.45     
Median ITL (ms):                         38.99     
P99 ITL (ms):                            46.28     
==================================================

# gsm8k
# acc result
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9462|±  |0.0062|
|     |       |strict-match    |     5|exact_match|↑  |0.9447|±  |0.0063|

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request successfully optimizes the deepseek MLA implementation by removing two redundant D2D copies, which is achieved by leveraging in-place writes within the kernels. The changes in _v_up_proj and _forward_prefill are well-aligned with the performance goals. I have one suggestion to correct a return type annotation that was missed during refactoring.

[ganyi1996ppo]

@HAIAI Please take a look

[HAIAI]

LGTM

[simon-mo]

@LucasWilkinson PTAL

[LucasWilkinson]

LGTM, thanks!

[sunway513]

@ganyi1996ppo - can you fix the failed CI signals?

[gshtras]

FAILED v1/attention/test_mla_backends.py::test_backend_correctness[1-deepseek-ai/DeepSeek-R1-small_prefill] - RuntimeError: prefix_output heads must be contiguous in memory
This failure looks relevant to the proposed change

[ganyi1996ppo]

FAILED v1/attention/test_mla_backends.py::test_backend_correctness[1-deepseek-ai/DeepSeek-R1-small_prefill] - RuntimeError: prefix_output heads must be contiguous in memory This failure looks relevant to the proposed change

Thanks @sunway513 @gshtras for point that out, I'll take a look.

[ganyi1996ppo]

hi @LucasWilkinson @HAIAI @gshtras , just notice the merge_attn_states didn't support strided load and store, and that triggers the ci failure of the pad_v case. I add the strided memory access pattern support on both cuda and triton version of this function in this PR, please take a look again.