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 12, 14:45 UTC
[ROCm][MLA] enable fp8 MLA decode on ROCm (#28032)
Signed-off-by: guanbao <gyu@amd.com> Signed-off-by: Guanbao Yu <gyu@amd.com> Signed-off-by: gbyu-amd <Guanbao.Yu@amd.com> Co-authored-by: guanbao <gyu@amd.com>
gbyu-amd committed
Nov 25, 2025 at 10:15 UTC
cb7214d8eaa231c67416282668f1ca274f8068ba
2 files changed
+21
-1
vllm/_aiter_ops.py
+10
@@ -294,6 +294,8 @@ def _rocm_aiter_mla_decode_fwd_impl(
294
kv_last_page_lens: torch.Tensor | None = None,
295
sm_scale: float = 1.0,
296
logit_cap: float = 0.0,
297
+ q_scale: torch.Tensor | None = None,
298
+ kv_scale: torch.Tensor | None = None,
299
) -> None:
300
from aiter.mla import mla_decode_fwd
301
@@ -308,6 +310,8 @@ def _rocm_aiter_mla_decode_fwd_impl(
310
max_seqlen_qo,
311
sm_scale=sm_scale,
312
logit_cap=logit_cap,
313
+ q_scale=q_scale,
314
+ kv_scale=kv_scale,
315
)
316
317
@@ -322,6 +326,8 @@ def _rocm_aiter_mla_decode_fwd_fake(
326
kv_last_page_lens: torch.Tensor | None = None,
327
sm_scale: float = 1.0,
328
logit_cap: float = 0.0,
329
+ q_scale: torch.Tensor | None = None,
330
+ kv_scale: torch.Tensor | None = None,
331
) -> None:
332
pass
333
@@ -806,6 +812,8 @@ class rocm_aiter_ops:
812
kv_indices: torch.Tensor | None = None,
813
kv_last_page_lens: torch.Tensor | None = None,
814
logit_cap: float = 0.0,
815
+ q_scale: torch.Tensor | None = None,
816
+ kv_scale: torch.Tensor | None = None,
817
):
818
torch.ops.vllm.rocm_aiter_mla_decode_fwd(
819
q,
@@ -818,6 +826,8 @@ class rocm_aiter_ops:
826
kv_last_page_lens,
827
sm_scale=sm_scale,
828
logit_cap=logit_cap,
829
+ q_scale=q_scale,
830
+ kv_scale=kv_scale,
831
)
832
833
@staticmethod
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+11
-1
@@ -49,6 +49,8 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
49
paged_kv_last_page_len: torch.Tensor | None = None
50
# The query indptr, shape : [num_decode + 1]
51
qo_indptr: torch.Tensor | None = None
52
+ # The dtype of MLA out tensor
53
+ attn_out_dtype: torch.dtype = torch.bfloat16
54
55
56
class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):
@@ -74,6 +76,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
76
)
77
78
self.compilation_config = vllm_config.compilation_config
79
+ self.decode_attn_out_dtype = vllm_config.model_config.dtype
80
# kernel block size is always 1.
81
max_num_pages_per_req = vllm_config.model_config.max_model_len
82
max_num_reqs = vllm_config.scheduler_config.max_num_seqs
@@ -162,6 +165,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
165
paged_kv_last_page_len=paged_kv_last_page_len,
166
qo_indptr=qo_indptr,
167
dcp_tot_seq_lens=dcp_tot_seq_lens_device,
168
+ attn_out_dtype=self.decode_attn_out_dtype,
169
)
170
171
return attn_metadata
@@ -242,7 +246,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
246
assert isinstance(q, torch.Tensor)
247
B = q.shape[0]
248
o = torch.zeros(
245
- B, self.num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device
249
+ B,
250
+ self.num_heads,
251
+ self.kv_lora_rank,
252
+ dtype=attn_metadata.decode.attn_out_dtype,
253
+ device=q.device,
254
)
255
256
kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
@@ -260,6 +268,8 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
268
attn_metadata.decode.paged_kv_indptr,
269
attn_metadata.decode.paged_kv_indices,
270
attn_metadata.decode.paged_kv_last_page_len,
271
+ q_scale=layer._q_scale,
272
+ kv_scale=layer._k_scale,
273
)
274
275
return o, None