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

[AMD][DSv4] Switch output projection gemm (oproj_a) to fp8 (#37423)

Browse Source 
Signed-off-by: Hemanth Acharya <heachary@amd.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: kk <43161300+kkHuang-amd@users.noreply.github.com>
This commit is contained in:
Hemanth Acharya
2026-09-03 23:56:07 -07:00
committed by GitHub
co-authored by Claude Opus 5 kk
parent b168f905c8
commit 1bda9694b7
4 changed files with 312 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/model_hook.py
+17 -1
View File
@@ -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)
python/sglang/srt/layers/quantization/fp8.py
+12
View File
@@ -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:
python/sglang/srt/models/deepseek_common/amd/deepseek_v4_wo_a_fp8.py
+223
View File
@@ -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",
]
python/sglang/srt/models/deepseek_v4.py
+60 -1
View File
@@ -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(