A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 20:15 UTC
[perf] Integrate flashinfer concat_mla_k (#31171)
jiahanc committed
Feb 5, 2026 at 18:23 UTC
59a5cb387ae4c11c73855d505adb0b2c7cd3861d
2 files changed
+64
-3
vllm/model_executor/layers/attention/mla_attention.py
+17
-3
@@ -1885,6 +1885,16 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1885
self.indexer = indexer
1886
self.q_pad_num_heads = q_pad_num_heads
1887
1888
+ # Use flashinfer's optimized concat_mla_k kernel when available.
1889
+ # The kernel is optimized for DeepSeek V3 dimensions:
1890
+ # num_heads=128, nope_dim=128, rope_dim=64
1891
+ self._use_flashinfer_concat_mla_k = (
1892
+ has_flashinfer()
1893
+ and (self.num_heads == 128)
1894
+ and (self.qk_nope_head_dim == 128)
1895
+ and (self.qk_rope_head_dim == 64)
1896
+ )
1897
+
1898
if use_trtllm_ragged_deepseek_prefill():
1899
logger.info_once(
1900
"Using TRT-LLM ragged DeepSeek prefill for MLA", scope="local"
@@ -2192,9 +2202,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2202
dtype=k_nope.dtype,
2203
device=k_nope.device,
2204
)
2195
- # Direct copies with efficient broadcasting
2196
- k[..., : k_nope.shape[-1]] = k_nope
2197
- k[..., k_nope.shape[-1] :] = k_pe
2205
+
2206
+ if self._use_flashinfer_concat_mla_k:
2207
+ torch.ops.vllm.flashinfer_concat_mla_k(k, k_nope, k_pe)
2208
+ else:
2209
+ # Fallback: Direct copies with efficient broadcasting
2210
+ k[..., : k_nope.shape[-1]] = k_nope
2211
+ k[..., k_nope.shape[-1] :] = k_pe
2212
return k
2213
2214
def _compute_prefill_context(
vllm/utils/flashinfer.py
+47
@@ -396,6 +396,53 @@ def use_trtllm_attention(
396
397
398
if has_flashinfer():
399
+ from vllm.utils.torch_utils import direct_register_custom_op
400
+
401
+ def _flashinfer_concat_mla_k(
402
+ k: torch.Tensor,
403
+ k_nope: torch.Tensor,
404
+ k_pe: torch.Tensor,
405
+ ) -> None:
406
+ """Custom op wrapper for flashinfer's concat_mla_k.
407
+
408
+ This is an in-place operation that concatenates k_nope and k_pe into k.
409
+
410
+ The kernel is optimized for DeepSeek V3 dimensions:
411
+ - num_heads=128
412
+ - nope_dim=128
413
+ - rope_dim=64
414
+
415
+ Key optimizations:
416
+ - Warp-based processing with software pipelining
417
+ - Vectorized memory access (int2 for nope, int for rope)
418
+ - L2 prefetching for next row while processing current
419
+ - Register reuse for rope values across all heads
420
+
421
+ Args:
422
+ k: Output tensor, shape [num_tokens, num_heads, nope_dim + rope_dim].
423
+ Modified in-place.
424
+ k_nope: The nope part of k, shape [num_tokens, num_heads, nope_dim].
425
+ k_pe: The rope part of k (shared), shape [num_tokens, 1, rope_dim].
426
+ This is broadcast to all heads.
427
+ """
428
+ from flashinfer.concat_ops import concat_mla_k
429
+
430
+ concat_mla_k(k, k_nope, k_pe)
431
+
432
+ def _flashinfer_concat_mla_k_fake(
433
+ k: torch.Tensor,
434
+ k_nope: torch.Tensor,
435
+ k_pe: torch.Tensor,
436
+ ) -> None:
437
+ return
438
+
439
+ # Register flashinfer concat_mla_k custom op
440
+ direct_register_custom_op(
441
+ op_name="flashinfer_concat_mla_k",
442
+ op_func=_flashinfer_concat_mla_k,
443
+ mutates_args=["k"], # k tensor is modified in-place
444
+ fake_impl=_flashinfer_concat_mla_k_fake,
445
+ )
446
447
@torch.library.custom_op(
448
"vllm::flashinfer_mm_fp4",