iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Bugfix] Revert #56433 + #51692 to fix accuracy breakdown for DeepSeek-V4 (#57132)

Browse Source 
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
This commit is contained in:
Shanshan Shen
2026-09-16 18:58:23 +00:00
committed by GitHub
parent 403f182c02
commit bc0f47cd03
8 changed files with 34 additions and 402 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/test_deepseek_v4_vl_rocm.py
-125
View File
@@ -15,131 +15,6 @@ pytestmark = pytest.mark.skipif(
)
def _preshuffled_fp8_linear(
holder: str = "quant_method", weight_shape: tuple[int, int] = (256, 128)
) -> nn.Module:
pytest.importorskip("aiter")
from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
AiterPreshuffledFp8BlockScaledMMKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
FP8ScaledMMLinearLayerConfig,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
kFp8Dynamic128Sym,
kFp8Static128BlockSym,
)
config = FP8ScaledMMLinearLayerConfig(
weight_quant_key=kFp8Static128BlockSym,
activation_quant_key=kFp8Dynamic128Sym,
input_dtype=torch.bfloat16,
out_dtype=torch.bfloat16,
weight_shape=weight_shape,
)
layer = nn.Module()
layer.weight = nn.Parameter(
torch.randn(*weight_shape, device="cuda").to(current_platform.fp8_dtype()),
requires_grad=False,
)
layer.register_parameter(
"weight_scale" if holder == "scheme" else "weight_scale_inv",
nn.Parameter(
torch.ones(*(dim // 128 for dim in weight_shape), device="cuda"),
requires_grad=False,
),
)
setattr(
layer,
holder,
SimpleNamespace(fp8_linear=AiterPreshuffledFp8BlockScaledMMKernel(config)),
)
return layer
@pytest.mark.parametrize("holder", ["quant_method", "scheme"])
def test_rocm_wo_a_keeps_row_major_weights(holder: str, default_vllm_config) -> None:
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
layer = _preshuffled_fp8_linear(holder)
layer.is_bmm = True
original = layer.weight.detach().clone()
kernel = getattr(layer, holder).fp8_linear
kernel.process_weights_after_loading(layer)
actual = _get_cached_wo_a_bf16(layer, 2, 128, 128)
torch.testing.assert_close(actual, original.to(torch.bfloat16).view(2, 128, 128))
@pytest.mark.parametrize("holder", ["quant_method", "scheme"])
def test_rocm_gateup_shuffles_once_and_down_proj_keeps_preshuffled_backend(
holder: str, monkeypatch: pytest.MonkeyPatch, default_vllm_config
) -> None:
from vllm._aiter_ops import rocm_aiter_ops
from vllm.models.deepseek_v4.amd import model as rocm_model
gate_up, down = _preshuffled_fp8_linear(holder), _preshuffled_fp8_linear(holder)
original_gate_up = gate_up.weight.detach().clone()
original_down = down.weight.detach().clone()
monkeypatch.setattr(
rocm_model, "MergedColumnParallelLinear", lambda *a, **k: gate_up
)
monkeypatch.setattr(rocm_model, "RowParallelLinear", lambda *a, **k: down)
monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True)
model = rocm_model.DeepseekV4MLP(128, 128, "silu")
for linear in (gate_up, down):
getattr(linear, holder).fp8_linear.process_weights_after_loading(linear)
model.prepare_gateup_preshuffle()
for linear, original in ((gate_up, original_gate_up), (down, original_down)):
expected = rocm_aiter_ops.shuffle_weight(original, layout=(16, 16))
torch.testing.assert_close(
linear.weight.view(torch.uint8), expected.view(torch.uint8)
)
def test_rocm_fused_qkv_quant_matches_preshuffled_gemm_scale_layout(
monkeypatch: pytest.MonkeyPatch, default_vllm_config
) -> None:
from vllm._aiter_ops import rocm_aiter_ops
from vllm.models.deepseek_v4.amd.rocm import (
DeepseekV4ROCMAiterMLAAttention,
apply_pre_quantized_block_scaled_mm,
)
if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(1024, 256):
pytest.skip("Requires a tuned AITER preshuffled FP8 GEMM")
linear = _preshuffled_fp8_linear(weight_shape=(1024, 256))
original = linear.weight.detach().clone()
linear.quant_method.fp8_linear.process_weights_after_loading(linear)
attention = DeepseekV4ROCMAiterMLAAttention.__new__(DeepseekV4ROCMAiterMLAAttention)
nn.Module.__init__(attention)
attention.wq_b = linear
attention.indexer = None
attention.q_lora_rank, attention.head_dim = 256, 128
attention.eps = 1e-5
attention.q_norm = SimpleNamespace(
weight=torch.ones(256, dtype=torch.bfloat16, device="cuda")
)
attention.kv_norm = SimpleNamespace(
weight=torch.ones(128, dtype=torch.bfloat16, device="cuda")
)
monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: True)
# Multiple rows and groups with different ranges expose byte-order errors.
inputs = torch.randn(16, 384, dtype=torch.bfloat16, device="cuda")
inputs[:, :128] *= 4
quantized, scales, _ = attention._split_qkv_and_norm(inputs)
assert scales is not None
output = apply_pre_quantized_block_scaled_mm(linear, quantized, scales)
q = inputs[:, :256].float()
normalized = q * torch.rsqrt(q.square().mean(dim=-1, keepdim=True) + attention.eps)
reference_scales = normalized.view(16, 2, 128).abs().amax(dim=-1)
reference_scales /= torch.finfo(current_platform.fp8_dtype()).max
dequantized = quantized.float() * reference_scales.repeat_interleave(128, dim=1)
expected = dequantized @ original.float().T
torch.testing.assert_close(output.float(), expected, atol=0.125, rtol=0.01)
def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None:
from vllm.config import CacheConfig
from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype
vllm/_aiter_ops.py
+5 -27
View File
@@ -175,26 +175,18 @@ def is_aiter_found_and_supported_on_rdna4() -> bool:
@functools.cache
def _load_gemm_tuned_configs(
csv_path: str,
filters: tuple[tuple[str, object], ...],
key_cols: tuple[str, ...] = ("N", "K", "M"),
) -> set[tuple[int, ...]]:
q_dtype_w: torch.dtype, csv_path: str
) -> set[tuple[int, int, int]]:
try:
df = pd.read_csv(csv_path).drop_duplicates()
for col, val in filters:
if col not in df.columns:
continue
if isinstance(val, int):
df = df[df[col].astype(int) == val]
else:
df = df[df[col].astype(str) == str(val)]
return set(zip(*(df[c].astype(int) for c in key_cols)))
df = df[df["q_dtype_w"] == str(q_dtype_w)]
return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int)))
except Exception:
return set()
def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool:
configs = _load_gemm_tuned_configs(csv_path, (("q_dtype_w", q_dtype_w),))
configs = _load_gemm_tuned_configs(q_dtype_w, csv_path)
l_m = (
[1, 2, 4]
+ list(range(8, 513, 8))
@@ -3329,20 +3321,6 @@ class rocm_aiter_ops:
csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE
return _check_kernel_tuned(N, K, q_dtype_w, csv_path)
@staticmethod
def is_blockscale_bpreshuffle_tuned(n: int, k: int) -> bool:
"""Whether (N, K) has a tuned aiter blockscale bpreshuffle config."""
if not current_platform.is_rocm():
return False
import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops
csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE
gfx = aiter_gemm_a8w8_ops.get_gfx()
cu_num = aiter_gemm_a8w8_ops.get_cu_num()
return (n, k) in _load_gemm_tuned_configs(
csv_path, (("gfx", gfx), ("cu_num", cu_num)), key_cols=("N", "K")
)
@staticmethod
def shuffle_weight(
tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
vllm/model_executor/kernels/linear/__init__.py
-4
View File
@@ -177,7 +177,6 @@ from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
AiterInt8ScaledMMLinearKernel,
AiterPerTokenFp8ScaledMMLinearKernel,
AiterPreshuffledFp8BlockScaledMMKernel,
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
)
from vllm.model_executor.kernels.linear.scaled_mm.b12x import (
@@ -322,7 +321,6 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
"aiter": {
AiterInt8ScaledMMLinearKernel,
AiterFp8BlockScaledMMKernel,
AiterPreshuffledFp8BlockScaledMMKernel,
AiterPerTokenFp8ScaledMMLinearKernel,
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
AiterMxfp4LinearKernel,
@@ -466,7 +464,6 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
BlockWiseTorchFP8ScaledMMLinearKernel,
],
PlatformEnum.ROCM: [
AiterPreshuffledFp8BlockScaledMMKernel,
AiterFp8BlockScaledMMKernel,
TritonFp8BlockScaledMMKernel,
],
@@ -1250,7 +1247,6 @@ __all__ = [
"ScaledMMLinearLayerConfig",
"AiterHipbMMPerTokenFp8ScaledMMLinearKernel",
"AiterPreshuffledPerTokenFp8ScaledMMLinearKernel",
"AiterPreshuffledFp8BlockScaledMMKernel",
"AiterPerTokenFp8ScaledMMLinearKernel",
"NvFp4LinearKernel",
"NvFp4LinearLayerConfig",
vllm/model_executor/kernels/linear/scaled_mm/aiter.py
-170
View File
@@ -2,8 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import ClassVar
import torch
from vllm import _custom_ops as ops
@@ -474,171 +472,3 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
return gemm_a8w8_blockscale_op(
A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype
)
class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
"""Aiter FP8 block-scaled GEMM using a pre-shuffled (bpreshuffle) weight."""
# gemm_a8w8_blockscale_bpreshuffle reads the activation scale column-major.
wants_transposed_act_scale: ClassVar[bool] = True
# process_weights_after_loading shuffles layer.weight to layout (16, 16).
preshuffles_weight: ClassVar[bool] = True
@classmethod
def is_supported(
cls, compute_capability: int | None = None
) -> tuple[bool, str | None]:
return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported(
compute_capability
)
@classmethod
def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
can_implement_base, reason = super().can_implement(config)
if not can_implement_base:
return can_implement_base, reason
act_quant_desc = config.activation_quant_key.scale
if act_quant_desc.group_shape != GroupShape(1, 128):
return (
False,
(
"Supports only dynamic per token group activation "
"quantization with group_shape=(1,128)."
),
)
# bpreshuffle GEMM requires aiter fp8 linear to be enabled and an fp8
# 2D weight with N and K divisible by 128. Require a tuned config as the
# fallback can trigger faults or numerical errors.
if not rocm_aiter_ops.is_linear_fp8_enabled():
return (
False,
(
"requires setting `VLLM_ROCM_USE_AITER=1` "
"and `VLLM_ROCM_USE_AITER_LINEAR=1`."
),
)
n, k = config.weight_shape
if not (n % 128 == 0 and k % 128 == 0):
return (
False,
(
f"requires N and K dimensions divisible by 128, received "
f"N={n} and K={k}."
),
)
if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(n, k):
return (
False,
(
f"requires a tuned aiter blockscale bpreshuffle config for "
f"N={n} and K={k}."
),
)
return True, None
@staticmethod
def _reads_weight_directly(layer: torch.nn.Module) -> bool:
"""True when something other than apply_weights consumes layer.weight.
Such a weight must stay in the plain layout: ``is_bmm`` marks a stack
of matrices (wo_a), ``skip_weight_relayout`` marks MLA's kv_b_proj.
Both are stamped after construction, so can_implement cannot see them.
"""
return bool(
getattr(layer, "is_bmm", False)
or getattr(layer, "skip_weight_relayout", False)
)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
super().process_weights_after_loading(layer)
if self._reads_weight_directly(layer):
return
params = FP8BlockParams.from_layer(layer)
if params.weight_scale_inv is not None:
ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV
else:
ws, attr = params.weight_scale, params.WEIGHT_SCALE
if ws is not None and ws.dtype == torch.float8_e8m0fnu:
replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous())
weight = params.weight
# runtime safety net
assert (
weight.dim() == 2
and weight.dtype == current_platform.fp8_dtype()
and weight.shape[0] % 128 == 0
and weight.shape[1] % 128 == 0
), (
"AiterPreshuffledFp8BlockScaledMMKernel requires a 2D fp8 weight "
"with N and K divisible by 128."
)
shuffled_weight = rocm_aiter_ops.shuffle_weight(
weight.contiguous(), layout=(16, 16)
)
replace_parameter(
layer,
params.WEIGHT,
torch.nn.Parameter(shuffled_weight.data, requires_grad=False),
)
def apply_weights(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
**kwargs,
) -> torch.Tensor:
params = self._get_layer_params(layer)
Bs = (
params.weight_scale
if params.weight_scale_inv is None
else params.weight_scale_inv
)
# Left unshuffled above; kv_b_proj still reaches here from MLA's
# prefill-context path.
plain = self._reads_weight_directly(layer)
x_2d = x.view(-1, x.shape[-1])
A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=not plain)
if plain:
output = rocm_aiter_ops.gemm_a8w8_blockscale(
A,
params.weight,
As,
Bs,
list(self.weight_group_shape),
output_dtype=self.config.out_dtype,
)
else:
output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
A, params.weight, As, Bs, output_dtype=self.config.out_dtype
)
if bias is not None:
output = output + bias
return output.view(*x.shape[:-1], params.weight.shape[0])
def apply_block_scaled_mm(
self,
A: torch.Tensor,
B: torch.Tensor,
As: torch.Tensor,
Bs: torch.Tensor,
) -> torch.Tensor:
"""Block-scaled GEMM for callers that pre-quantize their activations.
``As`` must be column-major; a row-major one will not raise, it just
returns wrong numbers for M > 1.
"""
return rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
A, B, As, Bs, output_dtype=self.config.out_dtype
)
vllm/model_executor/layers/attention/mla_attention.py
+4 -3
View File
@@ -582,9 +582,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):
)
self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False)
# MLA reads this weight directly to build W_UK/W_UV, so a backend must
# not relayout it at load time.
kv_b_proj.skip_weight_relayout = True
# AMX reads kv_b_proj's weight directly and never calls it live; the
# reference CPU MLA backend calls it but isn't perf-critical. Skip
# the packed-kernel dispatch either way.
kv_b_proj._cpu_skip_gemm_dispatch = True
self.use_direct_call = not current_platform.opaque_attention_op()
vllm_config = get_current_vllm_config()
vllm/model_executor/layers/linear.py
+3 -3
View File
@@ -205,9 +205,9 @@ class UnquantizedLinearMethod(LinearMethodBase):
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
if current_platform.is_cpu():
# MLA's kv_b_proj (see `skip_weight_relayout`): not perf-critical,
# so skip packing and use a plain fallback.
if getattr(layer, "skip_weight_relayout", False):
# MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not
# perf-critical, so skip packing and use a plain fallback.
if getattr(layer, "_cpu_skip_gemm_dispatch", False):
layer.cpu_linear = torch.nn.functional.linear
return
vllm/models/deepseek_v4/amd/model.py
+6 -11
View File
@@ -71,10 +71,7 @@ from vllm.model_executor.models.utils import (
make_layers,
maybe_prefix,
)
from vllm.models.deepseek_v4.amd.rocm import (
DeepseekV4ROCMAiterMLAAttention,
weight_already_preshuffled,
)
from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
from vllm.platforms import current_platform
from vllm.platforms.rocm import on_gfx950
from vllm.sequence import IntermediateTensors
@@ -152,13 +149,11 @@ class DeepseekV4MLP(nn.Module):
return
if ws.dtype == torch.float8_e8m0fnu:
ws = _upcast_e8m0_to_fp32(ws).contiguous()
# Skip if the linear's kernel already shuffled it.
if not weight_already_preshuffled(self.gate_up_proj):
replace_parameter(
self.gate_up_proj,
"weight",
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
replace_parameter(
self.gate_up_proj,
"weight",
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
self._gateup_scale = ws
def forward(self, x):
vllm/models/deepseek_v4/amd/rocm.py
+16 -59
View File
@@ -62,21 +62,6 @@ def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
return indptr
def weight_already_preshuffled(linear: torch.nn.Module) -> bool:
"""True when the linear's kernel already B-preshuffled ``weight``.
The hand-shuffles below (fused_wqa_wkv, wo_b, gate_up_proj) must be skipped
for those, since shuffle_weight is a permutation rather than an involution.
"""
return any(
getattr(getattr(method, "fp8_linear", None), "preshuffles_weight", False)
for method in (
getattr(linear, "quant_method", None),
getattr(linear, "scheme", None),
)
)
def apply_pre_quantized_block_scaled_mm(
linear: torch.nn.Module,
x_fp8: torch.Tensor,
@@ -792,13 +777,12 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
return None
if ws.dtype == torch.float8_e8m0fnu:
ws = _upcast_e8m0_to_fp32(ws).contiguous()
# Skip if the linear's kernel already shuffled it.
if not weight_already_preshuffled(linear):
replace_parameter(
linear,
"weight",
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
# Shuffle the weight in place (single weight, no unshuffled copy).
replace_parameter(
linear,
"weight",
rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
)
return ws
self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
@@ -896,26 +880,16 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
return qr_kv, kv_score, indexer_kv_score, indexer_weights
@functools.cached_property
def _wq_b_act_scale_transpose(self) -> bool | None:
"""Activation-scale byte order every wq_b consumer agrees on, else None.
None means "do not take the fused norm+quant path". Otherwise the value
is the ``transpose_scale`` the producer must pass so that the scale it
writes matches what the selected GEMM reads.
def _wq_b_uses_aiter_block_scaled(self) -> bool:
"""True when both wq_b GEMMs run the aiter block-scaled fp8 kernel.
Cached: the linear kernels and the aiter env gates are fixed once
the model is built, so this is evaluated at the first forward
only.
Two conditions, both necessary:
* The fused norm+quant path is only valid if the quant and GEMM it
replaces are exactly the aiter ones; otherwise fall back to the
shared path.
* One qr_scale feeds both self.wq_b and self.indexer.wq_b, which have
different (N, K) and so can resolve to kernels wanting opposite byte
orders. A single producer cannot serve both, so refuse the fast path
when they disagree rather than guessing.
The fused norm+quant path is only valid if the quant and GEMM it
replaces are exactly the aiter ones; otherwise fall back to the
shared path.
"""
from vllm._aiter_ops import rocm_aiter_ops
from vllm.model_executor.kernels.linear.scaled_mm import (
@@ -923,31 +897,16 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
)
if not rocm_aiter_ops.is_linear_fp8_enabled():
return None
return False
linears = [self.wq_b]
if self.indexer is not None:
linears.append(self.indexer.wq_b)
layouts: set[bool] = set()
for linear in linears:
kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None)
if not isinstance(kernel, Fp8BlockScaledMMLinearKernel):
return None
layouts.add(bool(getattr(kernel, "wants_transposed_act_scale", False)))
if len(layouts) != 1:
logger.warning_once(
"DeepSeek-V4 wq_b consumers disagree on activation-scale layout; "
"disabling the fused q/kv norm+quant path.",
scope="global",
)
return None
transpose_scale = layouts.pop()
logger.debug_once(
"DeepSeek-V4 wq_b: emitting %s-major activation scales",
"column" if transpose_scale else "row",
scope="global",
)
return transpose_scale
return False
return True
def _split_qkv_and_norm(
self, qr_kv: torch.Tensor
@@ -964,12 +923,11 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
when the aiter linear path is not active.
"""
qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
transpose_scale = self._wq_b_act_scale_transpose
if not (
qr.dim() == 2
and qr.shape[0] > 0
and self.q_lora_rank % 128 == 0
and transpose_scale is not None
and self._wq_b_uses_aiter_block_scaled
):
return super()._split_qkv_and_norm(qr_kv)
@@ -983,8 +941,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
kv_weight=self.kv_norm.weight.data,
kv_epsilon=self.eps,
group_size=128,
# Emit the byte order the wq_b GEMMs read.
transpose_scale=transpose_scale,
transpose_scale=False,
)
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: