From 1bda9694b7bb15a0fabe0377be25360c484e4f28 Mon Sep 17 00:00:00 2001 From: Hemanth Acharya Date: Fri, 4 Sep 2026 12:26:07 +0530 Subject: [PATCH] [AMD][DSv4] Switch output projection gemm (oproj_a) to fp8 (#37423) Signed-off-by: Hemanth Acharya Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com> --- python/sglang/srt/arg_groups/model_hook.py | 18 +- python/sglang/srt/layers/quantization/fp8.py | 12 + .../amd/deepseek_v4_wo_a_fp8.py | 223 ++++++++++++++++++ python/sglang/srt/models/deepseek_v4.py | 61 ++++- 4 files changed, 312 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/models/deepseek_common/amd/deepseek_v4_wo_a_fp8.py diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index a56fe3f31..6f66b0dec 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -44,6 +44,18 @@ from sglang.srt.utils.common import ( logger = logging.getLogger(__name__) +def _rocm_fp8_wo_a_supported() -> bool: + """True when ROCm can run the DeepSeek-V4 fp8 wo_a GEMM (gfx950 + aiter).""" + try: + from sglang.srt.models.deepseek_common.amd.deepseek_v4_wo_a_fp8 import ( + is_wo_a_fp8_mxscale_supported, + ) + + return is_wo_a_fp8_mxscale_supported() + except Exception: # pragma: no cover - env-dependent + return False + + def handle_model_specific_adjustments(server_args: Any): cfg = resolving_view(server_args) @@ -361,7 +373,11 @@ def handle_model_specific_adjustments(server_args: Any): envs.SGLANG_OPT_USE_TILELANG_INDEXER.set(True) elif get_platform().is_hip: envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False) - envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False) + # The fp8 wo_a GEMM is DeepGEMM-based on CUDA. ROCm has an aiter + # e8m0 block-scale equivalent, but only on gfx950 -- everywhere else + # keeps the bf16 absorb GEMM. + if not _rocm_fp8_wo_a_supported(): + envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False) envs.SGLANG_OPT_USE_JIT_INDEXER_METADATA.set(False) envs.SGLANG_OPT_USE_TOPK_V2.set(True) envs.SGLANG_OPT_USE_AITER_INDEXER.set(True) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 1ed1d21c8..406ce76ec 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -724,9 +724,16 @@ class Fp8LinearMethod(LinearMethodBase): layer.weight.data = weight.data layer.weight_scale_inv.data = weight_scale.data + # The preshuffle rewrites the weight into a layout only + # aiter_w8a8_block_fp8_linear can read, so it is correct exactly when + # this quant method is what consumes the weight. A layer whose weight is + # read directly by the model (DeepSeek-V4 wo_a, whose absorb GEMM takes + # .weight/.weight_scale_inv and runs its own batched kernel) sets + # skip_aiter_bpreshuffle and keeps the plain row-major layout. if ( _use_aiter_bpreshuffle_gfx95 and self.w8a8_block_fp8_linear is aiter_w8a8_block_fp8_linear + and not getattr(layer, "skip_aiter_bpreshuffle", False) ): n, k = layer.weight.shape if not use_aiter_triton_gemm_w8a8_tuned_gfx950(n, k): @@ -735,6 +742,11 @@ class Fp8LinearMethod(LinearMethodBase): t = shuffle_weight(layer.weight, (16, 16)) layer.weight.copy_(t) del t + # The shuffle is in place and preserves shape, dtype and + # strides, so nothing downstream can tell it happened. Record + # it so a consumer that needs the row-major layout can assert + # instead of silently reading a permuted weight. + layer.aiter_bpreshuffled = True def _process_mxfp8_linear_weight_scale(self, layer: Module) -> None: if not self.use_mxfp8: diff --git a/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_wo_a_fp8.py b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_wo_a_fp8.py new file mode 100644 index 000000000..9e6860ff1 --- /dev/null +++ b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_wo_a_fp8.py @@ -0,0 +1,223 @@ +"""DeepSeek-V4 ``wo_a`` (MLA output-absorb) fp8 GEMM for AMD gfx950. + +``wo_a`` is the first half of the o_proj: it absorbs the attention output into +the low-rank o space before ``wo_b`` projects back to the hidden size. On CUDA +it runs in fp8 via DeepGEMM's ``fp8_einsum`` (``SGLANG_OPT_FP8_WO_A_GEMM``), but +DeepGEMM is CUDA-only, so ROCm falls back to a bf16 batched GEMM and pays for +loading bf16 weights on a decode step that is bound by weight traffic. + +This module is the ROCm equivalent, built on aiter's microscaling (e8m0) +block-scale batched GEMM ``batched_gemm_a8w8_mxscale``: + +* a Triton quantizer for the post-inverse-RoPE attention output, emitting the + per-token-group fp8 codes and uint8 e8m0 scales that kernel expects, and +* the load-time conversion of the checkpoint's fp32 block scales to e8m0. + +Shapes per attention-TP rank, with ``G`` local o-groups and 128-wide blocks:: + + activation o [T, G, D] bf16 -> fp8 e4m3 + act scale o_s [T, G, D/128] uint8 e8m0 (per token-group) + weight wo_a [G, R, D] fp8 e4m3 (from the checkpoint) + wgt scale wo_a_s [G, R/128, D/128] uint8 e8m0 + output z [T, G, R] bf16 + +The e8m0 scale is a bare power-of-two exponent, so a block's scale must be +rounded *up* to a power of two before its values are divided by it. The +quantizer does that for activations; ``wo_a_weight_scale_to_e8m0`` does it for +weights, requantizing the weight itself if the checkpoint's scales are not +already powers of two (otherwise the exponent round-up would silently rescale +the weights). +""" + +from __future__ import annotations + +import logging +from typing import Tuple + +import torch +import triton +import triton.language as tl + +from sglang.srt.utils import is_gfx95_supported, is_hip + +logger = logging.getLogger(__name__) + +# Block size of the microscaling scales, on both operands and both axes. +WO_A_MXFP8_GROUP_SIZE = 128 +_FP8_MAX = 448.0 +_ABSMAX_EPS = 1e-10 + +_is_hip = is_hip() +_is_gfx95_supported = is_gfx95_supported() + +# The mxscale flatmm BMM is gfx950-only, so resolve availability once at import +# rather than per decode step. +_batched_gemm_a8w8_mxscale = None +if _is_hip and _is_gfx95_supported: + try: + from aiter.ops.batched_gemm_op_a8w8 import ( + batched_gemm_a8w8_mxscale as _batched_gemm_a8w8_mxscale, + ) + except Exception as err: # pragma: no cover - env-dependent + logger.warning( + "aiter batched_gemm_a8w8_mxscale import failed; the DSV4 wo_a fp8 " + "path is unavailable on this build: %s", + err, + ) + + +def is_wo_a_fp8_mxscale_supported() -> bool: + """True when the ROCm fp8 ``wo_a`` path can run on this build/arch.""" + return _batched_gemm_a8w8_mxscale is not None + + +@triton.jit +def _wo_a_quant_mxfp8_kernel( + x_ptr, + xq_ptr, + xs_ptr, + D, + NUM_GROUPS, + NUM_GROUPS_PADDED: tl.constexpr, + GROUP_SIZE: tl.constexpr, + FP8_MAX: tl.constexpr, + EPS: tl.constexpr, +): + """One program per token-group row of the [T*G, D] view. + + The row is tiled as ``[NUM_GROUPS, GROUP_SIZE]`` so a single program reduces + every 128-wide block of the row at once. Splitting the row across programs + instead leaves each one with a single 128-element load, which on a decode + step is all launch latency and no work. + """ + row = tl.program_id(0) + grp = tl.arange(0, NUM_GROUPS_PADDED) + offs = row * D + grp[:, None] * GROUP_SIZE + tl.arange(0, GROUP_SIZE)[None, :] + mask = grp[:, None] < NUM_GROUPS + x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) + + # Smallest power-of-two scale that keeps a block inside the fp8 range: take + # ceil(log2(absmax / FP8_MAX)) by bumping the float32 exponent whenever the + # mantissa is non-zero (deep_gemm's ceil_to_ue8m0 convention). + raw = tl.maximum(tl.max(tl.abs(x), axis=1) / FP8_MAX, EPS) + bits = raw.to(tl.int32, bitcast=True) + exp = (bits >> 23) & 0xFF + exp = exp + tl.where((bits & 0x7FFFFF) != 0, 1, 0) + exp = tl.minimum(tl.maximum(exp, 1), 254) + scale = (exp << 23).to(tl.float32, bitcast=True) + + q = tl.minimum(tl.maximum(x / scale[:, None], -FP8_MAX), FP8_MAX) + tl.store(xq_ptr + offs, q.to(xq_ptr.dtype.element_ty), mask=mask) + # e8m0 stores the biased exponent itself (127 == scale 1.0). + tl.store(xs_ptr + row * NUM_GROUPS + grp, exp.to(tl.uint8), mask=grp < NUM_GROUPS) + + +def quant_wo_a_act_mxfp8(o: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize the [T, G, D] attention output to fp8 + e8m0 block scales. + + Returns ``(o_fp8 [T, G, D], o_scale [T, G, D/128] uint8)`` in the layout + ``batched_gemm_a8w8_mxscale`` consumes directly. + """ + T, G, D = o.shape + assert D % WO_A_MXFP8_GROUP_SIZE == 0, ( + f"wo_a in-features ({D}) must be divisible by {WO_A_MXFP8_GROUP_SIZE}" + ) + if not o.is_contiguous(): + o = o.contiguous() + + num_groups = D // WO_A_MXFP8_GROUP_SIZE + o_fp8 = torch.empty((T, G, D), device=o.device, dtype=torch.float8_e4m3fn) + o_scale = torch.empty((T, G, num_groups), device=o.device, dtype=torch.uint8) + + _wo_a_quant_mxfp8_kernel[(T * G,)]( + o, + o_fp8, + o_scale, + D, + num_groups, + NUM_GROUPS_PADDED=triton.next_power_of_2(num_groups), + GROUP_SIZE=WO_A_MXFP8_GROUP_SIZE, + FP8_MAX=_FP8_MAX, + EPS=_ABSMAX_EPS, + num_warps=8, + ) + return o_fp8, o_scale + + +def apply_wo_a_fp8_mxscale( + o: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, +) -> torch.Tensor: + """fp8 ``wo_a``: quantize [T, G, D] then batched-GEMM against [G, R, D]. + + ``weight_scale`` is the uint8 e8m0 [G, R/128, D/128] tensor produced at load + time by ``wo_a_weight_scale_to_e8m0``. Returns bf16 [T, G, R]. + """ + o_fp8, o_scale = quant_wo_a_act_mxfp8(o) + return _batched_gemm_a8w8_mxscale( + o_fp8, weight, o_scale, weight_scale, dtype=torch.bfloat16 + ) + + +def _is_power_of_two(scale: torch.Tensor) -> bool: + """True when every fp32 scale is an exact power of two (zero mantissa).""" + bits = scale.detach().float().contiguous().view(torch.int32) + return bool(((bits & 0x7FFFFF) == 0).all().item()) + + +def wo_a_weight_scale_to_e8m0( + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + num_groups: int, + o_lora_rank: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Convert checkpoint ``wo_a`` fp8 weight + fp32 block scales to e8m0. + + Returns ``(weight [G, R, D] fp8, scale [G, R/128, D/128] uint8)``. + + DeepSeek-V4 checkpoints that carry fp4 experts already store their linear + block scales as powers of two, so the conversion is exact and the weight + passes through untouched. Otherwise the weight is dequantized and requantized + against power-of-two scales, because rounding the exponent up on its own + would rescale every value in the block. + """ + G, R = num_groups, o_lora_rank + D = weight.shape[-1] + block = WO_A_MXFP8_GROUP_SIZE + + scale = weight_scale_inv.detach().float() + if not _is_power_of_two(scale): + from sglang.srt.layers.quantization.fp8_utils import ( + block_quant_dequant, + quant_weight_ue8m0, + ) + + logger.info_once( + "DSV4 wo_a block scales are not power-of-two; requantizing the " + "weight to ue8m0 for the aiter mxscale GEMM." + ) + dequant = block_quant_dequant( + weight.view(G * R, D), + scale.view(-1, D // block), + [block, block], + torch.bfloat16, + ) + weight, scale = quant_weight_ue8m0( + weight_dequant=dequant, weight_block_size=[block, block] + ) + + exponent = (scale.contiguous().view(torch.int32) >> 23).to(torch.uint8) + return ( + weight.contiguous().view(G, R, D), + exponent.contiguous().view(G, R // block, D // block), + ) + + +__all__ = [ + "WO_A_MXFP8_GROUP_SIZE", + "apply_wo_a_fp8_mxscale", + "is_wo_a_fp8_mxscale_supported", + "quant_wo_a_act_mxfp8", + "wo_a_weight_scale_to_e8m0", +] diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 7a564232f..3ec618f60 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -387,6 +387,23 @@ if _wo_a_aiter_batched_gemm_enabled: # instead of re-raising (and re-logging) on every layer/token. _wo_a_aiter_batched_gemm_disabled = False +# ROCm fp8 wo_a. The CUDA fp8 path below is built on DeepGEMM's fp8_einsum, so +# gfx950 runs the equivalent aiter e8m0 block-scale batched GEMM instead. Both +# the kernel availability and the weight-scale converter resolve once at import; +# ``None`` here means the platform keeps the bf16 absorb GEMM. +_wo_a_fp8_mxscale = None +_wo_a_weight_scale_to_e8m0 = None +if _is_hip: + from sglang.srt.models.deepseek_common.amd.deepseek_v4_wo_a_fp8 import ( + apply_wo_a_fp8_mxscale, + is_wo_a_fp8_mxscale_supported, + wo_a_weight_scale_to_e8m0, + ) + + if is_wo_a_fp8_mxscale_supported(): + _wo_a_fp8_mxscale = apply_wo_a_fp8_mxscale + _wo_a_weight_scale_to_e8m0 = wo_a_weight_scale_to_e8m0 + def _apply_wo_a_bf16_matmul( o: torch.Tensor, wo_a: torch.Tensor, is_decode: bool @@ -745,6 +762,16 @@ class MqaAttentionBase(nn.Module): self.wo_a.weight_scale_inv.format_ue8m0 = ( deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 ) + # wo_a is quantized but never *applied* through its quant method: + # the absorb GEMM in forward() reads .weight / .weight_scale_inv and + # runs its own batched kernel (DeepGEMM fp8_einsum on CUDA, aiter + # mxscale BMM on gfx950), both of which want the plain row-major + # [G, R, D] weight. Opt out of any backend-private weight layout the + # linear method would otherwise install for its own GEMM -- on ROCm + # that is aiter's B-preshuffle, which silently permutes the weight + # in place (same shape, dtype and strides) and makes this GEMM + # return noise. + self.wo_a.skip_aiter_bpreshuffle = True self.wo_b = RowParallelLinear( self.n_groups * self.o_lora_rank, self.hidden_size, @@ -1752,7 +1779,17 @@ class MQALayer(MqaAttentionBase): o = o.view(o.shape[0], self.n_local_groups, -1) - if _FP8_WO_A_GEMM: + if _FP8_WO_A_GEMM and _wo_a_fp8_mxscale is not None: + # ROCm gfx950: same fp8 absorb GEMM as the DeepGEMM path below, but + # through aiter's e8m0 block-scale batched GEMM. The activation is + # quantized per token-group inside the helper. + T, G, D = o.shape + o = _wo_a_fp8_mxscale( + o, + self.wo_a.weight.view(G, self.o_lora_rank, D), + self.wo_a.weight_scale_inv.data, + ) + elif _FP8_WO_A_GEMM: import deep_gemm from sglang.srt.layers import deep_gemm_wrapper @@ -3480,6 +3517,28 @@ class DeepseekV4ForCausalLM(nn.Module): R = attn.o_lora_rank D = attn.wo_a.weight.shape[1] + if _wo_a_weight_scale_to_e8m0 is not None: + # ROCm: aiter's mxscale GEMM reads uint8 e8m0 block scales, and + # requantizes the weight when the checkpoint's scales are not + # already powers of two. It also needs the weight row-major, so + # check the linear method honoured skip_aiter_bpreshuffle: a + # preshuffled weight has the same shape, dtype and strides and + # would only show up as garbage output. + assert not getattr(attn.wo_a, "aiter_bpreshuffled", False), ( + "DSV4 wo_a was B-preshuffled by the fp8 linear method; the " + "aiter mxscale absorb GEMM needs the row-major weight" + ) + weight, scale = _wo_a_weight_scale_to_e8m0( + attn.wo_a.weight.data, + attn.wo_a.weight_scale_inv.data, + G, + R, + ) + attn.wo_a.weight.data = weight.view(G * R, D) + attn.wo_a.weight_scale_inv.data = scale + attn.wo_a.weight_scale_inv.format_ue8m0 = True + continue + raw_scale = attn.wo_a.weight_scale_inv.data.view(G, R // 128, D // 128) if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: attn.wo_a.weight_scale_inv.data = transform_sf_into_required_layout(