opt flashinfer mla cat

[xu-yfei]

Motivation

Base on #5748 and #5638 , for flashinfer mla, remove q and k cat.

Accuracy

Accuracy: 0.951
Invalid: 0.000
Latency: 228.672 s
Output throughput: 554.173 token/s

Performance

main branch:

{"run_name": "default", "batch_size": 1, "input_len": 1024, "output_len": 1024, "latency": 14.4687, "output_throughput": 70.77, "overall_throughput": 141.55}

{"run_name": "default", "batch_size": 16, "input_len": 1024, "output_len": 1024, "latency": 28.8723, "output_throughput": 567.47, "overall_throughput": 1134.93}

{"run_name": "default", "batch_size": 32, "input_len": 1024, "output_len": 1024, "latency": 38.0349, "output_throughput": 861.52, "overall_throughput": 1723.05}

this PR:

{"run_name": "default", "batch_size": 1, "input_len": 1024, "output_len": 1024, "latency": 14.5066, "output_throughput": 70.59, "overall_throughput": 141.18}

{"run_name": "default", "batch_size": 16, "input_len": 1024, "output_len": 1024, "latency": 28.4372, "output_throughput": 576.15, "overall_throughput": 1152.29}

{"run_name": "default", "batch_size": 32, "input_len": 1024, "output_len": 1024, "latency": 37.2972, "output_throughput": 878.57, "overall_throughput": 1757.13}

Profile

Prefill

main branch:

this PR:

47 us to 3us

Decode

main branch bs=1 cuda graph+torch compile:

this PR bs=1 cuda graph+torch compile:

bs=1, cuda graph+torch compile, almost the same

main branch bs=1 cuda graph+ without torch compile fused with other ops:

this PR bs=1 cuda graph+ without torch compile fused with other ops:

bs=1, cuda graph, without torch compile, 6~7 us -> 1us

Modifications

• Update deepseek_v2 code, remove q and k cat.
• In flashinfer_mla_backend:

Cat when ragged, no cat in other scenes.
Use set_mla_kv_buffer when k_rope is not empty

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[lambert0312]

I pulled the latest commit and did some experiments, and it seems to be consistent with the optimizations mentioned above.