diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 21e6e6369..928eb6352 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -203,17 +203,24 @@ def rocm_absorb_v_bmm( else: _bmm_buf = None if _use_aiter_gfx95 and attn.w_kc.dtype == torch.float8_e4m3fn: - attn_bmm_output = ( - batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( - X=attn_output, - WQ=attn.w_vc.transpose(-1, -2), - w_scale=attn.w_scale, - group_size=128, - YQ=None, - transpose_bm=False, - transpose_bm_in=True, - dtype=torch.bfloat16, - ) + # As in the mxfp4 path above, write (batch, heads, dim) so the + # post-GEMM flatten is a free view instead of a copy. + _bmm_buf = torch.empty( + attn_output.shape[0], + attn.num_local_heads, + attn.w_vc.shape[-1], + device=attn_output.device, + dtype=torch.bfloat16, + ) + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( + X=attn_output, + WQ=attn.w_vc.transpose(-1, -2), + w_scale=attn.w_scale, + group_size=128, + YQ=_bmm_buf, + transpose_bm=True, + transpose_bm_in=True, + dtype=torch.bfloat16, ) else: attn_bmm_output = torch.bmm( diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index b4557f795..9899f9fb8 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -31,6 +31,7 @@ from sglang.srt.layers.quantization.fp8_utils import ( block_quant_dequant, block_quant_to_tensor_quant, channel_quant_to_tensor_quant, + input_to_float8, inverse_transform_scale_ue8m0, normalize_e4m3fn_to_e4m3fnuz, quant_weight_ue8m0, @@ -639,6 +640,16 @@ class DeepseekV2WeightLoaderMixin: torch.bfloat16 ) + # GLM ships kv_b_proj as bf16, which falls back to torch.bmm. Quantize to + # per-tensor e4m3fn (not fnuz) to match forward_mla_rocm's dtype gate. + if ( + _use_aiter_gfx95 + and self.config.architectures + and self.config.architectures[0] == "GlmMoeDsaForCausalLM" + and w.dtype == torch.bfloat16 + ): + w, self_attn.w_scale = input_to_float8(w, dtype=torch.float8_e4m3fn) + w_kc, w_vc = w.unflatten( 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)