[perf][rocm]: accelerate strict VIME rollout with direct paged CK and graph-safe execution

[Flink-ddd]

Context

PR #388 established bitwise Qwen3/VIME train-rollout alignment on ROCm, but the strict R/R rollout path still carried substantial runtime overhead:

• paged KV was reconstructed into dense logical rows;
• GPU metadata was synchronized to the host in every decoder layer;
• Attention executed multiple Python-level row/KV-group dispatches;
• RoPE tables and weight layouts were repeatedly materialized;
• Attention and GEMM outputs passed through temporary buffers;
• deterministic collectives introduced staging copies and graph breaks.

PR #390 moves the strict ROCm rollout path from a materialization-heavy adapter to a direct-paged, graph-safe execution pipeline while preserving the train/rollout arithmetic contract.

Optimization architecture

flowchart TB
    INPUT(["VIME strict R/R rollout"]) --> META["Forward-scoped metadata cache"]

    subgraph HOTPATH["Strict ROCm compute hot path"]
        direction LR
        QKV["Deterministic QKV GEMM"] --> ROPE["Token-major deterministic RoPE"]
        ROPE --> ATTN["Direct paged AITER/CK"]
        ATTN --> OPROJ["Direct-output O-projection"]
    end

    META --> QKV
    OPROJ --> IPC["Stable-buffer IPC fixed-tree reduce"]
    IPC --> HEAD["Prepared LM-head and strict logprob"]
    HEAD --> GATE{"Exactness and provenance gate"}
    GATE --> OUTPUT(["Aligned rollout output"])

    subgraph CONTRACT["Strict arithmetic contract"]
        direction LR
        SPLIT["Split-KV disabled"] --> ACCUM["FP32 accumulation"]
        ACCUM --> TREE["Fixed reduction order"]
        TREE --> FALLBACK["No fallback"]
    end

    SPLIT -. "constrains" .-> ATTN
    TREE -. "constrains" .-> IPC
    FALLBACK -. "verified by" .-> GATE

    classDef runtime fill:#ddf4ff,stroke:#0969da,color:#24292f,stroke-width:2px
    classDef compute fill:#fff8c5,stroke:#bf8700,color:#24292f,stroke-width:2px
    classDef distributed fill:#dafbe1,stroke:#1a7f37,color:#24292f,stroke-width:2px
    classDef guard fill:#ffebe9,stroke:#cf222e,color:#24292f,stroke-width:2px

    class INPUT,META,OUTPUT runtime
    class QKV,ROPE,ATTN,OPROJ compute
    class IPC,HEAD distributed
    class SPLIT,ACCUM,TREE,FALLBACK,GATE guard

    style HOTPATH fill:#fffdf2,stroke:#bf8700,stroke-width:2px
    style CONTRACT fill:#fff5f5,stroke:#cf222e,stroke-width:2px

Loading

The solid arrows represent the steady-state rollout data path. The dashed arrows represent the strict arithmetic constraints applied to kernel dispatch, accumulation, communication, and runtime acceptance.

The optimized route keeps tensors on the GPU, preserves framework-owned storage, and removes redundant synchronization, layout conversion, temporary allocation, and output-copy operations from steady-state decode.

Important

PR #390 changes execution and data movement while keeping the strict arithmetic schedule fail-closed, deterministic, and auditable.

Key optimizations

