fix(rocm): stabilize direct paged CK attention for bitwise R/R

csrc/rocm/attention/fixed_paged_ck.hip

@@ -0,0 +1,116 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+// Instantiate the installed CK implementation with a fixed arithmetic schedule.
+// CK headers remain external; no AITER/vLLM source or dispatcher is patched.
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <array>
+#include "fmha_fwd.hpp"
+
+#ifndef RLK_CK_TILE_M
+#define RLK_CK_TILE_M 128
+#endif
+
+namespace {
+template <bool Causal, bool HasLSE, ck_tile::BlockAttentionKVCacheLoadModeEnum LoadMode>
+void launch_fixed(fmha_batch_prefill_args a, hipStream_t stream) {
+    using Config = FmhaFwdTypeConfig<FmhaFwdBf16>;
+    constexpr int M = RLK_CK_TILE_M;
+    constexpr int K = M == 64 ? 64 : 32;
+    constexpr int W = M == 64 ? 16 : 32;
+    using Shape = ck_tile::TileFmhaShape<
+        ck_tile::sequence<M, 128, K, 128, K, 128>,
+        ck_tile::sequence<4, 1, 1>, ck_tile::sequence<W, W, 16>,
+        ck_tile::sequence<4, 1, 1>, ck_tile::sequence<W, W, 16>, true>;
+    using Traits = ck_tile::TileFmhaBatchPrefillTraits<
+        true, true, true, true, false, ck_tile::BlockAttentionBiasEnum::NO_BIAS,
+        false, HasLSE, false, ck_tile::BlockAttentionQuantScaleEnum::NO_SCALE,
+        -1, false, false, 16,
+        ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::LINEAR_LAYOUT,
+        ck_tile::BlockAttentionKVCacheLookupTableEnum::VLLM_BLOCK_TABLE_2D, LoadMode>;
+    using Problem = ck_tile::BlockFmhaBatchPrefillPipelineProblem<
+        Config::QDataType, Config::KDataType, Config::VDataType, Config::SaccDataType,
+        Config::SMPLComputeDataType, Config::BiasDataType, Config::RandValOutputDataType,
+        Config::LSEDataType, Config::PDataType, Config::OaccDataType, Config::ODataType,
+        Shape, true, ck_tile::ComposedAttention<0, CK_TILE_FMHA_FWD_FAST_EXP2>,
+        ck_tile::SimplifiedGenericAttentionMask<Causal>, false, 16, Traits>;
+    using Pipeline = ck_tile::BlockFmhaBatchPrefillPipelineQRKSVSAsync<Problem>;
+    using Epilogue = ck_tile::Default2DEpilogue<ck_tile::Default2DEpilogueProblem<
+        Config::OaccDataType, Config::ODataType, true, true>>;
+    using Kernel = ck_tile::FmhaBatchPrefillWithPagedKVCacheKernel<Pipeline, Epilogue>;
+    auto [kargs, grid] = fmha_batch_prefill_create_kargs_and_grids<Kernel>(a);
+    ck_tile::launch_kernel(ck_tile::stream_config{stream},
+        ck_tile::make_kernel<Kernel::kBlockPerCu>(Kernel{}, grid, Kernel::BlockSize(), 0, kargs));
+}
+
+template <bool Causal, bool HasLSE>
+void dispatch_load(fmha_batch_prefill_args a, hipStream_t stream) {
+    using Mode = ck_tile::BlockAttentionKVCacheLoadModeEnum;
+    // Keep 64-bit global addressing when either interleaved view exceeds 2GB.
+    const auto k_bytes = int64_t(a.num_total_pages) * a.batch_stride_k * 2;
+    const auto v_bytes = int64_t(a.num_total_pages) * a.batch_stride_v * 2;
+    if (k_bytes > INT32_MAX || v_bytes > INT32_MAX)
+        launch_fixed<Causal, HasLSE, Mode::GLOBAL_LOAD_LDS>(a, stream);
+    else
+        launch_fixed<Causal, HasLSE, Mode::BUFFER_LOAD>(a, stream);
+}
+} // namespace
+
+// Pointer-only binding avoids importing the Torch C++ ABI into the CK build.
+// Shape, dtype, device and metadata checks live at the RL-Kernel entry point.
+void forward(const std::array<uint64_t, 9>& ptr,
+             const std::array<int, 16>& dims,
+             float scale, bool causal, bool has_lse) {
+    fmha_batch_prefill_args a{};
+    a.q_ptr = reinterpret_cast<void*>(ptr[0]);
+    a.k_ptr = reinterpret_cast<void*>(ptr[1]);
+    a.v_ptr = reinterpret_cast<void*>(ptr[2]);
+    a.o_ptr = reinterpret_cast<void*>(ptr[3]);
+    a.lse_ptr = has_lse ? reinterpret_cast<void*>(ptr[4]) : nullptr;
+    a.seqstart_q_ptr = reinterpret_cast<void*>(ptr[5]);
+    a.kv_page_indices = reinterpret_cast<void*>(ptr[6]);
+    a.seqlen_k_ptr = reinterpret_cast<void*>(ptr[7]);
+    auto stream = reinterpret_cast<hipStream_t>(ptr[8]);
+    a.seqlen_q = dims[0];
+    a.batch = dims[1];
+    a.max_seqlen_q = dims[2];
+    a.nhead_q = dims[3];
+    a.nhead_k = dims[4];
+    a.num_total_pages = dims[5];
+    a.batch_stride_block_table = dims[6];
+    a.stride_q = dims[7];
+    a.stride_o = dims[7];
+    a.stride_k = dims[8];
+    a.stride_v = dims[9];
+    a.nhead_stride_k = dims[10];
+    a.nhead_stride_v = dims[11];
+    a.batch_stride_k = dims[12];
+    a.batch_stride_v = dims[13];
+    a.nhead_stride_q = dims[14];
+    a.nhead_stride_o = dims[14];
+    a.nhead_stride_lse = dims[0];
+    a.hdim_q = 128;
+    a.hdim_v = 128;
+    a.page_block_size = 16;
+    a.kv_memory_layout = ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::LINEAR_LAYOUT;
+    a.kv_lookup_table = ck_tile::BlockAttentionKVCacheLookupTableEnum::VLLM_BLOCK_TABLE_2D;
+    a.scale_s = scale;
+    a.scale_p = 1;
+    a.scale_o = 1;
+    a.window_size_left = -1;
+    a.window_size_right = causal ? 0 : -1;
+    a.mask_type = static_cast<int>(causal ? ck_tile::GenericAttentionMaskEnum::MASK_FROM_BOTTOM_RIGHT
+                                         : ck_tile::GenericAttentionMaskEnum::NO_MASK);
+    if (causal) {
+        if (has_lse) dispatch_load<true, true>(a, stream);
+        else dispatch_load<true, false>(a, stream);
+    } else {
+        if (has_lse) dispatch_load<false, true>(a, stream);
+        else dispatch_load<false, false>(a, stream);
+    }
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &forward);
+    m.attr("tile_m") = RLK_CK_TILE_M;
+}

