[ROCm][Bugfix] Revert #56433 + #51692 to fix accuracy breakdown for DeepSeek-V4 (#57132)
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Shanshan Shen committed
Sep 17, 2026 at 02:58 UTC
bc0f47cd03d6ae99f9f217f096684423bb4ebc2b
8 files changed
+34
-402
tests/models/test_deepseek_v4_vl_rocm.py
-125
@@ -15,131 +15,6 @@ pytestmark = pytest.mark.skipif(
15
)
16
17
18
-def _preshuffled_fp8_linear(
19
- holder: str = "quant_method", weight_shape: tuple[int, int] = (256, 128)
20
-) -> nn.Module:
21
- pytest.importorskip("aiter")
22
- from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
23
- AiterPreshuffledFp8BlockScaledMMKernel,
24
- )
25
- from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
26
- FP8ScaledMMLinearLayerConfig,
27
- )
28
- from vllm.model_executor.layers.quantization.utils.quant_utils import (
29
- kFp8Dynamic128Sym,
30
- kFp8Static128BlockSym,
31
- )
32
-
33
- config = FP8ScaledMMLinearLayerConfig(
34
- weight_quant_key=kFp8Static128BlockSym,
35
- activation_quant_key=kFp8Dynamic128Sym,
36
- input_dtype=torch.bfloat16,
37
- out_dtype=torch.bfloat16,
38
- weight_shape=weight_shape,
39
- )
40
- layer = nn.Module()
41
- layer.weight = nn.Parameter(
42
- torch.randn(*weight_shape, device="cuda").to(current_platform.fp8_dtype()),
43
- requires_grad=False,
44
- )
45
- layer.register_parameter(
46
- "weight_scale" if holder == "scheme" else "weight_scale_inv",
47
- nn.Parameter(
48
- torch.ones(*(dim // 128 for dim in weight_shape), device="cuda"),
49
- requires_grad=False,
50
- ),
51
- )
52
- setattr(
53
- layer,
54
- holder,
55
- SimpleNamespace(fp8_linear=AiterPreshuffledFp8BlockScaledMMKernel(config)),
56
- )
57
- return layer
58
-
59
-
60
-@pytest.mark.parametrize("holder", ["quant_method", "scheme"])
61
-def test_rocm_wo_a_keeps_row_major_weights(holder: str, default_vllm_config) -> None:
62
- from vllm.v1.attention.ops.rocm_aiter_mla_sparse import _get_cached_wo_a_bf16
63
-
64
- layer = _preshuffled_fp8_linear(holder)
65
- layer.is_bmm = True
66
- original = layer.weight.detach().clone()
67
- kernel = getattr(layer, holder).fp8_linear
68
- kernel.process_weights_after_loading(layer)
69
-
70
- actual = _get_cached_wo_a_bf16(layer, 2, 128, 128)
71
- torch.testing.assert_close(actual, original.to(torch.bfloat16).view(2, 128, 128))
72
-
73
-
74
-@pytest.mark.parametrize("holder", ["quant_method", "scheme"])
75
-def test_rocm_gateup_shuffles_once_and_down_proj_keeps_preshuffled_backend(
76
- holder: str, monkeypatch: pytest.MonkeyPatch, default_vllm_config
77
-) -> None:
78
- from vllm._aiter_ops import rocm_aiter_ops
79
- from vllm.models.deepseek_v4.amd import model as rocm_model
80
-
81
- gate_up, down = _preshuffled_fp8_linear(holder), _preshuffled_fp8_linear(holder)
82
- original_gate_up = gate_up.weight.detach().clone()
83
- original_down = down.weight.detach().clone()
84
- monkeypatch.setattr(
85
- rocm_model, "MergedColumnParallelLinear", lambda *a, **k: gate_up
86
- )
87
- monkeypatch.setattr(rocm_model, "RowParallelLinear", lambda *a, **k: down)
88
- monkeypatch.setattr(rocm_aiter_ops, "is_enabled", lambda: True)
89
- model = rocm_model.DeepseekV4MLP(128, 128, "silu")
90
- for linear in (gate_up, down):
91
- getattr(linear, holder).fp8_linear.process_weights_after_loading(linear)
92
- model.prepare_gateup_preshuffle()
93
-
94
- for linear, original in ((gate_up, original_gate_up), (down, original_down)):
95
- expected = rocm_aiter_ops.shuffle_weight(original, layout=(16, 16))
96
- torch.testing.assert_close(
97
- linear.weight.view(torch.uint8), expected.view(torch.uint8)
98
- )
99
-
100
-
101
-def test_rocm_fused_qkv_quant_matches_preshuffled_gemm_scale_layout(
102
- monkeypatch: pytest.MonkeyPatch, default_vllm_config
103
-) -> None:
104
- from vllm._aiter_ops import rocm_aiter_ops
105
- from vllm.models.deepseek_v4.amd.rocm import (
106
- DeepseekV4ROCMAiterMLAAttention,
107
- apply_pre_quantized_block_scaled_mm,
108
- )
109
-
110
- if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(1024, 256):
111
- pytest.skip("Requires a tuned AITER preshuffled FP8 GEMM")
112
- linear = _preshuffled_fp8_linear(weight_shape=(1024, 256))
113
- original = linear.weight.detach().clone()
114
- linear.quant_method.fp8_linear.process_weights_after_loading(linear)
115
- attention = DeepseekV4ROCMAiterMLAAttention.__new__(DeepseekV4ROCMAiterMLAAttention)
116
- nn.Module.__init__(attention)
117
- attention.wq_b = linear
118
- attention.indexer = None
119
- attention.q_lora_rank, attention.head_dim = 256, 128
120
- attention.eps = 1e-5
121
- attention.q_norm = SimpleNamespace(
122
- weight=torch.ones(256, dtype=torch.bfloat16, device="cuda")
123
- )
124
- attention.kv_norm = SimpleNamespace(
125
- weight=torch.ones(128, dtype=torch.bfloat16, device="cuda")
126
- )
127
- monkeypatch.setattr(rocm_aiter_ops, "is_linear_fp8_enabled", lambda: True)
128
- # Multiple rows and groups with different ranges expose byte-order errors.
129
- inputs = torch.randn(16, 384, dtype=torch.bfloat16, device="cuda")
130
- inputs[:, :128] *= 4
131
- quantized, scales, _ = attention._split_qkv_and_norm(inputs)
132
- assert scales is not None
133
- output = apply_pre_quantized_block_scaled_mm(linear, quantized, scales)
134
- q = inputs[:, :256].float()
135
- normalized = q * torch.rsqrt(q.square().mean(dim=-1, keepdim=True) + attention.eps)
136
- reference_scales = normalized.view(16, 2, 128).abs().amax(dim=-1)
137
- reference_scales /= torch.finfo(current_platform.fp8_dtype()).max
138
- dequantized = quantized.float() * reference_scales.repeat_interleave(128, dim=1)
139
- expected = dequantized @ original.float().T
140
- torch.testing.assert_close(output.float(), expected, atol=0.125, rtol=0.01)
141
-
142
-
18
def test_rocm_packed_kv_cache_auto_uses_ds_mla_layout() -> None:
19
from vllm.config import CacheConfig
20
from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype
vllm/_aiter_ops.py
+5
-27
@@ -175,26 +175,18 @@ def is_aiter_found_and_supported_on_rdna4() -> bool:
175
176
@functools.cache
177
def _load_gemm_tuned_configs(
178
- csv_path: str,
179
- filters: tuple[tuple[str, object], ...],
180
- key_cols: tuple[str, ...] = ("N", "K", "M"),
181
-) -> set[tuple[int, ...]]:
178
+ q_dtype_w: torch.dtype, csv_path: str
179
+) -> set[tuple[int, int, int]]:
180
try:
181
df = pd.read_csv(csv_path).drop_duplicates()
184
- for col, val in filters:
185
- if col not in df.columns:
186
- continue
187
- if isinstance(val, int):
188
- df = df[df[col].astype(int) == val]
189
- else:
190
- df = df[df[col].astype(str) == str(val)]
191
- return set(zip(*(df[c].astype(int) for c in key_cols)))
182
+ df = df[df["q_dtype_w"] == str(q_dtype_w)]
183
+ return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int)))
184
except Exception:
185
return set()
186
187
188
def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool:
197
- configs = _load_gemm_tuned_configs(csv_path, (("q_dtype_w", q_dtype_w),))
189
+ configs = _load_gemm_tuned_configs(q_dtype_w, csv_path)
190
l_m = (
191
[1, 2, 4]
192
+ list(range(8, 513, 8))
@@ -3329,20 +3321,6 @@ class rocm_aiter_ops:
3321
csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE
3322
return _check_kernel_tuned(N, K, q_dtype_w, csv_path)
3323
3332
- @staticmethod
3333
- def is_blockscale_bpreshuffle_tuned(n: int, k: int) -> bool:
3334
- """Whether (N, K) has a tuned aiter blockscale bpreshuffle config."""
3335
- if not current_platform.is_rocm():
3336
- return False
3337
- import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops
3338
-
3339
- csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE
3340
- gfx = aiter_gemm_a8w8_ops.get_gfx()
3341
- cu_num = aiter_gemm_a8w8_ops.get_cu_num()
3342
- return (n, k) in _load_gemm_tuned_configs(
3343
- csv_path, (("gfx", gfx), ("cu_num", cu_num)), key_cols=("N", "K")
3344
- )
3345
-
3324
@staticmethod
3325
def shuffle_weight(
3326
tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
vllm/model_executor/kernels/linear/__init__.py
-4
@@ -177,7 +177,6 @@ from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
177
AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
178
AiterInt8ScaledMMLinearKernel,
179
AiterPerTokenFp8ScaledMMLinearKernel,
180
- AiterPreshuffledFp8BlockScaledMMKernel,
180
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
181
)
182
from vllm.model_executor.kernels.linear.scaled_mm.b12x import (
@@ -322,7 +321,6 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
321
"aiter": {
322
AiterInt8ScaledMMLinearKernel,
323
AiterFp8BlockScaledMMKernel,
325
- AiterPreshuffledFp8BlockScaledMMKernel,
324
AiterPerTokenFp8ScaledMMLinearKernel,
325
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
326
AiterMxfp4LinearKernel,
@@ -466,7 +464,6 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
464
BlockWiseTorchFP8ScaledMMLinearKernel,
465
],
466
PlatformEnum.ROCM: [
469
- AiterPreshuffledFp8BlockScaledMMKernel,
467
AiterFp8BlockScaledMMKernel,
468
TritonFp8BlockScaledMMKernel,
469
],
@@ -1250,7 +1247,6 @@ __all__ = [
1247
"ScaledMMLinearLayerConfig",
1248
"AiterHipbMMPerTokenFp8ScaledMMLinearKernel",
1249
"AiterPreshuffledPerTokenFp8ScaledMMLinearKernel",
1253
- "AiterPreshuffledFp8BlockScaledMMKernel",
1250
"AiterPerTokenFp8ScaledMMLinearKernel",
1251
"NvFp4LinearKernel",
1252
"NvFp4LinearLayerConfig",
vllm/model_executor/kernels/linear/scaled_mm/aiter.py
-170
@@ -2,8 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
-from typing import ClassVar
6
-
5
import torch
6
7
from vllm import _custom_ops as ops
@@ -474,171 +472,3 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
472
return gemm_a8w8_blockscale_op(
473
A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype
474
)
477
-
478
-
479
-class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
480
- """Aiter FP8 block-scaled GEMM using a pre-shuffled (bpreshuffle) weight."""
481
-
482
- # gemm_a8w8_blockscale_bpreshuffle reads the activation scale column-major.
483
- wants_transposed_act_scale: ClassVar[bool] = True
484
-
485
- # process_weights_after_loading shuffles layer.weight to layout (16, 16).
486
- preshuffles_weight: ClassVar[bool] = True
487
-
488
- @classmethod
489
- def is_supported(
490
- cls, compute_capability: int | None = None
491
- ) -> tuple[bool, str | None]:
492
- return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported(
493
- compute_capability
494
- )
495
-
496
- @classmethod
497
- def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
498
- can_implement_base, reason = super().can_implement(config)
499
- if not can_implement_base:
500
- return can_implement_base, reason
501
-
502
- act_quant_desc = config.activation_quant_key.scale
503
- if act_quant_desc.group_shape != GroupShape(1, 128):
504
- return (
505
- False,
506
- (
507
- "Supports only dynamic per token group activation "
508
- "quantization with group_shape=(1,128)."
509
- ),
510
- )
511
-
512
- # bpreshuffle GEMM requires aiter fp8 linear to be enabled and an fp8
513
- # 2D weight with N and K divisible by 128. Require a tuned config as the
514
- # fallback can trigger faults or numerical errors.
515
- if not rocm_aiter_ops.is_linear_fp8_enabled():
516
- return (
517
- False,
518
- (
519
- "requires setting `VLLM_ROCM_USE_AITER=1` "
520
- "and `VLLM_ROCM_USE_AITER_LINEAR=1`."
521
- ),
522
- )
523
-
524
- n, k = config.weight_shape
525
- if not (n % 128 == 0 and k % 128 == 0):
526
- return (
527
- False,
528
- (
529
- f"requires N and K dimensions divisible by 128, received "
530
- f"N={n} and K={k}."
531
- ),
532
- )
533
-
534
- if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(n, k):
535
- return (
536
- False,
537
- (
538
- f"requires a tuned aiter blockscale bpreshuffle config for "
539
- f"N={n} and K={k}."
540
- ),
541
- )
542
-
543
- return True, None
544
-
545
- @staticmethod
546
- def _reads_weight_directly(layer: torch.nn.Module) -> bool:
547
- """True when something other than apply_weights consumes layer.weight.
548
-
549
- Such a weight must stay in the plain layout: ``is_bmm`` marks a stack
550
- of matrices (wo_a), ``skip_weight_relayout`` marks MLA's kv_b_proj.
551
- Both are stamped after construction, so can_implement cannot see them.
552
- """
553
- return bool(
554
- getattr(layer, "is_bmm", False)
555
- or getattr(layer, "skip_weight_relayout", False)
556
- )
557
-
558
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
559
- super().process_weights_after_loading(layer)
560
-
561
- if self._reads_weight_directly(layer):
562
- return
563
-
564
- params = FP8BlockParams.from_layer(layer)
565
- if params.weight_scale_inv is not None:
566
- ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV
567
- else:
568
- ws, attr = params.weight_scale, params.WEIGHT_SCALE
569
- if ws is not None and ws.dtype == torch.float8_e8m0fnu:
570
- replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous())
571
-
572
- weight = params.weight
573
- # runtime safety net
574
- assert (
575
- weight.dim() == 2
576
- and weight.dtype == current_platform.fp8_dtype()
577
- and weight.shape[0] % 128 == 0
578
- and weight.shape[1] % 128 == 0
579
- ), (
580
- "AiterPreshuffledFp8BlockScaledMMKernel requires a 2D fp8 weight "
581
- "with N and K divisible by 128."
582
- )
583
-
584
- shuffled_weight = rocm_aiter_ops.shuffle_weight(
585
- weight.contiguous(), layout=(16, 16)
586
- )
587
- replace_parameter(
588
- layer,
589
- params.WEIGHT,
590
- torch.nn.Parameter(shuffled_weight.data, requires_grad=False),
591
- )
592
-
593
- def apply_weights(
594
- self,
595
- layer: torch.nn.Module,
596
- x: torch.Tensor,
597
- bias: torch.Tensor | None = None,
598
- **kwargs,
599
- ) -> torch.Tensor:
600
- params = self._get_layer_params(layer)
601
- Bs = (
602
- params.weight_scale
603
- if params.weight_scale_inv is None
604
- else params.weight_scale_inv
605
- )
606
-
607
- # Left unshuffled above; kv_b_proj still reaches here from MLA's
608
- # prefill-context path.
609
- plain = self._reads_weight_directly(layer)
610
-
611
- x_2d = x.view(-1, x.shape[-1])
612
- A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=not plain)
613
- if plain:
614
- output = rocm_aiter_ops.gemm_a8w8_blockscale(
615
- A,
616
- params.weight,
617
- As,
618
- Bs,
619
- list(self.weight_group_shape),
620
- output_dtype=self.config.out_dtype,
621
- )
622
- else:
623
- output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
624
- A, params.weight, As, Bs, output_dtype=self.config.out_dtype
625
- )
626
- if bias is not None:
627
- output = output + bias
628
- return output.view(*x.shape[:-1], params.weight.shape[0])
629
-
630
- def apply_block_scaled_mm(
631
- self,
632
- A: torch.Tensor,
633
- B: torch.Tensor,
634
- As: torch.Tensor,
635
- Bs: torch.Tensor,
636
- ) -> torch.Tensor:
637
- """Block-scaled GEMM for callers that pre-quantize their activations.
638
-
639
- ``As`` must be column-major; a row-major one will not raise, it just
640
- returns wrong numbers for M > 1.
641
- """
642
- return rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
643
- A, B, As, Bs, output_dtype=self.config.out_dtype
644
- )
vllm/model_executor/layers/attention/mla_attention.py
+4
-3
@@ -582,9 +582,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):
582
)
583
self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
584
self.is_amx_bmm_enabled = getattr(self.impl, "uses_amx_bmm", False)
585
- # MLA reads this weight directly to build W_UK/W_UV, so a backend must
586
- # not relayout it at load time.
587
- kv_b_proj.skip_weight_relayout = True
585
+ # AMX reads kv_b_proj's weight directly and never calls it live; the
586
+ # reference CPU MLA backend calls it but isn't perf-critical. Skip
587
+ # the packed-kernel dispatch either way.
588
+ kv_b_proj._cpu_skip_gemm_dispatch = True
589
self.use_direct_call = not current_platform.opaque_attention_op()
590
591
vllm_config = get_current_vllm_config()
vllm/model_executor/layers/linear.py
+3
-3
@@ -205,9 +205,9 @@ class UnquantizedLinearMethod(LinearMethodBase):
205
206
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
207
if current_platform.is_cpu():
208
- # MLA's kv_b_proj (see `skip_weight_relayout`): not perf-critical,
209
- # so skip packing and use a plain fallback.
210
- if getattr(layer, "skip_weight_relayout", False):
208
+ # MLA's kv_b_proj (see `_cpu_skip_gemm_dispatch`): not
209
+ # perf-critical, so skip packing and use a plain fallback.
210
+ if getattr(layer, "_cpu_skip_gemm_dispatch", False):
211
layer.cpu_linear = torch.nn.functional.linear
212
return
213
vllm/models/deepseek_v4/amd/model.py
+6
-11
@@ -71,10 +71,7 @@ from vllm.model_executor.models.utils import (
71
make_layers,
72
maybe_prefix,
73
)
74
-from vllm.models.deepseek_v4.amd.rocm import (
75
- DeepseekV4ROCMAiterMLAAttention,
76
- weight_already_preshuffled,
77
-)
74
+from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
75
from vllm.platforms import current_platform
76
from vllm.platforms.rocm import on_gfx950
77
from vllm.sequence import IntermediateTensors
@@ -152,13 +149,11 @@ class DeepseekV4MLP(nn.Module):
149
return
150
if ws.dtype == torch.float8_e8m0fnu:
151
ws = _upcast_e8m0_to_fp32(ws).contiguous()
155
- # Skip if the linear's kernel already shuffled it.
156
- if not weight_already_preshuffled(self.gate_up_proj):
157
- replace_parameter(
158
- self.gate_up_proj,
159
- "weight",
160
- rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
161
- )
152
+ replace_parameter(
153
+ self.gate_up_proj,
154
+ "weight",
155
+ rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
156
+ )
157
self._gateup_scale = ws
158
159
def forward(self, x):
vllm/models/deepseek_v4/amd/rocm.py
+16
-59
@@ -62,21 +62,6 @@ def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
62
return indptr
63
64
65
-def weight_already_preshuffled(linear: torch.nn.Module) -> bool:
66
- """True when the linear's kernel already B-preshuffled ``weight``.
67
-
68
- The hand-shuffles below (fused_wqa_wkv, wo_b, gate_up_proj) must be skipped
69
- for those, since shuffle_weight is a permutation rather than an involution.
70
- """
71
- return any(
72
- getattr(getattr(method, "fp8_linear", None), "preshuffles_weight", False)
73
- for method in (
74
- getattr(linear, "quant_method", None),
75
- getattr(linear, "scheme", None),
76
- )
77
- )
78
-
79
-
65
def apply_pre_quantized_block_scaled_mm(
66
linear: torch.nn.Module,
67
x_fp8: torch.Tensor,
@@ -792,13 +777,12 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
777
return None
778
if ws.dtype == torch.float8_e8m0fnu:
779
ws = _upcast_e8m0_to_fp32(ws).contiguous()
795
- # Skip if the linear's kernel already shuffled it.
796
- if not weight_already_preshuffled(linear):
797
- replace_parameter(
798
- linear,
799
- "weight",
800
- rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
801
- )
780
+ # Shuffle the weight in place (single weight, no unshuffled copy).
781
+ replace_parameter(
782
+ linear,
783
+ "weight",
784
+ rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
785
+ )
786
return ws
787
788
self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
@@ -896,26 +880,16 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
880
return qr_kv, kv_score, indexer_kv_score, indexer_weights
881
882
@functools.cached_property
899
- def _wq_b_act_scale_transpose(self) -> bool | None:
900
- """Activation-scale byte order every wq_b consumer agrees on, else None.
901
-
902
- None means "do not take the fused norm+quant path". Otherwise the value
903
- is the ``transpose_scale`` the producer must pass so that the scale it
904
- writes matches what the selected GEMM reads.
883
+ def _wq_b_uses_aiter_block_scaled(self) -> bool:
884
+ """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel.
885
886
Cached: the linear kernels and the aiter env gates are fixed once
887
the model is built, so this is evaluated at the first forward
888
only.
889
910
- Two conditions, both necessary:
911
-
912
- * The fused norm+quant path is only valid if the quant and GEMM it
913
- replaces are exactly the aiter ones; otherwise fall back to the
914
- shared path.
915
- * One qr_scale feeds both self.wq_b and self.indexer.wq_b, which have
916
- different (N, K) and so can resolve to kernels wanting opposite byte
917
- orders. A single producer cannot serve both, so refuse the fast path
918
- when they disagree rather than guessing.
890
+ The fused norm+quant path is only valid if the quant and GEMM it
891
+ replaces are exactly the aiter ones; otherwise fall back to the
892
+ shared path.
893
"""
894
from vllm._aiter_ops import rocm_aiter_ops
895
from vllm.model_executor.kernels.linear.scaled_mm import (
@@ -923,31 +897,16 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
897
)
898
899
if not rocm_aiter_ops.is_linear_fp8_enabled():
926
- return None
900
+ return False
901
902
linears = [self.wq_b]
903
if self.indexer is not None:
904
linears.append(self.indexer.wq_b)
931
- layouts: set[bool] = set()
905
for linear in linears:
906
kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None)
907
if not isinstance(kernel, Fp8BlockScaledMMLinearKernel):
935
- return None
936
- layouts.add(bool(getattr(kernel, "wants_transposed_act_scale", False)))
937
- if len(layouts) != 1:
938
- logger.warning_once(
939
- "DeepSeek-V4 wq_b consumers disagree on activation-scale layout; "
940
- "disabling the fused q/kv norm+quant path.",
941
- scope="global",
942
- )
943
- return None
944
- transpose_scale = layouts.pop()
945
- logger.debug_once(
946
- "DeepSeek-V4 wq_b: emitting %s-major activation scales",
947
- "column" if transpose_scale else "row",
948
- scope="global",
949
- )
950
- return transpose_scale
908
+ return False
909
+ return True
910
911
def _split_qkv_and_norm(
912
self, qr_kv: torch.Tensor
@@ -964,12 +923,11 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
923
when the aiter linear path is not active.
924
"""
925
qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
967
- transpose_scale = self._wq_b_act_scale_transpose
926
if not (
927
qr.dim() == 2
928
and qr.shape[0] > 0
929
and self.q_lora_rank % 128 == 0
972
- and transpose_scale is not None
930
+ and self._wq_b_uses_aiter_block_scaled
931
):
932
return super()._split_qkv_and_norm(qr_kv)
933
@@ -983,8 +941,7 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
941
kv_weight=self.kv_norm.weight.data,
942
kv_epsilon=self.eps,
943
group_size=128,
986
- # Emit the byte order the wq_b GEMMs read.
987
- transpose_scale=transpose_scale,
944
+ transpose_scale=False,
945
)
946
947
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: