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 1, 18:45 UTC
[ROCm][MLA] Replace torch.cat in sparse-MLA forward_mqa with fused concat_mla_q (#42838)
Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Signed-off-by: Frida Andersson <fanderss@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Frida Andersson <fanderss@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Markus Hartikainen committed
Jun 7, 2026 at 02:20 UTC
bc5745a00f58eb55b7c53078995edcfbefc02928
2 files changed
+12
-5
csrc/libtorch_stable/concat_mla_q.cuh
-3
@@ -1,9 +1,6 @@
1
#ifndef CONCAT_MLA_Q_CUH_
2
#define CONCAT_MLA_Q_CUH_
3
4
-#include <cuda_bf16.h>
5
-#include <cuda_fp16.h>
6
-
4
#include "cuda_vec_utils.cuh"
5
6
namespace vllm {
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+12
-2
@@ -9,7 +9,7 @@ import torch
9
10
from vllm import _custom_ops as ops
11
from vllm._aiter_ops import rocm_aiter_ops
12
-from vllm.config import VllmConfig
12
+from vllm.config import VllmConfig, get_current_vllm_config
13
from vllm.config.cache import CacheDType
14
from vllm.logger import init_logger
15
from vllm.model_executor.layers.attention.mla_attention import (
@@ -31,6 +31,7 @@ from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
31
AiterMLAHelper,
32
)
33
from vllm.v1.kv_cache_interface import AttentionSpec
34
+from vllm.v1.worker.workspace import current_workspace_manager
35
36
if TYPE_CHECKING:
37
from vllm.model_executor.models.deepseek_v2 import Indexer
@@ -644,6 +645,13 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
645
assert indexer is not None
646
self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
647
648
+ vllm_config = get_current_vllm_config()
649
+ max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
650
+ q_concat_shape = (max_tokens, num_heads, head_size)
651
+ (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous(
652
+ (q_concat_shape, vllm_config.model_config.dtype),
653
+ )
654
+
655
def _forward_mla(
656
self,
657
layer: AttentionLayer,
@@ -703,7 +711,9 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
711
712
# Concatenate q if it's a tuple (ql_nope, q_pe)
713
if isinstance(q, tuple):
706
- q = torch.cat(q, dim=-1)
714
+ ql_nope, q_pe = q
715
+ q = self.q_concat_buffer[: ql_nope.shape[0]]
716
+ ops.concat_mla_q(ql_nope, q_pe, q)
717
718
num_actual_toks = attn_metadata.num_actual_tokens
719