[ROCm][MLA] Add guarded gfx942 FP8 context prefill by maeehart · Pull Request #51551 · vllm-project/vllm

maeehart · 2026-08-09T10:16:12Z

Summary

Adds an opt-in gfx942 FP8 context-prefill path for Kimi-K3 MLA using the D192/V128 kernel from ROCm/aiter#4645. It also supports checkpoint-specific static per-head scales through a versioned external artifact, removing dynamic amax reductions while preserving the existing attention boundary.

Background

The existing AITER FP8 ASM kernels do not cover Kimi-K3's asymmetric D192 Q/K and D128 V shape on gfx942. ROCm/aiter#4645 adds a packed variable-length Gluon implementation and static per-head quantization for this shape.

What this PR does

Enables the AITER FP8 path only with VLLM_ROCM_KIMI_K3_FP8_PREFILL=1.
Restricts dispatch to BF16 Kimi-K3 MLA on gfx942 with 12 heads, D192 Q/K, and D128 V.
Quantizes Q once per context pass and reuses its per-head scale across context segments.
Collects layer-owned Q/K/V maxima after warmup in eager calibration mode.
Writes atomic TP-rank shards and merges them into a versioned safetensors artifact.
Validates model revision, TP/PP topology, cache mode, dimensions, dtype, layer coverage, and finite positive descales before compilation.
Uses a one-kernel static quantizer when an artifact is configured. Missing or incompatible artifacts fail instead of silently falling back.
Keeps the BF16 path for unsupported shapes and short contexts.

Validation

Model: moonshotai/Kimi-K3, TP8, MI325X.

Full five-shot gsm8k: 0.9606 exact match with 0.0054 stderr.
Needle retrieval: 9 of 9 passed at 128K, 512K, and 1,047,552 tokens with needle depths 10%, 50%, and 90%.
Deterministic completion sanity checks: 3 of 3 passed.
Matched 1,047,552-token dynamic TTFT: 434.65 s, 404.80 s, and 404.83 s.
Matched static TTFT: 431.38 s, 402.44 s, and 402.44 s.
Mean static TTFT improved by 0.64%. The two warm repeats improved by 0.59%.
AITER synthetic gfx942 tests: 9 passed.
vLLM artifact and configuration tests: 8 passed.

Notes / scope

This PR depends on ROCm/aiter#4645. The feature remains disabled by default. Calibration and validation prompt sets are separate, and calibration is eager-only.

Test plan

Full gsm8k, five-shot, 1319 samples
128K, 512K, and one-million-token needle retrieval
Three matched one-million-token dynamic and static repeats
API sanity checks
Artifact compatibility and merge tests

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.

maeehart · 2026-08-09T15:01:12Z

I have read the DCO document and hereby sign off past commits made by me.

maeehart · 2026-08-11T20:18:23Z

I have read the DCO and I hereby sign off the commits in this pull request.

maeehart · 2026-08-11T20:19:10Z

I have read the DCO and I hereby sign off on the commits in this pull request.

maeehart · 2026-08-11T21:04:58Z

Follow-up commit 22ab1fa46 makes static validation fail closed and hardens artifacts. Rank shards now require immutable checkpoint and calibration-run identities, the merger rejects malformed maxima, runtime cache and backend compatibility are checked, empty PP stages emit shards, worker cleanup survives save errors, and older dynamic-only AITER builds keep working. The artifact and configuration suite passes 9 tests.

mergify · 2026-08-13T17:10:49Z

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

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

maeehart · 2026-08-14T15:53:14Z

Rebased onto current main. The conflict resolution preserves current main's context-chunk output-buffer contract and adds the calibration parameters as keyword-only arguments.

Focused validation passes 11 tests. The Python 3.10 mypy hook passes for the changed modules, and Ruff check and formatting pass. The new head is cf2fa9addc. The PR remains draft until ROCm/aiter#4645 is available.

maeehart · 2026-08-14T16:03:44Z

Follow-up a59030238f addresses the pre-commit findings from the rebased branch: the merge tool now has the required SPDX header, calibration shutdown uses torch.accelerator.synchronize(), and the unit test mocks the portable API. The focused suite passes 11 tests, and both previously failing portability checks pass locally.

mergify · 2026-08-17T04:06:27Z

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

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

Use the AITER D192/V128 kernel for BF16 Kimi-K3 requests once the cached context reaches 786432 tokens. Keep the path opt-in, quantize Q once per context pass, and fall back to the existing BF16 implementation when the shape or dependency is unsupported. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com>

The FP8 V path and 786432-token threshold are validated behavior, not independent configuration. Keep only the dependency opt-in until ROCm/aiter#4645 lands, and cover the fixed dispatch choices with focused tests. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Dynamic per-head reductions add work to every long-context prefill step. Collect checkpoint-specific maxima after warmup, validate and load an external scale artifact before compilation, and use a one-kernel static quantizer while keeping the existing attention boundary. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Require immutable checkpoint and calibration identities, validate cache and backend compatibility, and reject malformed rank maxima before merging. Preserve worker cleanup when calibration saving fails and keep dynamic FP8 available with older AITER builds. Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Current main types MLAAttention.impl as a generic AttentionImpl, so forwarding calibration state through ad-hoc attributes fails mypy. Narrow to MLACommonBaseImpl, declare the shared state, and update focused tests for AttentionConfig and the calibrated quantizer call. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Add the required SPDX header to the calibration merge tool and use the accelerator-neutral synchronization API when saving rank shards. This keeps the ROCm path compatible with the repository-wide portability checks. Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: Markus Hartikainen <markus.hartikainen@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Co-authored-by: Cursor <cursoragent@cursor.com>

mergify · 2026-08-28T12:35:06Z

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

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

maeehart · 2026-09-29T14:02:13Z

I am closing this. It depends on ROCm/aiter#4645, which is now closed.

FMHA v3 BF16 already covers D192/V128 on gfx942. The FP8 kernel was faster only on the non-causal 4K/65K chunk, 3.498 ms versus 4.357 ms on MI325X, and slower on both causal shapes. Those times are in the close comment on ROCm/aiter#4645.

maeehart requested review from AndreasKaratzas and tjtanaa as code owners August 9, 2026 10:16

claude Bot reviewed Aug 9, 2026

View reviewed changes

mergify Bot added the rocm Related to AMD ROCm label Aug 9, 2026

github-project-automation Bot added this to AMD Aug 9, 2026

github-project-automation Bot moved this to Todo in AMD Aug 9, 2026

maeehart force-pushed the maeehart/kimi-k3-gfx942-fp8-prefill-vllm branch from 8439431 to 490340e Compare August 9, 2026 15:11

tjtanaa reviewed Aug 10, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py Outdated

tjtanaa reviewed Aug 10, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py Outdated

tjtanaa requested changes Aug 10, 2026

View reviewed changes

Comment thread vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py

maeehart requested review from LucasWilkinson, MatthewBonanni, ProExpertProg, WoosukKwon, houseroad, mgoin, njhill, pavanimajety, robertgshaw2-redhat, tlrmchlsmth, yewentao256 and youkaichao as code owners August 11, 2026 20:16

mergify Bot added kimi k3 nvidia labels Aug 11, 2026

github-project-automation Bot added this to NVIDIA Aug 11, 2026

github-project-automation Bot moved this to In review in NVIDIA Aug 11, 2026

maeehart force-pushed the maeehart/kimi-k3-gfx942-fp8-prefill-vllm branch from d521ac8 to f57de88 Compare August 11, 2026 20:52

maeehart marked this pull request as draft August 11, 2026 21:57

mergify Bot added the needs-rebase label Aug 13, 2026

RerankerGuo mentioned this pull request Aug 14, 2026

[Bugfix][Model] Validate Kimi-K3 plain FP8 prefill support #51394

Closed

4 tasks done

maeehart force-pushed the maeehart/kimi-k3-gfx942-fp8-prefill-vllm branch from 22ab1fa to cf2fa9a Compare August 14, 2026 15:52

mergify Bot removed the needs-rebase label Aug 14, 2026

mergify Bot added the needs-rebase label Aug 17, 2026

maeehart mentioned this pull request Aug 18, 2026

[ROCm][AMD] Kimi-K3 gfx942 / MI325X Gap and Roadmap #52803

Open

maeehart and others added 6 commits August 24, 2026 13:25

maeehart force-pushed the maeehart/kimi-k3-gfx942-fp8-prefill-vllm branch from a590302 to 2d13454 Compare August 24, 2026 20:27

Fix cpu_native prefill signature after rebase

5be5507

Co-authored-by: Cursor <cursoragent@cursor.com>

mergify Bot added cpu Related to CPU backends and removed needs-rebase labels Aug 24, 2026

mergify Bot added the needs-rebase label Aug 28, 2026

maeehart mentioned this pull request Sep 28, 2026

[Triton/Gluon] [MHA][gfx942] Add FP8 D192/V128 prefill ROCm/aiter#4645

Closed

4 tasks done

gbyu-amd mentioned this pull request Sep 29, 2026

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

Open

maeehart closed this Sep 29, 2026

Navigation Menu

Repository navigation

[ROCm][MLA] Add guarded gfx942 FP8 context prefill - #51551

[ROCm][MLA] Add guarded gfx942 FP8 context prefill#51551
maeehart wants to merge 7 commits into
vllm-project:mainfrom
maeehart:maeehart/kimi-k3-gfx942-fp8-prefill-vllm

maeehart commented Aug 9, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

maeehart commented Aug 9, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

maeehart commented Aug 11, 2026

Uh oh!

maeehart commented Aug 11, 2026

Uh oh!

maeehart commented Aug 11, 2026

Uh oh!

mergify Bot commented Aug 13, 2026

Uh oh!

maeehart commented Aug 14, 2026

Uh oh!

maeehart commented Aug 14, 2026

Uh oh!

mergify Bot commented Aug 17, 2026

Uh oh!

mergify Bot commented Aug 28, 2026

Uh oh!

maeehart commented Sep 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

maeehart commented Aug 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Background

What this PR does

Validation

Notes / scope

Test plan

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

maeehart commented Aug 9, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

maeehart commented Aug 11, 2026

Uh oh!

maeehart commented Aug 11, 2026

Uh oh!

maeehart commented Aug 11, 2026

Uh oh!

mergify Bot commented Aug 13, 2026

Uh oh!

maeehart commented Aug 14, 2026

Uh oh!

maeehart commented Aug 14, 2026

Uh oh!

mergify Bot commented Aug 17, 2026

Uh oh!

mergify Bot commented Aug 28, 2026

Uh oh!

maeehart commented Sep 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

maeehart commented Aug 9, 2026 •

edited

Loading