[ROCm][Perf] Add bpreshuffled blockscaled fp8 GEMM (#51692)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Simon Danielsson committed
Sep 10, 2026 at 18:18 UTC
7470082f57a0245703be7482544f4b0fd410e1f2
3 files changed
+159
-5
vllm/_aiter_ops.py
+27
-5
@@ -175,18 +175,26 @@ def is_aiter_found_and_supported_on_rdna4() -> bool:
175
176
@functools.cache
177
def _load_gemm_tuned_configs(
178
- q_dtype_w: torch.dtype, csv_path: str
179
-) -> set[tuple[int, int, int]]:
178
+ csv_path: str,
179
+ filters: tuple[tuple[str, object], ...],
180
+ key_cols: tuple[str, ...] = ("N", "K", "M"),
181
+) -> set[tuple[int, ...]]:
182
try:
183
df = pd.read_csv(csv_path).drop_duplicates()
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
+ 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)))
192
except Exception:
193
return set()
194
195
196
def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool:
189
- configs = _load_gemm_tuned_configs(q_dtype_w, csv_path)
197
+ configs = _load_gemm_tuned_configs(csv_path, (("q_dtype_w", q_dtype_w),))
198
l_m = (
199
[1, 2, 4]
200
+ list(range(8, 513, 8))
@@ -3200,6 +3208,20 @@ class rocm_aiter_ops:
3208
csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE
3209
return _check_kernel_tuned(N, K, q_dtype_w, csv_path)
3210
3211
+ @staticmethod
3212
+ def is_blockscale_bpreshuffle_tuned(n: int, k: int) -> bool:
3213
+ """Whether (N, K) has a tuned aiter blockscale bpreshuffle config."""
3214
+ if not current_platform.is_rocm():
3215
+ return False
3216
+ import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops
3217
+
3218
+ csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BLOCKSCALE_BPRESHUFFLE_FILE
3219
+ gfx = aiter_gemm_a8w8_ops.get_gfx()
3220
+ cu_num = aiter_gemm_a8w8_ops.get_cu_num()
3221
+ return (n, k) in _load_gemm_tuned_configs(
3222
+ csv_path, (("gfx", gfx), ("cu_num", cu_num)), key_cols=("N", "K")
3223
+ )
3224
+
3225
@staticmethod
3226
def shuffle_weight(
3227
tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)
vllm/model_executor/kernels/linear/__init__.py
+4
@@ -178,6 +178,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
178
AiterHipbMMPerTokenFp8ScaledMMLinearKernel,
179
AiterInt8ScaledMMLinearKernel,
180
AiterPerTokenFp8ScaledMMLinearKernel,
181
+ AiterPreshuffledFp8BlockScaledMMKernel,
182
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
183
)
184
from vllm.model_executor.kernels.linear.scaled_mm.b12x import (
@@ -316,6 +317,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
317
"aiter": {
318
AiterInt8ScaledMMLinearKernel,
319
AiterFp8BlockScaledMMKernel,
320
+ AiterPreshuffledFp8BlockScaledMMKernel,
321
AiterPerTokenFp8ScaledMMLinearKernel,
322
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
323
AiterMxfp4LinearKernel,
@@ -456,6 +458,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
458
BlockWiseTorchFP8ScaledMMLinearKernel,
459
],
460
PlatformEnum.ROCM: [
461
+ AiterPreshuffledFp8BlockScaledMMKernel,
462
AiterFp8BlockScaledMMKernel,
463
TritonFp8BlockScaledMMKernel,
464
],
@@ -1210,6 +1213,7 @@ __all__ = [
1213
"ScaledMMLinearLayerConfig",
1214
"AiterHipbMMPerTokenFp8ScaledMMLinearKernel",
1215
"AiterPreshuffledPerTokenFp8ScaledMMLinearKernel",
1216
+ "AiterPreshuffledFp8BlockScaledMMKernel",
1217
"AiterPerTokenFp8ScaledMMLinearKernel",
1218
"NvFp4LinearKernel",
1219
"NvFp4LinearLayerConfig",
vllm/model_executor/kernels/linear/scaled_mm/aiter.py
+128
@@ -457,3 +457,131 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
457
return gemm_a8w8_blockscale_op(
458
A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype
459
)
460
+
461
+
462
+class AiterPreshuffledFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
463
+ """Aiter FP8 block-scaled GEMM using a pre-shuffled (bpreshuffle) weight."""
464
+
465
+ @classmethod
466
+ def is_supported(
467
+ cls, compute_capability: int | None = None
468
+ ) -> tuple[bool, str | None]:
469
+ return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported(
470
+ compute_capability
471
+ )
472
+
473
+ @classmethod
474
+ def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
475
+ can_implement_base, reason = super().can_implement(config)
476
+ if not can_implement_base:
477
+ return can_implement_base, reason
478
+
479
+ act_quant_desc = config.activation_quant_key.scale
480
+ if act_quant_desc.group_shape != GroupShape(1, 128):
481
+ return (
482
+ False,
483
+ (
484
+ "Supports only dynamic per token group activation "
485
+ "quantization with group_shape=(1,128)."
486
+ ),
487
+ )
488
+
489
+ # bpreshuffle GEMM requires aiter fp8 linear to be enabled and an fp8
490
+ # 2D weight with N and K divisible by 128. Require a tuned config as the
491
+ # fallback can trigger faults or numerical errors.
492
+ if not rocm_aiter_ops.is_linear_fp8_enabled():
493
+ return (
494
+ False,
495
+ (
496
+ "requires setting `VLLM_ROCM_USE_AITER=1` "
497
+ "and `VLLM_ROCM_USE_AITER_LINEAR=1`."
498
+ ),
499
+ )
500
+
501
+ n, k = config.weight_shape
502
+ if not (n % 128 == 0 and k % 128 == 0):
503
+ return (
504
+ False,
505
+ (
506
+ f"requires N and K dimensions divisible by 128, received "
507
+ f"N={n} and K={k}."
508
+ ),
509
+ )
510
+
511
+ if not rocm_aiter_ops.is_blockscale_bpreshuffle_tuned(n, k):
512
+ return (
513
+ False,
514
+ (
515
+ f"requires a tuned aiter blockscale bpreshuffle config for "
516
+ f"N={n} and K={k}."
517
+ ),
518
+ )
519
+
520
+ return True, None
521
+
522
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
523
+ super().process_weights_after_loading(layer)
524
+
525
+ params = FP8BlockParams.from_layer(layer)
526
+ if params.weight_scale_inv is not None:
527
+ ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV
528
+ else:
529
+ ws, attr = params.weight_scale, params.WEIGHT_SCALE
530
+ if ws is not None and ws.dtype == torch.float8_e8m0fnu:
531
+ replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous())
532
+
533
+ weight = params.weight
534
+ # runtime safety net
535
+ assert (
536
+ weight.dim() == 2
537
+ and weight.dtype == current_platform.fp8_dtype()
538
+ and weight.shape[0] % 128 == 0
539
+ and weight.shape[1] % 128 == 0
540
+ ), (
541
+ "AiterPreshuffledFp8BlockScaledMMKernel requires a 2D fp8 weight "
542
+ "with N and K divisible by 128."
543
+ )
544
+
545
+ shuffled_weight = rocm_aiter_ops.shuffle_weight(
546
+ weight.contiguous(), layout=(16, 16)
547
+ )
548
+ replace_parameter(
549
+ layer,
550
+ params.WEIGHT,
551
+ torch.nn.Parameter(shuffled_weight.data, requires_grad=False),
552
+ )
553
+
554
+ def apply_weights(
555
+ self,
556
+ layer: torch.nn.Module,
557
+ x: torch.Tensor,
558
+ bias: torch.Tensor | None = None,
559
+ **kwargs,
560
+ ) -> torch.Tensor:
561
+ params = self._get_layer_params(layer)
562
+ Bs = (
563
+ params.weight_scale
564
+ if params.weight_scale_inv is None
565
+ else params.weight_scale_inv
566
+ )
567
+
568
+ x_2d = x.view(-1, x.shape[-1])
569
+ A, As = rocm_aiter_ops.group_fp8_quant(x_2d, transpose_scale=True)
570
+ output = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle(
571
+ A, params.weight, As, Bs, output_dtype=self.config.out_dtype
572
+ )
573
+ if bias is not None:
574
+ output = output + bias
575
+ return output.view(*x.shape[:-1], params.weight.shape[0])
576
+
577
+ def apply_block_scaled_mm(
578
+ self,
579
+ A: torch.Tensor,
580
+ B: torch.Tensor,
581
+ As: torch.Tensor,
582
+ Bs: torch.Tensor,
583
+ ) -> torch.Tensor:
584
+ raise NotImplementedError(
585
+ "AiterPreshuffledFp8BlockScaledMMKernel overrides apply_weights and "
586
+ "does not use apply_block_scaled_mm."
587
+ )