minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[ROCm] Eliminate redundant contiguous copy in MLA attention on ROCm MXFP4 (#25463)

Browse Source 
Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
This commit is contained in:
Rita Brugarolas
2026-05-29 01:28:32 -07:00
committed by GitHub
parent 05bef58297
commit 9062f583db
1 changed files with 21 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+21 -5
View File
@@ -575,13 +575,18 @@ class DeepseekMLAForwardMixin:
# TODO(haishaw): add bmm_fp8 to ROCm # TODO(haishaw): add bmm_fp8 to ROCm
if _use_aiter_gfx95 and self.w_vc.dtype == torch.uint8: if _use_aiter_gfx95 and self.w_vc.dtype == torch.uint8:
x = attn_output.transpose(0, 1) x = attn_output.transpose(0, 1)
attn_bmm_output = torch.empty( B_heads, M_batch = x.shape[0], x.shape[1]
x.shape[0], N_vdim = self.w_vc.shape[2]
x.shape[1], # Allocate in (batch, heads, dim) so the post-GEMM
self.w_vc.shape[2], # transpose+flatten is a free view instead of a copy.
_bmm_buf = torch.empty(
M_batch,
B_heads,
N_vdim,
device=x.device, device=x.device,
dtype=torch.bfloat16, dtype=torch.bfloat16,
) )
attn_bmm_output = _bmm_buf.transpose(0, 1)
batched_gemm_afp4wfp4_pre_quant( batched_gemm_afp4wfp4_pre_quant(
x, x,
self.w_vc.transpose(-2, -1), self.w_vc.transpose(-2, -1),
@@ -590,6 +595,7 @@ class DeepseekMLAForwardMixin:
attn_bmm_output, attn_bmm_output,
) )
else: else:
_bmm_buf = None
if _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn: if _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn:
attn_bmm_output = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( attn_bmm_output = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
X=attn_output, X=attn_output,
@@ -607,7 +613,17 @@ class DeepseekMLAForwardMixin:
self.w_vc.to(torch.bfloat16) * self.w_scale, self.w_vc.to(torch.bfloat16) * self.w_scale,
) )
if self.o_proj.weight.dtype == torch.uint8: if _bmm_buf is not None:
# _bmm_buf is already (batch, heads, dim) contiguous
if self.o_proj.weight.dtype == torch.uint8:
attn_bmm_output = fused_flatten_mxfp4_quant(_bmm_buf)
elif self.o_proj.weight.dtype == torch.float8_e4m3fn:
attn_bmm_output = fused_flatten_fp8_group_quant(
_bmm_buf, group_size=128, dtype_quant=torch.float8_e4m3fn
)
else:
attn_bmm_output = _bmm_buf.flatten(1, 2)
elif self.o_proj.weight.dtype == torch.uint8:
attn_bmm_output = attn_bmm_output.transpose(0, 1) attn_bmm_output = attn_bmm_output.transpose(0, 1)
attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output) attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output)
elif self.o_proj.weight.dtype == torch.float8_e4m3fn: elif self.o_proj.weight.dtype == torch.float8_e4m3fn: