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

[ROCm/gfx95] Fix fp8 per-channel attention for Kimi-K2.7-code-mxfp4 o… (#31105)

Browse Source 
Co-authored-by: Hung <Emmanuel0612@users.noreply.github.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
This commit is contained in:
Hsiu-Chun, Hung
2026-08-11 00:51:57 -07:00
committed by GitHub
co-authored by Hung HaiShaw
parent 1c06c160f9
commit e74ea5b1d7
7 changed files with 314 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha_rocm.py
+5 -7
View File
@@ -28,6 +28,7 @@ from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mha imp
resolve_attn_backend,
)
from sglang.srt.models.deepseek_common.utils import (
_is_block_scale_fp8,
_use_aiter_bpreshuffle_gfx95,
_use_aiter_gfx95,
)
@@ -74,10 +75,7 @@ class DeepseekMHARocmForwardMixin:
# on gfx95, we can still use fused RMSNorm+FP8 quant, but MUST request
# the unquantized output for q_lora; otherwise q_lora becomes the (fp8,scale)
# tuple.
if (
_use_aiter_gfx95
and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
):
if _use_aiter_gfx95 and _is_block_scale_fp8(self.q_b_proj):
q_quanted, q_lora, _, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
@@ -121,7 +119,7 @@ class DeepseekMHARocmForwardMixin:
None,
)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
elif _use_aiter_gfx95 and _is_block_scale_fp8(self.q_b_proj):
q, _, _, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
@@ -152,7 +150,7 @@ class DeepseekMHARocmForwardMixin:
kv_a, _ = latent_cache.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
latent_cache = latent_cache.unsqueeze(1)
if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
if _use_aiter_gfx95 and _is_block_scale_fp8(self.kv_b_proj):
kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
kv_a,
self.kv_a_layernorm.weight,
@@ -243,7 +241,7 @@ class DeepseekMHARocmForwardMixin:
)
)[0]
else:
if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
if _use_aiter_gfx95 and _is_block_scale_fp8(self.kv_b_proj):
kv = self.kv_b_proj(kv_a_quanted)[0]
else:
kv = self.kv_b_proj(kv_a)[0]
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py
+4 -3
View File
@@ -53,6 +53,7 @@ from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla imp
)
from sglang.srt.models.deepseek_common.utils import (
FORWARD_ABSORB_CORE_ATTENTION_BACKENDS,
_is_block_scale_fp8,
_is_gfx95_supported,
_use_aiter,
_use_aiter_bpreshuffle_gfx95,
@@ -224,7 +225,7 @@ def rocm_absorb_v_bmm(
# _bmm_buf is already (batch, heads, dim) contiguous
if attn.o_proj.weight.dtype == torch.uint8:
attn_bmm_output = fused_flatten_mxfp4_quant(_bmm_buf)
elif attn.o_proj.weight.dtype == torch.float8_e4m3fn:
elif _is_block_scale_fp8(attn.o_proj):
attn_bmm_output = fused_flatten_fp8_group_quant(
_bmm_buf,
group_size=128,
@@ -240,7 +241,7 @@ def rocm_absorb_v_bmm(
elif attn.o_proj.weight.dtype == torch.uint8:
attn_bmm_output = attn_bmm_output.transpose(0, 1)
attn_bmm_output = fused_flatten_mxfp4_quant(attn_bmm_output)
elif attn.o_proj.weight.dtype == torch.float8_e4m3fn:
elif _is_block_scale_fp8(attn.o_proj):
attn_bmm_output = attn_bmm_output.transpose(0, 1)
attn_bmm_output = fused_flatten_fp8_group_quant(
attn_bmm_output,
@@ -335,7 +336,7 @@ class DeepseekMLARocmForwardMixin:
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
)
elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
elif _use_aiter_gfx95 and _is_block_scale_fp8(self.q_b_proj):
if self.use_dsa:
q_quanted, q_lora, k_nope, _ = fused_rms_fp8_group_quant(
q,
python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
+5
View File
@@ -606,6 +606,10 @@ class DeepseekV2WeightLoaderMixin:
else:
weight = w
weight_scale = self_attn.kv_b_proj.weight_scale
# Per-channel scale is 1D [out]; reshape to [out, 1] so it
# broadcasts correctly against weight [out, in].
if weight_scale.dim() == 1:
weight_scale = weight_scale.view(-1, 1)
w, scale = channel_quant_to_tensor_quant(weight, weight_scale)
self_attn.w_scale = scale
@@ -638,6 +642,7 @@ class DeepseekV2WeightLoaderMixin:
and self.config.architectures
and self.config.architectures[0]
== "DeepseekV3ForCausalLM" # Avoid processing other models like GlmMoeDsaForCausalLM
and w.dtype not in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
):
w_kc, self_attn.w_scale_k, w_vc, self_attn.w_scale_v = (
quark_post_load_weights(self_attn, w, "mxfp4")
python/sglang/srt/models/deepseek_common/utils.py
+17
View File
@@ -72,6 +72,23 @@ FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [
]
def _is_block_scale_fp8(proj: torch.nn.Module) -> bool:
"""Return True if proj uses block-scale fp8 quantization.
Per-channel fp8 has weight_scale shape [N, 1] (one scale per output row).
Block-scale fp8 has weight_scale shape [N, K/block_size] (multiple columns).
The fused gfx95 kernels (fused_rms_fp8_group_quant, fused_flatten_fp8_group_quant)
are only compatible with block-scale layouts — per-channel layers must fall
through to the plain bf16 path instead.
"""
if not hasattr(proj, "weight") or proj.weight.dtype != torch.float8_e4m3fn:
return False
weight_scale = getattr(proj, "weight_scale", None)
if weight_scale is None or weight_scale.dim() != 2:
return False
return weight_scale.shape[-1] > 1
def awq_dequantize_func():
"""
Get the AWQ dequantize function for the current device
python/sglang/srt/models/deepseek_v2.py
+24 -5
View File
@@ -179,6 +179,7 @@ from sglang.srt.models.deepseek_common.deepseek_weight_loader import (
from sglang.srt.models.deepseek_common.utils import (
_device_sm,
_get_llama_4_scaling,
_is_block_scale_fp8,
_is_cpu,
_is_cpu_amx_available,
_is_cuda,
@@ -2402,17 +2403,35 @@ class DeepseekV2DecoderLayer(nn.Module):
def _detect_gfx95_quant_format(self) -> str:
if not _is_gfx95_supported:
return ""
weight = getattr(
getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None), "weight", None
)
proj = getattr(self.self_attn, "fused_qkv_a_proj_with_mqa", None)
weight = getattr(proj, "weight", None)
if weight is None:
return ""
if weight.dtype == torch.uint8:
return "mxfp4"
if weight.dtype == getattr(torch, "float8_e4m3fn", None):
return "fp8"
# Use _is_block_scale_fp8 to distinguish block-scale fp8 (K/128 scale
# cols, compatible with fused_rms_fp8_group_quant) from per-channel fp8
# ([N, 1] scale, must use the plain bf16 path).
# weight_scale may not be reshaped yet at __init__ time — return
# "fp8_pending" so _resolve_gfx95_quant_format re-checks on first forward.
weight_scale = getattr(proj, "weight_scale", None)
if weight_scale is None:
return "fp8_pending"
return "fp8" if _is_block_scale_fp8(proj) else ""
return ""
def _resolve_gfx95_quant_format(self) -> str:
"""Re-evaluate after weights are loaded if still pending."""
fmt = getattr(self, "_gfx95_quant_format", "")
if fmt == "fp8_pending":
fmt = self._detect_gfx95_quant_format()
if fmt == "fp8_pending":
# weight_scale still unavailable — default to bf16 (safe fallback).
fmt = ""
self._gfx95_quant_format = fmt
return fmt
def _is_layer_sparse(self, layer_id: int, is_nextn: bool) -> bool:
return is_nextn or (
self.config.n_routed_experts is not None
@@ -2440,7 +2459,7 @@ class DeepseekV2DecoderLayer(nn.Module):
residual,
forward_batch,
captured_last_layer_outputs=captured_last_layer_outputs,
quant_format=getattr(self, "_gfx95_quant_format", ""),
quant_format=self._resolve_gfx95_quant_format(),
)
)