[RFC][ROCm] End-to-end optimization of bitwise VIME R/R beyond native P/P

[Flink-ddd]

Status

• State: Proposed
• Target: Qwen3-8B BF16 on one node with 8 AMD Instinct MI300X GPUs
• Reference topology: Megatron TP4/CP2/PP1; two vLLM rollout engines with TP4 each
• Reference result: #400
• Feedback period: seven calendar days after this issue is opened

Summary

RL-Kernel now provides a fail-closed ROCm path in which VIME training and vLLM rollout execute the same strict Attention, FFN, collective, and linear-logp contracts. The 200-step G11 run in #400 preserved bitwise train-rollout equality over all 9,400,614 compared log-probability elements from 1,600 samples.

The strict route is not yet faster than VIME's native production route. Across the 200-step run, G11 R/R averaged 99.21 seconds per step versus 84.52 seconds for G10 P/P. After excluding rollout warmup and normalizing by generated tokens, end-to-end throughput was 59.28 versus 80.78 tokens/GPU/s.

This RFC proposes a profile-driven ROCm fast path with four primary changes:

1. a graph-captured, allocation-stable rollout execution plan for small-M decode;
2. a local-query CP Attention algorithm that removes duplicated full-sequence training work;
3. a fused MFMA LM-head plus log-probability path that avoids materializing full local-vocabulary logits when they are not otherwise required;
4. topology-aware deterministic HIP IPC collectives with fixed-order overlap for the dependency patterns that permit it.

The numerical contract remains fail-closed. Performance tuning may change launch geometry, storage placement, and data movement, but it may not change a contract's reduction tree, accumulation precision, rounding boundaries, transcendental implementation, or rank order without introducing a new versioned contract that passes the full train-rollout exactness ladder.

The exit condition is an end-to-end strict R/R path that is at least 5% faster than native P/P on the frozen matched-work benchmark, while retaining zero mismatches across a fresh 200-step VIME run.

Motivation

Established correctness

The G11 result in #400 is the baseline that this RFC must preserve:

Check	G11 R/R result
Rollout/training steps	200/200
Samples	1,600
Compared elements	9,400,614
Mismatch count	0
Maximum absolute difference	0
torch.equal	true
Arithmetic fallback	none

Measured performance gap

The following values come from the complete paired CSV published by #400. The timers inside train_time are nested and must not be added together.

Metric	G10 native P/P	G11 strict R/R	R/R relative to P/P
Rollout time	73.77 s	81.43 s	10.4% slower
Rollout throughput	93.79 tok/GPU/s	72.60 tok/GPU/s	22.6% lower
Reference log-probability time	1.56 s	3.33 s	113.4% slower
Actor train time	4.88 s	8.75 s	79.3% slower
Train time	8.54 s	14.93 s	74.7% slower
Weight update	1.113 s	1.121 s	0.7% slower
End-to-end step	84.52 s	99.21 s	17.4% slower

The paired mean step-time delta is +14.68 seconds, with a 95% paired-bootstrap interval of +11.28 to +17.87 seconds. Rollout contributes +7.66 seconds and train time contributes +6.38 seconds to the mean gap. Weight update is already at parity and is not a first-order target.

Because G10 and G11 are free-running training arms, their sampled response lengths differ. Two normalized views are therefore required:

View	Native P/P	Strict R/R	Remaining gap
Steps 1-199, end-to-end	80.78 tok/GPU/s	59.28 tok/GPU/s	26.6% lower
Steps 1-199, rollout	92.43 tok/GPU/s	72.20 tok/GPU/s	21.9% lower
Steps 1-199, actor train	11,305 tok/s	5,387 tok/s	52.3% lower
Steps 1-199, reference logp	4,500 tok/GPU/s	1,767 tok/GPU/s	60.7% lower
Equal 7,168-token subset (75 steps), end-to-end	81.89 tok/GPU/s	62.99 tok/GPU/s	23.1% lower

Raw free-running time remains useful for operational reporting, but it is not sufficient for accepting a kernel optimization. The optimization program needs a frozen-token replay in addition to the full training run.