examples/vime_rocm_attention_ablation/README.md

@@ -33,14 +33,12 @@ Those rows describe one-at-a-time root-cause probes; most still need concrete
 runtime mutation and restoration hooks. No row label is treated as execution
 evidence here.
 
-## Why every arm uses eager vLLM
+## HIP Graph execution
 
-All four arms pass `--vllm-enforce-eager`. The current strict ROCm QKV/O
-projection uses a fixed-tree collective with Python/lock bookkeeping, so it is
-not a valid vLLM fullgraph capture target. Freezing eager mode across `P` and
-`R` arms makes this a correctness-first, one-factor experiment. A future graph
-path should be introduced as a separate controlled variable, not enabled only
-for some arms.
+The launcher uses `FULL_AND_PIECEWISE` graph mode for every arm, with a maximum
+capture size of 32. Stateful strict collectives remain at piecewise graph
+boundaries. ROCm uses PyTorch's `torch.cuda` graph API to capture HIP work.
+Attention routes and fixed-paged settings participate in the graph cache key.
 
 ## Prerequisites
 
@@ -63,9 +61,11 @@ The default formal topology reuses PR #377's colocated eight-GPU schedule:
 - strict Logp vocabulary layout: 151936 real rows, padded to 152064 rows for TP4.
 
 The batch defaults are deliberately small correctness settings for the strict
