Skip to content

[ROCm][Perf][MLA] Add opt-in FlyDSL FP8 prefill attention for AITER MLA - #59124

Open
gbyu-amd wants to merge 6 commits into
vllm-project:mainfrom
gbyu-amd:rocm-k3-flydsl-fp8-prefill
Open

gbyu-amd wants to merge 6 commits into
vllm-project:mainfrom
gbyu-amd:rocm-k3-flydsl-fp8-prefill

Conversation

@gbyu-amd

@gbyu-amd gbyu-amd commented Sep 29, 2026 •

Copy link
Copy Markdown
Contributor

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 subclasses ROCM_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:

  • quantizes Q/K/V to FP8 with dynamic per-tensor scales and runs AITER's flydsl_flash_attn_fp8_func for the new tokens;
  • keeps the new-token FP8 Q for the context chunks. Each chunk quantizes only its own K/V and calls the kernel non-causally with return_lse=True. The LSE layout ([H, T], natural log) is unchanged, so merge_attn_states works as-is. The layer's softmax scale is passed directly.
  • takes over context-free prefills too, in place of the ASM FP8 prefill (mla_prefill_ps_asm_fwd). The metadata builder stops building ASM PS metadata when this backend is selected, and AiterMLAImpl.forward_mha then 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_capability requires gfx950;
  • is_available requires AITER and the FlyDSL FP8 API;
  • validate_configuration defers to AITER's flydsl_flash_attn_fp8_supported for 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.py is 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 the kv_b_proj output. 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 behind QuantFP8) gives the same result: descales within 1e-6 relative, values within 1 e4m3 ulp. It is too slow for this path, for three reasons:

  • It takes one 2D tensor per call, so Q/K/V need three calls (memset, amax, quant each).
  • The strided V (row stride 256) takes its per-row path: one 256-thread block and one atomicMax per 128-element row, i.e. T×H rows.
  • An all-zero input gets descale 0. ATOM uses a positive descale in that case to avoid FMHA NaNs, and this kernel does the same.

rocm_aiter_ops.per_tensor_quant returns wrong results for the strided V. GPU time on MI355X, 12 heads, q/k 192 and v 128 (v a strided view):

Tokens This kernel (Q/K/V, one call) scaled_fp8_quant ×3 scaled_fp8_quant ×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

At 65,536 tokens, V alone takes 9,284 µs with scaled_fp8_quant and 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:

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-rocm nightly (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_cross layout, 92 kernels) during AITER source builds, so serving no longer JIT-compiles them. The new-token call therefore passes cross_seqlen=True, as ATOM does.

  • With cross_seqlen=False, the new-token call would use the varlen layout, which is not pre-compiled.
  • For equal q/kv lengths, the two layouts give bit-identical output and LSE. Checked 8 batch shapes.
  • Kernel time with cross_seqlen=True is 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_fp8 from AITER main) were put into an empty FLYDSL_RUNTIME_CACHE_DIR, and FLYDSL_RUNTIME_RUN_ONLY=1 was set, which makes any cache miss raise.

  • Through this backend, all 12 new-token and context-chunk calls were served from the AOT cache, and the cache did not grow. The calls covered 1–32 sequences, 17–8192 new tokens and contexts of 0–131072 tokens.
  • With cross_seqlen=False, all 20 new-token calls missed.

Not a duplicate. I searched open PRs for FlyDSL/FP8 MLA prefill:

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_FP8 is added to PREFILL_BACKENDS_TO_TEST in test_mla_backends.py, so test_backend_correctness and test_chunked_context_backend_correctness check it against the SDPA reference.

  • GLM dims (v=256) are skipped at collection through validate_configuration.
  • For this backend only, the tolerance is 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.
  • A new context_free_prefill spec, run for this backend, covers the builder change that keeps ROCM_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:

  • the quant kernel against a torch reference: at most 1 e4m3 ulp, with strided V, empty and all-zero inputs;
  • the backend's FP8 new-token and context-chunk outputs against its BF16 path, including short-q/long-kv chunks and sequences with empty context;
  • FP8 Q reuse by the context chunk;
  • validate_configuration rejecting 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 inside ROCM_AITER_FA. Apart from cross_seqlen=True (bit-identical outputs), the kernel calls are unchanged.