What the existing ROCm work established

The proposal builds on the existing stack rather than replacing it:

• #316 added the ROCm-aware build and registration path.
• #319, #325, #328, and #347 established strict Attention, FFN, vocab-parallel logp, and linear-logp operators.
• #356, #357, and #380 established deterministic RCCL transport and the single-node HIP IPC fixed-rank tree.
• #390, #393, #394, and #396 removed dense-KV reconstruction, host metadata synchronization, temporary projection copies, stale graph-cache reuse, scalar-FMA GEMM leaves, and shape-dependent CK dispatch.
• #400 published the 200-step reference result and made the remaining gap measurable.
• [vllm-project/vime: 423] supplied the narrow provider ABI while leaving token layout, masking, loss composition, and orchestration under VIME ownership.

These results also narrow the remaining work:

1. perf(rocm): MFMA batch-invariant GEMM and chunked Triton attention for strict R/R #396 reduced large-M deterministic GEMM to roughly 1.2-1.4x hipBLASLt on the measured Qwen3 projections, but small decode QKV and down projections remain more than 2x the vendor GEMM latency.
2. The perf(rocm): MFMA batch-invariant GEMM and chunked Triton attention for strict R/R #396 Triton Attention core measured 68 us for four 7,168-token decode requests, compared with 316 us for the fixed-M128 CK path. The remaining rollout gap therefore cannot be assigned to the Attention core alone; the complete captured layer stack must be profiled.
3. Training CP2 still gathers Q/K/V and computes the complete sequence on both CP ranks. The strict Attention result is row-invariant, so duplicated query work is not required for exactness.
4. ROCm linear logp either consumes materialized local logits or creates them through a deterministic LM-head GEMM before a separate fixed-tree logp reduction. Both actor and reference scoring show a material end-to-end deficit.
5. The HIP IPC collective is already faster than native RCCL for the measured 4 KiB and 64 KiB messages, but 16 MiB deterministic AllReduce and ReduceScatter remain slower. The current API also reports no asynchronous overlap or compute/communication fusion.

Goals

1. Preserve bitwise train-rollout selected-token log-probability equality under the tested Qwen3-8B TP4/CP2 contract.
2. Make strict R/R at least 5% faster than native P/P on a frozen matched-work end-to-end benchmark.
3. Make the performance attribution reviewable: every claimed improvement must name the affected kernels, data movement, graph boundaries, and phase timer.
4. Keep VIME and vLLM changes limited to stable provider/runtime hooks. Kernels, arithmetic contracts, routing, and provenance remain owned by RL-Kernel.
5. Keep native execution unchanged when RL-Kernel is disabled.

Non-goals

• Relaxing equality to a tolerance-based criterion.
• Reusing native RCCL floating-point reduction when its parenthesization is not contractually fixed.
• Dynamic Split-KV or autotuning over bit-relevant reduction axes.
• Multi-node or RDMA performance in this RFC.
• FP8, quantized LM head, MoE, MLA, or non-Qwen3 model coverage.
• Changing reward computation, optimizer behavior, loss composition, sampling policy, or VIME scheduling ownership.
• Treating a single short run, a degraded node, or unmatched generated-token counts as the final performance claim.

Contract boundary

Each accepted optimization belongs to one of two categories:

Contract-preserving implementation change

The following may vary after exact-byte qualification:

• CTA/wave allocation along independent rows;
• workgroup count, pipeline stages, and prefetch distance;
• graph partitioning and launch aggregation;
• scratch-buffer ownership and output placement;
• transport mechanism for exact copies;
• overlap between independent work;
• offline selection among implementations that already produce identical bytes.

Contract-versioned arithmetic change

The following require a new contract ID and complete requalification on both training and rollout:

• K or KV chunk size when it changes the reduction grouping;
• split count or split boundaries;
• MFMA instruction form or kpack;
• FP32/BF16 rounding points;
• online-softmax merge tree;
• rank-reduction order;
• exp/log implementation;
• masked and padded-lane contribution rules.