Area	Previous overhead	PR #390 optimization	Exactness protection
Paged metadata	seqused_k, page bounds, and index metadata were reconstructed or validated in every decoder layer	Materialize metadata once per vLLM forward and reuse it through a forward-scoped cache	Cache ownership is limited to one immutable forward epoch
Training Attention positions	CP position IDs were transported and sorted again in every layer and recompute	Cache the immutable CP position transport, sort order, and inverse order across layers	Q/K/V transport and Attention arithmetic are unchanged; cache entries retain their source tensors
Paged KV	Physical pages were gathered into dense logical K/V rows	Pass vLLM-owned paged K/V directly to AITER mha_batch_prefill	Page order, sequence lengths, shape, dtype, and device remain validated
Attention dispatch	One Python/AITER launch per batch row and KV group	Execute one packed paged-varlen CK launch for the active query batch	Split-KV remains disabled and num_splits=1
Attention output	Temporary group outputs were assembled through cat and adapter copies	Write AITER output directly into the caller-owned vLLM output slice	Output shape, dtype, storage, and backend identity are checked
Decode LSE	Rollout produced LSE even when no downstream consumer required it	Disable LSE materialization on the inference-only direct-paged path	Training still returns FP32 LSE where required
RoPE	Q and K rebuilt equivalent position tables and required layout conversions	Prebuild the FP32 table, reuse it across Q/K, cache inv_freq, and execute token-major HIP RoPE	Q/K retain the same deterministic HIP arithmetic and positions
Deterministic GEMM	The fixed-tree root was written to workspace and copied to the destination	Write the final reduction node directly into the destination buffer	Leaf kernels and the canonical BF16 reduction tree are unchanged
LM head	A contiguous transposed weight was materialized for every projection	Maintain a prepared [K,N] weight cache and refresh it after IPC weight updates	Storage pointer, tensor version, source identity, and refresh generation are validated
TP O-projection	GEMM output was copied into collective staging	Write deterministic GEMM output directly into ROCm IPC staging	The fixed deterministic reduction tree is preserved
HIP Graph execution	Replaying stateful IPC inside a full graph could consume capture-time peer payloads	Force piecewise HIP Graphs on ROCm, split only the stateful TP reductions, and feed adjacent graph partitions through stable input/output buffers	IPC sequence state advances once per eager boundary; CUDA retains the PR #377 full-graph path
Provenance	Immutable device and Split-KV metadata was repeatedly resolved	Cache immutable provenance while retaining runtime readback	Backend, schedule, source fingerprint, and fallback status remain visible

Direct-paged Attention path

The optimized rollout path executes:

vLLM paged KV
    -> packed query metadata
    -> AITER mha_batch_prefill
    -> caller-owned output buffer

The direct path exposes the following runtime evidence:

paged_execution = direct_vllm_pages_to_aiter_batch_prefill_ck
dense_kv_materialized = false
query_schedule = paged_varlen_batch
core_launch_count = 1
split_kv = disabled
fallback = false

This removes dense KV reconstruction from the primary rollout path without introducing a moving Split-KV reduction schedule.

Exactness-preserving contract

The optimizations change data movement, dispatch, allocation, and storage placement. They do not relax the strict arithmetic contract:

• AITER/CK remains the production Attention arithmetic backend.
• Split-KV is disabled and num_splits is fixed at 1.
• Attention accumulation remains FP32.
• Output downcast occurs only at the final write.
• Dropout remains disabled.
• Deterministic GEMM retains its canonical reduction tree.
• Tensor-parallel collectives retain a fixed reduction order.
• Runtime fallback is forbidden.
• Backend identity, execution schedule, source fingerprint, and fallback state remain visible through provenance.
• Optimized results must remain byte-identical to the strict reference route.

Post-merge follow-up

PR #390 merged into test at fed8362. The subsequent full-Graph, route-specific AOT cache, position-plan optimization, and matched three-round P/P vs R/R results are tracked in #393 so results from different commit sets are not mixed in this merged PR.

Scope boundary

This PR optimizes the existing strict ROCm execution contract. It does not claim that every stage is faster than VIME production P/P, and it does not change arithmetic ordering solely for performance.

Potential follow-up work includes:

• profiling small-query mha_batch_prefill occupancy on gfx942;
• evaluating additional AITER paged kernels under a fixed non-Split-K contract;
• qualifying MFMA-specific deterministic GEMM schedules;
• reducing the remaining QKV, O-projection, and LM-head layout conversions;
• publishing a final 30-step or 200-step matched-workload comparison.

Related: #377, #385, #388.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: 80617107-478e-4700-aa3f-42ba462a3411

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[maxiaosong1124]

LGTM!