diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 149597cf9..d9bb49559 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1770,13 +1770,9 @@ class DeepseekV2AttentionMLA(nn.Module): res1=None, output_unquantized_inp1=True, # return unqaunt kv_a ) - kv = self.kv_b_proj( - kv_a_quanted, - )[0] else: kv_a = self.kv_a_layernorm(kv_a) - kv = self.kv_b_proj(kv_a)[0] # kv_a = self.kv_a_layernorm(kv_a) @@ -1800,7 +1796,12 @@ class DeepseekV2AttentionMLA(nn.Module): q.dtype, forward_batch, ) - kv = self.kv_b_proj(kv_a)[0] + if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn: + kv = self.kv_b_proj( + kv_a_quanted, + )[0] + else: + kv = self.kv_b_proj(kv_a)[0] kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) k_nope = kv[..., : self.qk_nope_head_dim] v = kv[..., self.qk_nope_head_dim :]