vllm serve moonshotai/Kimi-K3 --trust-remote-code --tensor-parallel-size 8 \
  --decode-context-parallel-size 8 --dcp-comm-backend a2a \
  --attention-backend ROCM_AITER_MLA --attention-config '{"mla_prefill_backend":"ROCM_AITER_FA"}' \
  --kv-cache-dtype fp8 --enable-prefix-caching --max-num-batched-tokens 8192 --max-num-seqs 96 \
  --max-model-len 1048576 --language-model-only --moe-backend auto --load-format fastsafetensors \
  --gpu-memory-utilization 0.9 --tool-call-parser kimi_k3 --reasoning-parser kimi_k3 --enable-auto-tool-choice \
  --compilation-config '{"mode":3,"cudagraph_mode":"FULL","max_cudagraph_capture_size":96,"custom_ops":["+fused_rms_norm_gated"],"cudagraph_capture_sizes":[2,...,96]}' \
  --kv-transfer-config '{"kv_connector":"SimpleCPUOffloadConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use_per_rank":224875000000,"lazy_offload":false}}'
# env: VLLM_ROCM_USE_AITER=1 VLLM_ROCM_AITER_MLA_ASM_PADDING=asm VLLM_ROCM_USE_AITER_MOE_SITUV2_A8W4=1
#      AITER_SITUV2_A8W4=1 AITER_BF16_FP8_MOE_BOUND=0 AITER_QUICK_REDUCE_QUANTIZATION=INT4
#      VLLM_USE_DIRECT_DCP_{A2A,Q_GATHER,KV_GATHER}=0
# FP8 arm: --attention-config '{"mla_prefill_backend":"ROCM_AITER_FLYDSL_FP8"}'
  • Accuracy:
    • lm-eval gsm8k, 5-shot, full 1319, via local-completions, num_concurrent=64, temperature 0.
    • InferenceX kimi_tool_call_schema_full (Kimi Vendor Verifier, 408 records, the CI eval for this job), run twice per arm.
  • Performance:
    • InferenceX agentic replay: 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.
    • Prefill TTFT: vllm bench serve --dataset-name random --random-output-len 1 --ignore-eos at 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. Three ROCM_AITER_FA cases (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 on main.
  • All 290 ROCM_AITER_FLYDSL_FP8 cases pass; the GLM-dim cases are skipped.
  • Backend, selector and registry tests: 48 passed.
  • pre-commit passes on the changed files. tools/pre_commit/mypy.py 3.12 reports Success: no issues found in 7 source files.

Accuracy (8×MI355X, TP8/DCP8):

Eval Baseline FlyDSL FP8
GSM8K 5-shot strict-match / flexible-extract 0.9621 / 0.9621 0.9621 / 0.9621
kimi_tool_call_schema_full, run 1 / run 2 0.9779 / 0.9877 0.9755 / 0.9951

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):

Metric Baseline FlyDSL FP8 Change
TTFT p50 / p90 / p99 (ms) 1176 / 3690 / 22443 1119 / 3106 / 16155 -4.8% / -15.8% / -28.0%
ITL p50 / p90 / p99 (ms) 51.1 / 74.6 / 188.0 49.2 / 68.2 / 130.5 -3.8% / -8.7% / -30.6%
P90 interactivity (1000 / ITL p90, tok/s) 13.40 14.67 +9.5%
Total token throughput (tok/s) 91,787 95,931 +4.5%
Output token throughput (tok/s) 609.2 638.5 +4.8%
Completed / errored requests 2653 / 8 2744 / 6

The errors are client-side aiperf InvalidInferenceResultError/ClientOSError records 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):

Case Baseline mean TTFT FlyDSL FP8 mean TTFT Change
8K fresh prefill ×128, c16 6676 ms 6659 ms -0.2%
32K (chunked, with context) ×32, c4 5528 ms 5441 ms -1.6%
128K ×8, c2 13252 ms 12283 ms -7.3%
64K with a 56K shared cached prefix ×32, c4 5071 ms* 4758 ms -6.2%

* 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:

