[Bugfix][ROCm] Stream-order AITER MLA FP8 prefill PS metadata updates

[amd-mghanimi]

Purpose

Fix a GPU memory fault (Memory access fault ... Write access to a read-only page, address
~0x6c000-0x6f000) in the ROCm AITER MLA backend with FP8 KV cache and async scheduling. It
also removes a likely source of silent prefill attention corruption on the same path.

aiter.get_ps_metadata_v1 plans the FP8 MLA prefill schedule on the host. It then writes the
plan into the caller's device tensors with blocking hipMemcpy that is not ordered on the
current stream
(aiter v1_2_host.cuh L395-L418).
AiterMLAMetadataBuilder._build_fp8_prefill_ps_metadata passes the persistent
fp8_ps_* device buffers directly, and every step's metadata aliases those buffers
(L901-L941).

With async scheduling, step N+1 is built while step N's prefill kernels
(mla_prefill_ps_asm_fwd + mla_reduce_v1, one launch per layer) are still queued. Step N
can then run with step N+1's work/reduce maps. The existing .item() sync on
reduce_indptr[-1] comes after the rewrite, so it does not help.

This PR plans into two pinned host staging slots and copies them to the device buffers with
copy_(non_blocking=True) on the current stream. That orders the update after every
previously launched kernel. A slot is reused only after its previous copy has completed
(CUDA event). num_partial_tiles is read from the host plan, which removes one GPU sync per
prefill step. work_metadata is not written by the host planner and stays on device.

Only rocm_aiter_mla.py calls get_ps_metadata_v1 in vLLM.

Test plan

Kimi-K3, TP8, MI355X, vllm/vllm-openai-rocm:nightly-rocm100-29468dde… (ROCm 10), dummy
weights, FP8 KV, max_num_batched_tokens=16384, max_num_seqs=256, FULL_AND_PIECEWISE,
async scheduling. sglang.bench_serving random 8k/1k, max concurrency 32, 1024 prompts.

1. Reproducer without the fix (deterministic).
2. Same with the fix (patched file bind-mounted, nothing else changed).
3. Diagnostics isolating the cause (below).
4. Plan equivalence: host-output plan equals device-output plan for 5 batch shapes.

Test result

Crash: unpatched vs this PR on the same request stream

Each row pairs runs from the same environment. Whether a given request stream triggers the
fault varies, so every fix run is paired with an unpatched run that faults on that stream.

Request stream (8k/1k random)	Unpatched	This PR
MC32, 1024 prompts, seed 4	fault at request 111 (x2, and with --enforce-eager)	1024/1024; second run 949/1024 with no fault (stopped early)
MC32, 2048 prompts	fault at request 79	2048/2048
MC64, 1024 prompts	fault at request 61	1024/1024
MC32, 1024 prompts, seed 5	fault at request 117	1024/1024

Diagnostics on the seed-4 stream, all unpatched otherwise:

Change	Result
FP8 ASM prefill disabled	1024/1024
torch.cuda.synchronize() before get_ps_metadata_v1	1024/1024 (x2)
--no-async-scheduling	1024/1024

Regression checks

Check	Unpatched	This PR
8k/1k MC32, output tok/s	1083.0	1082.6
8k/1k MC32, median ITL	18.14 ms	18.14 ms
8k/1k MC32, mean TTFT	577 ms	574 ms
8k/1k MC4, output tok/s	236.05	236.05
8k/1k MC4, median ITL	15.49 ms	15.50 ms
8k/1k MC4, mean TTFT	371 ms	369 ms
8k/1k MC32, --enforce-eager, 1024 prompts	n/a	1024/1024
ROCm 7.2 image, 8k/1k MC32, 1024 prompts	n/a	1024/1024

Accuracy (real weights, niah_single_2, 4k, a single run on each build)

num_concurrent	Unpatched	This PR
4	0.974	0.974
32	0.954	0.978
64	0.970	0.984

Suggestive, not conclusive: single runs.

Long context (real weights, niah_single_2, num_concurrent=32, a single run on each build):

Context	Unpatched	This PR
64k	0.922	0.928
128k	0.888	0.888

No accuracy change at long context (within single-run noise).

Full sweep with real weights (this PR)

12-cell serving sweep plus GPQA and NIAH, real weights, default max_num_seqs, FP8 KV,
FULL_AND_PIECEWISE, async scheduling. All 12 cells completed with no fault.

Cell (random, sglang.bench_serving)	Output tok/s
8k/1k, concurrency 4 / 8 / 16 / 32	216.9 / 353.3 / 542.7 / 799.6
8k/1k, concurrency 64 / 128 / 256	1097.8 / 1383.2 / 1648.9
60k/600, concurrency 4 / 16 / 64	110.1 / 161.5 / 189.5
128k/1k, concurrency 4 / 16	91.6 / 123.8

Eval	Score
GPQA (sglang.test.run_eval, thinking on, 32k max tokens)	0.889
niah_single_2, 64k, num_concurrent=32	0.928
niah_single_2, 128k, num_concurrent=32	0.908

Related: DSpark startup fault (vllm-project#58814)

vllm-project#58814 reports a GPU memory fault in the Kimi-K3 DSpark draft (stale KV cache
pointers after the startup profiling pass). This A/B checks whether this PR also covers it.
Kimi-K3 TP8 + DSpark (K=5, TRITON_MLA draft), ROCm 10 nightly 29468dde, real weights,
FULL_AND_PIECEWISE, async scheduling.

Arm	Setup	Patch	Result
d3	vllm-project#58814 recipe: FP8 KV, max_num_seqs=8, 8k/1k MC4	none	fault at startup, right after JIT kernel warmup: Memory access fault ... Reason: Unknown at a high address (matches vllm-project#58814)
d4	same	this PR	fault at startup, same signature as d3: this PR does not fix vllm-project#58814
d5	same	vllm-project#58814	64/64, no fault (positive control; DSpark mean acceptance length 4.4-5.9 at K=5)
d1	DSpark OOB recipe (BF16 KV, max_num_seqs=128), 8k/1k MC4/8/16/32	none	fault at startup on all ranks, same signature as d3, then collective timeout (matches vllm-project#58814's log)
d2	same	this PR	fault at startup on all ranks, same as d1

Conclusion: the two faults are independent. vllm-project#58814's fault is fixed by vllm-project#58814 and not by this
PR. It happens at startup, at a high address, with reason "Unknown". This PR's fault is a write
to a read-only page near address 0x6c000, while serving. Also, the code this PR changes only
runs with an FP8 KV cache, so it cannot affect d1/d2.

Standalone reproducer (no server, one GPU)

Plans batch A (one 8175-token prefill) into the persistent buffers, queues A's
mla_prefill_ps_asm_fwd + mla_reduce_v1 behind a GPU delay on a non-blocking stream, then
immediately plans batch B (three prefills) into the same buffers from the host.

Mode	ROCm 7.2	ROCm 10
unpatched pattern	GPU memory fault (write to a read-only page, near-null address)	GPU memory fault (same signature)
stream sync before planning B	output matches reference	output matches reference
this PR's copy (pinned staging + copy_ on the stream)	output matches reference	output matches reference

Plan equivalence

Host-vs-device plan equivalence was verified on MI355X with a standalone script (all 5 shapes
byte-identical). No unit test is added: the check exercises AITER rather than vLLM code.

Still to run before marking ready

• Repeat 4k NIAH (more runs).

Notes for reviewers

• ROCm 7.2 has the same race: vLLM runs these kernels on a non-blocking stream on both
  runtimes, the stream-less blocking hipMemcpy in get_ps_metadata_v1 does not wait for it
  on either, and the standalone reproducer above faults on both. In serving, the overlap
  (next plan built while the previous prefill is still pending) happens about equally often on
  both, about 9% of prefill steps. The difference is what the copy does then. On ROCm 7.2 most
  of those plan updates (about 55 of 80 per rank) only return after the in-flight prefill has
  finished, so the runtime happens to order them. On ROCm 10 they return in under 1 ms, while
  the prefill is still reading. Which runtime change causes this is not pinned down yet (the
  two images ship very different HIP runtime builds). Either way it is accidental ordering,
  not a guarantee.
• Runtime workarounds without this PR, each 1024/1024 on a ROCm 10 stream that faulted at
  request 78: AMD_SERIALIZE_COPY=1 (copies wait for already-queued work) and
  GPU_MAX_HW_QUEUES=1 (every stream shares one in-order hardware queue).
• AITER-side follow-up (stream-ordered host planner, or a device-side planner like the decode
  v1_2_device path): ROCm/aiter#5879.
• Related: [Bugfix][ROCm] Fix startup OOM in AITER MLA FP8 prefill workspace sizing vllm-project/vllm#57923 (FP8 prefill workspace sizing) touched the same function. This PR does not
  change workspace sizing.

[amd-mghanimi]

Opened upstream as vllm-project#58887 (fixes vllm-project#58886).