-path and may be overridden consistently on the CLI. On the rollout side, the current strict route logically gathers the
-vLLM paged KV layout before invoking dense AITER Attention; it does not claim a
-native paged Attention kernel or publishable performance numbers.
+path and may be overridden consistently on the CLI. Supported strict rollout
+calls consume the vLLM paged KV cache directly. Readbacks report
+`dense_kv_materialized` so a materialized fallback can be distinguished from
+this direct path. See [fixed paged CK attention](fixed_paged_ck.md) for the
+opt-in arithmetic schedule, supported shapes, and one-round long-workload runner.
 
 The command refuses to reuse a non-empty run directory or an already-running
 Ray cluster. Each arm receives its own dump, readback, and log directory. A
@@ -188,3 +188,15 @@ python examples/vime_rocm_attention_ablation/validate_artifacts.py \
 Do not add `matrix-plan.json`, validation JSON, rollout dumps, mismatch
 sidecars, checkpoints, or MI300X result files to the PR. Publish them as CI/job
 artifacts when needed.
+
+## Full-native PR377 workload
+
+The standalone P/P runner selects production attention, FFN, and logp on both
+the Megatron and vLLM sides. It uses actor TP4/CP2, two TP4 rollout engines,
+round-robin routing, eight samples, a 7168-token response limit, and Vime's
+rollout-logprob framework consistency mode. Three rounds are the default:
+
+```bash
+python -m examples.vime_rocm_attention_ablation.run_full_pp_pr377_workload \
+  --run-dir /app/model/vime-runs/pr394-full-native-pp-vime-tis
+```

examples/vime_rocm_attention_ablation/bench_paged_dispatch.py

@@ -0,0 +1,104 @@
+"""Compare current AITER and fixed CK launches on identical paged attention inputs."""
+
+from __future__ import annotations
+
+import json
+import statistics
+import time
+from functools import partial
+
+import torch
+from examples.vime_rocm_attention_ablation.probe_paged_dispatch import cache_for, packed_forward
+from rl_engine.kernels.ops.rocm.attention.fixed_paged_ck import fixed_paged_prefill
+from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore
+
+
+def measure(fn):
+    for _ in range(5):
+        fn()
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    for _ in range(50):
+        fn()
+    torch.cuda.synchronize()
+    eager_ms = (time.perf_counter() - start) * 1000 / 50
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            fn()
+    torch.cuda.current_stream().wait_stream(stream)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=stream):
+        for _ in range(20):
+            out = fn()
+    measurements = []
+    for _ in range(7):
+        begin, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+        begin.record()
+        for _ in range(10):
+            graph.replay()
+        end.record()
+        end.synchronize()
+        measurements.append(begin.elapsed_time(end) / 200)
+    assert bool(torch.isfinite(out).all())
+    return dict(
+        eager_ms=eager_ms,
+        graph_ms=statistics.median(measurements),
+        graph_min_ms=min(measurements),
+        graph_max_ms=max(measurements),
+    )
+
+
+@torch.inference_mode()
+def main():
+    for length, qlen, batch in [(4096, 4096, 1), (7168, 128, 4), (7168, 1, 4)]:
+        torch.manual_seed(1234)
+        q = torch.randn((batch * qlen, 8, 128), dtype=torch.bfloat16, device="cuda")
+        k = torch.randn((batch, 2, length, 128), dtype=torch.bfloat16, device="cuda")
+        v = torch.randn_like(k)
+        for layout in ("interleaved", "large"):
+            kc, vc, table = cache_for(k, v, layout, True)
+            cuq = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * qlen
+            indptr = torch.arange(batch + 1, device="cuda", dtype=torch.int32) * table.shape[1]
+            seqs = torch.full((batch,), length, device="cuda", dtype=torch.int32)
+            core = StrictRocmAiterCKAttentionCore()
+            original = core._mha_batch_prefill
+            for name, entry in [
+                ("aiter_dynamic", original),
+                ("fixed_m128", partial(fixed_paged_prefill, tile_m=128)),
+                ("fixed_m64", partial(fixed_paged_prefill, tile_m=64)),
+            ]:
+                core._mha_batch_prefill = entry
+                call = partial(
+                    packed_forward,
+                    core,
+                    q,
+                    kc,
+                    vc,
+                    table,
+                    seqs,
+                    cuq,
+                    indptr,
+                    qlen,
+                    length,
+                    causal=qlen > 1,
+                    lse=qlen > 1,
+                )
+                print(
+                    json.dumps(
+                        dict(
+                            route=name,
+                            length=length,
+                            qlen=qlen,
+                            batch=batch,
+                            layout=layout,
+                            **measure(call),
+                        )
+                    ),
+                    flush=True,
+                )
+
+
+if __name__ == "__main__":
+    main()