Shape BF16 flash_attn_varlen_func FlyDSL FP8 Speedup
new tokens 1×8192, causal 313 µs 229 µs 1.37×
new q 8192 + context kv 64K 3821 µs 2469 µs 1.55×
new q 1024 + context kv 64K 1278 µs 549 µs 2.33×
new q 128 + context kv 64K 1251 µs 351 µs 3.56×

Short 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.

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>
@mergify mergify Bot added the rocm Related to AMD ROCm label Sep 29, 2026
@github-project-automation github-project-automation Bot moved this to Todo in AMD Sep 29, 2026
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>
@gbyu-amd
gbyu-amd marked this pull request as ready for review September 29, 2026 05:50

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@simondanielsson simondanielsson left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work!

Suggestions:

  1. Can we create a standalone mla prefill backend for this? there's a lot of useful existing infra for this
  2. Can we plug that new backend into the existing attention correctness tests?
  3. 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(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question: what's the rationale behind implementing this here rather than re-using an existing quant kernel?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_quant returns wrong results for the strided V.
  • ops.scaled_fp8_quant matches 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 atomicMax per 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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question: does this work even though this backend does not supports_out?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
@gbyu-amd

Copy link
Copy Markdown
Contributor Author

Nice work!

Suggestions:

  1. Can we create a standalone mla prefill backend for this? there's a lot of useful existing infra for this
  2. Can we plug that new backend into the existing attention correctness tests?
  3. 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

hi @simondanielsson , thanks for your review! Here are the updates regarding your concerns:

  1. Standalone backend. The FlyDSL FP8 path now lives in its own prefill backend, ROCM_AITER_FLYDSL_FP8 (vllm/v1/attention/backends/mla/prefill/aiter_flydsl_fp8.py). It subclasses AiterFlashAttnPrefillBackend, so the few inputs FlyDSL can't take fall back to the parent's BF16 path. ROCM_AITER_FA is back to its main version.

    • Select it with --attention-config '{"mla_prefill_backend":"ROCM_AITER_FLYDSL_FP8"}'. The env var is gone.
    • Gating goes through the standard hooks:
      • supported_dtypes = [bf16].
      • supports_compute_capability requires gfx950.
      • is_available requires AITER and the FlyDSL FP8 API.
      • validate_configuration defers to AITER's flydsl_flash_attn_fp8_supported for the head dims.
    • An explicit but unsupported selection therefore fails at startup, instead of silently running BF16 as before.
    • The backend is not in the ROCm auto-selection priorities.
  2. Correctness tests. The backend is added to PREFILL_BACKENDS_TO_TEST in tests/v1/attention/test_mla_backends.py, so test_backend_correctness and test_chunked_context_backend_correctness run it against the SDPA reference.

    • GLM dims (v=256) are skipped at collection through validate_configuration.
    • The tolerance is loosened for this backend only: rtol=2**-4, atol=0.3. With per-tensor e4m3 Q/K/V, a row that attends to only a few tokens is off by up to about 2⁻⁴·max|V|. The largest observed diff was about 0.25.
    • I also added a context_free_prefill spec to test_chunked_context_backend_correctness, which this backend runs. It covers the builder change that stops ROCM_AITER_MLA's ASM FP8 prefill from bypassing the selected prefill backend.
    • Registry and selector unit tests are added as well.
    • Result on MI355X with VLLM_ROCM_USE_AITER=1: pytest tests/v1/attention/test_mla_backends.py -k ROCM_AITER → 869 passed and 296 skipped. Three ROCM_AITER_FA cases hit CUDA OOM because another job filled the shared GPU; all three pass individually on this branch and on main. All ROCM_AITER_FLYDSL_FP8 cases pass.
  3. Default for FP8 KV. Agreed, this should stay opt-in until there is a broader sweep. It also overlaps with your [ROCm][MLA] AITER FP8 persistent ASM prefill backend #44544 (ROCM_AITER_ASM as the gfx950 default for FP8 KV). I think the two can coexist as separate backends:

    • ROCM_AITER_ASM consumes the FP8 KV cache directly.
    • This one quantizes the BF16 Q/K/V dynamically per tensor, so it doesn't depend on the KV cache dtype.

    If it helps decide the default, I can run a same-image sweep of FA, ASM and FlyDSL-FP8 (ISL 1k/8k/32k+, concurrency 8–256, plus the Kimi-K3 agentic replay) once [ROCm][MLA] AITER FP8 persistent ASM prefill backend #44544 is rebased.

@ChuanLi1101

Copy link
Copy Markdown
Collaborator

Three technical suggestions:

  1. The quantizer grid uses the largest tensor size. V is smaller, so masked threads can still form out-of-range pointers. Please clamp indices or use separate grids.

  2. The quantizer allocates temporary buffers on every call. Reusing fixed buffers could reduce the ~50 µs overhead for short inputs.

  3. Please disable the ASM FP8 prefill path directly when this backend is selected, instead of relying only on missing metadata to bypass it.

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>
@gbyu-amd

Copy link
Copy Markdown
Contributor Author

@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:

  • _partial_amax derives its trip count per program, tl.cdiv(N - start, PARTS * BLOCK), which is 0 once the program's stripe starts past the end of this input. It still stores its (zero) partial, which the reduce needs.
  • _quant guards the body with if start < N or start == 0. Program 0 stays alive unconditionally so an empty input still gets its descale written.

I added test_fused_per_tensor_fp8_quant_covers_every_element to the kernel test: it plants the global max at {0, BLOCK-1, BLOCK, STRIPE-1, STRIPE, numel-1} for sizes straddling the 4096-element block and 1048576-element stripe boundaries, up to 3,145,735 elements. It fails if a single element is dropped (verified by shrinking the amax loop to one stripe). The change is perf-neutral: host 46.3–46.8 µs, e2e 349.2 µs at 65K tokens, same as before.

2. Per-call temporary buffers

I measured this on MI355X before changing it, since the buffers are tiny (partial is 3×256 f32, scales is 3 f32). Host cost of fused_per_tensor_fp8_quant(q, k, v):

Variant Host Δ
current (5 allocations) 46.6 µs —
reuse partial only 44.2 µs −2.4 µs
reuse partial + scales 42.7 µs −3.9 µs
one merged scratch allocation 46.7 µs +0.1 µs

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 torch.empty ×5, and the rest is the Python wrapper. So the ceiling on reuse is ~8% of host cost and ~1% of the 349 µs e2e at 65K tokens. The three FP8 outputs can't be reused at all — they are returned and stay live for the whole attention call.

Reusing scales specifically has a correctness trap: the backend keeps self._q_fp8 = (q, q_fp8, q_descale) across the forward, where q_descale is a view into that call's scales, and each context chunk then calls the quantizer again for its own K/V. A shared per-instance scales would be overwritten under the still-live q_descale unless keyed by tensor count — exactly the kind of implicit invariant that breaks later.

Given ~2 µs for the only unambiguously safe part (partial, which never escapes the function) against making a pure function stateful, I left it as is. Happy to add the partial cache if you'd still prefer it — it's a five-line change.

3. Disabling the ASM FP8 prefill directly

Agreed, keying off absent PS metadata was too indirect. There's now a single predicate in rocm_aiter_mla.py:

def _asm_fp8_prefill_superseded(prefill_backend: object) -> bool:
    return isinstance(prefill_backend, AiterFlyDSLFP8PrefillBackend)

used in both places so they can't drift:

  • AiterMLAMetadataBuilder.__init__ folds it into _fp8_prefill_enabled, so the PS workspace buffers aren't allocated either.
  • AiterMLAImpl.forward_mha checks the actual prefill_metadata.prefill_backend instance and defers to the parent, instead of inferring it from fp8_prefill_qo_indptr is None.

test_mla_backends.py::test_chunked_context_backend_correctness covers this through the context_free_prefill spec.

4. Opt-in

Agreed — it stays out of the ROCm auto-selection priorities and is only reachable via --attention-config '{"mla_prefill_backend":"ROCM_AITER_FLYDSL_FP8"}'.

@mergify

mergify Bot commented Oct 8, 2026

Copy link
Copy Markdown
Contributor

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

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

@mergify mergify Bot added the needs-rebase label Oct 8, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

needs-rebase rocm Related to AMD ROCm

Projects

Status: Todo

Development

Successfully merging this pull request may close these issues.

4 participants