diff --git a/tests/kernels/quantization/test_rocm_mxfp8_linear.py b/tests/kernels/quantization/test_rocm_mxfp8_linear.py index b784551928a6..5411caf55fa1 100644 --- a/tests/kernels/quantization/test_rocm_mxfp8_linear.py +++ b/tests/kernels/quantization/test_rocm_mxfp8_linear.py @@ -17,6 +17,8 @@ QDQ fp32 reference, and bit-identical fallbacks for ineligible inputs. """ +import copy + import pytest import torch @@ -27,6 +29,7 @@ mxfp8_e4m3_quantize, ) from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_default_torch_dtype pytestmark = pytest.mark.skipif( not current_platform.is_rocm() or not current_platform.supports_mx(), @@ -606,3 +609,165 @@ def test_rocm_mxfp8_fused_quant_cuda_graph_capture(warm_counters): graph.replay() torch.accelerator.synchronize() assert torch.equal(out, expected), "replay differs from stock" + + +# With AITER enabled, 32x32 block-scaled layers (DeepSeek V4.1) run AITER's +# FlyDSL GEMM on a preshuffled weight; every other layer keeps the paths above. + + +@pytest.fixture +def flydsl(monkeypatch: pytest.MonkeyPatch): + """Enables AITER; skips the test unless AITER's FlyDSL GEMM can run.""" + from vllm._aiter_ops import rocm_aiter_ops + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + _flydsl_enabled, + ) + + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + rocm_aiter_ops.refresh_env_variables() + try: + if not _flydsl_enabled(): + pytest.skip("requires AITER on gfx950") + yield + finally: + monkeypatch.undo() + rocm_aiter_ops.refresh_env_variables() + + +def _kernel(n: int, k: int, dtype: torch.dtype = torch.bfloat16): + """The kernel as a model of this dtype builds it.""" + from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( + Mxfp8LinearLayerConfig, + ) + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + RocmDotScaledMxfp8LinearKernel, + ) + + with set_default_torch_dtype(dtype): + return RocmDotScaledMxfp8LinearKernel( + Mxfp8LinearLayerConfig(weight_shape=(n, k)) + ) + + +@pytest.mark.parametrize("name,n,k", V41_BLOCK32_SHAPES, ids=lambda v: str(v)) +@pytest.mark.parametrize("num_tokens", [1, 7, 32, 33, 192, 1024]) +@torch.inference_mode() +def test_rocm_mxfp8_linear_flydsl_matches_block32_gemm(flydsl, name, n, k, num_tokens): + """The FlyDSL GEMM agrees with the block-scaled Triton GEMM it replaces.""" + kernel = _kernel(n, k) + ref = _kernel(n, k, torch.float32) # not a bf16 model: keeps the Triton GEMM + torch.manual_seed(num_tokens) + layer = _make_block32_layer(n, k, "cuda") + ref_layer = copy.deepcopy(layer) + kernel.process_weights_after_loading(layer) + ref.process_weights_after_loading(ref_layer) + assert kernel._use_flydsl(layer) and not ref._use_flydsl(ref_layer) + + x = torch.randn(num_tokens, k, device="cuda", dtype=torch.bfloat16) * 0.5 + out = kernel.apply_weights(layer, x) + assert out.shape == (num_tokens, n) and out.dtype == torch.bfloat16 + rel = _rel_err(out, ref.apply_weights(ref_layer, x).float()) + assert rel < 1e-3, f"{name}: relative error {rel:.2e}" + + +@torch.inference_mode() +def test_rocm_mxfp8_linear_flydsl_leaves_other_layers_alone(flydsl): + """Only 32x32 block-scaled layers get a preshuffled weight: 1x32 layers + (every other MXFP8 model), K % 64 != 0 and batched layers (DeepSeek V4.1's + wo_a, which the model runs itself) are untouched.""" + torch.manual_seed(0) + batched = _make_block32_layer(256, 512, "cuda") + batched.bmm_batch_size = 2 + layers = { + "per_row": _make_layer(256, 512, "cuda"), + "k_tail": _make_block32_layer(256, 288, "cuda"), + "batched": batched, + } + for name, layer in layers.items(): + kernel = _kernel(*layer.weight.shape) + loaded = layer.weight.data.view(torch.uint8).clone() + kernel.process_weights_after_loading(layer) + assert not kernel._use_flydsl(layer), name + assert torch.equal(layer.weight.view(torch.uint8), loaded), name + + +_WARMUP_SWEEP = """ +import glob, os, sys +from types import SimpleNamespace +import torch +from vllm._aiter_ops import rocm_aiter_ops +from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( + Mxfp8LinearLayerConfig, +) +from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + RocmDotScaledMxfp8LinearKernel, + _flydsl_enabled, +) +from vllm.model_executor.warmup.rocm_flydsl_mxfp8_warmup import ( + rocm_flydsl_mxfp8_warmup, +) +{make_layer} +rocm_aiter_ops.refresh_env_variables() +if not _flydsl_enabled(): + print("SKIP"); sys.exit(0) +n, k, max_tokens = {n}, {k}, {max_tokens} +torch.set_default_dtype(torch.bfloat16) +kernel = RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig(weight_shape=(n, k))) +model = torch.nn.Module() +model.proj = _make_block32_layer(n, k, "cuda") +model.proj.quant_method = SimpleNamespace(kernel=kernel) +kernel.process_weights_after_loading(model.proj) + + +def compiled(): + return len(glob.glob(os.environ["FLYDSL_RUNTIME_CACHE_DIR"] + "/**/*.pkl", + recursive=True)) + + +if {warm}: + rocm_flydsl_mxfp8_warmup(model, max_tokens) +warmed = compiled() +with torch.inference_mode(): + for m in range(1, max_tokens + 1): + kernel.apply_weights(model.proj, torch.randn(m, k, device="cuda")) +torch.accelerator.synchronize() +print("COMPILED", warmed, compiled()) +""" + + +@pytest.mark.parametrize("warm", [True, False], ids=["warmup", "no_warmup"]) +def test_rocm_flydsl_mxfp8_warmup_leaves_nothing_to_compile(tmp_path, warm): + """After the warmup, serving every token count up to max_num_batched_tokens + compiles no FlyDSL kernel; without it, serving compiles them (the cold-start + stall). Runs in a subprocess: FlyDSL fixes its cache directory at import.""" + import inspect + import os + import subprocess + import sys + + script = _WARMUP_SWEEP.format( + make_layer=inspect.getsource(_make_block32_layer), + n=5120, # DeepSeek V4.1 shared expert down at TP1 + k=2304, + max_tokens=512, + warm=warm, + ) + env = dict( + os.environ, VLLM_ROCM_USE_AITER="1", FLYDSL_RUNTIME_CACHE_DIR=str(tmp_path) + ) + proc = subprocess.run( + [sys.executable, "-c", script], + env=env, + capture_output=True, + text=True, + timeout=1800, + check=False, + ) + assert proc.returncode == 0, proc.stderr[-3000:] + if "SKIP" in proc.stdout: + pytest.skip("requires AITER on gfx950") + warmed, served = map(int, proc.stdout.split("COMPILED")[1].split()) + if warm: + assert warmed > 0 and served == warmed, (warmed, served) + else: + assert served > warmed == 0, (warmed, served) diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py index 878a9ee1cd0e..61bf036169c7 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py +++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py @@ -9,11 +9,16 @@ tiles K by 128; a weight whose K is not a multiple of that is dequantized to BF16 once at load and served by a plain linear instead, since ``can_implement`` does not filter on K. + +Weights with 32x32 scale blocks (DeepSeek V4.1) run a block-scaled Triton GEMM, +or AITER's FlyDSL GEMM on gfx950 when AITER is enabled. FlyDSL reads the weight +(16, 16)-preshuffled, so those weights are shuffled once at load. """ import torch import vllm.envs as envs +from vllm._aiter_ops import rocm_aiter_ops from vllm.model_executor.layers.fusion.quant_activation import ( QuantizedActivation, as_quantized_activation, @@ -360,11 +365,27 @@ def _as_block32_scale(weight_scale: torch.Tensor) -> torch.Tensor | None: return blocks[:, 0].contiguous() +_FLYDSL_K_ALIGN = 64 # AITER's FlyDSL GEMM needs K % 64 == 0 + + +def _flydsl_enabled() -> bool: + """AITER is enabled for linears; its FlyDSL MXFP8 GEMM is gfx950-only.""" + from vllm.platforms.rocm import on_gfx950 + + return bool(rocm_aiter_ops.is_linear_enabled()) and on_gfx950() + + class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel): """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``.""" supports_pre_processed_weights = True + def __init__(self, c: Mxfp8LinearLayerConfig) -> None: + super().__init__(c) + self.flydsl_enabled = ( + torch.get_default_dtype() == torch.bfloat16 and _flydsl_enabled() + ) + @classmethod def is_supported( cls, compute_capability: int | None = None @@ -382,6 +403,18 @@ def is_supported( def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: return True, None + def _use_flydsl(self, layer: torch.nn.Module) -> bool: + """Whether a loaded layer runs AITER's FlyDSL GEMM: 32x32 scale blocks + and K % 64 == 0.""" + return ( + self.flydsl_enabled + and layer.weight_scale.shape[0] != layer.weight.shape[0] + and layer.weight.shape[1] % _FLYDSL_K_ALIGN == 0 + # A batched layer (DeepSeek V4.1's wo_a) runs on rocm_mxfp8_wo_a_bmm, + # which reads the loaded weight. + and getattr(layer, "bmm_batch_size", None) is None + ) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: weight = layer.weight.data # [N, K] fp8 N, K = weight.shape @@ -400,6 +433,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale) replace_parameter(layer, "weight", weight.contiguous()) replace_parameter(layer, "weight_scale", weight_scale) + if self._use_flydsl(layer): + replace_parameter( + layer, + "weight", + rocm_aiter_ops.shuffle_weight(layer.weight.data, layout=(16, 16)), + ) def input_quant_key(self) -> QuantKey: # Row-major [M, K // 32] E8M0 scales, unlike FlashInfer's swizzled ones @@ -441,7 +480,19 @@ def apply_weights( out_shape = (*in_shape[:-1], layer.weight.shape[0]) if layer.weight_scale.shape[0] != layer.weight.shape[0]: # One scale row per 32 weight rows (see process_weights_after_loading). - if qa is not None: + if self._use_flydsl(layer): + if qa is None: + x_q, x_scale = mxfp8_e4m3_quantize(x2d) + # MXFP8 scales are stored as uint8; AITER takes its FlyDSL path + # only for E8M0-typed scales. + out = rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( + x_q, + layer.weight, + x_scale.view(torch.float8_e8m0fnu), + layer.weight_scale.view(torch.float8_e8m0fnu), + output_dtype=out_dtype, + ) + elif qa is not None: out = rocm_mxfp8_block32_gemm( x_q, x_scale, layer.weight, layer.weight_scale, out_dtype ) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 0e706a390541..d8588020c183 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -230,6 +230,14 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): ) rocm_segmented_attn_autotune_warmup(worker) + if worker.vllm_config.kernel_config.enable_jit_warmup: + from vllm.model_executor.warmup.rocm_flydsl_mxfp8_warmup import ( + rocm_flydsl_mxfp8_warmup, + ) + + rocm_flydsl_mxfp8_warmup( + worker.get_model(), worker.scheduler_config.max_num_batched_tokens + ) flashinfer_sparse_mla_decode_autotune_warmup(worker) deepseek_v4_sparse_mla_attention_warmup(worker) diff --git a/vllm/model_executor/warmup/rocm_flydsl_mxfp8_warmup.py b/vllm/model_executor/warmup/rocm_flydsl_mxfp8_warmup.py new file mode 100644 index 000000000000..58ecc9c9e8f6 --- /dev/null +++ b/vllm/model_executor/warmup/rocm_flydsl_mxfp8_warmup.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up AITER's FlyDSL GEMM for 32x32-block MXFP8 linears on gfx950. + +AITER picks a FlyDSL kernel per (padded) token count and JIT-compiles it on +first use. Without a warmup, a cold server compiles several of them inside its +first long prefill step, which can exceed the execute-model RPC timeout. +""" + +from __future__ import annotations + +import time + +import torch + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def _flydsl_layers(model: torch.nn.Module): + from vllm.model_executor.kernels.linear.mxfp8.rocm_native import ( + RocmDotScaledMxfp8LinearKernel, + ) + + for module in model.modules(): + # Quark keeps the kernel on the layer's scheme, ModelOpt on the method. + kernel = getattr(getattr(module, "scheme", None), "ocp_mx_linear", None) + if kernel is None: + kernel = getattr(getattr(module, "quant_method", None), "kernel", None) + if isinstance(kernel, RocmDotScaledMxfp8LinearKernel) and kernel._use_flydsl( + module + ): + yield module + + +def rocm_flydsl_mxfp8_warmup(model: torch.nn.Module, max_tokens: int) -> None: + """Run each FlyDSL MXFP8 linear shape once per distinct kernel AITER can + pick for 1..max_tokens tokens, so serving compiles none of them.""" + from aiter.ops.gemm_op_a8w8 import get_mxscale_bpreshuffle_config + + from vllm._aiter_ops import rocm_aiter_ops + + layers = {tuple(m.weight.shape): m for m in _flydsl_layers(model)} + if not layers: + return + start, num_kernels = time.perf_counter(), 0 + device = next(iter(layers.values())).weight.device + max_k = max(k for _, k in layers) + x_buf = torch.zeros(max_tokens * max_k, dtype=torch.float8_e4m3fn, device=device) + scale_buf = torch.full( + (max_tokens * max_k // 32,), 127, dtype=torch.uint8, device=device + ).view(torch.float8_e8m0fnu) + for (n, k), layer in layers.items(): + x_q = x_buf[: max_tokens * k].view(max_tokens, k) + x_scale = scale_buf[: max_tokens * k // 32].view(max_tokens, k // 32) + kernels = set() + for m in range(1, max_tokens + 1): + row = get_mxscale_bpreshuffle_config(m, n, k, "32x32") + if row is not None: + name = row["kernelName"] + elif m & (m - 1) == 0 or m == max_tokens: + # Untuned: AITER's fallback heuristic, sampled at powers of two. + name = ("untuned", m) + else: + continue + if name in kernels: + continue + kernels.add(name) + rocm_aiter_ops.gemm_a8w8_blockscale_bpreshuffle( + x_q[:m], + layer.weight, + x_scale[:m], + layer.weight_scale.view(torch.float8_e8m0fnu), + output_dtype=torch.bfloat16, + ) + num_kernels += len(kernels) + torch.accelerator.synchronize() + logger.info( + "Warmed up %d FlyDSL MXFP8 GEMM kernels for %d shapes in %.1fs.", + num_kernels, + len(layers), + time.perf_counter() - start, + )