[Feature][Kernel]: DeepSeek-R1 KV Proj is Too Slow for TP

[robertgshaw2-redhat]

🚀 The feature, motivation and pitch

TRT-LLM has a SwapAB kernel for KV proj for DSR1. We should integrate this by collaborating with the FlashInfer team

Current situation: we run CUTLASS Block Fp8 for KV proj because DeepGEMM upstream does not support it

• the CUTLASS Block Fp8 kernels are ~1/2 the speed of DeepGEMM for other Linear layers
• the CUTLASS Block Fp8 kernels have padding overhead

Example Trace:

note: in this example, we are using DeepGEMM for the Routed Expert. This will run with Triton / DeepGEMM Swap AB TEP in the final implementation

Alternatives

• improve cutlass impl. IUUC in recent versions there is swap AB and an option for no padding

Additional context

cc @mgoin @yewentao256 @LucasWilkinson @MatthewBonanni @alexm-redhat

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[tunglinwood]

@robertgshaw2-redhat Hi Robert, just for my curiosity, how do I get the kernel progress chart like yours?

[eicherseiji]

@tunglinwood This is the torch profiler, which can be enabled by setting VLLM_TORCH_PROFILER_DIR and sending requests via vllm bench serve --profile, or manually starting/stopping profiling with curl -X POST localhost:8000/(start|stop)_profile. The resulting trace can be loaded via https://ui.perfetto.dev/#!/viewer

[jhaotingc]

Hi @robertgshaw2-redhat, we have a duplicated #28316.

@xuanzic has a workable implementation of exposing deepgemm swapAB to linear gemm. We haven't started micro benchmarking on DSR1 KV proj gemm shape yet.

Previous nsys showed that per attention layer may be improved from 117 us (VLLM) to 98 us (TRTLLM) where the first GEMM (should be qkv) may be improved from 23us (VLLM cutlass) to 11.4us (TRTLLM DG swapAB).

(VLLM, TEP8, 8xH200, 1600/600, FLASH_ATTN_MLA, concurrency=4)

(TRTLLM, TEP8, 8xH200, 1600/600, concurrency=4)

[jhaotingc]

We have a FlashInfer PR flashinfer-ai/flashinfer#2101 and we're trying to include this in VLLM. cc @katec846

[robertgshaw2-redhat]

very nice. Note that we also have #28687 which improves vs CUTLASS

[jhaotingc]

Thanks. We'll take #28687 into account.

If it meets my expectation, the swapAB in the flashinfer PR should give better perf at low batch size, we'll use https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/bench_block_fp8_gemm.py#L147 for micro benchmarking these kernels (DeepGEMM, DeepGEMM from FlashInfer (TRTLLM), original CUTLASS).