[ROCm][DSV4] B-preshuffle the attention fp8 projections (#46720)
Signed-off-by: Mehmet Cagri Kaymak <mehmet.kaymak@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Mehmet Cagri committed
Jul 30, 2026 at 13:59 UTC
12a34a6bc794135e3b93dd721a6bdcdb1bf734b9
5 files changed
+187
-9
vllm/_aiter_ops.py
+53
-2
@@ -687,6 +687,32 @@ def _rocm_aiter_gemm_a8w8_blockscale_fake(
687
return Y
688
689
690
+def _rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_impl(
691
+ A: torch.Tensor,
692
+ B: torch.Tensor,
693
+ As: torch.Tensor,
694
+ Bs: torch.Tensor,
695
+ output_dtype: torch.dtype = torch.float16,
696
+) -> torch.Tensor:
697
+ from aiter import gemm_a8w8_blockscale_bpreshuffle
698
+
699
+ # B preshuffled (shuffle_weight (16,16)); As column-major group scale.
700
+ return gemm_a8w8_blockscale_bpreshuffle(A, B, As, Bs, dtype=output_dtype)
701
+
702
+
703
+def _rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_fake(
704
+ A: torch.Tensor,
705
+ B: torch.Tensor,
706
+ As: torch.Tensor,
707
+ Bs: torch.Tensor,
708
+ output_dtype: torch.dtype = torch.float16,
709
+) -> torch.Tensor:
710
+ m = A.shape[0]
711
+ n = B.shape[0]
712
+ Y = torch.empty(m, n, dtype=output_dtype, device=A.device)
713
+ return Y
714
+
715
+
716
def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl(
717
x: torch.Tensor,
718
residual: torch.Tensor,
@@ -1145,17 +1171,21 @@ def _rocm_aiter_fused_rms_gated_fp8_group_quant_fake(
1171
def _rocm_aiter_group_fp8_quant_impl(
1172
x: torch.Tensor,
1173
group_size: int,
1174
+ transpose_scale: bool = False,
1175
) -> tuple[torch.Tensor, torch.Tensor]:
1176
assert x.shape[-1] % group_size == 0, "Input shape must be divisible by group size"
1177
from aiter import QuantType, get_hip_quant
1178
1179
aiter_per1x128_quant = get_hip_quant(QuantType.per_1x128)
1153
- return aiter_per1x128_quant(x.contiguous(), quant_dtype=FP8_DTYPE)
1180
+ return aiter_per1x128_quant(
1181
+ x.contiguous(), quant_dtype=FP8_DTYPE, transpose_scale=transpose_scale
1182
+ )
1183
1184
1185
def _rocm_aiter_group_fp8_quant_fake(
1186
x: torch.Tensor,
1187
group_size: int,
1188
+ transpose_scale: bool = False,
1189
) -> tuple[torch.Tensor, torch.Tensor]:
1190
M, N = x.shape
1191
x_fp8 = torch.empty((M, N), dtype=FP8_DTYPE, device=x.device)
@@ -1955,6 +1985,12 @@ class rocm_aiter_ops:
1985
fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake,
1986
)
1987
1988
+ direct_register_custom_op(
1989
+ op_name="rocm_aiter_gemm_a8w8_blockscale_bpreshuffle",
1990
+ op_func=_rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_impl,
1991
+ fake_impl=_rocm_aiter_gemm_a8w8_blockscale_bpreshuffle_fake,
1992
+ )
1993
+
1994
direct_register_custom_op(
1995
op_name="rocm_aiter_rmsnorm_fused_dynamic_quant",
1996
op_func=_rocm_aiter_rmsnorm_fused_dynamic_quant_impl,
@@ -2209,6 +2245,18 @@ class rocm_aiter_ops:
2245
A, B, As, Bs, output_dtype
2246
)
2247
2248
+ @staticmethod
2249
+ def gemm_a8w8_blockscale_bpreshuffle(
2250
+ A: torch.Tensor,
2251
+ B: torch.Tensor,
2252
+ As: torch.Tensor,
2253
+ Bs: torch.Tensor,
2254
+ output_dtype: torch.dtype = torch.float16,
2255
+ ) -> torch.Tensor:
2256
+ return torch.ops.vllm.rocm_aiter_gemm_a8w8_blockscale_bpreshuffle(
2257
+ A, B, As, Bs, output_dtype
2258
+ )
2259
+
2260
@staticmethod
2261
def fused_moe(
2262
hidden_states: torch.Tensor,
@@ -2702,9 +2750,12 @@ class rocm_aiter_ops:
2750
def group_fp8_quant(
2751
input_2d: torch.Tensor,
2752
group_size: int = 128,
2753
+ transpose_scale: bool = False,
2754
) -> tuple[torch.Tensor, torch.Tensor]:
2755
assert group_size == 128, "Group size must be 128"
2707
- return torch.ops.vllm.rocm_aiter_group_fp8_quant(input_2d, group_size)
2756
+ return torch.ops.vllm.rocm_aiter_group_fp8_quant(
2757
+ input_2d, group_size, transpose_scale
2758
+ )
2759
2760
@staticmethod
2761
def is_triton_gemm_w8a8_tuned(n: int, k: int) -> bool:
vllm/model_executor/model_loader/utils.py
+4
@@ -142,6 +142,10 @@ def process_weights_after_loading(
142
if isinstance(module, HpcModule):
143
module.process_weights_after_loading(model)
144
145
+ # Model-level post-load hook, after the per-layer quant finalize.
146
+ if hasattr(model, "process_weights_after_loading"):
147
+ model.process_weights_after_loading()
148
+
149
# Needed for torchao model reloading via model.reload_weights
150
# @kylesayrs @jerryzh168 this can be removed if callers move to `reload_weights`
151
if model_config.quantization == "torchao":
vllm/models/deepseek_v4/amd/model.py
+53
-3
@@ -9,6 +9,7 @@ import torch
9
import torch.nn as nn
10
11
import vllm.envs as envs
12
+from vllm._aiter_ops import rocm_aiter_ops
13
from vllm.config import VllmConfig, get_current_vllm_config
14
from vllm.distributed import (
15
get_pp_group,
@@ -104,8 +105,50 @@ class DeepseekV4MLP(nn.Module):
105
else:
106
self.act_fn = SiluAndMul()
107
108
+ # gate_up_proj B-preshuffle (ColumnParallel -> no all-reduce); set at load.
109
+ self._gateup = rocm_aiter_ops.is_enabled()
110
+ # Block scale for the preshuffled gate_up weight; None = not preshuffled.
111
+ self._gateup_scale: torch.Tensor | None = None
112
+
113
+ def prepare_gateup_preshuffle(self) -> None:
114
+ # B-preshuffle the gate_up_proj weight in place (single weight).
115
+ if not self._gateup:
116
+ return
117
+ from vllm.model_executor.utils import replace_parameter
118
+
119
+ w = getattr(self.gate_up_proj, "weight", None)
120
+ ws = getattr(self.gate_up_proj, "weight_scale_inv", None) # per-block scale
121
+ if w is None or ws is None or w.dim() != 2:
122
+ return
123
+ # K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold.
124
+ if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0:
125
+ return
126
+ if ws.dtype == torch.float8_e8m0fnu:
127
+ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
128
+ _upcast_e8m0_to_fp32,
129
+ )
130
+
131
+ ws = _upcast_e8m0_to_fp32(ws).contiguous()
132
+ replace_parameter(
133
+ self.gate_up_proj,
134
+ "weight",
135
+ rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
136
+ )
137
+ self._gateup_scale = ws
138
+
139
def forward(self, x):
108
- gate_up, _ = self.gate_up_proj(x)
140
+ if self._gateup_scale is not None and x.dim() == 2:
141
+ # gate_up via fp8 group-quant (col-major) + B-preshuffle GEMM.
142
+ x_fp8, x_scale = rocm_aiter_ops.group_fp8_quant(x, transpose_scale=True)
143
+ gate_up = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
144
+ x_fp8,
145
+ self.gate_up_proj.weight,
146
+ x_scale,
147
+ self._gateup_scale,
148
+ output_dtype=x.dtype,
149
+ )
150
+ else:
151
+ gate_up, _ = self.gate_up_proj(x)
152
x = self.act_fn(gate_up)
153
x, _ = self.down_proj(x)
154
return x
@@ -961,8 +1004,15 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
1004
1005
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1006
loader = AutoWeightsLoader(self, skip_substrs=["mtp."])
964
- loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
965
- return loaded_params
1007
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1008
+
1009
+ def process_weights_after_loading(self) -> None:
1010
+ # After per-layer quant finalize, so we preshuffle the final fp8 weights.
1011
+ for module in self.modules():
1012
+ if isinstance(module, DeepseekV4ROCMAiterMLAAttention):
1013
+ module.prepare_attn_preshuffle()
1014
+ elif isinstance(module, DeepseekV4MLP):
1015
+ module.prepare_gateup_preshuffle()
1016
1017
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1018
return self.model.get_expert_mapping()
vllm/models/deepseek_v4/amd/rocm.py
+71
-1
@@ -6,6 +6,10 @@ from typing import cast
6
7
import torch
8
9
+from vllm.distributed import (
10
+ get_tensor_model_parallel_world_size,
11
+ tensor_model_parallel_all_reduce,
12
+)
13
from vllm.forward_context import get_forward_context
14
from vllm.models.deepseek_v4.attention import DeepseekV4Attention
15
from vllm.models.deepseek_v4.common.ops import dequantize_and_gather_k_cache
@@ -442,10 +446,73 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
446
447
backend_cls = DeepseekV4ROCMAiterMLASparseBackend
448
449
+ def __init__(self, *args, **kwargs):
450
+ super().__init__(*args, **kwargs)
451
+ # Block scale for the preshuffled weight; None = not preshuffled.
452
+ self._wqa_wkv_scale: torch.Tensor | None = None
453
+ self._wo_b_scale: torch.Tensor | None = None
454
+
455
@classmethod
456
def get_padded_num_q_heads(cls, num_heads: int) -> int:
457
return num_heads
458
459
+ def prepare_attn_preshuffle(self) -> None:
460
+ from vllm._aiter_ops import rocm_aiter_ops
461
+
462
+ if not rocm_aiter_ops.is_enabled():
463
+ return
464
+ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
465
+ _upcast_e8m0_to_fp32,
466
+ )
467
+ from vllm.model_executor.utils import replace_parameter
468
+
469
+ def _prep(linear) -> torch.Tensor | None:
470
+ w = getattr(linear, "weight", None)
471
+ if w is None or w.dim() != 2:
472
+ return None
473
+ # K % 128 (group-128 quant) and N % 16 (shuffle_weight) must hold.
474
+ if w.shape[-1] % 128 != 0 or w.shape[0] % 16 != 0:
475
+ return None
476
+ ws = getattr(linear, "weight_scale_inv", None) # per-block scale
477
+ if ws is None:
478
+ return None
479
+ if ws.dtype == torch.float8_e8m0fnu:
480
+ ws = _upcast_e8m0_to_fp32(ws).contiguous()
481
+ # Shuffle the weight in place (single weight, no unshuffled copy).
482
+ replace_parameter(
483
+ linear,
484
+ "weight",
485
+ rocm_aiter_ops.shuffle_weight(w.data, layout=(16, 16)),
486
+ )
487
+ return ws
488
+
489
+ self._wqa_wkv_scale = _prep(self.fused_wqa_wkv)
490
+ self._wo_b_scale = _prep(self.wo_b)
491
+
492
+ def _bpre_attn_gemm(
493
+ self,
494
+ weight: torch.Tensor,
495
+ scale: torch.Tensor,
496
+ x: torch.Tensor,
497
+ reduce_tp: bool,
498
+ ) -> torch.Tensor:
499
+ from vllm._aiter_ops import rocm_aiter_ops
500
+
501
+ x_fp8, x_scale = rocm_aiter_ops.group_fp8_quant(x, transpose_scale=True)
502
+ out = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
503
+ x_fp8, weight, x_scale, scale, output_dtype=x.dtype
504
+ )
505
+ if reduce_tp and get_tensor_model_parallel_world_size() > 1:
506
+ out = tensor_model_parallel_all_reduce(out)
507
+ return out
508
+
509
+ def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
510
+ if self._wqa_wkv_scale is not None and hidden_states.dim() == 2:
511
+ return self._bpre_attn_gemm(
512
+ self.fused_wqa_wkv.weight, self._wqa_wkv_scale, hidden_states, False
513
+ )
514
+ return super()._fused_wqa_wkv_gemm(hidden_states)
515
+
516
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
517
# ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
518
z = rocm_inv_rope_einsum(
@@ -457,7 +524,10 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
524
self.o_lora_rank,
525
self.wo_a,
526
)
460
- return self.wo_b(z.flatten(1))
527
+ zf = z.flatten(1)
528
+ if self._wo_b_scale is not None and zf.dim() == 2:
529
+ return self._bpre_attn_gemm(self.wo_b.weight, self._wo_b_scale, zf, True)
530
+ return self.wo_b(zf)
531
532
def forward_mqa(
533
self,
vllm/models/deepseek_v4/attention.py
+6
-3
@@ -390,6 +390,11 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
390
# Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
391
return self._o_proj(o, positions)
392
393
+ def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
394
+ # MergedColumnParallelLinear returns (output, bias); bias is None.
395
+ qr_kv, _ = self.fused_wqa_wkv(hidden_states)
396
+ return qr_kv
397
+
398
def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]:
399
aux_streams = self.aux_stream_list
400
if aux_streams is not None:
@@ -434,9 +439,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
439
aux_fns[2] = indexer_compressor_kv_score
440
441
def fused_wqa_wkv() -> torch.Tensor:
437
- # MergedColumnParallelLinear returns (output, bias); bias is None.
438
- qr_kv, _ = self.fused_wqa_wkv(hidden_states)
439
- return qr_kv
442
+ return self._fused_wqa_wkv_gemm(hidden_states)
443
444
qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel(
445
fused_wqa_wkv,