examples/vime_rocm_attention_ablation/fixed_paged_ck.md

@@ -0,0 +1,62 @@
+# Fixed paged CK attention on ROCm
+
+Set `RL_KERNEL_ROCM_FIXED_PAGED_TILE=128` before starting training and rollout
+workers to select the RL-Kernel-owned CK instantiation. `64` selects the other
+fixed schedule; `0` (default) retains the AITER entrypoint. This path reuses
+installed `aiter_meta`/Composable Kernel headers and requires a ROCm C++ build
+toolchain. The extension builds lazily and must be warmed before HIP Graph
+capture. It was validated on MI300X VF (`gfx942`).
+
+The supported specialization is BF16, head dimension 128, page size 16, with
+no dropout, softcap, bias, or sliding window. Both training and rollout use
+the same fixed tile. Implicit scalar FMA contraction is disabled in this
+extension to preserve rounding across differently sized query chunks;
+explicit CK MFMA instructions remain enabled. Cache addressing stays 64-bit
+when either K/V view spans more than 2 GB.
+
+The adapter pads page tables to full 128-token KV tiles and sanitizes unused
+columns and inactive graph rows before CK can load their page IDs. Decode
+query rows retain their request mapping as the active batch shrinks. These
+operations prepare metadata without copying the dense KV tensor. A
+materialized fallback remains available for other adapter call layouts and
+uses the fixed arithmetic when enabled.
+
+`RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS` optionally bounds the scheduling length.
+It must cover the complete prompt-plus-response KV length; an asynchronous
+device assertion rejects lengths above the bound. Leave it unset for general
+workloads. Changing either setting changes the graph cache identity; restart
+workers when changing settings.
+
+## One-round R/R reproduction
+
+In the validated AMD container, install this checkout at
+`/workspace/RL-Kernel-pr390`, with Vime at `/workspace/vime`, Megatron at
+`/workspace/Megatron-LM-vime`, and the model/data paths used by
+`run_full_rr_single_arm_v90.py`. From the repository root, run:
+
+```bash
+python -m examples.vime_rocm_attention_ablation.run_full_rr_fixed_paged_ck \
+  --tile 128 --run-dir /app/model/vime-runs/fixed-paged-ck-rr
+```
+
+Use a new output directory. This runs one R/R rollout/training iteration,
+with 8 GPUs, training TP4/CP2, two TP4 rollout engines, round-robin routing,
+Qwen3-8B BF16, batch 1 with 8 samples, global batch 8, response limit 7168,
+4096 tokens per training GPU, seed 1234, and KV scheduling bound 8192.
+FFN and logp remain R/R. The helper is also the shared dependency of the
+existing 200-round R/R runner; adding it does not enable fixed tiles in that
+runner by default.
+
+## Diagnostic probes
+
+```bash
+python -m examples.vime_rocm_attention_ablation.probe_paged_dispatch \
+  --fixed-tile 128 --guard-pages --graph
+python -m examples.vime_rocm_attention_ablation.bench_paged_dispatch
+```
+
+The probe prints raw BF16 bit comparisons on identical input data. It covers
+full/suffix queries, physical KV layouts, and dynamic graph replay. The
+benchmark compares AITER dynamic dispatch against fixed CK on the same
+inputs. Neither replaces full training/rollout logprob validation. Performance
+depends on cache layout and query shape; fixed M128 is not uniformly faster.

