[ahmed-bsod]

Overview

On gfx950, DeepSeek V4.1's MXFP8 linears with 32x32 weight-scale blocks run on vLLM's Triton block-scaled GEMM (rocm_mxfp8_block32_gemm). With AITER enabled, this PR runs them on AITER's FlyDSL MXFP8 GEMM instead. No other model or layer changes path.

Claims

• +3.7% total token throughput at ISL 350k / OSL 350 / conc 32 and +2.6% at ISL 32k / OSL 32 / conc 32
• Accuracy is unchanged: GSM8K 0.9735 vs 0.9727 (strict match, ±0.0045).

ISL 350k / OSL 350 / conc 32:

metric	baseline	this PR	change
total tok/s	44,011	45,641	+3.70%
TTFT mean	165.9 s	159.9 s	−3.63%
TPOT mean	210.9 ms	203.4 ms	−3.55%

ISL 32k / OSL 32 / conc 32:

metric	baseline	this PR	change
total tok/s	49,488	50,765	+2.58%
TTFT mean	14.14 s	13.70 s	−3.09%
TPOT mean	119.8 ms	119.0 ms	−0.68%

Validation

Unit Tests

pytest -q tests/kernels/quantization/test_rocm_mxfp8_linear.py \
  "tests/kernels/test_minimax_m3_amd_ops.py::test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load"

• test_rocm_mxfp8_linear_flydsl_matches_block32_gemm checks FlyDSL against the Triton block32 GEMM. It covers DeepSeek V4.1's seven 32x32 shapes (TP2/TP4) at M = 1, 7, 32, 33, 192 and 1024, with relative error < 1e-3.
• test_rocm_mxfp8_linear_flydsl_leaves_other_layers_alone checks that 1x32-scaled layers (every other MXFP8 model), K % 64 != 0 and batched layers keep their path and their weight bytes.

E2E
image used: vllm/vllm-openai-rocm:nightly-rocm100

serve command:

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_MOE=1
export VLLM_USE_BREAKABLE_CUDAGRAPH=1
export VLLM_USE_RUST_FRONTEND=1
export AITER_TRITON_LOG_LEVEL=ERROR
export VLLM_ENGINE_READY_TIMEOUT_S=3600
export OMP_NUM_THREADS=1
export VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1
export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4
vllm serve /data/models/DeepSeek-V4.1-Flash \
  --host 0.0.0.0 --port 8000 \
  --tensor-parallel-size 4 \
  --moe-backend aiter \
  --gpu-memory-utilization 0.80 \
  --language-model-only \
  --max-model-len 360448 \
  --max-num-seqs 128 \
  --max-num-batched-tokens 16384 \
  --no-async-scheduling \
  --no-enable-prefix-caching \
  --attention-config '{"indexer_kv_dtype":"mxfp4","indexer_sparse_logits":true}' \
  --speculative-config '{"method":"dspark","num_speculative_tokens":5,"draft_sample_method":"probabilistic","rejection_sample_method":"block","enable_adaptive_verification":false}'

bench command:

vllm bench serve \
    --model="$MODEL_PATH" \
    --backend=vllm \
    --tokenizer="$MODEL_PATH" \
    --dataset-name=random \
    --random-input-len=32768 \
    --random-output-len=32 \
    --random-range-ratio 0 \
    --num-prompts=64 \
    --max-concurrency=32 \
    --ignore-eos \
    --temperature=0 \
    --save-result --result-dir . \
    --metric-percentiles="50,90,99" \
    --percentile-metrics="ttft,tpot,itl,e2el"  2>&1 | tee -a "$LOG_FILE"

vllm bench serve \
    --model="$MODEL_PATH" \
    --backend=vllm \
    --tokenizer="$MODEL_PATH" \
    --dataset-name=random \
    --random-input-len=350000 \
    --random-output-len=350 \
    --random-range-ratio 0 \
    --num-prompts=128 \
    --max-concurrency=32 \
    --ignore-eos \
    --temperature=0 \
    --save-result --result-dir . \
    --metric-percentiles="50,90,99" \
    --percentile-metrics="ttft,tpot,itl,e2el"  2>&1 | tee -a "$LOG_FILE"

lm-eval command:

MODEL_PATH="${1:-/data/models/DeepSeek-V4.1-Flash}"
lm_eval --model local-chat-completions --apply_chat_template \
    --tasks /app/scripts/infx/gsm8k.yaml \
    --output_path eval_infx_out --log_samples \
    --model_args "model=${MODEL_PATH},base_url=http://0.0.0.0:8000/v1/chat/completions,api_key=EMPTY,eos_string=</s>,max_retries=5,num_concurrent=16,timeout=1800,tokenized_requests=False,max_length=32768" \
    --gen_kwargs "max_tokens=8192,temperature=0,top_p=1,thinking=false"

Details