Strict mode rejects an unrecognized contract, mixed contract IDs between engines, missing provenance, or any fallback.

Proposed architecture

Optimization map

flowchart TB
    B["#400 baseline: exact R/R, slower than P/P"] --> P["Paired rocprof + frozen-token replay"]
    P --> E["Forward-scoped execution plan"]

    E --> R["Rollout decode"]
    E --> T["Actor and reference scoring"]
    E --> C["Training compute"]

    R --> RG["Small-M MFMA schedules"]
    R --> RA["Fixed-chunk paged Attention"]
    R --> RH["Full HIP Graph with stable IPC state"]

    T --> LF["Fused LM-head + local logp statistics"]
    LF --> LM["Rank-ordered TP merge"]
    LM --> LO["Selected logp / optional entropy"]

    C --> CQ["Keep CP-local query rows"]
    CQ --> CK["Gather K/V in fixed global order"]
    CK --> CA["Local-row Attention forward/backward"]

    RG --> D["Direct-output buffers"]
    RA --> D
    RH --> D
    CA --> X["Chunked deterministic HIP IPC"]
    D --> X
    X --> G["Exactness and provenance gate"]
    LO --> G
    G --> O["R/R performance acceptance"]

Loading

Steady-state execution

flowchart TB
    I["Frozen inputs, weights, topology, contract IDs"] --> B["Build one execution plan per shape bucket"]
    B --> W["Warm kernels, allocate stable buffers, capture graphs"]

    W --> RD["Rollout: RMSNorm -> MFMA projections -> RoPE"]
    RD --> PA["Paged fixed-chunk Attention -> direct O-projection"]
    PA --> AR["Device-resident fixed-tree TP reduction"]

    W --> TD["Training: CP-local Q plus rank-ordered K/V"]
    TD --> AT["Local-row Attention and packed FFN"]
    AT --> BW["Backward with fixed-order chunk overlap"]

    W --> LS["Actor/ref: streaming LM-head tiles"]
    LS --> ST["Local max/sumexp/target statistics"]
    ST --> TM["Rank-ordered TP merge"]

    AR --> V["Runtime readback outside the hot path"]
    BW --> V
    TM --> V
    V --> Q{"Contract and exactness valid?"}
    Q -->|yes| N["Advance step"]
    Q -->|no| F["Fail closed; no native fallback"]

Loading

Detailed design

Workstream 0: reproducible attribution and performance budget

Before changing kernels, add a profiler mode to the #400 harness.

• Record rocprofv3 HIP API, kernel, and memory-copy traces for every TP rank during one representative prefill, a bounded decode window, actor scoring, reference scoring, actor forward/backward, and weight update.
• Add NVTX/ROCTX ranges with stable names for Attention, QKV, O-projection, gate/up, down, LM head, logp merge, CP transport, TP reduction, graph replay, and framework wait.
• Export per-shape kernel counts, total device time, launch gaps, achieved occupancy, LDS use, HBM traffic, and XGMI traffic.
• Alternate P/P and R/R process pairs on the same node. Reject a pair when the health probe, copy bandwidth, collective latency, clock state, or KFD process inventory moves outside the recorded healthy envelope.
• Add a frozen-token replay containing the same prompt tokens, response tokens, active masks, page tables, sequence lengths, and pre-update weights for both routes.
• Keep the free-running 200-step experiment for realism and exactness, but use frozen replay for performance acceptance.

Deliverable: one checked-in Markdown/JSON report that attributes at least 95% of GPU device time and reconciles phase timers with wall-clock step time.

Workstream 1: small-M rollout execution plan

Target files:

• rl_engine/integrations/vllm_runtime.py
• rl_engine/kernels/ops/rocm/attention/strict_runtime.py
• rl_engine/kernels/ops/rocm/matmul/det_gemm.py
• rl_engine/kernels/ops/triton/matmul/mfma_gemm.py
• rl_engine/kernels/ops/triton/attention/chunked_flash_attn.py

Changes:

1. Build an immutable RocmStrictExecutionPlan per graph bucket. It owns prepared weights, paged metadata, RoPE tables, output slices, collective staging, scratch buffers, and provenance constants.
2. Resolve backend selection and validate static invariants once at plan construction. Dynamic validation remains device-side and is summarized after replay; it must not synchronize the host per layer.
3. Specialize bit-neutral MFMA launch axes for the measured decode row buckets. Preserve the current K walk, MFMA form, chunk tree, and BF16 boundary.
4. Route QKV, O-projection, gate/up, and down results directly into the next contract-owned buffer. A fused epilogue must perform any required BF16 round explicitly before the next operation observes the value.
5. Promote the fixed-chunk Triton Attention path only after long-run qualification. Its monolithic and split schedules may be selected by shape only if they remain byte-identical under the same contract.
6. Capture the full decoder-layer sequence, including stateful HIP IPC operations, using device-visible generation counters and stable addresses. An eager boundary remains only where a true framework dependency requires it.
7. Report graph-break count, kernel-launch count per generated token, and device idle time as acceptance metrics.

Workstream 2: CP-local-row training Attention

The current CP2 route gathers Q/K/V and computes full-sequence Attention on every CP rank. Replace it with a row-owned algorithm:

1. Keep Q in its local zigzag ownership order.
2. Transport K/V chunks in fixed global block order. Exact copies may use HIP IPC or RCCL transport; no floating-point reduction occurs during the gather.
3. Compute output and LSE only for locally owned query rows. Consume KV blocks in the same ascending contract order used by rollout.
4. In backward, compute local dQ and rank-local dK/dV partials from the owned query rows.
5. Merge dK/dV partials with the deterministic fixed-rank tree and return only the owner shard.
6. Pipeline transport of KV chunk n+1 with computation of chunk n on a second stream. The merge order remains fixed even if transport arrival order differs.
7. Preserve final-write-only downcast and the existing Attention contract readback.

At CP2, this removes duplicated query work and the Q all-gather. It also reduces the O(S^2) workspace pressure observed in deterministic Attention backward. The exact speedup is a measurement result, not assumed by this RFC.

Workstream 3: fused ROCm LM-head and linear logp

The 200-step result shows a 66.8% actor logp time increase and a 113.4% reference logp time increase. The strict ROCm path should own projection and vocabulary statistics as one operation, following the same boundary already established for the CUDA linear-logp design in #332.

Proposed operator:

rocm_fused_linear_logp(
    hidden,
    lm_head_shard,
    targets,
    vocab_partition,
    temperature,
    return_entropy,
    return_logits,
    requires_grad
) -> logp[, lse, entropy, logits]

Implementation requirements:

• Use the qualified MFMA K order for each local vocabulary tile.
• Produce local (max, sumexp, target_logit[, entropy_stat]) without storing [tokens, vocab_local] when no downstream consumer requires logits.
• Keep a serving entry that optionally stores the unscaled logits required by sampling. logp and lse must be byte-identical whether return_logits is enabled or disabled.
• Merge vocabulary tiles and TP ranks with an explicit fixed tree owned by RL-Kernel, not a framework sum, logsumexp, or native reduction.
• Apply temperature, padded-vocabulary masking, and the final clamp at the same contract-defined points on training and rollout.
• Provide a forward-only reference path and an autograd actor path. Backward may recompute fixed tiles from saved LSE to avoid retaining full logits; its correctness and memory tradeoff must be benchmarked separately.
• Reuse prepared [K,N] weights and refresh them by weight-generation ID after IPC update.
• Extend the VIME provider hook from 423 at the pre-projection boundary so that strict scoring can skip native logits materialization. This is a companion VIME change, not a transfer of training-loop ownership.

Workstream 4: deterministic communication and overlap

Target files:

• csrc/rocm/distributed/deterministic_collective.hip
• rl_engine/distributed/rocm_collectives.py
• rl_engine/distributed/collectives.py
• rl_engine/kernels/ops/pytorch/ffn/ffn.py

Changes:

1. Add a vectorized large-payload HIP IPC path that reduces only the destination shard and uses a topology-aware but fixed rank tree.
2. Use multiple persistent staging slots so the producer can publish tile n+1 while the consumer reduces tile n.
3. Preserve system-scope release/acquire generations and add an explicit completion event before a buffer can be reused.
4. Tune size thresholds offline. Runtime selection may depend on byte size, alignment, world size, and topology, but every candidate must implement the same rank parenthesization.
5. Overlap only independent dependencies:
   • gate-lane reduction with up-lane backward compute;
   • CP K/V chunk transport with local-row Attention;
   • completed projection tiles with fixed-tree reduction when each element's GEMM and rank-reduction boundaries remain unchanged.
6. Keep transport-only RCCL as the multi-node and unsupported-shape fallback. Strict mode must never use RCCL's floating-point reduction.

The weight-update path remains unchanged until profiling identifies a regression greater than noise; #400 shows it is already within 0.7% of native.

Workstream 5: integration, dispatch, and provenance

• Introduce one versioned capability record for the combined ROCm fast path: architecture, operator contract IDs, graph policy, collective tree ID, supported shapes, and fallback status.
• Cache only immutable or generation-keyed state. Weight, page-table, sequence-length, or topology changes invalidate the affected plan.
• Keep tuning tables checked in and reviewable. Runtime autotuning is prohibited in strict mode.
• Add a shadow mode that runs the new and current strict implementation on sampled layers or recorded fixtures and compares exact bytes before the new path becomes default.
• Maintain an explicit current-strict fallback for debugging. A fallback may be selected before a run starts; it may not occur silently inside an accepted strict run.

Implementation sequence

flowchart TB
    P0["PR0: profiler and frozen replay"] --> P1["PR1: rollout plan and small-M MFMA"]
    P0 --> P2["PR2: CP-local-row Attention"]
    P0 --> P3["PR3: fused linear logp"]
    P0 --> P4["PR4: large-message IPC and overlap"]
    P1 --> P5["PR5: integrated 30-step gate"]
    P2 --> P5
    P3 --> P5
    P4 --> P5
    P5 --> P6["PR6: 200-step release evidence"]

Loading

Each implementation PR must contain an isolated base/head benchmark, contract tests, runtime readback, and a statement of which final RFC gate it advances. The 30-step integration PR is the first point at which combined end-to-end speed may be claimed. The final 200-step PR is the release decision.

Validation ladder

Level	Required check	Acceptance
L0	Static contract and build	gfx942 build; no unversioned bit-relevant knob; no CUDA behavior change
L1	Repeated operator execution	exact bytes across repetitions and graph/eager execution
L2	Shape invariance	exact rows across M buckets, padding, mixed batches, page layouts, and allowed launch schedules
L3	Distributed operator	exact outputs at TP/CP 1/2/4/8 where supported; fixed rank order proved by readback
L4	Engine replay	same frozen tokens and weights through Megatron and vLLM; mismatch_count=0, max_abs_diff=0
L5	30-step VIME integration	all strict readbacks present; no fallback; phase performance gates pass
L6	200-step release run	zero mismatches for every step and final end-to-end performance target passes

For a new arithmetic contract version, L1-L4 must compare both engines against frozen golden fingerprints. Equality to the previous R/R contract is not required, but mixing old and new contract IDs is forbidden.

Performance methodology and acceptance criteria

Hardware and software controls

• One healthy 8xMI300X node; no unrelated GPU processes or retained KFD queues.
• Fixed GPU clocks/power policy where the environment permits it.
• Identical RL-Kernel, VIME, vLLM, Megatron, PyTorch, ROCm, AITER/CK, model, tokenizer, dataset, and launcher revisions for paired arms.
• Same TP4/CP2/PP1 topology, graph mode, memory utilization, prompt/token fixture, masks, and weights.
• Warmup is excluded by a declared rule applied equally to both arms.
• At least three independently launched paired repetitions for the frozen replay; P/P and R/R order alternates.

Required performance gates