examples/vime_rocm_attention_ablation/launch_arm.sh

@@ -53,8 +53,9 @@ if [[ "${RL_KERNEL_ATTENTION_CASE:-}" != "${RLK_ABLATION_CASE_ID}" ]]; then
   echo "RL_KERNEL_ATTENTION_CASE disagrees with the arm ID" >&2
   exit 2
 fi
-if [[ "${RL_KERNEL_FFN_CASE:-}" != "R/R" || "${RL_KERNEL_LOGP_CASE:-}" != "R/R" ]]; then
-  echo "the strict dense matrix requires FFN=R/R and Logp=R/R" >&2
+if [[ "${RL_KERNEL_FFN_CASE:-}" != "${RL_KERNEL_LOGP_CASE:-}" ]] ||
+   [[ "${RL_KERNEL_FFN_CASE:-}" != "R/R" && "${RL_KERNEL_FFN_CASE:-}" != "P/P" ]]; then
+  echo "FFN and Logp must both use P/P or both use R/R" >&2
   exit 2
 fi
 if [[ "${RL_KERNEL_VLLM_INTEGRATION:-}" != "1" ]]; then
@@ -201,7 +202,11 @@ names = [
     "RL_KERNEL_READBACK_DIR",
     "RL_KERNEL_MISMATCH_SIDECAR_DIR",
 ]
-print(json.dumps({"env_vars": {name: os.environ[name] for name in names}}))
+env_vars = {name: os.environ[name] for name in names}
+for name in ("RL_KERNEL_ROCM_FIXED_PAGED_TILE", "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS"):
+    if name in os.environ:
+        env_vars[name] = os.environ[name]
+print(json.dumps({"env_vars": env_vars}))
 PY
 )"
 
@@ -247,6 +252,18 @@ sleep 10
 
 # The RL-Kernel paged CK path and device-sequenced IPC collectives are captured
 # by vLLM's HIP graph runtime after adapter-owned warmup.
+LINEAR_LOGP_ARGS=()
+if [[ "${RL_KERNEL_LOGP_CASE%%/*}" == "R" ]]; then
+  LINEAR_LOGP_ARGS+=(
+    --linear-logp-provider
+    rl_engine.integrations.vime.linear_logp_provider.provider
+    --linear-logp-provider-mode strict
+  )
+fi
+ROLLOUT_LOGPROBS_ARGS=()
+if [[ "${RLK_ABLATION_USE_ROLLOUT_LOGPROBS:-0}" == "1" ]]; then
+  ROLLOUT_LOGPROBS_ARGS+=(--use-rollout-logprobs)
+fi
 ray job submit \
   --address="${ray_job_address}" \
   --runtime-env-json="${RUNTIME_ENV_JSON}" \
@@ -313,9 +330,8 @@ ray job submit \
   --attention-backend flash \
   --train-memory-margin-bytes 2147483648 \
   --no-gradient-accumulation-fusion \
-  --linear-logp-provider \
-  rl_engine.integrations.vime.linear_logp_provider.provider \
-  --linear-logp-provider-mode strict \
+  "${LINEAR_LOGP_ARGS[@]}" \
+  "${ROLLOUT_LOGPROBS_ARGS[@]}" \
   --get-mismatch-metrics \
   --custom-tis-function-path \
   vime_rocm_attention_ablation.tis_metrics.metrics_only_tis \