• Only the GEMM call in RocmDotScaledMxfp8LinearKernel's existing 32x32 branch is changed, other branches as is, and there is no new env var.
• FlyDSL runs when: AITER linear ops on, gfx950, BF16 model, 32x32 scale blocks, K % 64 == 0.
• The weights are preshuffled once at load.
• Batched gemm for attention output projection wo_a stays on the Triton BMM path (the perf gain in e2e wasn't worth the code footprint).

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ahmed-bsod.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Fangzhou-Ai]

hi @ahmed-bsod thanks for the PR, can you also add TP1 support here (for DP attention)?

[Fangzhou-Ai]

Suggestion from running this PR's FlyDSL MXFP8 path in DeepSeek-V4.1-Flash DPATP2 serving (TP1 × DP2, MI355X, AgentX c128): warm the FlyDSL GEMM kernels at weight load.

aiter picks the FlyDSL kernel per padded M and JIT-compiles it on first use. On a cold FlyDSL cache, the server compiled several kernels inside its first long prefill step and exceeded the 300 s execute-model RPC timeout. We saw this as a startup death on first traffic.

The fix below runs the GEMM once per distinct kernel when the preshuffled weight is loaded, once per (N, K) per process. That covers every tuned M row for the shape up to max_num_batched_tokens, plus powers of two for the untuned fallback. With a cold cache, a sweep of serving-sized M compiles nothing afterwards. It adds 57 lines in rocm_block32_gemm.py / rocm_native.py. Feel free to fold it into this PR if useful.

patch

diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
index 6739372e47..22e9c88631 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_block32_gemm.py
@@ -887,6 +887,52 @@ def shuffle_bpreshuffle_weight(weight: torch.Tensor) -> torch.Tensor:
     return shuffle_weight(weight, layout=(16, 16))
 
 
+_warmed_bpreshuffle: set[tuple[int, int]] = set()
+
+
+def warm_up_bpreshuffle_gemm(
+    weight: torch.Tensor, weight_scale: torch.Tensor, max_tokens: int
+) -> None:
+    """Compile every FlyDSL kernel the GEMM can pick for this weight shape.
+
+    aiter picks the kernel per (padded) M, and each first use JIT-compiles,
+    so a cold first prefill would otherwise compile several kernels inside
+    one step. Runs once per (N, K) in the process.
+    """
+    from aiter.ops.gemm_op_a8w8 import (
+        _load_mxscale_bpreshuffle_tuned,
+        get_mxscale_bpreshuffle_config,
+    )
+
+    N, K = weight.shape
+    if (N, K) in _warmed_bpreshuffle:
+        return
+    _warmed_bpreshuffle.add((N, K))
+    # Tuned rows are keyed (gfx, cu_num, M, N, K, w_scale_block); untuned M
+    # fall back to a heuristic, sampled at powers of two.
+    tuned = {
+        key[2]
+        for key in _load_mxscale_bpreshuffle_tuned()
+        if key[3:] == (N, K, "32x32") and key[2] <= max_tokens
+    }
+    powers = {1 << i for i in range(max_tokens.bit_length())}
+    sizes = tuned | powers | {max_tokens}
+    seen = set()
+    for m in sorted(sizes):
+        row = get_mxscale_bpreshuffle_config(m, N, K, "32x32")
+        kernel = ("untuned", m) if row is None else row["kernelName"]
+        if kernel in seen:
+            continue
+        seen.add(kernel)
+        x = torch.zeros((m, K), dtype=torch.float8_e4m3fn, device=weight.device)
+        x_scale = torch.full(
+            (m, K // BLOCK_ROWS), 127, dtype=torch.uint8, device=weight.device
+        )
+        _rocm_mxfp8_block32_bpreshuffle_gemm_impl(
+            x, x_scale, weight, weight_scale, torch.bfloat16
+        )
+
+
 def _rocm_mxfp8_block32_bpreshuffle_gemm_impl(
     x: torch.Tensor,
     x_scale: torch.Tensor,
diff --git a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
index 7c751139c9..ac15e43d42 100644
--- a/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+++ b/vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
@@ -39,6 +39,7 @@ from .rocm_block32_gemm import (
     rocm_mxfp8_block32_gemm,
     shuffle_bpreshuffle_weight,
     use_bpreshuffle,
+    warm_up_bpreshuffle_gemm,
 )
 
 
@@ -301,6 +302,16 @@ class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
             weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
         replace_parameter(layer, "weight", weight.contiguous())
         replace_parameter(layer, "weight_scale", weight_scale)
+        if block_scale is not None and self._bpreshuffle(layer.weight):
+            from vllm.config import get_current_vllm_config_or_none
+
+            config = get_current_vllm_config_or_none()
+            if config is not None:
+                warm_up_bpreshuffle_gemm(
+                    layer.weight,
+                    layer.weight_scale,
+                    config.scheduler_config.max_num_batched_tokens,
+                )
 
     def input_quant_key(self) -> QuantKey:
         # Row-major [M, K // 32] E8M0 scales, unlike FlashInfer's swizzled ones

(AI-assisted; reviewed and tested by me.)

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Fangzhou-Ai]

need to have ROCm/aiter#7126 landed first