Repository navigation
Conversation
Add VLLM_ROCM_USE_AITER_FLYDSL_FP8_PREFILL (default off). When set, the ROCM_AITER_FA MLA prefill backend runs AITER's FlyDSL FP8 attention on dynamically per-tensor quantized Q/K/V for new tokens and context chunks, reusing the new-token FP8 Q across chunks. Support is decided by AITER's flydsl_flash_attn_fp8_supported; unsupported devices, head dims, dtypes or AITER builds fall back to BF16 attention. When active, context-free prefills also use it instead of the ASM FP8 prefill. Adds a Triton kernel that quantizes up to three strided tensors with independent per-tensor scales in two launches. Signed-off-by: Guanbao Yu <Guanbao.Yu@amd.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
AITER's AOT build of the FlyDSL FP8 FMHA (ROCm/aiter#5796, vllm-project#5901) covers the varlen_cross layout for Kimi-K3 TP8. With cross_seqlen=False the new-token call used the varlen layout, which is not pre-compiled, so it still JIT-compiled on live requests. For equal q/kv lengths the two layouts are bit-identical. Signed-off-by: Guanbao Yu <Guanbao.Yu@amd.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
simondanielsson
left a comment
There was a problem hiding this comment.
Nice work!
Suggestions:
- Can we create a standalone mla prefill backend for this? there's a lot of useful existing infra for this
- Can we plug that new backend into the existing attention correctness tests?
- Optional: if we'd like to make this default for fp8 kv, It'd be great if we could have a wider benchmark sweep to see in what domains this outperforms the otherwise default backend
| ) | ||
|
|
||
|
|
||
| def fused_per_tensor_fp8_quant( |
There was a problem hiding this comment.
Question: what's the rationale behind implementing this here rather than re-using an existing quant kernel?
There was a problem hiding this comment.
I compared against the existing kernels first. On MI355X, with 12 heads, q/k 192 and v 128, and V a strided view of the kv_b_proj output:
rocm_aiter_ops.per_tensor_quantreturns wrong results for the strided V.ops.scaled_fp8_quantmatches this kernel: descales within 1e-6, values within 1 e4m3 ulp. It is much slower here, though:
| Tokens | This kernel (Q/K/V, one call) | scaled_fp8_quant ×3 |
×3, V made contiguous first |
|---|---|---|---|
| 1,024 | 50 µs | 189 µs | 59 µs |
| 8,192 | 50 µs | 1,318 µs | 237 µs |
| 65,536 | 352 µs | 10,165 µs | 1,306 µs |
Why it is slower:
- It takes one 2D tensor per call, so Q/K/V need 3 × (memset + amax + quant).
- V's row stride sends it down the per-row path: one 256-thread block and one
atomicMaxper 128-element row, and there are T×H rows. At 65K tokens, V alone takes 9.3 ms there versus 110 µs here. - An all-zero input gets descale 0.
This kernel is adapted from ATOM's fused QKV quantizer. The one change is that token counts are runtime args, so varying batch shapes don't recompile.
I'd be happy to move it somewhere more general, or to fix the strided path of dynamic_scaled_fp8_quant in a follow-up, if you prefer.
| q_descale=q_descale, | ||
| k_descale=k_descale, | ||
| v_descale=v_descale, | ||
| out=out, |
There was a problem hiding this comment.
Question: does this work even though this backend does not supports_out?
There was a problem hiding this comment.
Yes. forward_mha passes out only together with output_scale, i.e. fused quantized output. That requires supports_quant_output(), which this backend doesn't claim (neither does ROCM_AITER_FA), so out is always None here and FlyDSL allocates the output itself.
If a caller did pass out, FlyDSL would write into it: it validates a contiguous BF16 [T, H, Dv] buffer.
I kept supports_out() False because nothing on the ROCm MLA path consumes it yet; only the Kimi-K3 NVIDIA module does. The context-chunk path asserts out is None.
Saving the extra output.copy_ for context-free prefills could be a follow-up with numbers.
Address review: move the FlyDSL FP8 path out of ROCM_AITER_FA into a new
ROCM_AITER_FLYDSL_FP8 prefill backend, selected with
--attention-config '{"mla_prefill_backend":"ROCM_AITER_FLYDSL_FP8"}', and
drop VLLM_ROCM_USE_AITER_FLYDSL_FP8_PREFILL. Support is checked by the
standard hooks (bf16 only, gfx950, AITER FlyDSL FP8 API, and AITER's
flydsl_flash_attn_fp8_supported for the head dims), so an unsupported
explicit selection fails at startup.
Run the backend in test_mla_backends.py (FP8 tolerance for this backend)
and add a context-free prefill case covering the ASM FP8 prefill bypass;
add registry and selector tests.
Signed-off-by: Guanbao Yu <Guanbao.Yu@amd.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
hi @simondanielsson , thanks for your review! Here are the updates regarding your concerns:
|
|
Three technical suggestions:
Keeping this opt-in makes sense: short prefills show little benefit, while FP8 accuracy needs a wider long-context evaluation. |
Programs that land past the end of a smaller input in the shared quantization grid now compute no address at all: `_partial_amax` derives its trip count per program and `_quant` guards its body, keeping program 0 alive so empty inputs still get a descale. Add a boundary-coverage test across the block and stripe boundaries. Replace the indirect "PS metadata is absent" check in `forward_mha` with `_asm_fp8_prefill_superseded`, shared with the metadata builder, so the ASM FP8 prefill cannot claim prefills owned by the FlyDSL FP8 backend. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: ganyi <ygan@amd.com>
|
@ChuanLi1101 Thanks for the careful read — points 1 and 3 are fixed in 19a77e3, and I have numbers for point 2. 1. Out-of-range pointers in the shared quant grid You're right that the grid is sized for the largest input, so programs can land past the end of a smaller one. Every load and store was already masked and the offset math is int64, so nothing was dereferenced or wrapped — but those programs still computed addresses they had no business computing. Both kernels now stop before that:
I added 2. Per-call temporary buffers I measured this on MI355X before changing it, since the buffers are tiny (
Where the 46.6 µs goes: 28 µs is the two Triton launches (an empty kernel with the same signature already costs 15.2 µs), 10.2 µs is Reusing Given ~2 µs for the only unambiguously safe part ( 3. Disabling the ASM FP8 prefill directly Agreed, keying off absent PS metadata was too indirect. There's now a single predicate in def _asm_fp8_prefill_superseded(prefill_backend: object) -> bool:
return isinstance(prefill_backend, AiterFlyDSLFP8PrefillBackend)used in both places so they can't drift:
4. Opt-in Agreed — it stays out of the ROCm auto-selection priorities and is only reachable via |
|
This pull request has merge conflicts that must be resolved before it can be |
Purpose
This PR adds an opt-in FlyDSL FP8 path for MLA prefill attention on ROCm, following the approach of ROCm/ATOM#2191. The target is Kimi-K3 on MI355X (gfx950).
It adds a new, opt-in MLA prefill backend,
ROCM_AITER_FLYDSL_FP8, which subclassesROCM_AITER_FA. Select it with--attention-config '{"mla_prefill_backend":"ROCM_AITER_FLYDSL_FP8"}'; it is not in the ROCm auto-selection priorities. The backend:flydsl_flash_attn_fp8_funcfor the new tokens;return_lse=True. The LSE layout ([H, T], natural log) is unchanged, somerge_attn_statesworks as-is. The layer's softmax scale is passed directly.mla_prefill_ps_asm_fwd). The metadata builder stops building ASM PS metadata when this backend is selected, andAiterMLAImpl.forward_mhathen falls back to the common path.Support is checked through the standard prefill-backend hooks. vLLM does not hard-code any shape or model:
supported_dtypes = [bf16], because the kernel writes BF16 output;supports_compute_capabilityrequires gfx950;is_availablerequires AITER and the FlyDSL FP8 API;validate_configurationdefers to AITER'sflydsl_flash_attn_fp8_supportedfor the head dims.An explicit selection that fails these checks raises at startup. At runtime, inputs the kernel cannot take (empty, or ≥2³¹ elements) use the parent's BF16 path.
vllm/v1/attention/ops/triton_per_tensor_fp8_quant.pyis a small Triton kernel adapted from ATOM#2191's fused QKV quantizer (MIT). It quantizes up to three strided[tokens, heads, dim]tensors in two launches, each with its own amax/descale. V is quantized directly from its strided view of thekv_b_projoutput. Compile-time specialization does not depend on token counts (only the grid size does), so varying batch shapes do not recompile, and no atomics are used. Unlike ATOM's KV quantizer, element counts are runtime arguments.Why not reuse
ops.scaled_fp8_quant. vLLM's dynamic per-tensor quantizer (dynamic_scaled_fp8_quant, also behindQuantFP8) gives the same result: descales within 1e-6 relative, values within 1 e4m3 ulp. It is too slow for this path, for three reasons:atomicMaxper 128-element row, i.e. T×H rows.rocm_aiter_ops.per_tensor_quantreturns wrong results for the strided V. GPU time on MI355X, 12 heads, q/k 192 and v 128 (v a strided view):scaled_fp8_quant×3scaled_fp8_quant×3, V made contiguous firstAt 65,536 tokens, V alone takes 9,284 µs with
scaled_fp8_quantand 110 µs here. The existing op is faster only for a single small contiguous tensor (Q alone at 1,024 tokens: 21 µs vs 44 µs), which this path never quantizes.AITER requirement. This needs an AITER build containing:
flydsl_flash_attn_fp8_supported)softmax_scaleinflydsl_flash_attn_fp8_func)All three are already on AITER main. #5718 is the newest of them, so a build that has the capability check also has the other two. #5508 only affects performance. The current
vllm/vllm-openai-rocmnightly (AITER 0.1.22.post1) predates them, so there this backend reports itself unavailable.AOT build. ROCm/aiter#5796 and #5901 pre-compile the FlyDSL FP8 FMHA for Kimi-K3 TP8 (
12:12:192:128,varlen_crosslayout, 92 kernels) during AITER source builds, so serving no longer JIT-compiles them. The new-token call therefore passescross_seqlen=True, as ATOM does.cross_seqlen=False, the new-token call would use thevarlenlayout, which is not pre-compiled.cross_seqlen=Trueis within ±2% for most shapes and 7% faster for 1×8192. Batches of several equal-length sequences are slower: 4×2048 +15 µs, 8×1024 +4 µs.Coverage check: the 92 AOT kernels (built with
aiter.aot.flydsl.fmha_fp8from AITER main) were put into an emptyFLYDSL_RUNTIME_CACHE_DIR, andFLYDSL_RUNTIME_RUN_ONLY=1was set, which makes any cache miss raise.cross_seqlen=False, all 20 new-token calls missed.Not a duplicate. I searched open PRs for FlyDSL/FP8 MLA prefill:
_fp8_prefill_enabledgate inAiterMLAMetadataBuilder.__init__, so whichever lands second needs a small rebase.AI assistance (Claude Code) was used to implement this change and to run the tests and benchmarks below.
Test Plan
Unit tests (MI355X, gfx950):
VLLM_ROCM_USE_AITER=1 python -m pytest tests/v1/attention/test_mla_backends.py -k ROCM_AITER python -m pytest -v tests/kernels/attention/test_rocm_aiter_mla_flydsl_fp8_prefill.py \ tests/v1/attention/test_mla_prefill_selector.py \ tests/v1/attention/test_mla_prefill_registry.py pre-commit run --files <changed files> python tools/pre_commit/mypy.py 3.12 <changed vllm files>ROCM_AITER_FLYDSL_FP8is added toPREFILL_BACKENDS_TO_TESTintest_mla_backends.py, sotest_backend_correctnessandtest_chunked_context_backend_correctnesscheck it against the SDPA reference.validate_configuration.rtol=2**-4, atol=0.3. With per-tensor e4m3 Q/K/V, a row that attends to a few tokens is off by up to about 2⁻⁴·max|V|. The largest observed diff was about 0.25.context_free_prefillspec, run for this backend, covers the builder change that keepsROCM_AITER_MLA's ASM FP8 prefill from bypassing the selected prefill backend.The registry and selector tests cover registration, bf16-only, gfx950 gating, availability, and deferring to AITER's check. The backend test file covers:
validate_configurationrejecting a head dim that AITER does not support.End-to-end. The server config is the InferenceX Kimi-K3 MI355X job (
kimik3_fp4_mi355x_mtp.sh, conc 48, no spec decode). A/B differs only in the MLA prefill backend, and both runs use the same AITER build. The end-to-end numbers were measured with the first revision of this PR, which enabled the same code path through an env flag insideROCM_AITER_FA. Apart fromcross_seqlen=True(bit-identical outputs), the kernel calls are unchanged.gsm8k, 5-shot, full 1319, vialocal-completions,num_concurrent=64, temperature 0.kimi_tool_call_schema_full(Kimi Vendor Verifier, 408 records, the CI eval for this job), run twice per arm.aiperf profile --scenario inferencex-agentx-mvp --concurrency 48 --benchmark-duration 3600 --warmup-grace-period 1800 --public-dataset semianalysis_cc_traces_weka_062126 ..., the same command and aiperf commit as CI. One run per arm.vllm bench serve --dataset-name random --random-output-len 1 --ignore-eosat 8K, 32K, 128K and 64K-with-56K-shared-prefix inputs. A warmup pass runs first, and each arm is run twice.Test Result
Unit tests (MI355X,
VLLM_ROCM_USE_AITER=1):test_mla_backends.py -k ROCM_AITER: 869 passed, 296 skipped. ThreeROCM_AITER_FAcases (fp8 KV, scales 2/3, TP4) failed with CUDA OOM because another job had filled the shared GPU. All three pass when rerun individually, both on this branch and onmain.ROCM_AITER_FLYDSL_FP8cases pass; the GLM-dim cases are skipped.pre-commitpasses on the changed files.tools/pre_commit/mypy.py 3.12reportsSuccess: no issues found in 7 source files.Accuracy (8×MI355X, TP8/DCP8):
The tool-call eval varies by about 2 points between runs of the same server. The two arms are within that noise.
Agentic replay (InferenceX agentx-mvp, conc 48, 3600 s):
The errors are client-side aiperf
InvalidInferenceResultError/ClientOSErrorrecords and appear in both arms. Neither server logged any errors. This workload has very long inputs (ISL p50 about 86K) and about 97% prefix-cache hits. Most prefill work is therefore short new-token segments against long cached context, which is where the FP8 kernel gains most. Shorter prefills also shorten the stalls of decodes batched with them, which lowers the ITL tail.Prefill TTFT (
vllm bench serve, output_len 1, mean of 2 runs):* Baseline run 2 only. Run 1 overlapped another client's load on the same server and was discarded.
Backend micro-benchmark (one GPU, 12 heads, qk 192 / v 128; FP8 times include quantization). The "context" rows time the new-token call plus one context-chunk call, with Q quantized once and reused, as in
forward_mha:flash_attn_varlen_funcShort new-token-only calls (≤2K tokens) carry about 50 µs of host-side launch overhead from the quantizer. This did not show up end-to-end; the 8K fresh-prefill case is flat.
JIT. With an AITER build that includes the AOT kernels (#5796/#5901), the FlyDSL FMHA is not JIT-compiled while serving. Without them, each new FMHA configuration compiles on first use (about 1 s) and is then cached on disk. The Triton quantizer compiles 3 variants (Q/K/V, K/V, Q) once, which Triton caches.