[ROCm][DSv4.1] Warm up AITER FlyDSL MXFP8 GEMM kernels before serving

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)

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
                 )

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)

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,
+    )