[ROCm][Perf] DSv3.2: fuse MLA Q concat+fp8-quant in forward_mqa (#43673)
Signed-off-by: Frida Andersson <fanderss@amd.com>
frida-andersson committed
Jun 23, 2026 at 12:12 UTC
20b5af55c1c9ad0d1a00dc45771997dd490c8461
1 file changed
+12
-7
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+12
-7
@@ -713,11 +713,16 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
713
# NOTE(lucas): for the sparse FlashMLA kernels the kernels want to use
714
# MQA 576/512 approach for both prefill and decode
715
716
- # Concatenate q if it's a tuple (ql_nope, q_pe)
716
+ fp8_attention = self.kv_cache_dtype.startswith("fp8")
717
if isinstance(q, tuple):
718
ql_nope, q_pe = q
719
- q = self.q_concat_buffer[: ql_nope.shape[0]]
720
- ops.concat_mla_q(ql_nope, q_pe, q)
719
+ if fp8_attention:
720
+ q = layer._decode_concat_quant_fp8_op( # type: ignore[attr-defined]
721
+ ql_nope, q_pe, layer._q_scale
722
+ )
723
+ else:
724
+ q = self.q_concat_buffer[: ql_nope.shape[0]]
725
+ ops.concat_mla_q(ql_nope, q_pe, q)
726
727
num_actual_toks = attn_metadata.num_actual_tokens
728
@@ -736,12 +741,12 @@ class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata])
741
)
742
743
# write the latent and rope to kv cache
739
- fp8_attention = self.kv_cache_dtype.startswith("fp8")
744
if fp8_attention:
741
- original_q_shape = q.shape
745
kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(current_platform.fp8_dtype())
743
- q, _ = ops.scaled_fp8_quant(q.view(q.shape[0], -1), layer._q_scale)
744
- q = q.view(original_q_shape)
746
+ if q.dtype != current_platform.fp8_dtype():
747
+ original_q_shape = q.shape
748
+ q, _ = ops.scaled_fp8_quant(q.view(q.shape[0], -1), layer._q_scale)
749
+ q = q.view(original_q_shape)
750
mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
751
attn_out = self._forward_mla(
752
layer, mla_padded_q, kv_c_and_k_pe_cache, attn_metadata