1. Primary end-to-end gate: strict R/R median step time is at most 0.95x native P/P on the frozen matched-work replay, and the paired 95% bootstrap interval for the R/R/P/P ratio remains below 1.0.
2. Rollout gate: strict rollout tokens/GPU/s is at least 1.05x native on the same frozen sequence-length distribution.
3. Training gate: actor-train tokens/s and reference-logp tokens/GPU/s are each at least at parity with native; no phase may hide a regression behind shorter generated responses.
4. Communication gate: no regression greater than 5% for 4 KiB-16 MiB AllReduce, AllGather, or ReduceScatter; the measured production message sizes must improve.
5. Tail gate: R/R step-time p95 is no worse than native P/P p95.
6. Correctness gate: all 200 R/R steps satisfy mismatch_count=0, max_abs_diff=0, and torch.equal=true, with complete route provenance and no fallback.

The final report must publish raw, warmup-excluded token-normalized, equal-length, and frozen-replay views. A free-running reward or response-length difference is reported but is not used as evidence of a kernel speedup.

Failure handling and fallback

• Unsupported dtype, head dimension, page size, topology, graph bucket, or contract combination fails during plan construction.
• A plan invalidated by a weight or metadata generation change is rebuilt before execution.
• A runtime provenance mismatch, stale buffer generation, collective timeout, or exactness failure aborts strict execution and preserves the diagnostic artifact.
• Production/native fallback remains available only when strict mode was not requested.
• During rollout, the new path remains opt-in until L6 passes. The current docs(rocm): publish G10/G11 200-step RL results #400 strict route is retained as the bisect reference for one release cycle.

Risks

1. Overfitting to one shape. Decode and training use different M/K/V regimes. The tuning table must cover the observed distribution and retain a safe strict implementation for qualified outliers.
2. Graph state races. Stateful IPC collectives inside HIP Graphs require device-visible generations, stable storage, and explicit reuse barriers. Stress tests must include graph replay after shrinking and growing batches.
3. Hidden arithmetic drift from fusion. Fusion may remove a BF16 store/load boundary or change an FP32 tree. Required boundaries must be represented explicitly and protected by golden fingerprints.
4. Profiler distortion. Kernel tracing changes launch timing. Claims use unprofiled synchronized wall time; traces are used only for attribution.
5. Vendor-version drift. AITER, CK, Triton, and ROCm upgrades may change code generation. Source/version fingerprints and post-upgrade replay qualification are mandatory.
6. Node degradation. Stale vLLM GPU contexts previously inflated all phases. Health gates must reject such runs before results are published.

Open questions for review

1. Should the first accepted fast path keep the current triton_mfma_rocm.v1 and triton_chunked_flash_attention.v1 arithmetic IDs, or introduce one combined ROCm v2 contract before operator fusion?
2. For CP-local-row backward, should dK/dV use fixed-tree ReduceScatter directly to owners, or a chunked AllReduce followed by a view when that wins on MI300X?
3. Can the VIME pre-projection provider hook bypass native local-logits materialization without changing any non-RL-Kernel path?
4. Which gfx942 performance counters and XGMI topology signals should AMD treat as authoritative for the rollout and large-message collective workstreams?
5. Should the final release gate require the 5% end-to-end win on both CK and Triton Attention routes, or only on the selected default strict route while the other remains a diagnostic backend?

Feedback period

The proposed review period is seven calendar days. Workstream 0 is non-contractual and may begin during review. Changes to arithmetic contracts, the VIME pre-projection hook, or the default strict route should wait for maintainer agreement on the questions above.

CC

@Flink-ddd @inaniloquentee @KJLdefeated @bitborne @EthanZero2Hero @frank-2077 @maxiaosong1124 @hihaluemen

References

• ROCm 200-step result and artifacts: #400
• ROCm performance stack: #390, #393, #394, #396
• ROCm ablation and integration: #385, #388
• ROCm deterministic communication: #356, #357, #380
• ROCm operators and platform support: #316, #319, #325, #328, #347
• VIME provider integration: [vllm-project/vime: 423]
• CUDA linear-logp contract reference: #332