From b2209b86b0aa31d9c3b5834f1e57ead789367431 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 9 Sep 2026 10:45:54 +0800 Subject: [PATCH 1/7] fix(rocm): stabilize direct paged CK attention for bitwise R/R --- csrc/rocm/attention/strict_paged_ck.cu | 116 ++++++++ .../vime_rocm_attention_ablation/README.md | 20 +- .../bench_paged_dispatch.py | 104 +++++++ .../fixed_paged_ck.md | 62 +++++ .../launch_arm.sh | 6 +- .../probe_paged_dispatch.py | 256 ++++++++++++++++++ examples/vime_rocm_attention_ablation/run.py | 9 +- .../run_full_rr_fixed_paged_ck.py | 39 +++ .../run_full_rr_single_arm_v90.py | 90 ++++++ rl_engine/integrations/framework_operators.py | 112 +++++++- rl_engine/integrations/vllm_runtime.py | 4 +- .../ops/rocm/attention/fixed_paged_ck.py | 197 ++++++++++++++ .../kernels/ops/rocm/attention/flash_attn.py | 93 +++++-- 13 files changed, 1065 insertions(+), 43 deletions(-) create mode 100644 csrc/rocm/attention/strict_paged_ck.cu create mode 100644 examples/vime_rocm_attention_ablation/bench_paged_dispatch.py create mode 100644 examples/vime_rocm_attention_ablation/fixed_paged_ck.md create mode 100644 examples/vime_rocm_attention_ablation/probe_paged_dispatch.py create mode 100644 examples/vime_rocm_attention_ablation/run_full_rr_fixed_paged_ck.py create mode 100644 examples/vime_rocm_attention_ablation/run_full_rr_single_arm_v90.py create mode 100644 rl_engine/kernels/ops/rocm/attention/fixed_paged_ck.py diff --git a/csrc/rocm/attention/strict_paged_ck.cu b/csrc/rocm/attention/strict_paged_ck.cu new file mode 100644 index 00000000..2f714034 --- /dev/null +++ b/csrc/rocm/attention/strict_paged_ck.cu @@ -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 +#include +#include +#include "fmha_fwd.hpp" + +#ifndef RLK_CK_TILE_M +#define RLK_CK_TILE_M 128 +#endif + +namespace { +template +void launch_fixed(fmha_batch_prefill_args a, hipStream_t stream) { + using Config = FmhaFwdTypeConfig; + 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, + ck_tile::sequence<4, 1, 1>, ck_tile::sequence, + ck_tile::sequence<4, 1, 1>, ck_tile::sequence, 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, false, 16, Traits>; + using Pipeline = ck_tile::BlockFmhaBatchPrefillPipelineQRKSVSAsync; + using Epilogue = ck_tile::Default2DEpilogue>; + using Kernel = ck_tile::FmhaBatchPrefillWithPagedKVCacheKernel; + auto [kargs, grid] = fmha_batch_prefill_create_kargs_and_grids(a); + ck_tile::launch_kernel(ck_tile::stream_config{stream}, + ck_tile::make_kernel(Kernel{}, grid, Kernel::BlockSize(), 0, kargs)); +} + +template +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(a, stream); + else + launch_fixed(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& ptr, + const std::array& dims, + float scale, bool causal, bool has_lse) { + fmha_batch_prefill_args a{}; + a.q_ptr = reinterpret_cast(ptr[0]); + a.k_ptr = reinterpret_cast(ptr[1]); + a.v_ptr = reinterpret_cast(ptr[2]); + a.o_ptr = reinterpret_cast(ptr[3]); + a.lse_ptr = has_lse ? reinterpret_cast(ptr[4]) : nullptr; + a.seqstart_q_ptr = reinterpret_cast(ptr[5]); + a.kv_page_indices = reinterpret_cast(ptr[6]); + a.seqlen_k_ptr = reinterpret_cast(ptr[7]); + auto stream = reinterpret_cast(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(causal ? ck_tile::GenericAttentionMaskEnum::MASK_FROM_BOTTOM_RIGHT + : ck_tile::GenericAttentionMaskEnum::NO_MASK); + if (causal) { + if (has_lse) dispatch_load(a, stream); + else dispatch_load(a, stream); + } else { + if (has_lse) dispatch_load(a, stream); + else dispatch_load(a, stream); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &forward); + m.attr("tile_m") = RLK_CK_TILE_M; +} diff --git a/examples/vime_rocm_attention_ablation/README.md b/examples/vime_rocm_attention_ablation/README.md index 77207679..1dfa6c30 100644 --- a/examples/vime_rocm_attention_ablation/README.md +++ b/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 diff --git a/examples/vime_rocm_attention_ablation/bench_paged_dispatch.py b/examples/vime_rocm_attention_ablation/bench_paged_dispatch.py new file mode 100644 index 00000000..dc7f8276 --- /dev/null +++ b/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() diff --git a/examples/vime_rocm_attention_ablation/fixed_paged_ck.md b/examples/vime_rocm_attention_ablation/fixed_paged_ck.md new file mode 100644 index 00000000..16d6ba54 --- /dev/null +++ b/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. diff --git a/examples/vime_rocm_attention_ablation/launch_arm.sh b/examples/vime_rocm_attention_ablation/launch_arm.sh index 7d675741..378fa512 100644 --- a/examples/vime_rocm_attention_ablation/launch_arm.sh +++ b/examples/vime_rocm_attention_ablation/launch_arm.sh @@ -201,7 +201,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 )" diff --git a/examples/vime_rocm_attention_ablation/probe_paged_dispatch.py b/examples/vime_rocm_attention_ablation/probe_paged_dispatch.py new file mode 100644 index 00000000..d1707fc6 --- /dev/null +++ b/examples/vime_rocm_attention_ablation/probe_paged_dispatch.py @@ -0,0 +1,256 @@ +"""Isolate CK dispatch, physical KV addressing and HIP replay on identical data. + +This is a diagnostic, not an end-to-end bitwise/performance acceptance test. +The inflated max-Q variant identifies dispatch effects; it is not a proposed +production optimization because it also launches unused query blocks. +""" + +from __future__ import annotations + +import argparse +import json +from functools import partial + +import torch +from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore + + +def report(label, actual, expected, **context): + actual = actual.contiguous() + expected = expected.contiguous() + finite = bool(torch.isfinite(actual).all() & torch.isfinite(expected).all()) + mismatch = int((actual.view(torch.int16) != expected.view(torch.int16)).sum()) + print( + json.dumps( + dict( + check=label, + finite=finite, + bitwise=mismatch == 0, + mismatch=mismatch, + elements=actual.numel(), + max_abs=float((actual.float() - expected.float()).abs().max()), + **context, + ) + ), + flush=True, + ) + + +def invoke(label, fn, profile): + print(json.dumps(dict(begin=label)), flush=True) + result = fn() + torch.cuda.synchronize() + if profile: + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) as prof: + result = fn() + torch.cuda.synchronize() + kernels = sorted( + { + e.name + for e in prof.events() + if e.device_type == torch.autograd.DeviceType.CUDA and "fmha" in e.name.lower() + } + ) + print(json.dumps(dict(profile=label, kernels=kernels)), flush=True) + return result + + +def cache_for(k, v, layout, guard_pages): + batch, heads, length, dim = k.shape + page_size = 16 + pages = (length + page_size - 1) // page_size + count = batch * pages + 1 + if layout == "large": + count = max(count, (2**31 // (page_size * heads * 2 * dim * 2)) + batch * pages + 1) + if layout == "separate": + kc = torch.empty((count, page_size, heads, dim), device=k.device, dtype=k.dtype) + vc = torch.empty_like(kc) + else: + cache = torch.empty((count, page_size, heads, 2 * dim), device=k.device, dtype=k.dtype) + kc, vc = cache.split(dim, dim=-1) + # Reverse high physical page IDs; all referenced pages and tail bytes are initialized. + ids = torch.arange( + count - 1, count - 1 - batch * pages, -1, device=k.device, dtype=torch.int64 + ).reshape(batch, pages) + kp = torch.zeros((batch, pages * page_size, heads, dim), device=k.device, dtype=k.dtype) + vp = torch.zeros_like(kp) + kp[:, :length].copy_(k.transpose(1, 2)) + vp[:, :length].copy_(v.transpose(1, 2)) + kc.index_copy_(0, ids.flatten(), kp.reshape(-1, page_size, heads, dim)) + vc.index_copy_(0, ids.flatten(), vp.reshape(-1, page_size, heads, dim)) + kc[0].zero_() + vc[0].zero_() + table = ids.to(torch.int32).contiguous() + assert bool(((table >= 0) & (table < count)).all()) + assert torch.equal(kc[ids].reshape_as(kp)[:, :length], k.transpose(1, 2)) + assert torch.equal(vc[ids].reshape_as(vp)[:, :length], v.transpose(1, 2)) + if guard_pages: + # CK reads physical page IDs for a whole 128-token KV tile before + # applying the logical sequence mask. Fill out that tile with page 0. + padded = torch.zeros((batch, ((pages + 7) // 8) * 8), dtype=torch.int32, device=k.device) + padded[:, :pages].copy_(table) + table = padded + return kc, vc, table + + +def packed_forward( + core, packed, kc, vc, table, seqs, cuq, indptr, maxq, length, causal=True, lse=True +): + return core.forward_paged_varlen_with_lse( + packed, + kc, + vc, + page_table=table, + seqused_k=seqs, + cu_seqlens_q=cuq, + kv_indptr=indptr, + max_seqlen_q=maxq, + max_seqlen_k=length, + causal=causal, + scale=128**-0.5, + return_lse=lse, + ).out + + +@torch.inference_mode() +def run_case(args, seed, length, batch): + torch.manual_seed(seed) + core = StrictRocmAiterCKAttentionCore() + if args.fixed_tile: + from rl_engine.kernels.ops.rocm.attention.fixed_paged_ck import fixed_paged_prefill + + core._mha_batch_prefill = partial(fixed_paged_prefill, tile_m=args.fixed_tile) + q = torch.randn((batch, 8, length, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn((batch, 2, length, 128), device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + positions = torch.arange(length, device="cuda").expand(batch, -1).contiguous() + # The training core executes one logical sequence per invocation, even + # when rollout packs several requests into a single launch. + ref = ( + torch.cat( + [ + invoke( + f"train/{row}", + lambda row=row: ( + core.forward_with_lse( + q[row : row + 1], + k[row : row + 1], + v[row : row + 1], + causal=True, + scale=128**-0.5, + query_position_ids=positions[row : row + 1], + key_position_ids=positions[row : row + 1], + ).out + ), + args.profile, + ) + for row in range(batch) + ] + ) + .transpose(1, 2) + .contiguous() + ) + for layout in args.layouts: + kc, vc, table = cache_for(k, v, layout, args.guard_pages) + context = dict( + seed=seed, + length=length, + batch=batch, + layout=layout, + pool_bytes=kc.shape[0] * kc.stride(0) * kc.element_size(), + ) + torch.cuda.synchronize() + print( + json.dumps( + dict( + cache_validated=True, + min_page=int(table.min()), + max_page=int(table.max()), + **context, + ) + ), + flush=True, + ) + seqs = torch.full((batch,), length, dtype=torch.int32, device="cuda") + indptr = torch.arange(batch + 1, dtype=torch.int32, device="cuda") * table.shape[1] + for qlen in dict.fromkeys((length, min(length, 128), min(length, 17), 1)): + packed = q[:, :, -qlen:].transpose(1, 2).contiguous().reshape(batch * qlen, 8, 128) + cuq = torch.arange(batch + 1, dtype=torch.int32, device="cuda") * qlen + expected = ref[:, -qlen:].reshape_as(packed) + for maxq in dict.fromkeys((qlen, max(4096, length))): + call = partial( + packed_forward, core, packed, kc, vc, table, seqs, cuq, indptr, maxq, length + ) + label = f"{layout}/q{qlen}/maxq{maxq}" + actual = invoke(label, call, args.profile) + report("full_vs_suffix", actual, expected, qlen=qlen, maxq=maxq, **context) + if qlen == 1: + decode = invoke( + label + "/decode", lambda call=call: call(False, False), args.profile + ) + report("causal_lse_vs_decode", decode, actual, qlen=qlen, maxq=maxq, **context) + if args.graph and maxq == 1: + # Retain graph addresses while lengths and active requests change. + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + call(False, False) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + graph_out = call(False, False) + original_table = table.clone() + for active in range(batch, 0, -1): + for new_length in (length, max(1, length - 1), max(1, length - 16)): + table.copy_(original_table) + table[active:].zero_() + seqs.fill_(new_length) + seqs[active:].fill_(1) + eager = call(False, False).clone() + graph.replay() + torch.cuda.synchronize() + report( + "graph_vs_eager", + graph_out, + eager, + active=active, + new_length=new_length, + **context, + ) + table.copy_(original_table) + seqs.fill_(length) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--seeds", type=int, nargs="+", default=[1234, 3]) + parser.add_argument("--seq-lens", type=int, nargs="+", default=[513, 4096, 4171, 7168]) + parser.add_argument("--batches", type=int, nargs="+", default=[1, 4]) + parser.add_argument( + "--layouts", + nargs="+", + choices=["separate", "interleaved", "large"], + default=["separate", "interleaved", "large"], + ) + parser.add_argument("--profile", action="store_true") + parser.add_argument("--graph", action="store_true") + parser.add_argument("--guard-pages", action="store_true") + parser.add_argument("--fixed-tile", type=int, choices=[64, 128]) + args = parser.parse_args() + print( + json.dumps(dict(gpu=str(torch.cuda.get_device_properties(0)), args=vars(args))), flush=True + ) + for seed in args.seeds: + for length in args.seq_lens: + for batch in args.batches: + run_case(args, seed, length, batch) + + +if __name__ == "__main__": + main() diff --git a/examples/vime_rocm_attention_ablation/run.py b/examples/vime_rocm_attention_ablation/run.py index 863489df..b3c31524 100644 --- a/examples/vime_rocm_attention_ablation/run.py +++ b/examples/vime_rocm_attention_ablation/run.py @@ -127,11 +127,12 @@ def validate(self, *, require_paths: bool) -> None: raise ValueError("rollout GPU count must be divisible by rollout TP") if self.router_policy != "round_robin": raise ValueError("the two-engine strict matrix requires round_robin routing") - if self.rollout_batch_size < self.rollout_engines: + generated = self.rollout_batch_size * self.samples_per_prompt + if generated < self.rollout_engines: raise ValueError( - "rollout_batch_size must issue at least one request per rollout engine" + "rollout_batch_size*samples_per_prompt must issue at least one " + "request per rollout engine" ) - generated = self.rollout_batch_size * self.samples_per_prompt if self.global_batch_size != generated: raise ValueError( "global_batch_size must equal rollout_batch_size*samples_per_prompt " @@ -558,6 +559,8 @@ def public_arm_environment(environment: Mapping[str, str]) -> dict[str, str]: "CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "RL_KERNEL_ATTENTION_CASE", + "RL_KERNEL_ROCM_FIXED_PAGED_TILE", + "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS", "RL_KERNEL_FFN_CASE", "RL_KERNEL_LOGP_CASE", "RL_KERNEL_VLLM_REAL_VOCAB_SIZE", diff --git a/examples/vime_rocm_attention_ablation/run_full_rr_fixed_paged_ck.py b/examples/vime_rocm_attention_ablation/run_full_rr_fixed_paged_ck.py new file mode 100644 index 00000000..3466cbb1 --- /dev/null +++ b/examples/vime_rocm_attention_ablation/run_full_rr_fixed_paged_ck.py @@ -0,0 +1,39 @@ +"""One PR377-workload R/R iteration with the fixed CK paged candidate.""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path + +from examples.vime_rocm_attention_ablation import run_full_rr_single_arm_v90 as base + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--tile", choices=["64", "128"], default="128") + args = parser.parse_args() + os.environ["RL_KERNEL_ROCM_FIXED_PAGED_TILE"] = args.tile + os.environ["RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS"] = "8192" + config_type = base.MatrixConfig + + def config(**kwargs): + kwargs.update( + num_rollout=1, + rollout_batch_size=1, + samples_per_prompt=8, + global_batch_size=8, + max_response_length=7168, + max_tokens_per_gpu=4096, + rollout_seed=1234, + ) + return config_type(**kwargs) + + base.MatrixConfig = config + base.RUN_DIR = args.run_dir + return base.main() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/vime_rocm_attention_ablation/run_full_rr_single_arm_v90.py b/examples/vime_rocm_attention_ablation/run_full_rr_single_arm_v90.py new file mode 100644 index 00000000..9d27b0d9 --- /dev/null +++ b/examples/vime_rocm_attention_ablation/run_full_rr_single_arm_v90.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import json +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +from examples.vime_rocm_attention_ablation.run import ( + MatrixConfig, + _prepare_run_dir, + build_arm_environment, + frozen_input_manifest, + public_arm_environment, +) +from examples.vime_rocm_attention_ablation.validate_artifacts import ( + CASE_IMPLEMENTATIONS, + validate_arm, + write_report, +) + +RUN_DIR = Path("/app/model/vime-runs/pr390-strict-direct-paged-hipgraph-rr-v90") + + +def main() -> int: + root = Path("/workspace/RL-Kernel-pr390") + config = MatrixConfig( + vime_root=Path("/workspace/vime"), + rl_kernel_root=root, + megatron_root=Path("/workspace/Megatron-LM-vime"), + model_root=Path("/app/model/Qwen3-8B"), + reference_checkpoint=Path("/app/model/Qwen3-8B_torch_dist"), + prompt_data=Path("/app/model/dapo-math-17k/dapo-math-17k.jsonl"), + run_dir=RUN_DIR, + launcher=root / "examples/vime_rocm_attention_ablation/launch_arm.sh", + ) + config.validate(require_paths=True) + _prepare_run_dir(RUN_DIR) + frozen_before = frozen_input_manifest(config) + write_report(RUN_DIR / "frozen-inputs.before.json", frozen_before) + + case_id = "R/R" + arm_dir = RUN_DIR / "arms/r-r" + for directory in ( + arm_dir / "readbacks", + arm_dir / "dump", + arm_dir / "checkpoint", + arm_dir / "mismatch_sidecars", + ): + directory.mkdir(parents=True, exist_ok=False) + environment = build_arm_environment(config, case_id, arm_dir, arm_index=3) + environment["VLLM_GPU_MEMORY_UTILIZATION"] = "0.38" + launch = { + "schema_version": "rlkernel.vime_rocm_attention_arm_launch.v1", + "case_id": case_id, + "expected_implementations": CASE_IMPLEMENTATIONS[case_id], + "frozen_input_fingerprint": frozen_before["fingerprint"], + "command": ["bash", str(config.launcher.resolve())], + "environment": public_arm_environment(environment), + "started_at": datetime.now(timezone.utc).isoformat(), + } + write_report(arm_dir / "launch.json", launch) + with (arm_dir / "launcher.log").open("w", encoding="utf-8") as log_handle: + process = subprocess.run( + ["bash", str(config.launcher.resolve())], + cwd=config.rl_kernel_root, + env=environment, + stdout=log_handle, + stderr=subprocess.STDOUT, + check=False, + ) + report = validate_arm(arm_dir, case_id, launcher_returncode=process.returncode) + write_report(arm_dir / "validation.json", report) + frozen_after = frozen_input_manifest(config) + write_report(RUN_DIR / "frozen-inputs.after.json", frozen_after) + summary = { + "run_dir": str(RUN_DIR), + "launcher_returncode": process.returncode, + "passed": report["passed"], + "errors": report["errors"], + "metrics": report["metrics"], + "frozen_sources_match": frozen_before["fingerprint"] == frozen_after["fingerprint"], + } + write_report(RUN_DIR / "single-arm-summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True), flush=True) + return 0 if report["passed"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/rl_engine/integrations/framework_operators.py b/rl_engine/integrations/framework_operators.py index 52ec4d88..425263c3 100644 --- a/rl_engine/integrations/framework_operators.py +++ b/rl_engine/integrations/framework_operators.py @@ -395,6 +395,24 @@ def _tensor_cache_token(tensor: torch.Tensor) -> tuple[Any, ...]: ) +@lru_cache(maxsize=1) +def _rocm_paged_kv_max_tokens() -> int | None: + """Return an optional workload bound for AITER paged scheduling.""" + + value = os.environ.get("RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS", "").strip() + if not value: + return None + try: + limit = int(value) + except ValueError as exc: + raise RuntimeError( + "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS must be an integer" + ) from exc + if limit <= 0: + raise RuntimeError("RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS must be positive") + return limit + + def _compact_attention_provenance(value: Mapping[str, Any]) -> dict[str, Any]: """Keep strict backend identity without retaining one record per token.""" @@ -1274,7 +1292,7 @@ def _rocm_direct_paged_metadata( num_actual: int, cache_owner: Any, ) -> tuple[dict[str, Any], bool] | None: - """Reuse vLLM's sequence-level GPU metadata without token expansion.""" + """Prepare graph-safe paged metadata once and share it across layers.""" num_decodes = int(getattr(attn_metadata, "num_decodes", 0)) num_prefills = int(getattr(attn_metadata, "num_prefills", 0)) @@ -1301,23 +1319,31 @@ def _rocm_direct_paged_metadata( causal = False if not isinstance(query_start_loc, torch.Tensor) or max_seqlen_q <= 0: return None - seq_lens = self._metadata_tensor(attn_metadata, "seq_lens") + query_starts_source = query_start_loc + seq_lens_source = self._metadata_tensor(attn_metadata, "seq_lens") max_seq_len = int( getattr(attn_metadata, "max_seq_len", block_table.size(1) * block_size) ) + configured_kv_limit = _rocm_paged_kv_max_tokens() + kernel_max_seqlen_k = ( + max_seq_len + if configured_kv_limit is None + else min(max_seq_len, configured_kv_limit) + ) page_count = min( block_table.size(1), (max_seq_len + block_size - 1) // block_size, ) key = ( mode, - _tensor_cache_token(query_start_loc), - _tensor_cache_token(seq_lens), + _tensor_cache_token(query_starts_source), + _tensor_cache_token(seq_lens_source), _tensor_cache_token(block_table), sequence_count, num_actual, max_seqlen_q, max_seq_len, + kernel_max_seqlen_k, page_count, ) owner_id = id(cache_owner) @@ -1329,19 +1355,80 @@ def _rocm_direct_paged_metadata( self._rocm_paged_metadata_owners.add(owner_id) return self._rocm_paged_metadata_value, True - query_start_loc = query_start_loc[: sequence_count + 1].to( + query_start_loc = query_starts_source.to( device=block_table.device, dtype=torch.int32 ) if not query_start_loc.is_contiguous(): query_start_loc = query_start_loc.contiguous() - seq_lens = seq_lens[:sequence_count].to( - device=block_table.device, dtype=torch.int32 - ) + seq_lens = seq_lens_source.to(device=block_table.device, dtype=torch.int32) if not seq_lens.is_contiguous(): seq_lens = seq_lens.contiguous() - pages = block_table[:sequence_count, :page_count].to(dtype=torch.int32) + # Graph metadata is request-level while packed Q is query-token-level. + # Expand decode page rows to query rows without materializing KV. + if mode == "decode": + query_starts = query_start_loc[: sequence_count + 1] + query_ends = query_starts[1:] + query_indices = torch.arange( + num_actual, dtype=torch.int32, device=block_table.device + ) + request_indices = torch.searchsorted( + query_ends, query_indices, right=True + ).to(dtype=torch.long) + request_indices = request_indices.clamp_max(sequence_count - 1) + request_query_ends = query_ends.index_select(0, request_indices) + request_seq_lens = seq_lens.index_select(0, request_indices) + active_queries = query_indices < query_starts[-1] + seqused_k = request_seq_lens - ( + request_query_ends - query_indices + ) + 1 + seqused_k = torch.where( + active_queries, seqused_k, torch.ones_like(seqused_k) + ) + pages = block_table.index_select(0, request_indices)[:, :page_count] + query_start_loc = torch.arange( + num_actual + 1, dtype=torch.int32, device=block_table.device + ) + sequence_count = num_actual + max_seqlen_q = 1 + active_rows = active_queries & (request_seq_lens > 0) + else: + query_start_loc = query_start_loc[: sequence_count + 1] + seq_lens = seq_lens[:sequence_count] + active_rows = seq_lens > 0 + seqused_k = seq_lens + pages = block_table[:sequence_count, :page_count] if not pages.is_contiguous(): pages = pages.contiguous() + if mode != "decode": + active_rows = seqused_k > 0 + if configured_kv_limit is not None: + # Zero-length rows are legal vLLM graph padding. They must not + # trip the bound assertion when a dynamic decode batch shrinks. + torch._assert_async( + torch.all((seq_lens >= 0) & (seq_lens <= kernel_max_seqlen_k)), + "vLLM sequence length exceeds RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS", + ) + pages = pages.to(dtype=torch.int32) + # vLLM CUDA/HIP graphs retain padded request rows after a request + # finishes. Their sequence length is zero and their page-table row + # is not guaranteed to contain a valid physical page. AITER CK does + # not accept either state; route inactive rows to page 0 and clamp + # their logical length. vLLM ignores outputs for graph-padding rows. + # CK looks up complete 128-token KV tiles before applying its logical + # mask. Unused columns must be valid even for active request rows. + # Reuse the row's first live page so masked loads see initialized KV. + safe_page = torch.where(active_rows, pages[:, 0], torch.zeros_like(seqused_k)) + columns = torch.arange(page_count, dtype=torch.int32, device=pages.device) + live_columns = active_rows[:, None] & ( + columns[None, :] * block_size < seqused_k[:, None] + ) + pages = torch.where(live_columns, pages, safe_page[:, None]) + tile_pages = max(1, 128 // block_size) + guard_columns = (-page_count) % tile_pages + if guard_columns: + pages = torch.cat((pages, safe_page[:, None].expand(-1, guard_columns)), dim=1) + page_count += guard_columns + seq_lens_for_kernel = seqused_k.clamp_min(1) indptr_key = ( block_table.device.type, block_table.device.index, @@ -1361,11 +1448,12 @@ def _rocm_direct_paged_metadata( "sequence_count": sequence_count, "page_count": page_count, "pages": pages, - "seqused_k": seq_lens, + "seqused_k": seq_lens_for_kernel, "cu_seqlens_q": query_start_loc, "kv_indptr": kv_indptr, "max_seqlen_q": max_seqlen_q, - "max_seqlen_k": page_count * block_size, + "max_seqlen_k": kernel_max_seqlen_k, + "configured_kv_limit": configured_kv_limit, "causal": causal, } self._rocm_paged_metadata_key = key @@ -1449,6 +1537,8 @@ def _rocm_direct_paged( "attention_phase": metadata["mode"], "sequence_count": metadata["sequence_count"], "query_token_count": num_actual, + "max_seqlen_k": metadata["max_seqlen_k"], + "configured_kv_limit": metadata["configured_kv_limit"], "launch_group_count": 1, "metadata_source": "vllm_gpu_sequence_level", "metadata_reused_across_layers": reused, diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index e032ea01..822ef41a 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -78,11 +78,13 @@ "rl_kernel::rocm_det_gemm_linear_all_reduce_inference", "rl_kernel::qwen3_ffn_packed_tp_inference_rocm", ) -_ROCM_FULL_GRAPH_CACHE_NAMESPACE = "rl_kernel_rocm_full_graph_v1" +_ROCM_FULL_GRAPH_CACHE_NAMESPACE = "rl_kernel_rocm_full_graph_v4" _ROCM_GRAPH_ROUTE_ENVIRONMENT = ( "RL_KERNEL_ATTENTION_CASE", "RL_KERNEL_FFN_CASE", "RL_KERNEL_LOGP_CASE", + "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS", + "RL_KERNEL_ROCM_FIXED_PAGED_TILE", ) diff --git a/rl_engine/kernels/ops/rocm/attention/fixed_paged_ck.py b/rl_engine/kernels/ops/rocm/attention/fixed_paged_ck.py new file mode 100644 index 00000000..c594a156 --- /dev/null +++ b/rl_engine/kernels/ops/rocm/attention/fixed_paged_ck.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""RL-Kernel-owned, fixed-schedule CK paged attention implementation.""" + +from __future__ import annotations + +import hashlib +import importlib.util +from functools import lru_cache +from pathlib import Path + +import torch + + +@lru_cache(maxsize=2) +def load_fixed_paged_ck(tile_m: int = 128): + from torch.utils.cpp_extension import load + + if tile_m not in (64, 128): + raise ValueError("fixed CK tile must be 64 or 128") + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("warm fixed CK attention before HIP Graph capture") + spec = importlib.util.find_spec("aiter_meta") + if spec is None or not spec.submodule_search_locations: + raise RuntimeError("fixed CK attention requires the installed aiter_meta headers") + meta = Path(next(iter(spec.submodule_search_locations))) + ck = meta / "3rdparty/composable_kernel" + source = Path(__file__).resolve().parents[5] / "csrc/rocm/attention/strict_paged_ck.cu" + fingerprint = hashlib.sha256(source.read_bytes()) + for header in ( + ck / "example/ck_tile/01_fmha/fmha_fwd.hpp", + ck + / "include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp", + ck / "include/ck_tile/core/tensor/tile_scatter_gather.hpp", + ): + fingerprint.update(header.read_bytes()) + arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.split(":")[0] + flags = [ + # In an all-masked KV tile, m_old == m must give an exact rescale of + # one. Contracting scale*m_old - rounded(scale*m) into FMA breaks that + # identity and makes results depend on the surrounding query chunk. + # Explicit CK MFMA instructions are unaffected by this scalar flag. + "-O3", + "-std=c++20", + "-ffp-contract=off", + f"--offload-arch={arch}", + f"-DRLK_CK_TILE_M={tile_m}", + "-DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=2", + "-DCK_TILE_FMHA_FWD_FAST_EXP2=1", + "-DCK_TILE_ATTENTION_LOGITS_SOFT_CAP_DEFAULT=0", + "-DCK_TILE_ATTENTION_USE_SOFTSIGN_ASM=1", + "-U__HIP_NO_HALF_CONVERSIONS__", + "-U__HIP_NO_HALF_OPERATORS__", + "-fbracket-depth=1024", + "-fgpu-flush-denormals-to-zero", + "-fno-offload-uniform-block", + "-mllvm", + "--amdgpu-kernarg-preload-count=16", + "-mllvm", + "--lsr-drop-solution=1", + "-mllvm", + "-amdgpu-early-inline-all=true", + "-mllvm", + "-amdgpu-function-calls=false", + "-mllvm", + "-enable-post-misched=0", + "-fno-gpu-rdc", + ] + fingerprint.update(" ".join(flags).encode()) + return load( + name=f"rlk_fixed_paged_ck_m{tile_m}_{fingerprint.hexdigest()[:16]}", + sources=[str(source)], + extra_include_paths=[ + str(ck / "include"), + str(ck / "library/include"), + str(ck / "example/ck_tile/01_fmha"), + str(meta / "3rdparty/ck_helper"), + ], + extra_cflags=["-O3", "-std=c++20"], + extra_cuda_cflags=flags, + with_cuda=True, + ) + + +def fixed_paged_prefill( + q, + k, + v, + cuq, + indptr, + flat_pages, + maxq, + maxk, + dropout, + scale, + softcap, + zero_tensors, + causal, + window_left, + window_right, + sink, + return_lse, + return_dropout, + *, + block_table, + seqlen_k, + out=None, + tile_m=128, +): + """Subset of the AITER entrypoint required by strict Qwen3 BF16 attention.""" + if ( + q.dtype != torch.bfloat16 + or q.shape[-1] != 128 + or k.shape[1] != 16 + or dropout + or softcap + or sink + or return_dropout + or zero_tensors + or window_left != -1 + or window_right != -1 + ): + raise ValueError( + "fixed CK entrypoint requires BF16/D128/page16 without dropout, bias or windows" + ) + tensors = (q, k, v, cuq, block_table, seqlen_k) + if any(t.device != q.device for t in tensors) or not q.is_cuda: + raise ValueError("fixed CK inputs must share one ROCm device") + if k.dtype != q.dtype or v.dtype != q.dtype or v.shape != k.shape: + raise ValueError("fixed CK K/V must match Q dtype and each other") + if not q.is_contiguous() or k.stride(-1) != 1 or v.stride(-1) != 1: + raise ValueError("fixed CK requires packed Q and contiguous K/V head dimensions") + batch = block_table.shape[0] + if ( + block_table.ndim != 2 + or block_table.stride(-1) != 1 + or cuq.shape != (batch + 1,) + or seqlen_k.shape != (batch,) + or any(t.dtype != torch.int32 for t in (cuq, block_table, seqlen_k)) + or not cuq.is_contiguous() + or not seqlen_k.is_contiguous() + ): + raise ValueError("fixed CK requires packed int32 query/length metadata and a 2D page table") + module = load_fixed_paged_ck(tile_m) + guard_columns = (-block_table.shape[1]) % 8 + if guard_columns: + block_table = torch.nn.functional.pad(block_table, (0, guard_columns), value=0) + if out is None: + out = torch.empty_like(q) + elif ( + out.shape != q.shape + or out.dtype != q.dtype + or out.device != q.device + or not out.is_contiguous() + ): + raise ValueError("fixed CK output must match packed Q") + lse = torch.empty( + (q.shape[1], q.shape[0]) if return_lse else (0,), dtype=torch.float32, device=q.device + ) + module.forward( + [ + q.data_ptr(), + k.data_ptr(), + v.data_ptr(), + out.data_ptr(), + lse.data_ptr(), + cuq.data_ptr(), + block_table.data_ptr(), + seqlen_k.data_ptr(), + torch.cuda.current_stream(q.device).cuda_stream, + ], + [ + q.shape[0], + block_table.shape[0], + maxq, + q.shape[1], + k.shape[2], + k.shape[0], + block_table.stride(0), + q.stride(0), + k.stride(1), + v.stride(1), + k.stride(2), + v.stride(2), + k.stride(0), + v.stride(0), + q.stride(1), + 0, + ], + scale, + causal, + return_lse, + ) + # Dropout is disabled; AITER backward accepts the unused, correctly typed + # state buffer just as for the native forward entrypoint. + rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) + return out, lse, torch.empty((0,), dtype=q.dtype, device=q.device), rng_state diff --git a/rl_engine/kernels/ops/rocm/attention/flash_attn.py b/rl_engine/kernels/ops/rocm/attention/flash_attn.py index 7a8e7296..6bb443cd 100644 --- a/rl_engine/kernels/ops/rocm/attention/flash_attn.py +++ b/rl_engine/kernels/ops/rocm/attention/flash_attn.py @@ -8,13 +8,11 @@ import inspect import math import os +from functools import partial from pathlib import Path from typing import Any, Callable import torch -from torch.autograd import Function -from torch.autograd.function import once_differentiable - from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, STRICT_ATTENTION_ROCM_SCHEDULE_ID, @@ -27,6 +25,8 @@ RLKernelDeterministicAttentionCore, ) from rl_engine.utils.logger import logger +from torch.autograd import Function +from torch.autograd.function import once_differentiable _MAX_TESTED_ROCM_TRITON_HEAD_DIM = 512 _AITER_API_SOURCE = "aiter.ops.mha" @@ -145,14 +145,12 @@ def _validate_aiter_schema( ) -def _load_aiter_ck_ops() -> tuple[ - Callable[..., Any], Callable[..., Any], Callable[..., Any], str -]: +def _load_aiter_ck_ops() -> tuple[Callable[..., Any], Callable[..., Any], Callable[..., Any], str]: try: module = importlib.import_module(_AITER_API_SOURCE) - mha_fwd = getattr(module, "mha_fwd") - mha_bwd = getattr(module, "mha_bwd") - mha_batch_prefill = getattr(module, "mha_batch_prefill") + mha_fwd = module.mha_fwd + mha_bwd = module.mha_bwd + mha_batch_prefill = module.mha_batch_prefill except (AttributeError, ImportError, OSError, RuntimeError) as exc: raise StrictRocmAttentionUnavailable( "strict ROCm Attention requires AITER dense backward and batch-prefill CK ops" @@ -301,9 +299,14 @@ def forward( ) v_cache = v_padded.reshape_as(k_cache) - block_table = torch.arange( - batch * page_count, dtype=torch.int32, device=q.device - ).reshape(batch, page_count) + block_table = torch.arange(batch * page_count, dtype=torch.int32, device=q.device).reshape( + batch, page_count + ) + # CK fetches page IDs for a complete 128-token tile before masking. + # Page 0 contains valid, initialized training KV for the unused columns. + guard_columns = (-page_count) % (128 // _AiterCKPagedAttentionFn.page_size) + if guard_columns: + block_table = torch.nn.functional.pad(block_table, (0, guard_columns), value=0) cu_seqlens_q = torch.arange(batch + 1, dtype=torch.int32, device=q.device) * q_len kv_indptr = torch.arange(batch + 1, dtype=torch.int32, device=q.device) * page_count seqlen_k = torch.full((batch,), kv_len, dtype=torch.int32, device=q.device) @@ -427,6 +430,24 @@ def __init__( mha_bwd = _mha_bwd mha_batch_prefill = _mha_batch_prefill source_sha256 = "test-double" if _source_sha256 is None else _source_sha256 + self._fixed_paged_tile = 0 + fixed_tile = os.environ.get("RL_KERNEL_ROCM_FIXED_PAGED_TILE", "0") + if _mha_fwd is None and fixed_tile != "0": + if fixed_tile not in ("64", "128"): + raise ValueError("RL_KERNEL_ROCM_FIXED_PAGED_TILE must be 0, 64 or 128") + from .fixed_paged_ck import fixed_paged_prefill + + self._fixed_paged_tile = int(fixed_tile) + mha_batch_prefill = partial(fixed_paged_prefill, tile_m=self._fixed_paged_tile) + source_digest = hashlib.sha256(source_sha256.encode()) + source_digest.update(Path(inspect.getsourcefile(fixed_paged_prefill)).read_bytes()) + source_digest.update( + ( + Path(__file__).resolve().parents[5] / "csrc/rocm/attention/strict_paged_ck.cu" + ).read_bytes() + ) + source_digest.update(fixed_tile.encode()) + source_sha256 = source_digest.hexdigest() if not callable(mha_fwd) or not callable(mha_bwd): raise StrictRocmAttentionUnavailable("AITER CK MHA entry points are not callable") self.split_kv = requested @@ -487,7 +508,11 @@ def forward_with_lse( self._mha_batch_prefill, self._mha_bwd, ) - forward_entrypoint = "mha_batch_prefill" + forward_entrypoint = ( + f"rl_kernel_fixed_paged_ck_m{self._fixed_paged_tile}" + if self._fixed_paged_tile + else "mha_batch_prefill" + ) kv_layout = "sequential_linear_pages" expected_lse_shape = (q.size(0), q.size(1), q.size(2)) if out.shape != q.shape or out.dtype != resolved_dtype: @@ -568,8 +593,7 @@ def forward_paged_varlen_with_lse( if cu_seqlens_q.shape != (batch + 1,) or kv_indptr.shape != (batch + 1,): raise ValueError("paged indptr tensors must carry batch + 1 entries") if any( - tensor.device != q.device - for tensor in (page_table, seqused_k, cu_seqlens_q, kv_indptr) + tensor.device != q.device for tensor in (page_table, seqused_k, cu_seqlens_q, kv_indptr) ): raise ValueError("paged metadata must be on the Q device") if not q.is_contiguous() or not page_table.is_contiguous(): @@ -644,7 +668,11 @@ def forward_paged_varlen_with_lse( "gpu_arch": gpu_arch, "aiter_api_source": self.api_source, "aiter_source_sha256": self.source_sha256, - "forward_entrypoint": "mha_batch_prefill", + "forward_entrypoint": ( + f"rl_kernel_fixed_paged_ck_m{self._fixed_paged_tile}" + if self._fixed_paged_tile + else "mha_batch_prefill" + ), "kv_layout": "vllm_linear_paged", "dense_kv_materialized": False, "num_splits": self.num_splits, @@ -740,7 +768,11 @@ def _validate_paged_inputs( raise ValueError("paged q/k/v layouts must be BHSD and page-linear BSHD") if q.size(1) % k.size(2) or q.size(3) != k.size(3): raise ValueError("paged Q/K head counts or dimensions are incompatible") - if q.dtype not in (torch.float16, torch.bfloat16) or k.dtype != q.dtype or v.dtype != q.dtype: + if ( + q.dtype not in (torch.float16, torch.bfloat16) + or k.dtype != q.dtype + or v.dtype != q.dtype + ): raise ValueError("paged q/k/v must share FP16 or BF16 dtype") if not (q.is_cuda and k.is_cuda and v.is_cuda) or not (q.device == k.device == v.device): raise ValueError("paged q/k/v must share one ROCm device") @@ -788,6 +820,33 @@ def forward_bshd_with_lse( if not out_bshd.is_contiguous(): raise ValueError("strict BSHD decode output must expose a contiguous AITER view") + if self._fixed_paged_tile: + # Keep the same arithmetic if a caller already materialized BSHD + # inputs (e.g. a mixed prefill/decode fallback). + fixed_out, fixed_lse = _AiterCKPagedAttentionFn.apply( + q_bshd.transpose(1, 2), + k_bshd.transpose(1, 2), + v_bshd.transpose(1, 2), + bool(causal), + resolved_scale, + self._mha_batch_prefill, + self._mha_bwd, + ) + if out is not None: + out.copy_(fixed_out) + fixed_out = out + return DeterministicAttentionCoreResult( + out=fixed_out, + lse=fixed_lse, + provenance={ + "actual_backend": self.backend_id, + "forward_entrypoint": f"rl_kernel_fixed_paged_ck_m{self._fixed_paged_tile}", + "aiter_source_sha256": self.source_sha256, + "dense_kv_materialized": True, + "fallback": False, + }, + ) + result = self._mha_fwd( q_bshd, k_bshd, From 95a98dfecb9accd0ee725a87b11415aeb01c1e5d Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 9 Sep 2026 13:09:42 +0800 Subject: [PATCH 2/7] perf(rocm): reuse dynamic decode GEMM reduction --- .../ops/rocm/loss/vocab_parallel_logp.py | 133 ++++++++++++++++-- .../kernels/ops/triton/matmul/det_gemm.py | 11 +- 2 files changed, 126 insertions(+), 18 deletions(-) diff --git a/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py index a1f3578f..7f58862f 100644 --- a/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py +++ b/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py @@ -21,6 +21,7 @@ from __future__ import annotations +from collections import OrderedDict from typing import Any import torch @@ -29,7 +30,6 @@ from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( DEFAULT_NUM_VOCAB_TILES, VocabParallelLogprobOp, - _gather_target_logit, _gather_tile_stats, _merge_tile_partials, _preflight_cross_rank_agreement, @@ -40,6 +40,102 @@ BACKEND_ID = "rocm-vocab-parallel-logp-ws2" +# These tensors are immutable metadata. Vime's logprob requests normally use +# an all-active mask and repeat the same token count/sharding on every forward; +# recreating them from Python tuples showed up as synchronous ``aten::to`` / +# ``aten::copy_`` activity in the ROCm trace. Keep the cache bounded because +# sequence packing can expose many token counts over a long run. +_METADATA_CACHE_LIMIT = 32 +_ACTIVE_MASK_CACHE: OrderedDict[tuple[Any, ...], torch.Tensor] = OrderedDict() +_SHARD_START_CACHE: OrderedDict[tuple[Any, ...], torch.Tensor] = OrderedDict() + + +def _device_key(device: torch.device) -> tuple[str, int | None]: + return device.type, device.index + + +def _cached_active_mask( + contract: LogprobContract, device: torch.device +) -> tuple[torch.Tensor, bool]: + values = contract.mask.active_mask + all_active = all(values) + # The all-active case is by far the common path; avoid hashing/copying the + # complete mask tuple for that case. + signature: Any = ("all", len(values)) if all_active else ("mask", values) + key = (_device_key(device), signature) + cached = _ACTIVE_MASK_CACHE.get(key) + if cached is None: + cached = ( + torch.ones((len(values),), dtype=torch.bool, device=device) + if all_active + else torch.tensor(values, dtype=torch.bool, device=device) + ) + _ACTIVE_MASK_CACHE[key] = cached + if len(_ACTIVE_MASK_CACHE) > _METADATA_CACHE_LIMIT: + _ACTIVE_MASK_CACHE.popitem(last=False) + else: + _ACTIVE_MASK_CACHE.move_to_end(key) + return cached, all_active + + +def _cached_shard_starts( + bounds: tuple[tuple[int, int], ...], device: torch.device +) -> torch.Tensor: + key = (_device_key(device), bounds) + cached = _SHARD_START_CACHE.get(key) + if cached is None: + cached = torch.tensor( + [start for start, _ in bounds], dtype=torch.long, device=device + ) + _SHARD_START_CACHE[key] = cached + if len(_SHARD_START_CACHE) > _METADATA_CACHE_LIMIT: + _SHARD_START_CACHE.popitem(last=False) + else: + _SHARD_START_CACHE.move_to_end(key) + return cached + + +def _gather_target_logit_cached( + z_masked: torch.Tensor, + safe_target: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> torch.Tensor: + """ROCm copy of the exact owner gather with cached immutable metadata.""" + + sharding = contract.sharding + n = z_masked.shape[0] + start = sharding.local_vocab_start + local_vocab = sharding.local_vocab_size + local_idx = (safe_target - start).clamp(0, max(local_vocab - 1, 0)) + owns = (safe_target >= start) & (safe_target < sharding.local_vocab_end) + rows = torch.arange(n, device=z_masked.device) + local_contrib = torch.where( + owns, + z_masked[rows, local_idx], + torch.zeros_like(safe_target, dtype=z_masked.dtype), + ).contiguous() + + if sharding.tp_world_size == 1: + stacked = local_contrib.unsqueeze(0) + else: + if ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + ): + raise LogprobContractError( + "vocab-parallel logprob requires initialized torch.distributed" + ) + gathered = [ + torch.empty_like(local_contrib) for _ in range(sharding.tp_world_size) + ] + torch.distributed.all_gather(gathered, local_contrib, group=tp_group) + stacked = torch.stack(gathered, dim=0) + + starts = _cached_shard_starts(sharding.vocab_shard_bounds, safe_target.device) + owner = torch.bucketize(safe_target, starts, right=True) - 1 + return stacked[owner, rows] + def _native_backward_available() -> bool: try: @@ -80,7 +176,9 @@ class _RocmVocabParallelLogprobFunction(torch.autograd.Function): """ROCm tile statistics and backward with the shared WS2 merge contract.""" @staticmethod - def forward(ctx, local_logits, target_1d, active_mask, contract, tp_group, tile): + def forward( + ctx, local_logits, target_1d, active_mask, contract, tp_group, tile, all_active + ): sharding = contract.sharding shard = local_logits.contiguous() local_tiles = sharding.local_vocab_size // tile @@ -102,18 +200,25 @@ def forward(ctx, local_logits, target_1d, active_mask, contract, tp_group, tile) tp_group, tile_counts, ) - safe_target = torch.where(active_mask, target_1d, torch.zeros_like(target_1d)) - target_logit = _gather_target_logit( + safe_target = ( + target_1d + if all_active + else torch.where(active_mask, target_1d, torch.zeros_like(target_1d)) + ) + target_logit = _gather_target_logit_cached( shard, safe_target, contract, tp_group ).float() lse = _merge_tile_partials(m_all, s_all) - selected_logp = torch.where( - active_mask, target_logit - lse, torch.zeros_like(lse) + selected_logp = ( + target_logit - lse + if all_active + else torch.where(active_mask, target_logit - lse, torch.zeros_like(lse)) ) ctx.save_for_backward(shard, lse, safe_target, active_mask) ctx.local_vocab_start = sharding.local_vocab_start ctx.real_vocab_size = sharding.real_vocab_size + ctx.all_active = all_active ctx.set_materialize_grads(False) return selected_logp, lse @@ -126,13 +231,15 @@ def backward(ctx, grad_logp, grad_lse): start = ctx.local_vocab_start if grad_logp is not None: coef_logp = ( - torch.where(active_mask, grad_logp, torch.zeros_like(grad_logp)) + grad_logp.float().contiguous() + if ctx.all_active + else torch.where(active_mask, grad_logp, torch.zeros_like(grad_logp)) .float() .contiguous() ) owns = (safe_target >= start) & (safe_target < start + local_vocab) target_local = torch.where( - owns & active_mask, + owns if ctx.all_active else owns & active_mask, safe_target - start, torch.full_like(safe_target, -1), ).contiguous() @@ -157,7 +264,7 @@ def backward(ctx, grad_logp, grad_lse): ctx.real_vocab_size, has_lse_grad, ) - return grad, None, None, None, None, None + return grad, None, None, None, None, None, None def _apply_with_kernels( @@ -177,11 +284,7 @@ def _apply_with_kernels( target_1d = target_ids.reshape(-1).to( device=local_logits.device, dtype=torch.long ) - active_mask = torch.tensor( - contract.mask.active_mask, - dtype=torch.bool, - device=local_logits.device, - ) + active_mask, all_active = _cached_active_mask(contract, local_logits.device) if validate: _validate_active_targets( target_1d, active_mask, contract.sharding.real_vocab_size @@ -192,7 +295,7 @@ def _apply_with_kernels( ) selected_logp, lse = _RocmVocabParallelLogprobFunction.apply( - local_logits, target_1d, active_mask, contract, tp_group, tile + local_logits, target_1d, active_mask, contract, tp_group, tile, all_active ) if validate and bool((~torch.isfinite(lse) & active_mask).any().item()): raise LogprobContractError( diff --git a/rl_engine/kernels/ops/triton/matmul/det_gemm.py b/rl_engine/kernels/ops/triton/matmul/det_gemm.py index a4fc1529..b614e7e2 100644 --- a/rl_engine/kernels/ops/triton/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/triton/matmul/det_gemm.py @@ -358,12 +358,15 @@ def _det_gemm_tree_reduce_kernel( ) @triton.jit - def _det_gemm_tree_reduce_to_output_kernel( + def _det_gemm_tree_reduce_to_output_rocm_kernel( workspace_ptr, output_ptr, lower_nodes_ptr, upper_nodes_ptr, - M: tl.constexpr, + # M is runtime-sized so decode batch changes reuse one compiled + # reduction kernel instead of triggering a Triton JIT per batch size. + # The reduction order and BF16 store boundary are unchanged. + M, N: tl.constexpr, BLOCK: tl.constexpr, ): @@ -644,7 +647,9 @@ def _triton_tree_gemm( raise RuntimeError("the final deterministic GEMM tree level must contain one root") if write_final_output: grid = (triton.cdiv(m_size * n_size, reduction_block),) - _det_gemm_tree_reduce_to_output_kernel[grid]( + # direct_root_output is true only for gfx942; CUDA retains its + # established constexpr-M reduction and root-copy path. + _det_gemm_tree_reduce_to_output_rocm_kernel[grid]( workspace, result, lower, From 409370807e53352a22b11a3728baf56b1daa502b Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 9 Sep 2026 14:01:15 +0800 Subject: [PATCH 3/7] perf(rocm): fuse inference GEMM tree levels --- .../kernels/ops/triton/matmul/det_gemm.py | 155 +++++++++++++++--- 1 file changed, 132 insertions(+), 23 deletions(-) diff --git a/rl_engine/kernels/ops/triton/matmul/det_gemm.py b/rl_engine/kernels/ops/triton/matmul/det_gemm.py index b614e7e2..c5af6967 100644 --- a/rl_engine/kernels/ops/triton/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/triton/matmul/det_gemm.py @@ -176,6 +176,9 @@ class _DeviceTreePlan: leaf_lengths: torch.Tensor leaf_nodes: torch.Tensor reduction_levels: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...] + rocm_fused_reduction_pairs: tuple[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ... + ] def _build_tree_plan(k_size: int) -> _TreePlan: @@ -239,12 +242,35 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: for operations in host.reduction_levels: lower, upper, output = zip(*operations, strict=True) levels.append((indices(lower), indices(upper), indices(output))) + fused_pairs = [] + if _device_arch(device_index) == "gfx942": + for level_index in range(0, len(host.reduction_levels) - 1, 2): + first = { + output: (lower, upper) + for lower, upper, output in host.reduction_levels[level_index] + } + grandchildren = [] + outputs = [] + for lower, upper, output in host.reduction_levels[level_index + 1]: + grandchildren.append((*first[lower], *first[upper])) + outputs.append(output) + node0, node1, node2, node3 = zip(*grandchildren, strict=True) + fused_pairs.append( + ( + indices(node0), + indices(node1), + indices(node2), + indices(node3), + indices(tuple(outputs)), + ) + ) result = _DeviceTreePlan( host=host, leaf_starts=indices(host.leaf_starts), leaf_lengths=indices(host.leaf_lengths), leaf_nodes=indices(host.leaf_nodes), reduction_levels=tuple(levels), + rocm_fused_reduction_pairs=tuple(fused_pairs), ) _TREE_PLANS[key] = result return result @@ -394,6 +420,60 @@ def _det_gemm_tree_reduce_to_output_rocm_kernel( mask=mask, ) + @triton.jit + def _det_gemm_tree_reduce_two_levels_rocm_kernel( + workspace_ptr, + output_ptr, + node0_ptr, + node1_ptr, + node2_ptr, + node3_ptr, + output_nodes_ptr, + M, + N: tl.constexpr, + BLOCK: tl.constexpr, + WRITE_OUTPUT: tl.constexpr, + ): + operation = tl.program_id(0) + block = tl.program_id(1) + offsets = (block * BLOCK + tl.arange(0, BLOCK)).to(tl.int64) + elements = M * N + mask = offsets < elements + node0 = tl.load(node0_ptr + operation).to(tl.int64) + node1 = tl.load(node1_ptr + operation).to(tl.int64) + node2 = tl.load(node2_ptr + operation).to(tl.int64) + node3 = tl.load(node3_ptr + operation).to(tl.int64) + value0 = tl.load( + workspace_ptr + node0 * elements + offsets, mask=mask, other=0.0 + ).to(tl.float32) + value1 = tl.load( + workspace_ptr + node1 * elements + offsets, mask=mask, other=0.0 + ).to(tl.float32) + value2 = tl.load( + workspace_ptr + node2 * elements + offsets, mask=mask, other=0.0 + ).to(tl.float32) + value3 = tl.load( + workspace_ptr + node3 * elements + offsets, mask=mask, other=0.0 + ).to(tl.float32) + # Match the two original kernel boundaries exactly: each first-level + # FP32 add is rounded to BF16 before the second-level FP32 add. + lower = (value0 + value1).to(workspace_ptr.dtype.element_ty).to(tl.float32) + upper = (value2 + value3).to(workspace_ptr.dtype.element_ty).to(tl.float32) + result = lower + upper + if WRITE_OUTPUT: + tl.store( + output_ptr + offsets, + result.to(output_ptr.dtype.element_ty), + mask=mask, + ) + else: + output_node = tl.load(output_nodes_ptr + operation).to(tl.int64) + tl.store( + workspace_ptr + output_node * elements + offsets, + result.to(workspace_ptr.dtype.element_ty), + mask=mask, + ) + @triton.jit def _copy_tree_root_kernel( workspace_ptr, @@ -633,42 +713,71 @@ def _triton_tree_gemm( reduction_block = 256 device_index = a.device.index if a.device.index is not None else torch.cuda.current_device() direct_root_output = not transpose_output and _device_arch(device_index) == "gfx942" - for level_index, (operations, (lower, upper, output)) in enumerate( - zip( - plan.host.reduction_levels, - plan.reduction_levels, - strict=True, - ) - ): - write_final_output = ( - direct_root_output and level_index == len(plan.host.reduction_levels) - 1 - ) - if write_final_output and len(operations) != 1: - raise RuntimeError("the final deterministic GEMM tree level must contain one root") - if write_final_output: - grid = (triton.cdiv(m_size * n_size, reduction_block),) - # direct_root_output is true only for gfx942; CUDA retains its - # established constexpr-M reduction and root-copy path. - _det_gemm_tree_reduce_to_output_rocm_kernel[grid]( + if direct_root_output and torch.is_inference_mode_enabled(): + blocks = triton.cdiv(m_size * n_size, reduction_block) + for pair_index, nodes in enumerate(plan.rocm_fused_reduction_pairs): + second_level = pair_index * 2 + 1 + operations = plan.host.reduction_levels[second_level] + write_final_output = second_level == len(plan.host.reduction_levels) - 1 + _det_gemm_tree_reduce_two_levels_rocm_kernel[(len(operations), blocks)]( workspace, result, - lower, - upper, + *nodes, M=m_size, N=n_size, BLOCK=reduction_block, + WRITE_OUTPUT=write_final_output, ) - else: - grid = (len(operations), triton.cdiv(m_size * n_size, reduction_block)) - _det_gemm_tree_reduce_kernel[grid]( + if len(plan.host.reduction_levels) % 2: + lower, upper, _output = plan.reduction_levels[-1] + _det_gemm_tree_reduce_to_output_rocm_kernel[(blocks,)]( workspace, + result, lower, upper, - output, M=m_size, N=n_size, BLOCK=reduction_block, ) + else: + for level_index, (operations, (lower, upper, output)) in enumerate( + zip( + plan.host.reduction_levels, + plan.reduction_levels, + strict=True, + ) + ): + write_final_output = ( + direct_root_output and level_index == len(plan.host.reduction_levels) - 1 + ) + if write_final_output and len(operations) != 1: + raise RuntimeError( + "the final deterministic GEMM tree level must contain one root" + ) + if write_final_output: + grid = (triton.cdiv(m_size * n_size, reduction_block),) + # direct_root_output is true only for gfx942; CUDA retains its + # established reduction and root-copy path. + _det_gemm_tree_reduce_to_output_rocm_kernel[grid]( + workspace, + result, + lower, + upper, + M=m_size, + N=n_size, + BLOCK=reduction_block, + ) + else: + grid = (len(operations), triton.cdiv(m_size * n_size, reduction_block)) + _det_gemm_tree_reduce_kernel[grid]( + workspace, + lower, + upper, + output, + M=m_size, + N=n_size, + BLOCK=reduction_block, + ) copy_block = 256 if transpose_output: From 2921b4a8d17f0410bb23872722086d8e99f70819 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 9 Sep 2026 14:24:30 +0800 Subject: [PATCH 4/7] refactor(rocm): name CK source as HIP --- csrc/rocm/attention/{strict_paged_ck.cu => fixed_paged_ck.hip} | 0 rl_engine/kernels/ops/rocm/attention/fixed_paged_ck.py | 2 +- rl_engine/kernels/ops/rocm/attention/flash_attn.py | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) rename csrc/rocm/attention/{strict_paged_ck.cu => fixed_paged_ck.hip} (100%) diff --git a/csrc/rocm/attention/strict_paged_ck.cu b/csrc/rocm/attention/fixed_paged_ck.hip similarity index 100% rename from csrc/rocm/attention/strict_paged_ck.cu rename to csrc/rocm/attention/fixed_paged_ck.hip diff --git a/rl_engine/kernels/ops/rocm/attention/fixed_paged_ck.py b/rl_engine/kernels/ops/rocm/attention/fixed_paged_ck.py index c594a156..8cee606a 100644 --- a/rl_engine/kernels/ops/rocm/attention/fixed_paged_ck.py +++ b/rl_engine/kernels/ops/rocm/attention/fixed_paged_ck.py @@ -25,7 +25,7 @@ def load_fixed_paged_ck(tile_m: int = 128): raise RuntimeError("fixed CK attention requires the installed aiter_meta headers") meta = Path(next(iter(spec.submodule_search_locations))) ck = meta / "3rdparty/composable_kernel" - source = Path(__file__).resolve().parents[5] / "csrc/rocm/attention/strict_paged_ck.cu" + source = Path(__file__).resolve().parents[5] / "csrc/rocm/attention/fixed_paged_ck.hip" fingerprint = hashlib.sha256(source.read_bytes()) for header in ( ck / "example/ck_tile/01_fmha/fmha_fwd.hpp", diff --git a/rl_engine/kernels/ops/rocm/attention/flash_attn.py b/rl_engine/kernels/ops/rocm/attention/flash_attn.py index 6bb443cd..3462a367 100644 --- a/rl_engine/kernels/ops/rocm/attention/flash_attn.py +++ b/rl_engine/kernels/ops/rocm/attention/flash_attn.py @@ -443,7 +443,7 @@ def __init__( source_digest.update(Path(inspect.getsourcefile(fixed_paged_prefill)).read_bytes()) source_digest.update( ( - Path(__file__).resolve().parents[5] / "csrc/rocm/attention/strict_paged_ck.cu" + Path(__file__).resolve().parents[5] / "csrc/rocm/attention/fixed_paged_ck.hip" ).read_bytes() ) source_digest.update(fixed_tile.encode()) From 512d69b4b34790f65fd00d784434c96a608bf123 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 9 Sep 2026 14:44:27 +0800 Subject: [PATCH 5/7] perf(rocm): fuse decode GEMM leaves --- .../kernels/ops/triton/matmul/det_gemm.py | 327 +++++++++++++++--- 1 file changed, 287 insertions(+), 40 deletions(-) diff --git a/rl_engine/kernels/ops/triton/matmul/det_gemm.py b/rl_engine/kernels/ops/triton/matmul/det_gemm.py index c5af6967..526a482c 100644 --- a/rl_engine/kernels/ops/triton/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/triton/matmul/det_gemm.py @@ -42,6 +42,8 @@ _MAX_TREE_WORKSPACE_ELEMENTS = 128 * 1024 * 1024 _TREE_PLANS: dict[tuple[int, int], "_DeviceTreePlan"] = {} _TREE_PLAN_LOCK = threading.Lock() +_ROCM_TUNE_DECODE_LEAF = True +_ROCM_FUSE_LEAF_REDUCTION = True @dataclass(frozen=True) @@ -88,6 +90,29 @@ class _TreeLeafConfig: (6144, 32, 4096): _TreeLeafConfig(128, 64, 2, True), } +_GFX942_QWEN_TP4_DECODE_SHAPES = { + (4096, 1536), # QKV projection + (1024, 4096), # attention output projection + (4096, 6144), # gate/up projection + (3072, 4096), # down projection + (4096, 37984), # vocabulary-parallel LM head +} + + +def _gfx942_qwen_tp4_decode_leaf_config( + m_size: int, + k_size: int, + n_size: int, +) -> _TreeLeafConfig | None: + """Return the gfx942 small-batch schedule for exact Qwen3-8B TP4 shapes.""" + + if not 1 <= m_size <= 32 or (k_size, n_size) not in _GFX942_QWEN_TP4_DECODE_SHAPES: + return None + block_m = 1 << (m_size - 1).bit_length() + block_n = 64 if m_size == 1 or (k_size, n_size) == (3072, 4096) else 128 + num_warps = 1 if block_m <= 2 else 2 if block_m <= 8 else 4 + return _TreeLeafConfig(block_m, block_n, num_warps, True) + def _gfx942_qwen_tree_leaf_config( m_size: int, @@ -179,6 +204,16 @@ class _DeviceTreePlan: rocm_fused_reduction_pairs: tuple[ tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ... ] + rocm_leaf_reduction: tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ] | None + rocm_fused_reduction_pairs_after_leaf: tuple[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ... + ] def _build_tree_plan(k_size: int) -> _TreePlan: @@ -243,27 +278,67 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: lower, upper, output = zip(*operations, strict=True) levels.append((indices(lower), indices(upper), indices(output))) fused_pairs = [] + leaf_reduction = None + fused_pairs_after_leaf = [] if _device_arch(device_index) == "gfx942": - for level_index in range(0, len(host.reduction_levels) - 1, 2): - first = { - output: (lower, upper) - for lower, upper, output in host.reduction_levels[level_index] - } - grandchildren = [] - outputs = [] - for lower, upper, output in host.reduction_levels[level_index + 1]: - grandchildren.append((*first[lower], *first[upper])) - outputs.append(output) - node0, node1, node2, node3 = zip(*grandchildren, strict=True) - fused_pairs.append( - ( - indices(node0), - indices(node1), - indices(node2), - indices(node3), + def build_fused_pairs(start_level: int): + result = [] + for level_index in range( + start_level, len(host.reduction_levels) - 1, 2 + ): + first = { + output: (lower, upper) + for lower, upper, output in host.reduction_levels[level_index] + } + grandchildren = [] + outputs = [] + for lower, upper, output in host.reduction_levels[level_index + 1]: + grandchildren.append((*first[lower], *first[upper])) + outputs.append(output) + node0, node1, node2, node3 = zip(*grandchildren, strict=True) + result.append( + ( + indices(node0), + indices(node1), + indices(node2), + indices(node3), + indices(tuple(outputs)), + ) + ) + return result + + fused_pairs = build_fused_pairs(0) + fused_pairs_after_leaf = build_fused_pairs(1) + if host.reduction_levels: + first_level = host.reduction_levels[0] + paired_nodes = {node for op in first_level for node in op[:2]} + if paired_nodes == set(host.leaf_nodes): + leaf_metadata = { + node: (start, length) + for start, length, node in zip( + host.leaf_starts, + host.leaf_lengths, + host.leaf_nodes, + strict=True, + ) + } + lower_starts, lower_lengths = [], [] + upper_starts, upper_lengths, outputs = [], [], [] + for lower, upper, output in first_level: + lower_start, lower_length = leaf_metadata[lower] + upper_start, upper_length = leaf_metadata[upper] + lower_starts.append(lower_start) + lower_lengths.append(lower_length) + upper_starts.append(upper_start) + upper_lengths.append(upper_length) + outputs.append(output) + leaf_reduction = ( + indices(tuple(lower_starts)), + indices(tuple(lower_lengths)), + indices(tuple(upper_starts)), + indices(tuple(upper_lengths)), indices(tuple(outputs)), ) - ) result = _DeviceTreePlan( host=host, leaf_starts=indices(host.leaf_starts), @@ -271,6 +346,8 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: leaf_nodes=indices(host.leaf_nodes), reduction_levels=tuple(levels), rocm_fused_reduction_pairs=tuple(fused_pairs), + rocm_leaf_reduction=leaf_reduction, + rocm_fused_reduction_pairs_after_leaf=tuple(fused_pairs_after_leaf), ) _TREE_PLANS[key] = result return result @@ -348,6 +425,101 @@ def _det_gemm_tree_leaf_kernel( # more descriptive implementation name used by the optimized tree path. _det_gemm_kernel = _det_gemm_tree_leaf_kernel + @triton.jit(do_not_specialize=["M"]) + def _det_gemm_tree_leaf_reduce_rocm_kernel( + a_ptr, + b_ptr, + workspace_ptr, + output_ptr, + lower_starts_ptr, + lower_lengths_ptr, + upper_starts_ptr, + upper_lengths_ptr, + output_nodes_ptr, + M, + N: tl.constexpr, + stride_am: tl.constexpr, + stride_ak: tl.constexpr, + stride_bk: tl.constexpr, + stride_bn: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + N_FASTEST: tl.constexpr, + WRITE_OUTPUT: tl.constexpr, + ): + if N_FASTEST: + pid_n = tl.program_id(0) + pid_m = tl.program_id(1) + pair = tl.program_id(2) + else: + pair = tl.program_id(0) + pid_m = tl.program_id(1) + pid_n = tl.program_id(2) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + lower_start = tl.load(lower_starts_ptr + pair).to(tl.int64) + lower_length = tl.load(lower_lengths_ptr + pair).to(tl.int64) + upper_start = tl.load(upper_starts_ptr + pair).to(tl.int64) + upper_length = tl.load(upper_lengths_ptr + pair).to(tl.int64) + + lower_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for offset in tl.static_range(0, BLOCK_K): + k_offset = lower_start + offset + active = offset < lower_length + a = tl.load( + a_ptr + offs_m * stride_am + k_offset * stride_ak, + mask=(offs_m < M) & active, + other=0.0, + ).to(tl.float32) + b = tl.load( + b_ptr + k_offset * stride_bk + offs_n * stride_bn, + mask=(offs_n < N) & active, + other=0.0, + ).to(tl.float32) + lower_acc += a[:, None] * b[None, :] + # This conversion is the original lower-leaf workspace store boundary. + lower = lower_acc.to(workspace_ptr.dtype.element_ty) + + upper_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for offset in tl.static_range(0, BLOCK_K): + k_offset = upper_start + offset + active = offset < upper_length + a = tl.load( + a_ptr + offs_m * stride_am + k_offset * stride_ak, + mask=(offs_m < M) & active, + other=0.0, + ).to(tl.float32) + b = tl.load( + b_ptr + k_offset * stride_bk + offs_n * stride_bn, + mask=(offs_n < N) & active, + other=0.0, + ).to(tl.float32) + upper_acc += a[:, None] * b[None, :] + # Match both leaf BF16 stores followed by the original first-level FP32 + # add and BF16 store. Only the intermediate global-memory trip is gone. + result = lower.to(tl.float32) + upper_acc.to( + workspace_ptr.dtype.element_ty + ).to(tl.float32) + output_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + offsets = ( + offs_m[:, None].to(tl.int64) * N + + offs_n[None, :].to(tl.int64) + ) + if WRITE_OUTPUT: + tl.store( + output_ptr + offsets, + result.to(output_ptr.dtype.element_ty), + mask=output_mask, + ) + else: + output_node = tl.load(output_nodes_ptr + pair).to(tl.int64) + tl.store( + workspace_ptr + output_node * (M * N) + offsets, + result.to(workspace_ptr.dtype.element_ty), + mask=output_mask, + ) + @triton.jit(do_not_specialize=["M"]) def _det_gemm_tree_reduce_kernel( workspace_ptr, @@ -675,6 +847,8 @@ def _triton_tree_gemm( dtype=torch.bfloat16, device=a.device, ) + device_index = a.device.index if a.device.index is not None else torch.cuda.current_device() + is_gfx942 = _device_arch(device_index) == "gfx942" leaf_config = _tree_leaf_config( a.device, m_size, @@ -683,6 +857,17 @@ def _triton_tree_gemm( transpose_output=transpose_output, preserve_a_strides=preserve_a_strides, ) + if ( + _ROCM_TUNE_DECODE_LEAF + and is_gfx942 + and torch.is_inference_mode_enabled() + and not transpose_output + and not preserve_a_strides + ): + leaf_config = ( + _gfx942_qwen_tp4_decode_leaf_config(m_size, k_size, n_size) + or leaf_config + ) tiles_m = triton.cdiv(m_size, leaf_config.block_m) tiles_n = triton.cdiv(n_size, leaf_config.block_n) leaf_grid = ( @@ -690,30 +875,92 @@ def _triton_tree_gemm( if leaf_config.n_fastest else (len(plan.host.leaf_nodes), tiles_m, tiles_n) ) - _det_gemm_kernel[leaf_grid]( - a, - b, - workspace, - plan.leaf_starts, - plan.leaf_lengths, - plan.leaf_nodes, - M=m_size, - N=n_size, - K=k_size, - stride_am=a.stride(0), - stride_ak=a.stride(1), - stride_bk=b.stride(0), - stride_bn=b.stride(1), - BLOCK_M=leaf_config.block_m, - BLOCK_N=leaf_config.block_n, - BLOCK_K=_BLOCK_K, - N_FASTEST=leaf_config.n_fastest, - num_warps=leaf_config.num_warps, + direct_root_output = not transpose_output and is_gfx942 + fuse_leaf_reduction = ( + _ROCM_FUSE_LEAF_REDUCTION + and direct_root_output + and torch.is_inference_mode_enabled() + and k_size >= 1536 + and plan.rocm_leaf_reduction is not None ) + if fuse_leaf_reduction: + operation_count = len(plan.host.reduction_levels[0]) + fused_leaf_grid = ( + (tiles_n, tiles_m, operation_count) + if leaf_config.n_fastest + else (operation_count, tiles_m, tiles_n) + ) + _det_gemm_tree_leaf_reduce_rocm_kernel[fused_leaf_grid]( + a, + b, + workspace, + result, + *plan.rocm_leaf_reduction, + M=m_size, + N=n_size, + stride_am=a.stride(0), + stride_ak=a.stride(1), + stride_bk=b.stride(0), + stride_bn=b.stride(1), + BLOCK_M=leaf_config.block_m, + BLOCK_N=leaf_config.block_n, + BLOCK_K=_BLOCK_K, + N_FASTEST=leaf_config.n_fastest, + WRITE_OUTPUT=len(plan.host.reduction_levels) == 1, + num_warps=leaf_config.num_warps, + ) + else: + _det_gemm_kernel[leaf_grid]( + a, + b, + workspace, + plan.leaf_starts, + plan.leaf_lengths, + plan.leaf_nodes, + M=m_size, + N=n_size, + K=k_size, + stride_am=a.stride(0), + stride_ak=a.stride(1), + stride_bk=b.stride(0), + stride_bn=b.stride(1), + BLOCK_M=leaf_config.block_m, + BLOCK_N=leaf_config.block_n, + BLOCK_K=_BLOCK_K, + N_FASTEST=leaf_config.n_fastest, + num_warps=leaf_config.num_warps, + ) reduction_block = 256 - device_index = a.device.index if a.device.index is not None else torch.cuda.current_device() - direct_root_output = not transpose_output and _device_arch(device_index) == "gfx942" - if direct_root_output and torch.is_inference_mode_enabled(): + if fuse_leaf_reduction: + blocks = triton.cdiv(m_size * n_size, reduction_block) + for pair_index, nodes in enumerate( + plan.rocm_fused_reduction_pairs_after_leaf + ): + second_level = pair_index * 2 + 2 + operations = plan.host.reduction_levels[second_level] + write_final_output = second_level == len(plan.host.reduction_levels) - 1 + _det_gemm_tree_reduce_two_levels_rocm_kernel[(len(operations), blocks)]( + workspace, + result, + *nodes, + M=m_size, + N=n_size, + BLOCK=reduction_block, + WRITE_OUTPUT=write_final_output, + ) + remaining_levels = len(plan.host.reduction_levels) - 1 + if remaining_levels % 2: + lower, upper, _output = plan.reduction_levels[-1] + _det_gemm_tree_reduce_to_output_rocm_kernel[(blocks,)]( + workspace, + result, + lower, + upper, + M=m_size, + N=n_size, + BLOCK=reduction_block, + ) + elif direct_root_output and torch.is_inference_mode_enabled(): blocks = triton.cdiv(m_size * n_size, reduction_block) for pair_index, nodes in enumerate(plan.rocm_fused_reduction_pairs): second_level = pair_index * 2 + 1 From 3d8764a06e749daa7c09037f8b824f46f43806aa Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 9 Sep 2026 15:16:39 +0800 Subject: [PATCH 6/7] perf(rocm): fuse no-grad GEMM reductions --- rl_engine/kernels/ops/rocm/matmul/det_gemm.py | 34 +++++++++++++++---- .../kernels/ops/triton/matmul/det_gemm.py | 11 ++++-- 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/rl_engine/kernels/ops/rocm/matmul/det_gemm.py b/rl_engine/kernels/ops/rocm/matmul/det_gemm.py index 06febba6..aed6c8f1 100644 --- a/rl_engine/kernels/ops/rocm/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/rocm/matmul/det_gemm.py @@ -136,11 +136,17 @@ def det_gemm_linear( *, native_op: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None, out: torch.Tensor | None = None, + inference_schedule: bool = False, ) -> torch.Tensor: """Apply a native [N,K] weight through the strict ROCm backend.""" del native_op - return _triton_tree_gemm(a, _cached_weight_transpose(weight), out=out) + return _triton_tree_gemm( + a, + _cached_weight_transpose(weight), + out=out, + inference_schedule=inference_schedule, + ) @torch.library.custom_op("rl_kernel::rocm_det_gemm_linear_inference", mutates_args=()) @@ -148,7 +154,7 @@ def _det_gemm_linear_inference( a: torch.Tensor, weight: torch.Tensor, ) -> torch.Tensor: - return det_gemm_linear(a, weight) + return det_gemm_linear(a, weight, inference_schedule=True) @_det_gemm_linear_inference.register_fake @@ -168,7 +174,7 @@ def _det_gemm_linear_inference_out( weight: torch.Tensor, out: torch.Tensor, ) -> None: - det_gemm_linear(a, weight, out=out) + det_gemm_linear(a, weight, out=out, inference_schedule=True) @_det_gemm_linear_inference_out.register_fake @@ -207,14 +213,14 @@ def _det_gemm_linear_all_reduce_inference( runtime_handle, direct_input, ) - det_gemm_linear(a, weight, out=direct_input) + det_gemm_linear(a, weight, out=direct_input, inference_schedule=True) _C.deterministic_collective_rocm_ipc_all_reduce_staged( runtime_handle, direct_input, output ) return output else: # Profiling and uncaptured prefill can exceed the decode capture bound. - output = det_gemm_linear(a, weight) + output = det_gemm_linear(a, weight, inference_schedule=True) _C.deterministic_collective_rocm_ipc_all_reduce_input( runtime_handle, output, @@ -317,6 +323,17 @@ def prepare_det_gemm_linear_weight( with torch.no_grad(): out.copy_(weight.t()) + # Layerwise checkpoint reloads update stable Parameters through + # ``.data.copy_()``, which intentionally does not advance the Parameter + # version counter. Keep any generic inference transpose for this same + # source coherent with the lifecycle-managed LM-head buffer. + key = (weight.data_ptr(), str(weight.device), tuple(weight.shape), weight.dtype) + cached = _WEIGHT_TRANSPOSE_CACHE.get(key) + if cached is not None: + cached_tensor = cached[1] + if cached_tensor is not out: + cached_tensor.copy_(weight.t()) + _WEIGHT_TRANSPOSE_CACHE[key] = (_tensor_version(weight), cached_tensor) return out @@ -334,7 +351,12 @@ def det_gemm_linear_prepared( raise ValueError("prepared deterministic linear weight must be contiguous") if out is not None and (torch._C._overlaps(out, a) or torch._C._overlaps(out, weight_t)): raise ValueError("prepared deterministic linear output must not alias its inputs") - return _triton_tree_gemm(a, weight_t, out=out) + return _triton_tree_gemm( + a, + weight_t, + out=out, + inference_schedule=True, + ) def det_gemm_linear_input_gradient( diff --git a/rl_engine/kernels/ops/triton/matmul/det_gemm.py b/rl_engine/kernels/ops/triton/matmul/det_gemm.py index 526a482c..1fd370d6 100644 --- a/rl_engine/kernels/ops/triton/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/triton/matmul/det_gemm.py @@ -784,6 +784,7 @@ def _triton_tree_gemm( transpose_output: bool = False, out: torch.Tensor | None = None, preserve_a_strides: bool = False, + inference_schedule: bool = False, ) -> torch.Tensor: if not _TRITON_AVAILABLE: raise RuntimeError("Triton is unavailable") @@ -819,6 +820,7 @@ def _triton_tree_gemm( b, transpose_output=transpose_output, preserve_a_strides=preserve_a_strides, + inference_schedule=inference_schedule, ) chunks.append(chunk) return torch.cat(chunks, dim=1 if transpose_output else 0) @@ -849,6 +851,9 @@ def _triton_tree_gemm( ) device_index = a.device.index if a.device.index is not None else torch.cuda.current_device() is_gfx942 = _device_arch(device_index) == "gfx942" + use_inference_schedule = ( + torch.is_inference_mode_enabled() or inference_schedule + ) leaf_config = _tree_leaf_config( a.device, m_size, @@ -860,7 +865,7 @@ def _triton_tree_gemm( if ( _ROCM_TUNE_DECODE_LEAF and is_gfx942 - and torch.is_inference_mode_enabled() + and use_inference_schedule and not transpose_output and not preserve_a_strides ): @@ -879,7 +884,7 @@ def _triton_tree_gemm( fuse_leaf_reduction = ( _ROCM_FUSE_LEAF_REDUCTION and direct_root_output - and torch.is_inference_mode_enabled() + and use_inference_schedule and k_size >= 1536 and plan.rocm_leaf_reduction is not None ) @@ -960,7 +965,7 @@ def _triton_tree_gemm( N=n_size, BLOCK=reduction_block, ) - elif direct_root_output and torch.is_inference_mode_enabled(): + elif direct_root_output and use_inference_schedule: blocks = triton.cdiv(m_size * n_size, reduction_block) for pair_index, nodes in enumerate(plan.rocm_fused_reduction_pairs): second_level = pair_index * 2 + 1 From e911eed4080812f7cc09a50b7c82e0027d6cf72d Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 9 Sep 2026 18:10:03 +0800 Subject: [PATCH 7/7] test(rocm): add full-native PR377 workload runner Signed-off-by: lamentropetion <3051000145@qq.com> --- .../vime_rocm_attention_ablation/README.md | 12 + .../launch_arm.sh | 22 +- .../run_full_pp_pr377_workload.py | 208 ++++++++++++++++++ 3 files changed, 237 insertions(+), 5 deletions(-) create mode 100644 examples/vime_rocm_attention_ablation/run_full_pp_pr377_workload.py diff --git a/examples/vime_rocm_attention_ablation/README.md b/examples/vime_rocm_attention_ablation/README.md index 1dfa6c30..c9ce8913 100644 --- a/examples/vime_rocm_attention_ablation/README.md +++ b/examples/vime_rocm_attention_ablation/README.md @@ -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 +``` diff --git a/examples/vime_rocm_attention_ablation/launch_arm.sh b/examples/vime_rocm_attention_ablation/launch_arm.sh index 378fa512..c51cc630 100644 --- a/examples/vime_rocm_attention_ablation/launch_arm.sh +++ b/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 @@ -251,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}" \ @@ -317,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 \ diff --git a/examples/vime_rocm_attention_ablation/run_full_pp_pr377_workload.py b/examples/vime_rocm_attention_ablation/run_full_pp_pr377_workload.py new file mode 100644 index 00000000..95bd56ec --- /dev/null +++ b/examples/vime_rocm_attention_ablation/run_full_pp_pr377_workload.py @@ -0,0 +1,208 @@ +"""Run the full-native P/P arm with the PR377 comparison workload.""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +from examples.vime_rocm_attention_ablation.run import ( + MatrixConfig, + _canonical_fingerprint, + _prepare_run_dir, + build_arm_environment, + frozen_input_manifest, + public_arm_environment, +) +from examples.vime_rocm_attention_ablation.validate_artifacts import ( + compare_train_rollout_logps, + load_readbacks, + load_rollout_identity, + write_report, +) + +CASE_ID = "P/P" + + +class FullStackConfig(MatrixConfig): + def frozen_parameters(self): + value = super().frozen_parameters() + value["ffn_case"] = CASE_ID + value["logp_case"] = CASE_ID + value["framework_consistency"] = { + "use_rollout_logprobs": True, + "get_mismatch_metrics": True, + "custom_tis_function": ("vime_rocm_attention_ablation.tis_metrics.metrics_only_tis"), + } + return value + + +def sealed_manifest(config: MatrixConfig): + value = frozen_input_manifest(config) + value["fingerprint"] = _canonical_fingerprint( + {key: item for key, item in value.items() if key != "fingerprint"} + ) + return value + + +def validate_native_readbacks(readback_dir: Path): + errors = [] + frameworks = set() + paths = [] + for record in load_readbacks(readback_dir): + paths.append(record["_path"]) + framework = record.get("framework") + if framework in {"megatron", "vllm"}: + frameworks.add(framework) + if record.get("fallbacks"): + errors.append(f"{record['_path']}: unexpected adapter fallback") + operators = record.get("operators", {}) + for module in ("attention", "ffn", "logp"): + operator = operators.get(module) + if not isinstance(operator, dict): + errors.append(f"{record['_path']}: missing {module} readback") + continue + if operator.get("case_id") != CASE_ID: + errors.append( + f"{record['_path']}: {module} case is " f"{operator.get('case_id')!r}" + ) + if operator.get("implementation") != "production": + errors.append( + f"{record['_path']}: {module} implementation is " + f"{operator.get('implementation')!r}" + ) + if frameworks != {"megatron", "vllm"}: + errors.append(f"readback frameworks are {sorted(frameworks)!r}") + return {"passed": not errors, "errors": errors, "paths": paths} + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--num-rollout", type=int, default=3) + args = parser.parse_args() + + root = Path("/workspace/RL-Kernel-pr390") + config = FullStackConfig( + vime_root=Path("/workspace/vime"), + rl_kernel_root=root, + megatron_root=Path("/workspace/Megatron-LM-vime"), + model_root=Path("/app/model/Qwen3-8B"), + reference_checkpoint=Path("/app/model/Qwen3-8B_torch_dist"), + prompt_data=Path("/app/model/dapo-math-17k/dapo-math-17k.jsonl"), + run_dir=args.run_dir, + launcher=root / "examples/vime_rocm_attention_ablation/launch_arm.sh", + num_rollout=args.num_rollout, + rollout_batch_size=1, + samples_per_prompt=8, + global_batch_size=8, + max_response_length=7168, + max_tokens_per_gpu=4096, + seed=1234, + rollout_seed=1234, + ) + config.validate(require_paths=True) + _prepare_run_dir(args.run_dir) + frozen_before = sealed_manifest(config) + write_report(args.run_dir / "frozen-inputs.before.json", frozen_before) + + arm_dir = args.run_dir / "arms/p-p" + for directory in ( + arm_dir / "readbacks", + arm_dir / "dump", + arm_dir / "checkpoint", + arm_dir / "mismatch_sidecars", + ): + directory.mkdir(parents=True, exist_ok=False) + + environment = build_arm_environment(config, CASE_ID, arm_dir, arm_index=0) + environment.update( + { + "RL_KERNEL_ATTENTION_CASE": CASE_ID, + "RL_KERNEL_FFN_CASE": CASE_ID, + "RL_KERNEL_LOGP_CASE": CASE_ID, + "RL_KERNEL_ROCM_FIXED_PAGED_TILE": "128", + "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS": "8192", + "RLK_ABLATION_USE_ROLLOUT_LOGPROBS": "1", + "VLLM_GPU_MEMORY_UTILIZATION": "0.38", + } + ) + launch = { + "schema_version": "rlkernel.vime_rocm_full_stack_arm_launch.v1", + "case_id": CASE_ID, + "expected_implementations": { + "attention": "production", + "ffn": "production", + "logp": "production", + }, + "framework_consistency": { + "use_rollout_logprobs": True, + "mismatch_metrics_recompute": True, + }, + "frozen_input_fingerprint": frozen_before["fingerprint"], + "command": ["bash", str(config.launcher.resolve())], + "environment": public_arm_environment(environment), + "started_at": datetime.now(timezone.utc).isoformat(), + } + write_report(arm_dir / "launch.json", launch) + + with (arm_dir / "launcher.log").open("w", encoding="utf-8") as log_handle: + process = subprocess.run( + ["bash", str(config.launcher.resolve())], + cwd=config.rl_kernel_root, + env=environment, + stdout=log_handle, + stderr=subprocess.STDOUT, + check=False, + ) + + errors = [] + if process.returncode: + errors.append(f"Vime launcher exited with status {process.returncode}") + try: + readbacks = validate_native_readbacks(arm_dir / "readbacks") + except Exception as exc: + readbacks = {"passed": False, "errors": [str(exc)], "paths": []} + errors.extend(readbacks["errors"]) + rollout_identity = load_rollout_identity(arm_dir / "dump" / "rollout_data") + errors.extend(rollout_identity["errors"]) + try: + metrics = compare_train_rollout_logps( + arm_dir / "mismatch_sidecars", + require_exact=False, + tensor_parallel_size=4, + context_parallel_size=2, + ) + except Exception as exc: + metrics = {"passed": False, "errors": [str(exc)]} + errors.extend(metrics["errors"]) + + frozen_after = sealed_manifest(config) + write_report(args.run_dir / "frozen-inputs.after.json", frozen_after) + frozen_match = frozen_before["fingerprint"] == frozen_after["fingerprint"] + if not frozen_match: + errors.append("frozen source fingerprint changed during the run") + report = { + "run_dir": str(args.run_dir), + "num_rollout": args.num_rollout, + "case_id": CASE_ID, + "framework_consistency": "use_rollout_logprobs", + "launcher_returncode": process.returncode, + "passed": not errors, + "errors": errors, + "readbacks": readbacks, + "rollout_identity": rollout_identity, + "metrics": metrics, + "frozen_sources_match": frozen_match, + } + write_report(arm_dir / "validation.json", report) + write_report(args.run_dir / "single-arm-summary.json", report) + print(json.dumps(report, indent=2, sort_keys=True), flush=True) + return 0 if report["passed"] else 1 + + +if __name__ == "__main__": + sys.exit(main())