diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 711b8df4..13f93722 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -102,7 +102,11 @@ jobs: run: python -m pytest tests/test_attention_contract.py -v - name: Run WS2 Logprob Contract Tests (CPU-safe) - run: python -m pytest tests/test_logprob_contract.py -v + run: | + python -m pytest \ + tests/test_logprob_contract.py \ + tests/test_rocm_lm_head_weight_cache.py \ + -v - name: Run WS2 Vocab-Parallel Logprob Tests (CPU-safe) run: python -m pytest tests/test_vocab_parallel_logp.py -v diff --git a/benchmarks/benchmark_rocm_aiter_direct_output.py b/benchmarks/benchmark_rocm_aiter_direct_output.py new file mode 100644 index 00000000..5fbecdcf --- /dev/null +++ b/benchmarks/benchmark_rocm_aiter_direct_output.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Benchmark direct caller-output writes in strict ROCm paged decode. + +Both arms run the full logical-page gather and the same one-row/one-KV-group +AITER CK schedule. The baseline hides the new direct-output entry point so the +runtime reproduces PR #390's staged group outputs plus ``torch.cat(..., out=)``. +The candidate lets AITER write each group directly into the caller's vLLM +output slice. Output and LSE must remain byte-identical. + +Example on one MI300X: + + HIP_VISIBLE_DEVICES=2 CUDA_VISIBLE_DEVICES=2 \ + python benchmarks/benchmark_rocm_aiter_direct_output.py \ + --cached-lengths 32,128,512,2048 --blocks 10 --iterations 100 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +from collections.abc import Callable + +import torch + +from rl_engine.kernels.ops.rocm.attention.flash_attn import ( + StrictRocmAiterCKAttentionCore, +) +from rl_engine.kernels.ops.rocm.attention.strict_runtime import StrictRocmAttentionRuntime + + +class _StagedCore: + """Expose only the pre-optimization core interface to the runtime.""" + + core_id = StrictRocmAiterCKAttentionCore.core_id + strict_schedule = StrictRocmAiterCKAttentionCore.strict_schedule + backend_id = StrictRocmAiterCKAttentionCore.backend_id + + def __init__(self, delegate: StrictRocmAiterCKAttentionCore) -> None: + self._delegate = delegate + + def forward_with_lse(self, *args, **kwargs): + return self._delegate.forward_with_lse(*args, **kwargs) + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _elapsed_ms(function: Callable[[], object], iterations: int) -> float: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def _benchmark_length( + *, + cached_length: int, + batch: int, + query_heads: int, + key_heads: int, + head_dim: int, + page_size: int, + warmup: int, + blocks: int, + iterations: int, +) -> dict[str, object]: + pages_per_row = (cached_length + page_size - 1) // page_size + total_pages = batch * pages_per_row + generator = torch.Generator(device="cpu").manual_seed(390_500 + cached_length) + q = torch.randn( + batch, + query_heads, + 1, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + k_cache = torch.randn( + total_pages, + page_size, + key_heads, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + v_cache = torch.randn( + total_pages, + page_size, + key_heads, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + page_table = torch.arange(total_pages, device="cuda", dtype=torch.int32).reshape( + batch, + pages_per_row, + ) + seqused_k = torch.full((batch,), cached_length, device="cuda", dtype=torch.int32) + cached_lengths = (cached_length,) * batch + core = StrictRocmAiterCKAttentionCore() + staged_runtime = StrictRocmAttentionRuntime(core=_StagedCore(core)) + direct_runtime = StrictRocmAttentionRuntime(core=core) + staged_out = torch.empty_like(q) + direct_out = torch.empty_like(q) + common = { + "page_table": page_table, + "seqused_k": seqused_k, + "max_seqlen_k": pages_per_row * page_size, + "scale": head_dim**-0.5, + "cached_lengths": cached_lengths, + } + + @torch.inference_mode() + def staged(): + return staged_runtime.forward_paged_with_lse( + q, + k_cache, + v_cache, + out=staged_out, + **common, + ) + + @torch.inference_mode() + def direct(): + return direct_runtime.forward_paged_with_lse( + q, + k_cache, + v_cache, + out=direct_out, + **common, + ) + + for _ in range(warmup): + staged() + direct() + torch.cuda.synchronize() + + staged_result = staged() + direct_result = direct() + torch.cuda.synchronize() + output_equal = torch.equal( + staged_result.out.view(torch.uint8), + direct_result.out.view(torch.uint8), + ) + lse_equal = torch.equal( + staged_result.lse.view(torch.uint8), + direct_result.lse.view(torch.uint8), + ) + if not output_equal or not lse_equal: + raise AssertionError("direct-output paged decode differs from the staged baseline") + if staged_result.out.data_ptr() != staged_out.data_ptr(): + raise AssertionError("baseline did not return its staged caller output buffer") + if staged_result.provenance["core_output_staging"] != "runtime_group_cat": + raise AssertionError("baseline no longer reproduces PR #390 output staging") + if direct_result.out.data_ptr() != direct_out.data_ptr(): + raise AssertionError("candidate did not return the caller's output buffer") + if direct_result.provenance["core_output_staging"] != "aiter_direct_caller_group": + raise AssertionError("candidate did not enter the direct AITER output path") + + samples: dict[str, list[float]] = {"staged": [], "direct": []} + functions = {"staged": staged, "direct": direct} + for block in range(blocks): + order = ("staged", "direct") if block % 2 == 0 else ("direct", "staged") + for name in order: + samples[name].append(_elapsed_ms(functions[name], iterations)) + + staged_ms = statistics.median(samples["staged"]) + direct_ms = statistics.median(samples["direct"]) + return { + "cached_length": cached_length, + "batch": batch, + "query_heads": query_heads, + "key_heads": key_heads, + "head_dim": head_dim, + "staged_ms": staged_ms, + "direct_ms": direct_ms, + "latency_reduction_percent": 100.0 * (staged_ms - direct_ms) / staged_ms, + "speedup": staged_ms / direct_ms, + "raw_output_equal": output_equal, + "raw_lse_equal": lse_equal, + "output_sha256": _digest(direct_result.out), + "lse_sha256": _digest(direct_result.lse), + "staged_samples_ms": samples["staged"], + "direct_samples_ms": samples["direct"], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cached-lengths", default="32,128,512,2048") + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--query-heads", type=int, default=8) + parser.add_argument("--key-heads", type=int, default=2) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--page-size", type=int, default=16) + parser.add_argument("--warmup", type=int, default=30) + parser.add_argument("--blocks", type=int, default=10) + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires ROCm PyTorch and a visible GPU") + results = [ + _benchmark_length( + cached_length=int(cached_length), + batch=args.batch, + query_heads=args.query_heads, + key_heads=args.key_heads, + head_dim=args.head_dim, + page_size=args.page_size, + warmup=args.warmup, + blocks=args.blocks, + iterations=args.iterations, + ) + for cached_length in args.cached_lengths.split(",") + ] + print( + json.dumps( + { + "device": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "blocks": args.blocks, + "iterations_per_block": args.iterations, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_det_gemm_root_output.py b/benchmarks/benchmark_rocm_det_gemm_root_output.py new file mode 100644 index 00000000..44ada2d6 --- /dev/null +++ b/benchmarks/benchmark_rocm_det_gemm_root_output.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""A/B the gfx942 deterministic-GEMM root copy against direct output. + +Both arms launch the same leaf kernels and canonical BF16 reduction tree with +preallocated storage. The legacy arm stores the final node in the workspace +and copies it to the output; the candidate stores that same final node directly +in the output. Blocks are run in alternating AB/BA order to limit clock drift. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import subprocess +import sys +from dataclasses import asdict, dataclass +from pathlib import Path + +import torch +import triton + +from rl_engine.kernels.ops.triton.matmul import det_gemm + + +@dataclass(frozen=True) +class Case: + name: str + m_size: int + k_size: int + n_size: int + + +_CASES = { + case.name: case + for case in ( + Case("qwen3_tp4_decode_qkv", 1, 4096, 1536), + Case("qwen3_tp4_decode_o_proj", 1, 1024, 4096), + Case("qwen3_tp4_decode_gate", 1, 4096, 3072), + Case("qwen3_tp4_decode_down", 1, 3072, 4096), + ) +} + + +@dataclass +class CaseState: + case: Case + a: torch.Tensor + b: torch.Tensor + workspace: torch.Tensor + output: torch.Tensor + plan: object + leaf_config: object + + +def _raw_sha256(tensor: torch.Tensor) -> str: + raw = tensor.contiguous().view(torch.uint8).cpu().numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _launch(state: CaseState, *, direct_root_output: bool) -> None: + case = state.case + config = state.leaf_config + plan = state.plan + tiles_m = triton.cdiv(case.m_size, config.block_m) + tiles_n = triton.cdiv(case.n_size, config.block_n) + leaf_grid = ( + (tiles_n, tiles_m, len(plan.host.leaf_nodes)) + if config.n_fastest + else (len(plan.host.leaf_nodes), tiles_m, tiles_n) + ) + det_gemm._det_gemm_kernel[leaf_grid]( + state.a, + state.b, + state.workspace, + plan.leaf_starts, + plan.leaf_lengths, + plan.leaf_nodes, + M=case.m_size, + N=case.n_size, + K=case.k_size, + stride_am=state.a.stride(0), + stride_ak=state.a.stride(1), + stride_bk=state.b.stride(0), + stride_bn=state.b.stride(1), + BLOCK_M=config.block_m, + BLOCK_N=config.block_n, + BLOCK_K=det_gemm._BLOCK_K, + N_FASTEST=config.n_fastest, + num_warps=config.num_warps, + ) + + reduction_block = 256 + for level_index, (operations, (lower, upper, output)) in enumerate( + zip( + plan.host.reduction_levels, + plan.reduction_levels, + strict=True, + ) + ): + final_level = level_index == len(plan.host.reduction_levels) - 1 + if direct_root_output and final_level: + if len(operations) != 1: + raise RuntimeError("final deterministic GEMM tree level must have one root") + grid = (triton.cdiv(case.m_size * case.n_size, reduction_block),) + det_gemm._det_gemm_tree_reduce_to_output_kernel[grid]( + state.workspace, + state.output, + lower, + upper, + M=case.m_size, + N=case.n_size, + BLOCK=reduction_block, + ) + else: + grid = ( + len(operations), + triton.cdiv(case.m_size * case.n_size, reduction_block), + ) + det_gemm._det_gemm_tree_reduce_kernel[grid]( + state.workspace, + lower, + upper, + output, + M=case.m_size, + N=case.n_size, + BLOCK=reduction_block, + ) + + if not direct_root_output: + det_gemm._copy_tree_root_kernel[ + (triton.cdiv(state.output.numel(), reduction_block),) + ]( + state.workspace, + state.output, + plan.host.root, + state.output.numel(), + BLOCK=reduction_block, + ) + + +def _block_median_ms(state: CaseState, direct: bool, samples: int) -> float: + starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples)] + for start, end in zip(starts, ends, strict=True): + start.record() + _launch(state, direct_root_output=direct) + end.record() + torch.cuda.synchronize() + return statistics.median(start.elapsed_time(end) for start, end in zip(starts, ends)) + + +def _run_case(case: Case, *, warmup: int, samples: int, blocks: int) -> dict[str, object]: + a = torch.randn((case.m_size, case.k_size), device="cuda", dtype=torch.bfloat16) + b = torch.randn((case.k_size, case.n_size), device="cuda", dtype=torch.bfloat16) + plan = det_gemm._device_tree_plan(case.k_size, a.device) + config = det_gemm._tree_leaf_config( + a.device, + case.m_size, + case.k_size, + case.n_size, + transpose_output=False, + preserve_a_strides=False, + ) + state = CaseState( + case=case, + a=a, + b=b, + workspace=torch.empty( + (plan.host.node_count, case.m_size, case.n_size), + device="cuda", + dtype=torch.bfloat16, + ), + output=torch.empty((case.m_size, case.n_size), device="cuda", dtype=torch.bfloat16), + plan=plan, + leaf_config=config, + ) + + for direct in (False, True): + for _ in range(warmup): + _launch(state, direct_root_output=direct) + torch.cuda.synchronize() + _launch(state, direct_root_output=False) + legacy = state.output.clone() + _launch(state, direct_root_output=True) + candidate = state.output.clone() + torch.cuda.synchronize() + raw_bytes_equal = torch.equal(legacy.view(torch.uint8), candidate.view(torch.uint8)) + if not raw_bytes_equal: + raise RuntimeError(f"root-output raw-byte mismatch for {case.name}") + + series = {"legacy_copy_ms": [], "direct_output_ms": []} + for block_index in range(blocks): + arms = ((False, "legacy_copy_ms"), (True, "direct_output_ms")) + if block_index % 2: + arms = tuple(reversed(arms)) + for direct, label in arms: + series[label].append(_block_median_ms(state, direct, samples)) + legacy_ms = statistics.median(series["legacy_copy_ms"]) + direct_ms = statistics.median(series["direct_output_ms"]) + return { + "case": asdict(case), + "tree": { + "leaf_count": len(plan.host.leaf_nodes), + "reduction_levels": len(plan.host.reduction_levels), + }, + "leaf_config": { + "block_m": config.block_m, + "block_n": config.block_n, + "num_warps": config.num_warps, + "n_fastest": config.n_fastest, + }, + "raw_bytes_equal": raw_bytes_equal, + "sha256": _raw_sha256(candidate), + "series": series, + "median_ms": {"legacy_copy": legacy_ms, "direct_output": direct_ms}, + "speedup_percent": (legacy_ms / direct_ms - 1.0) * 100.0, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--cases", default="qwen3_tp4_decode_qkv,qwen3_tp4_decode_o_proj") + parser.add_argument("--warmup", type=int, default=60) + parser.add_argument("--samples", type=int, default=400) + parser.add_argument("--blocks", type=int, default=10) + parser.add_argument("--seed", type=int, default=20260906) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires a ROCm GPU") + torch.cuda.set_device(args.device) + properties = torch.cuda.get_device_properties(args.device) + arch = str(getattr(properties, "gcnArchName", "")).partition(":")[0] + if arch != "gfx942": + raise RuntimeError(f"root-output promotion is qualified only on gfx942, got {arch!r}") + selected = [name.strip() for name in args.cases.split(",") if name.strip()] + unknown = sorted(set(selected) - set(_CASES)) + if unknown: + raise ValueError(f"unknown cases {unknown}; choices are {sorted(_CASES)}") + if args.warmup < 0 or args.samples <= 0 or args.blocks <= 0: + raise ValueError("warmup must be non-negative; samples and blocks must be positive") + + torch.manual_seed(args.seed) + results = [ + _run_case(_CASES[name], warmup=args.warmup, samples=args.samples, blocks=args.blocks) + for name in selected + ] + legacy_sum = sum(item["median_ms"]["legacy_copy"] for item in results) + direct_sum = sum(item["median_ms"]["direct_output"] for item in results) + payload = { + "environment": { + "device": properties.name, + "architecture": str(getattr(properties, "gcnArchName", "")), + "torch": torch.__version__, + "hip": torch.version.hip, + "triton": triton.__version__, + "git_commit": subprocess.run( + ["git", "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip(), + "command": sys.argv, + }, + "methodology": { + "timing": "GPU events around preallocated complete tree launches", + "order": "AB/BA alternating blocks", + "warmup": args.warmup, + "samples_per_block": args.samples, + "blocks": args.blocks, + "seed": args.seed, + }, + "results": results, + "combined_projection_medians_ms": { + "legacy_copy": legacy_sum, + "direct_output": direct_sum, + "speedup_percent": (legacy_sum / direct_sum - 1.0) * 100.0, + }, + } + rendered = json.dumps(payload, indent=2, sort_keys=True) + print(rendered) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_head_major_page_gather.py b/benchmarks/benchmark_rocm_head_major_page_gather.py new file mode 100644 index 00000000..6491e521 --- /dev/null +++ b/benchmarks/benchmark_rocm_head_major_page_gather.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Benchmark one-pass head-major page gathers for ROCm paged decode. + +By default both arms consume the pair-axis layout used by the vLLM revision in +the ROCm VIME image: ``[blocks, 2, block, kv_heads, head_dim]``. The adapter's +packed-last compatibility layout can be selected explicitly. The legacy arm +first copies pages in token-major order and then copies again to transpose +them; the candidate indexes a head-major view and materializes the final order +in one pass. AITER CK launch count/order and all arithmetic remain unchanged. +This is an isolated strict-runtime timing, not an end-to-end VIME result. + +The candidate is active only for gradient-free ``key_heads > 1`` and +``cached_length > 1``. Singleton cases intentionally retain the legacy path +and are covered as fallback controls in the regression tests, not timed here. + +Example on one MI300X: + + HIP_VISIBLE_DEVICES=2 CUDA_VISIBLE_DEVICES=2 \ + python benchmarks/benchmark_rocm_head_major_page_gather.py \ + --cached-lengths 32,127,512,2048 --batch 1 --layers 36 \ + --blocks 10 --iterations 10 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import time +from collections.abc import Callable + +import torch + +from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore +from rl_engine.kernels.ops.rocm.attention.strict_runtime import StrictRocmAttentionRuntime + + +class _LegacyGatherRuntime(StrictRocmAttentionRuntime): + """Restore the two-pass page-major materialization for the baseline.""" + + @staticmethod + def _gather_paged_row( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_row: torch.Tensor, + cached_length: int, + *, + validate_bounds: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + page_size = k_cache.size(1) + page_count = (cached_length + page_size - 1) // page_size + if page_count > page_row.numel(): + raise ValueError("page_table row is shorter than the cached length requires") + pages = page_row[:page_count] + if validate_bounds: + bounds_ok = torch.all((pages >= 0) & (pages < k_cache.size(0))) + if pages.is_cuda: + torch._assert_async(bounds_ok, "page_table entries are outside the KV cache") + elif not bool(bounds_ok.item()): + raise ValueError("page_table entries are outside the KV cache") + + def gather(cache: torch.Tensor) -> torch.Tensor: + selected = cache.index_select(0, pages) + flat = selected.reshape(page_count * page_size, cache.size(2), cache.size(3)) + return flat[:cached_length].permute(1, 0, 2).unsqueeze(0).contiguous() + + return gather(k_cache), gather(v_cache) + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _elapsed_ms(function: Callable[[], object], iterations: int) -> tuple[float, float]: + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + wall_start = time.perf_counter() + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + wall_ms = (time.perf_counter() - wall_start) * 1_000.0 / iterations + return start.elapsed_time(end) / iterations, wall_ms + + +def _benchmark_length( + *, + cached_length: int, + batch: int, + layers: int, + query_heads: int, + key_heads: int, + head_dim: int, + page_size: int, + kv_layout: str, + dtype: torch.dtype, + warmup: int, + blocks: int, + iterations: int, +) -> dict[str, object]: + if key_heads <= 1 or cached_length <= 1: + raise ValueError("head-major gather benchmark requires key_heads > 1 and length > 1") + pages_per_row = (cached_length + page_size - 1) // page_size + total_pages = batch * pages_per_row + generator = torch.Generator(device="cpu").manual_seed(390_900 + cached_length + batch) + q = torch.randn( + batch, + query_heads, + 1, + head_dim, + dtype=dtype, + generator=generator, + ).cuda() + if kv_layout == "vllm-rocm-5d": + # vLLM ROCM_AITER_FA: [blocks, 2, block, kv_heads, head_dim]. Keep a + # leading simulated-layer axis while preserving each layer's strides. + packed_kv = torch.randn( + layers, + total_pages, + 2, + page_size, + key_heads, + head_dim, + dtype=dtype, + generator=generator, + ).cuda() + k_caches, v_caches = packed_kv.unbind(2) + elif kv_layout == "packed-last": + # Adapter compatibility layout: [blocks, kv_heads, block, 2 * head_dim]. + packed_kv = torch.randn( + layers, + total_pages, + key_heads, + page_size, + 2 * head_dim, + dtype=dtype, + generator=generator, + ).cuda() + k_caches, v_caches = packed_kv.transpose(2, 3).split(head_dim, dim=-1) + else: # pragma: no cover - argparse owns the public validation. + raise ValueError(f"unknown KV layout {kv_layout!r}") + page_table = ( + torch.arange(total_pages, device="cuda", dtype=torch.int32) + .reshape(batch, pages_per_row) + .flip(1) + .contiguous() + ) + seqused_k = torch.full((batch,), cached_length, device="cuda", dtype=torch.int32) + cached_lengths = (cached_length,) * batch + core = StrictRocmAiterCKAttentionCore() + baseline_runtime = _LegacyGatherRuntime(core=core) + candidate_runtime = StrictRocmAttentionRuntime(core=core) + baseline_out = torch.empty_like(q) + candidate_out = torch.empty_like(q) + common = { + "page_table": page_table, + "seqused_k": seqused_k, + "max_seqlen_k": pages_per_row * page_size, + "scale": head_dim**-0.5, + "cached_lengths": cached_lengths, + } + + with torch.inference_mode(): + legacy_k, legacy_v = baseline_runtime._gather_paged_row( + k_caches[0], + v_caches[0], + page_table[0], + cached_length, + ) + candidate_k, candidate_v = candidate_runtime._gather_paged_row( + k_caches[0], + v_caches[0], + page_table[0], + cached_length, + ) + gather_k_equal = torch.equal( + legacy_k.view(torch.uint8), candidate_k.contiguous().view(torch.uint8) + ) + gather_v_equal = torch.equal( + legacy_v.view(torch.uint8), candidate_v.contiguous().view(torch.uint8) + ) + if not gather_k_equal or not gather_v_equal: + raise AssertionError("head-major gather changed materialized K/V bytes") + for group in range(key_heads): + if not candidate_k[:, group : group + 1].transpose(1, 2).is_contiguous(): + raise AssertionError("candidate K group would add a pre-AITER copy") + if not candidate_v[:, group : group + 1].transpose(1, 2).is_contiguous(): + raise AssertionError("candidate V group would add a pre-AITER copy") + + @torch.inference_mode() + def run_layers(runtime, out): + epoch = runtime.new_page_bounds_epoch() + result = None + for layer in range(layers): + result = runtime.forward_paged_with_lse( + q, + k_caches[layer], + v_caches[layer], + out=out, + page_bounds_epoch=epoch, + **common, + ) + if result is None: + raise AssertionError("layers must be positive") + return result + + def baseline(): + return run_layers(baseline_runtime, baseline_out) + + def candidate(): + return run_layers(candidate_runtime, candidate_out) + + for _ in range(warmup): + baseline() + candidate() + torch.cuda.synchronize() + + baseline_result = baseline() + candidate_result = candidate() + candidate_output_snapshot = candidate_result.out.clone() + candidate_lse_snapshot = candidate_result.lse.clone() + replay_result = candidate() + torch.cuda.synchronize() + output_equal = torch.equal( + baseline_result.out.view(torch.uint8), + candidate_result.out.view(torch.uint8), + ) and torch.equal( + candidate_output_snapshot.view(torch.uint8), + replay_result.out.view(torch.uint8), + ) + lse_equal = torch.equal( + baseline_result.lse.view(torch.uint8), + candidate_result.lse.view(torch.uint8), + ) and torch.equal( + candidate_lse_snapshot.view(torch.uint8), + replay_result.lse.view(torch.uint8), + ) + if not output_equal or not lse_equal: + raise AssertionError("head-major gather changed output or LSE bytes") + expected_launches = batch * key_heads + if baseline_result.provenance["core_launch_count"] != expected_launches: + raise AssertionError("baseline AITER launch count changed") + if candidate_result.provenance["core_launch_count"] != expected_launches: + raise AssertionError("candidate AITER launch count changed") + + functions = {"page_major_two_pass": baseline, "head_major_one_pass": candidate} + gpu_samples: dict[str, list[float]] = {name: [] for name in functions} + wall_samples: dict[str, list[float]] = {name: [] for name in functions} + for block in range(blocks): + order = ( + ("page_major_two_pass", "head_major_one_pass") + if block % 2 == 0 + else ("head_major_one_pass", "page_major_two_pass") + ) + for name in order: + gpu_ms, wall_ms = _elapsed_ms(functions[name], iterations) + gpu_samples[name].append(gpu_ms) + wall_samples[name].append(wall_ms) + + baseline_ms = statistics.median(gpu_samples["page_major_two_pass"]) + candidate_ms = statistics.median(gpu_samples["head_major_one_pass"]) + baseline_wall_ms = statistics.median(wall_samples["page_major_two_pass"]) + candidate_wall_ms = statistics.median(wall_samples["head_major_one_pass"]) + return { + "cached_length": cached_length, + "batch": batch, + "layers_per_forward": layers, + "query_heads": query_heads, + "key_heads": key_heads, + "head_dim": head_dim, + "dtype": str(dtype).removeprefix("torch."), + "kv_layout": kv_layout, + "kv_cache_stride": list(k_caches[0].stride()), + "head_major_active": key_heads > 1 and cached_length > 1, + "distinct_kv_cache_per_layer": True, + "benchmark_scope": "strict_runtime_native_vllm_kv_simulated_layers", + "aiter_core_launches_per_layer": expected_launches, + "median_ms": { + "gpu": { + "page_major_two_pass": baseline_ms, + "head_major_one_pass": candidate_ms, + }, + "wall": { + "page_major_two_pass": baseline_wall_ms, + "head_major_one_pass": candidate_wall_ms, + }, + }, + "latency_reduction_percent": 100.0 * (baseline_ms - candidate_ms) / baseline_ms, + "wall_latency_reduction_percent": ( + 100.0 * (baseline_wall_ms - candidate_wall_ms) / baseline_wall_ms + ), + "saved_ms_per_forward": baseline_ms - candidate_ms, + "raw_gather_k_equal": gather_k_equal, + "raw_gather_v_equal": gather_v_equal, + "raw_output_equal": output_equal, + "raw_lse_equal": lse_equal, + "gather_k_sha256": _digest(candidate_k), + "gather_v_sha256": _digest(candidate_v), + "output_sha256": _digest(candidate_result.out), + "lse_sha256": _digest(candidate_result.lse), + "samples_ms": {"gpu": gpu_samples, "wall": wall_samples}, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cached-lengths", default="32,127,512,2048") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--layers", type=int, default=36) + parser.add_argument("--query-heads", type=int, default=8) + parser.add_argument("--key-heads", type=int, default=2) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--page-size", type=int, default=16) + parser.add_argument( + "--kv-layout", + choices=("vllm-rocm-5d", "packed-last"), + default="vllm-rocm-5d", + ) + parser.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16") + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--blocks", type=int, default=10) + parser.add_argument("--iterations", type=int, default=10) + args = parser.parse_args() + + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires ROCm PyTorch and a visible GPU") + if ( + min( + args.batch, + args.layers, + args.query_heads, + args.key_heads, + args.head_dim, + args.page_size, + args.blocks, + args.iterations, + ) + <= 0 + ): + raise ValueError("benchmark dimensions, blocks, and iterations must be positive") + if args.query_heads % args.key_heads: + raise ValueError("--query-heads must be divisible by --key-heads") + dtype = getattr(torch, args.dtype) + results = [ + _benchmark_length( + cached_length=int(cached_length), + batch=args.batch, + layers=args.layers, + query_heads=args.query_heads, + key_heads=args.key_heads, + head_dim=args.head_dim, + page_size=args.page_size, + kv_layout=args.kv_layout, + dtype=dtype, + warmup=args.warmup, + blocks=args.blocks, + iterations=args.iterations, + ) + for cached_length in args.cached_lengths.split(",") + ] + print( + json.dumps( + { + "device": torch.cuda.get_device_name(), + "torch": torch.__version__, + "hip": torch.version.hip, + "dtype": args.dtype, + "kv_layout": args.kv_layout, + "blocks": args.blocks, + "iterations_per_block": args.iterations, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_lm_head_weight_cache.py b/benchmarks/benchmark_rocm_lm_head_weight_cache.py new file mode 100644 index 00000000..8ac8cbfc --- /dev/null +++ b/benchmarks/benchmark_rocm_lm_head_weight_cache.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""A/B the ROCm LM-head transpose hot path against a prepared weight cache. + +The baseline follows the current ``RocmDetGemmOp.linear`` path and therefore +materializes ``weight.T`` for every projection. The candidate prepares that +same contiguous logical ``[K, N]`` right-hand side once, outside the timed +region, and passes it to ``linear_prepared``. Both arms execute the same +deterministic Triton GEMM tree. Blocks alternate AB/BA order to limit clock +drift. + +The defaults reproduce one Qwen3-8B TP4 decode LM-head projection on MI300X:: + + HIP_VISIBLE_DEVICES=2 CUDA_VISIBLE_DEVICES=2 \ + python benchmarks/benchmark_rocm_lm_head_weight_cache.py \ + --blocks 10 --iterations 400 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import subprocess +import sys +import time +from collections.abc import Callable +from pathlib import Path + +import torch + +from rl_engine.integrations.vllm_runtime import ( + _refresh_lm_head_weight_cache, + _validated_lm_head_weight_cache, +) +from rl_engine.kernels.ops.rocm.matmul.det_gemm import RocmDetGemmOp, prepare_det_gemm_linear_weight + + +def _raw_sha256(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _elapsed_ms(function: Callable[[], object], iterations: int) -> tuple[float, float]: + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + wall_start = time.perf_counter() + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + wall_ms = (time.perf_counter() - wall_start) * 1_000.0 / iterations + return start.elapsed_time(end) / iterations, wall_ms + + +def _benchmark( + *, + m_size: int, + n_size: int, + k_size: int, + warmup: int, + iterations: int, + blocks: int, +) -> dict[str, object]: + activation = torch.randn( + (m_size, k_size), + device="cuda", + dtype=torch.bfloat16, + ) + + class LmHeadLayer(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter( + torch.randn( + (n_size, k_size), + device="cuda", + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + + layer = LmHeadLayer() + weight = layer.weight + operator = RocmDetGemmOp() + + # Preparation is deliberately outside the timed steady-state projection. + # Weight synchronization refreshes this same storage before the next + # projection. Track the actual persistent allocator cost as well. + allocated_before = torch.cuda.memory_allocated() + reserved_before = torch.cuda.memory_reserved() + state = _refresh_lm_head_weight_cache( + layer, + prepare_det_gemm_linear_weight, + ) + prepared_weight_t = state.weight_t + allocated_after = torch.cuda.memory_allocated() + reserved_after = torch.cuda.memory_reserved() + + @torch.inference_mode() + def materialize_each_call() -> torch.Tensor: + return operator.linear(activation, weight) + + @torch.inference_mode() + def prepared_cache() -> torch.Tensor: + return operator.linear_prepared( + activation, + _validated_lm_head_weight_cache( + layer, + prepare_det_gemm_linear_weight, + ), + ) + + functions = { + "materialize_each_call": materialize_each_call, + "prepared_cache": prepared_cache, + } + for _ in range(warmup): + materialize_each_call() + prepared_cache() + torch.cuda.synchronize() + + baseline_output = materialize_each_call() + candidate_output = prepared_cache() + torch.cuda.synchronize() + raw_bytes_equal = torch.equal( + baseline_output.view(torch.uint8), + candidate_output.view(torch.uint8), + ) + hashes = { + "materialize_each_call": _raw_sha256(baseline_output), + "prepared_cache": _raw_sha256(candidate_output), + } + if not raw_bytes_equal or len(set(hashes.values())) != 1: + raise RuntimeError("prepared LM-head weight changed deterministic output bytes") + + gpu_samples: dict[str, list[float]] = {name: [] for name in functions} + wall_samples: dict[str, list[float]] = {name: [] for name in functions} + for block_index in range(blocks): + order = ( + ("materialize_each_call", "prepared_cache") + if block_index % 2 == 0 + else ("prepared_cache", "materialize_each_call") + ) + for name in order: + gpu_ms, wall_ms = _elapsed_ms(functions[name], iterations) + gpu_samples[name].append(gpu_ms) + wall_samples[name].append(wall_ms) + + refresh_gpu_samples = [] + refresh_wall_samples = [] + for _ in range(blocks): + gpu_ms, wall_ms = _elapsed_ms( + lambda: _refresh_lm_head_weight_cache( + layer, + prepare_det_gemm_linear_weight, + ), + 1, + ) + refresh_gpu_samples.append(gpu_ms) + refresh_wall_samples.append(wall_ms) + + baseline_ms = statistics.median(gpu_samples["materialize_each_call"]) + candidate_ms = statistics.median(gpu_samples["prepared_cache"]) + baseline_wall_ms = statistics.median(wall_samples["materialize_each_call"]) + candidate_wall_ms = statistics.median(wall_samples["prepared_cache"]) + return { + "shape": { + "activation_mk": [m_size, k_size], + "weight_nk": [n_size, k_size], + "output_mn": [m_size, n_size], + }, + "prepared_weight": { + "shape": list(prepared_weight_t.shape), + "stride": list(prepared_weight_t.stride()), + "contiguous": prepared_weight_t.is_contiguous(), + "bytes": prepared_weight_t.numel() * prepared_weight_t.element_size(), + "allocated_delta_bytes": allocated_after - allocated_before, + "reserved_delta_bytes": reserved_after - reserved_before, + }, + "raw_bytes_equal": raw_bytes_equal, + "sha256": hashes, + "samples_ms": { + "gpu": gpu_samples, + "wall": wall_samples, + "refresh_gpu": refresh_gpu_samples, + "refresh_wall": refresh_wall_samples, + }, + "median_ms": { + "gpu": { + "materialize_each_call": baseline_ms, + "prepared_cache": candidate_ms, + "refresh": statistics.median(refresh_gpu_samples), + }, + "wall": { + "materialize_each_call": baseline_wall_ms, + "prepared_cache": candidate_wall_ms, + "refresh": statistics.median(refresh_wall_samples), + }, + }, + "latency_reduction_percent": 100.0 * (baseline_ms - candidate_ms) / baseline_ms, + "wall_latency_reduction_percent": ( + 100.0 * (baseline_wall_ms - candidate_wall_ms) / baseline_wall_ms + ), + "speedup": baseline_ms / candidate_ms, + "wall_speedup": baseline_wall_ms / candidate_wall_ms, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--m-size", type=int, default=1) + parser.add_argument("--n-size", type=int, default=38_016) + parser.add_argument("--k-size", type=int, default=4_096) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=400) + parser.add_argument("--blocks", type=int, default=10) + parser.add_argument("--seed", type=int, default=20260906) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires ROCm PyTorch and a visible GPU") + if min(args.m_size, args.n_size, args.k_size) <= 0: + raise ValueError("M, N, and K must be positive") + if args.warmup < 0 or args.iterations <= 0 or args.blocks <= 0: + raise ValueError("warmup must be non-negative; iterations and blocks must be positive") + + torch.cuda.set_device(args.device) + torch.manual_seed(args.seed) + properties = torch.cuda.get_device_properties(args.device) + result = _benchmark( + m_size=args.m_size, + n_size=args.n_size, + k_size=args.k_size, + warmup=args.warmup, + iterations=args.iterations, + blocks=args.blocks, + ) + payload = { + "environment": { + "device": properties.name, + "architecture": str(getattr(properties, "gcnArchName", "")), + "torch": torch.__version__, + "hip": torch.version.hip, + "git_commit": subprocess.run( + ["git", "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip(), + "command": sys.argv, + }, + "methodology": { + "timing": ( + "GPU events and synchronized host wall time around complete " + "projection calls; initial preparation is outside the steady-state " + "region; refresh is reported separately" + ), + "order": "AB/BA alternating blocks", + "warmup": args.warmup, + "iterations_per_block": args.iterations, + "blocks": args.blocks, + "seed": args.seed, + }, + "result": result, + } + rendered = json.dumps(payload, indent=2, sort_keys=True) + print(rendered) + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_paged_bounds_cache.py b/benchmarks/benchmark_rocm_paged_bounds_cache.py new file mode 100644 index 00000000..acbc25c5 --- /dev/null +++ b/benchmarks/benchmark_rocm_paged_bounds_cache.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Benchmark one page-bounds check per simulated decoder forward on ROCm. + +This is an isolated strict-runtime benchmark, not an end-to-end vLLM timing. +Each timed iteration simulates every Qwen3 decoder layer over the same paged +metadata. The baseline validates physical page indices in every layer; the +candidate validates the first layer and reuses that runtime-owned proof for +the remaining layers. Page gathers, AITER calls, outputs, and LSE are otherwise +identical and must remain byte-exact. End-to-end benefit must still be measured +with the VIME rollout A/B. + +Example on one MI300X: + + HIP_VISIBLE_DEVICES=2 CUDA_VISIBLE_DEVICES=2 \ + python benchmarks/benchmark_rocm_paged_bounds_cache.py \ + --cached-lengths 32,128,512,2048 --layers 36 \ + --blocks 10 --iterations 10 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import time +from collections.abc import Callable + +import torch + +from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore +from rl_engine.kernels.ops.rocm.attention.strict_runtime import StrictRocmAttentionRuntime + + +class _CountingRuntime(StrictRocmAttentionRuntime): + """Count validation rows outside the timed benchmark.""" + + def __init__(self, *, core: StrictRocmAiterCKAttentionCore) -> None: + super().__init__(core=core) + self.bounds_validation_rows = 0 + self.core_launches = 0 + + def _gather_paged_row( + self, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_row: torch.Tensor, + cached_length: int, + *, + validate_bounds: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + if validate_bounds: + self.bounds_validation_rows += 1 + return StrictRocmAttentionRuntime._gather_paged_row( + k_cache, + v_cache, + page_row, + cached_length, + validate_bounds=validate_bounds, + ) + + def _run_core(self, *args, **kwargs): + result = super()._run_core(*args, **kwargs) + self.core_launches += result[3] + return result + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _elapsed_ms(function: Callable[[], object], iterations: int) -> tuple[float, float]: + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + wall_start = time.perf_counter() + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + wall_ms = (time.perf_counter() - wall_start) * 1_000.0 / iterations + return start.elapsed_time(end) / iterations, wall_ms + + +def _benchmark_length( + *, + cached_length: int, + batch: int, + layers: int, + query_heads: int, + key_heads: int, + head_dim: int, + page_size: int, + warmup: int, + blocks: int, + iterations: int, +) -> dict[str, object]: + pages_per_row = (cached_length + page_size - 1) // page_size + total_pages = batch * pages_per_row + generator = torch.Generator(device="cpu").manual_seed(390_700 + cached_length + batch) + q = torch.randn( + batch, + query_heads, + 1, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + k_cache = torch.randn( + total_pages, + page_size, + key_heads, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + v_cache = torch.randn( + total_pages, + page_size, + key_heads, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + page_table = torch.arange(total_pages, device="cuda", dtype=torch.int32).reshape( + batch, + pages_per_row, + ) + seqused_k = torch.full((batch,), cached_length, device="cuda", dtype=torch.int32) + cached_lengths = (cached_length,) * batch + core = StrictRocmAiterCKAttentionCore() + counting_runtime = _CountingRuntime(core=core) + baseline_runtime = StrictRocmAttentionRuntime(core=core) + candidate_runtime = StrictRocmAttentionRuntime(core=core) + counting_out = torch.empty_like(q) + baseline_out = torch.empty_like(q) + candidate_out = torch.empty_like(q) + common = { + "page_table": page_table, + "seqused_k": seqused_k, + "max_seqlen_k": pages_per_row * page_size, + "scale": head_dim**-0.5, + "cached_lengths": cached_lengths, + } + + @torch.inference_mode() + def run_layers(runtime, out, *, scoped): + epoch = runtime.new_page_bounds_epoch() if scoped else None + result = None + for _ in range(layers): + result = runtime.forward_paged_with_lse( + q, + k_cache, + v_cache, + out=out, + page_bounds_epoch=epoch, + **common, + ) + if result is None: + raise AssertionError("layers must be positive") + return result + + counting_runtime.bounds_validation_rows = 0 + counting_runtime.core_launches = 0 + run_layers(counting_runtime, counting_out, scoped=False) + baseline_validation_rows = counting_runtime.bounds_validation_rows + baseline_core_launches = counting_runtime.core_launches + counting_runtime.bounds_validation_rows = 0 + counting_runtime.core_launches = 0 + run_layers(counting_runtime, counting_out, scoped=True) + candidate_validation_rows = counting_runtime.bounds_validation_rows + candidate_core_launches = counting_runtime.core_launches + expected_validation_rows = {"baseline": batch * layers, "candidate": batch} + if { + "baseline": baseline_validation_rows, + "candidate": candidate_validation_rows, + } != expected_validation_rows: + raise AssertionError("page-bounds validation count changed") + expected_core_launches = batch * key_heads * layers + if baseline_core_launches != expected_core_launches: + raise AssertionError("baseline AITER core launch count changed") + if candidate_core_launches != expected_core_launches: + raise AssertionError("candidate AITER core launch count changed") + + def baseline(): + return run_layers(baseline_runtime, baseline_out, scoped=False) + + def candidate(): + return run_layers(candidate_runtime, candidate_out, scoped=True) + + for _ in range(warmup): + baseline() + candidate() + torch.cuda.synchronize() + + baseline_result = baseline() + candidate_result = candidate() + candidate_output_snapshot = candidate_result.out.clone() + repeated_result = candidate() + torch.cuda.synchronize() + output_equal = torch.equal( + baseline_result.out.view(torch.uint8), + candidate_result.out.view(torch.uint8), + ) and torch.equal( + candidate_output_snapshot.view(torch.uint8), + repeated_result.out.view(torch.uint8), + ) + lse_equal = torch.equal( + baseline_result.lse.view(torch.uint8), + candidate_result.lse.view(torch.uint8), + ) and torch.equal( + candidate_result.lse.view(torch.uint8), + repeated_result.lse.view(torch.uint8), + ) + if not output_equal or not lse_equal: + raise AssertionError("cached page-bounds validation changed output or LSE bytes") + if baseline_result.provenance["page_bounds_validation_reused"]: + raise AssertionError("baseline unexpectedly reused page validation") + if not candidate_result.provenance["page_bounds_validation_reused"]: + raise AssertionError("candidate did not reuse page validation") + + functions = {"every_layer": baseline, "once_per_forward": candidate} + gpu_samples: dict[str, list[float]] = {name: [] for name in functions} + wall_samples: dict[str, list[float]] = {name: [] for name in functions} + for block in range(blocks): + order = ( + ("every_layer", "once_per_forward") + if block % 2 == 0 + else ("once_per_forward", "every_layer") + ) + for name in order: + gpu_ms, wall_ms = _elapsed_ms(functions[name], iterations) + gpu_samples[name].append(gpu_ms) + wall_samples[name].append(wall_ms) + + baseline_ms = statistics.median(gpu_samples["every_layer"]) + candidate_ms = statistics.median(gpu_samples["once_per_forward"]) + baseline_wall_ms = statistics.median(wall_samples["every_layer"]) + candidate_wall_ms = statistics.median(wall_samples["once_per_forward"]) + return { + "cached_length": cached_length, + "batch": batch, + "layers_per_forward": layers, + "query_heads": query_heads, + "key_heads": key_heads, + "head_dim": head_dim, + "benchmark_scope": "strict_runtime_simulated_decoder_layers", + "bounds_validation_rows": expected_validation_rows, + "core_launches": { + "every_layer": baseline_core_launches, + "once_per_forward": candidate_core_launches, + }, + "median_ms": { + "gpu": { + "every_layer": baseline_ms, + "once_per_forward": candidate_ms, + }, + "wall": { + "every_layer": baseline_wall_ms, + "once_per_forward": candidate_wall_ms, + }, + }, + "latency_reduction_percent": 100.0 * (baseline_ms - candidate_ms) / baseline_ms, + "wall_latency_reduction_percent": ( + 100.0 * (baseline_wall_ms - candidate_wall_ms) / baseline_wall_ms + ), + "saved_ms_per_forward": baseline_ms - candidate_ms, + "raw_output_equal": output_equal, + "raw_lse_equal": lse_equal, + "output_sha256": _digest(candidate_result.out), + "lse_sha256": _digest(candidate_result.lse), + "samples_ms": {"gpu": gpu_samples, "wall": wall_samples}, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cached-lengths", default="32,128,512,2048") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--layers", type=int, default=36) + parser.add_argument("--query-heads", type=int, default=8) + parser.add_argument("--key-heads", type=int, default=2) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--page-size", type=int, default=16) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--blocks", type=int, default=10) + parser.add_argument("--iterations", type=int, default=10) + args = parser.parse_args() + + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires ROCm PyTorch and a visible GPU") + if ( + min( + args.batch, + args.layers, + args.query_heads, + args.key_heads, + args.head_dim, + args.page_size, + args.blocks, + args.iterations, + ) + <= 0 + ): + raise ValueError("benchmark dimensions, blocks, and iterations must be positive") + if args.layers < 2: + raise ValueError("--layers must be at least 2 to measure cross-layer reuse") + results = [ + _benchmark_length( + cached_length=int(cached_length), + batch=args.batch, + layers=args.layers, + query_heads=args.query_heads, + key_heads=args.key_heads, + head_dim=args.head_dim, + page_size=args.page_size, + warmup=args.warmup, + blocks=args.blocks, + iterations=args.iterations, + ) + for cached_length in args.cached_lengths.split(",") + ] + print( + json.dumps( + { + "device": torch.cuda.get_device_name(), + "torch": torch.__version__, + "hip": torch.version.hip, + "dtype": "bfloat16", + "blocks": args.blocks, + "iterations_per_block": args.iterations, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_provenance_cache.py b/benchmarks/benchmark_rocm_provenance_cache.py new file mode 100644 index 00000000..cf2cd059 --- /dev/null +++ b/benchmarks/benchmark_rocm_provenance_cache.py @@ -0,0 +1,518 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Benchmark cached ROCm Attention provenance metadata lookups. + +The baseline restores the pre-cache behavior by resolving the GPU description +and immutable Split-KV execution plan for every AITER KV-group launch. The +candidate uses the production core caches, but clears only its one-entry plan +cache at the start of every simulated model forward. This models a new decode +length while retaining the process-stable device description. + +Both arms run 36 distinct layer inputs backed by vLLM's ROCm 5-D KV layout, +``[blocks, 2, block_size, kv_heads, head_dim]``. Page gathers, caller-output +writes, AITER launch count/order, and tensor arithmetic are otherwise +identical. This is an isolated strict-runtime timing, not an end-to-end VIME +result. + +Example on one MI300X:: + + HIP_VISIBLE_DEVICES=2 CUDA_VISIBLE_DEVICES=2 \ + python benchmarks/benchmark_rocm_provenance_cache.py \ + --cached-lengths 32,127,512,2048 --blocks 12 --iterations 20 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import time +from collections.abc import Callable, Sequence + +import torch + +from rl_engine.kernels.attention_contract import SplitKVExecutionPlan +from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore +from rl_engine.kernels.ops.rocm.attention.strict_runtime import StrictRocmAttentionRuntime + +_BATCH = 1 +_QUERY_HEADS = 8 +_KV_HEADS = 2 +_HEAD_DIM = 128 +_PAGE_SIZE = 16 +_LAYERS = 36 + + +class _LegacyProvenanceCore(StrictRocmAiterCKAttentionCore): + """Restore uncached host metadata construction for the baseline.""" + + def _device_description(self, device: torch.device) -> tuple[str, str]: + properties = torch.cuda.get_device_properties(device) + return properties.name, getattr(properties, "gcnArchName", "unknown") + + def _resolve_split_kv_plan(self, total_kv_tokens: int) -> SplitKVExecutionPlan: + return self.split_kv.resolve(total_kv_tokens, backend=self.backend_id) + + +class _NewDecodeCandidateCore(StrictRocmAiterCKAttentionCore): + """Use production caches while modeling one new sequence length per forward.""" + + def begin_simulated_forward(self) -> None: + # Decode normally advances to a new cached length. Clear only the + # one-entry plan cache so the first KV group resolves that new length; + # the process-stable device description deliberately remains warm. + self._split_kv_plan_cache = None + + +class _CountingLegacyCore(_LegacyProvenanceCore): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.device_helper_calls = 0 + self.device_property_lookups = 0 + self.plan_helper_calls = 0 + self.split_plan_resolutions = 0 + + def _device_description(self, device: torch.device) -> tuple[str, str]: + self.device_helper_calls += 1 + self.device_property_lookups += 1 + return super()._device_description(device) + + def _resolve_split_kv_plan(self, total_kv_tokens: int) -> SplitKVExecutionPlan: + self.plan_helper_calls += 1 + self.split_plan_resolutions += 1 + return super()._resolve_split_kv_plan(total_kv_tokens) + + +class _CountingCandidateCore(_NewDecodeCandidateCore): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.device_helper_calls = 0 + self.device_property_lookups = 0 + self.plan_helper_calls = 0 + self.split_plan_resolutions = 0 + + def _device_description(self, device: torch.device) -> tuple[str, str]: + self.device_helper_calls += 1 + cached = self._device_description_cache + if cached is None or cached[0] != device: + self.device_property_lookups += 1 + return super()._device_description(device) + + def _resolve_split_kv_plan(self, total_kv_tokens: int) -> SplitKVExecutionPlan: + self.plan_helper_calls += 1 + key = (self.split_kv, total_kv_tokens, self.backend_id) + cached = self._split_kv_plan_cache + if cached is None or cached[0] != key: + self.split_plan_resolutions += 1 + return super()._resolve_split_kv_plan(total_kv_tokens) + + +def _digest(tensors: Sequence[torch.Tensor]) -> str: + digest = hashlib.sha256() + for tensor in tensors: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + digest.update(raw) + return digest.hexdigest() + + +def _raw_equal(left: Sequence[torch.Tensor], right: Sequence[torch.Tensor]) -> bool: + if len(left) != len(right): + return False + return all( + torch.equal( + left_tensor.detach().contiguous().view(torch.uint8), + right_tensor.detach().contiguous().view(torch.uint8), + ) + for left_tensor, right_tensor in zip(left, right, strict=True) + ) + + +def _elapsed_ms(function: Callable[[], object], iterations: int) -> tuple[float, float]: + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + wall_start = time.perf_counter() + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + wall_ms = (time.perf_counter() - wall_start) * 1_000.0 / iterations + return start.elapsed_time(end) / iterations, wall_ms + + +def _counter_snapshot(core) -> dict[str, int]: + return { + "device_helper_calls": core.device_helper_calls, + "device_property_lookups": core.device_property_lookups, + "plan_helper_calls": core.plan_helper_calls, + "split_plan_resolutions": core.split_plan_resolutions, + } + + +def _counter_delta(after: dict[str, int], before: dict[str, int]) -> dict[str, int]: + return {name: after[name] - before[name] for name in after} + + +def _assert_nested_provenance_is_disjoint(left: dict, right: dict) -> None: + if left is right: + raise AssertionError("separate Attention results shared their provenance dictionary") + left_core = left["core"] + right_core = right["core"] + if left_core is right_core: + raise AssertionError("separate Attention results shared core provenance") + left_split = left_core["split_kv"] + right_split = right_core["split_kv"] + if left_split is right_split: + raise AssertionError("cached Split-KV plans leaked a mutable provenance dictionary") + left_boundaries = left_split["actual_split_boundaries"] + right_boundaries = right_split["actual_split_boundaries"] + if left_boundaries is right_boundaries: + raise AssertionError("cached Split-KV plans leaked a mutable boundaries list") + if any( + left_boundary is right_boundary + for left_boundary, right_boundary in zip( + left_boundaries, + right_boundaries, + strict=True, + ) + ): + raise AssertionError("cached Split-KV plans leaked a mutable boundary row") + + +def _benchmark_length( + *, + cached_length: int, + dtype: torch.dtype, + warmup: int, + blocks: int, + iterations: int, +) -> dict[str, object]: + pages_per_row = (cached_length + _PAGE_SIZE - 1) // _PAGE_SIZE + generator = torch.Generator(device="cpu").manual_seed(391_000 + cached_length) + layer_inputs: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + for _ in range(_LAYERS): + query = torch.randn( + _BATCH, + _QUERY_HEADS, + 1, + _HEAD_DIM, + dtype=dtype, + generator=generator, + ).cuda() + # vLLM ROCm 5-D cache: [blocks, K/V, block, KV heads, head dim]. + packed_cache = torch.randn( + pages_per_row, + 2, + _PAGE_SIZE, + _KV_HEADS, + _HEAD_DIM, + dtype=dtype, + generator=generator, + ).cuda() + key_cache, value_cache = packed_cache.unbind(1) + layer_inputs.append((query, key_cache, value_cache)) + + layer_cache_storages = { + key_cache.untyped_storage().data_ptr() for _query, key_cache, _value_cache in layer_inputs + } + if len(layer_cache_storages) != _LAYERS: + raise AssertionError("simulated decoder layers must own distinct KV cache storage") + + page_table = ( + torch.arange(pages_per_row, device="cuda", dtype=torch.int32) + .reshape(_BATCH, pages_per_row) + .flip(1) + .contiguous() + ) + seqused_k = torch.full((_BATCH,), cached_length, device="cuda", dtype=torch.int32) + cached_lengths = (cached_length,) + + prototype = StrictRocmAiterCKAttentionCore() + core_kwargs = { + "_mha_fwd": prototype._mha_fwd, + "_mha_bwd": prototype._mha_bwd, + "_source_sha256": prototype.source_sha256, + } + baseline_core = _LegacyProvenanceCore(**core_kwargs) + candidate_core = _NewDecodeCandidateCore(**core_kwargs) + baseline_runtime = StrictRocmAttentionRuntime(core=baseline_core) + candidate_runtime = StrictRocmAttentionRuntime(core=candidate_core) + baseline_outputs = [torch.empty_like(query) for query, _key, _value in layer_inputs] + candidate_outputs = [torch.empty_like(query) for query, _key, _value in layer_inputs] + + @torch.inference_mode() + def run_layers(runtime, outputs, *, capture_lse): + epoch = runtime.new_page_bounds_epoch() + last_result = None + captured_lse = [] + launches = 0 + for (query, key_cache, value_cache), output in zip( + layer_inputs, + outputs, + strict=True, + ): + last_result = runtime.forward_paged_with_lse( + query, + key_cache, + value_cache, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_k=pages_per_row * _PAGE_SIZE, + scale=_HEAD_DIM**-0.5, + out=output, + cached_lengths=cached_lengths, + page_bounds_epoch=epoch, + ) + if capture_lse: + captured_lse.append(last_result.lse) + launches += int(last_result.provenance["core_launch_count"]) + if last_result is None: + raise AssertionError("the simulated decoder must contain at least one layer") + return last_result, tuple(captured_lse), launches + + def baseline(): + return run_layers(baseline_runtime, baseline_outputs, capture_lse=False)[0] + + def candidate(): + candidate_core.begin_simulated_forward() + return run_layers(candidate_runtime, candidate_outputs, capture_lse=False)[0] + + counting_baseline_core = _CountingLegacyCore(**core_kwargs) + counting_candidate_core = _CountingCandidateCore(**core_kwargs) + counting_baseline_runtime = StrictRocmAttentionRuntime(core=counting_baseline_core) + counting_candidate_runtime = StrictRocmAttentionRuntime(core=counting_candidate_core) + counting_baseline_outputs = [torch.empty_like(query) for query, _key, _value in layer_inputs] + counting_candidate_outputs = [torch.empty_like(query) for query, _key, _value in layer_inputs] + + def counted_baseline_forward(): + before = _counter_snapshot(counting_baseline_core) + result = run_layers( + counting_baseline_runtime, + counting_baseline_outputs, + capture_lse=True, + ) + return result, _counter_delta(_counter_snapshot(counting_baseline_core), before) + + def counted_candidate_forward(): + counting_candidate_core.begin_simulated_forward() + before = _counter_snapshot(counting_candidate_core) + result = run_layers( + counting_candidate_runtime, + counting_candidate_outputs, + capture_lse=True, + ) + return result, _counter_delta(_counter_snapshot(counting_candidate_core), before) + + (baseline_counted_first, baseline_counts_first) = counted_baseline_forward() + (baseline_counted_next, baseline_counts_next) = counted_baseline_forward() + (candidate_counted_first, candidate_counts_first) = counted_candidate_forward() + (candidate_counted_next, candidate_counts_next) = counted_candidate_forward() + expected_launches = _LAYERS * _BATCH * _KV_HEADS + for result, _captured_lse, launches in ( + baseline_counted_first, + baseline_counted_next, + candidate_counted_first, + candidate_counted_next, + ): + if launches != expected_launches: + raise AssertionError("AITER launch count changed while caching host provenance") + if result.provenance["core_launch_count"] != _BATCH * _KV_HEADS: + raise AssertionError("per-layer AITER launch count changed") + + expected_helper_calls = expected_launches + expected_counts = { + "baseline_first_forward": { + "device_helper_calls": expected_helper_calls, + "device_property_lookups": expected_helper_calls, + "plan_helper_calls": expected_helper_calls, + "split_plan_resolutions": expected_helper_calls, + }, + "baseline_next_forward": { + "device_helper_calls": expected_helper_calls, + "device_property_lookups": expected_helper_calls, + "plan_helper_calls": expected_helper_calls, + "split_plan_resolutions": expected_helper_calls, + }, + "candidate_first_forward": { + "device_helper_calls": expected_helper_calls, + "device_property_lookups": 1, + "plan_helper_calls": expected_helper_calls, + "split_plan_resolutions": 1, + }, + "candidate_next_forward": { + "device_helper_calls": expected_helper_calls, + "device_property_lookups": 0, + "plan_helper_calls": expected_helper_calls, + "split_plan_resolutions": 1, + }, + } + observed_counts = { + "baseline_first_forward": baseline_counts_first, + "baseline_next_forward": baseline_counts_next, + "candidate_first_forward": candidate_counts_first, + "candidate_next_forward": candidate_counts_next, + } + if observed_counts != expected_counts: + raise AssertionError( + f"provenance cache lookup counts changed: {observed_counts} != {expected_counts}" + ) + + for _ in range(warmup): + baseline() + candidate() + torch.cuda.synchronize() + + baseline_result, baseline_lses, baseline_launches = run_layers( + baseline_runtime, + baseline_outputs, + capture_lse=True, + ) + candidate_core.begin_simulated_forward() + candidate_result, candidate_lses, candidate_launches = run_layers( + candidate_runtime, + candidate_outputs, + capture_lse=True, + ) + candidate_output_snapshot = tuple(output.clone() for output in candidate_outputs) + candidate_lse_snapshot = tuple(lse.clone() for lse in candidate_lses) + candidate_core.begin_simulated_forward() + replay_result, replay_lses, replay_launches = run_layers( + candidate_runtime, + candidate_outputs, + capture_lse=True, + ) + torch.cuda.synchronize() + + output_equal = _raw_equal(baseline_outputs, candidate_output_snapshot) and _raw_equal( + candidate_output_snapshot, + candidate_outputs, + ) + lse_equal = _raw_equal(baseline_lses, candidate_lse_snapshot) and _raw_equal( + candidate_lse_snapshot, + replay_lses, + ) + provenance_equal = ( + baseline_result.provenance == candidate_result.provenance == replay_result.provenance + ) + if not output_equal or not lse_equal: + raise AssertionError("provenance caching changed output or LSE bytes") + if not provenance_equal: + raise AssertionError("provenance caching changed the reported strict contract") + _assert_nested_provenance_is_disjoint( + baseline_result.provenance, + candidate_result.provenance, + ) + _assert_nested_provenance_is_disjoint( + candidate_result.provenance, + replay_result.provenance, + ) + if {baseline_launches, candidate_launches, replay_launches} != {expected_launches}: + raise AssertionError("AITER launch count changed in the exactness replay") + + functions = {"uncached_provenance": baseline, "cached_provenance": candidate} + gpu_samples: dict[str, list[float]] = {name: [] for name in functions} + wall_samples: dict[str, list[float]] = {name: [] for name in functions} + for block in range(blocks): + order = ( + ("uncached_provenance", "cached_provenance") + if block % 2 == 0 + else ("cached_provenance", "uncached_provenance") + ) + for name in order: + gpu_ms, wall_ms = _elapsed_ms(functions[name], iterations) + gpu_samples[name].append(gpu_ms) + wall_samples[name].append(wall_ms) + + baseline_ms = statistics.median(gpu_samples["uncached_provenance"]) + candidate_ms = statistics.median(gpu_samples["cached_provenance"]) + baseline_wall_ms = statistics.median(wall_samples["uncached_provenance"]) + candidate_wall_ms = statistics.median(wall_samples["cached_provenance"]) + return { + "cached_length": cached_length, + "dtype": str(dtype).removeprefix("torch."), + "layout": "vllm_rocm_[blocks,2,block,kv_heads,head_dim]", + "batch": _BATCH, + "query_heads": _QUERY_HEADS, + "kv_heads": _KV_HEADS, + "head_dim": _HEAD_DIM, + "page_size": _PAGE_SIZE, + "distinct_layers_per_forward": _LAYERS, + "distinct_layer_kv_storages": True, + "candidate_cache_model": { + "split_plan": "cleared_once_per_simulated_forward", + "device_description": "retained_across_simulated_forwards", + }, + "host_metadata_counts": observed_counts, + "aiter_core_launches_per_forward": expected_launches, + "median_ms": { + "gpu": { + "uncached_provenance": baseline_ms, + "cached_provenance": candidate_ms, + }, + "wall": { + "uncached_provenance": baseline_wall_ms, + "cached_provenance": candidate_wall_ms, + }, + }, + "latency_reduction_percent": 100.0 * (baseline_ms - candidate_ms) / baseline_ms, + "wall_latency_reduction_percent": ( + 100.0 * (baseline_wall_ms - candidate_wall_ms) / baseline_wall_ms + ), + "saved_ms_per_forward": baseline_ms - candidate_ms, + "raw_output_equal": output_equal, + "raw_lse_equal": lse_equal, + "full_provenance_equal": provenance_equal, + "nested_provenance_disjoint": True, + "output_sha256_36_layers": _digest(candidate_output_snapshot), + "lse_sha256_36_layers": _digest(candidate_lse_snapshot), + "samples_ms": {"gpu": gpu_samples, "wall": wall_samples}, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cached-lengths", default="32,127,512,2048") + parser.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16") + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--blocks", type=int, default=12) + parser.add_argument("--iterations", type=int, default=20) + args = parser.parse_args() + + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires ROCm PyTorch and a visible GPU") + if args.warmup < 0 or args.blocks <= 0 or args.iterations <= 0: + raise ValueError("warmup must be non-negative; blocks and iterations must be positive") + cached_lengths = tuple(int(value.strip()) for value in args.cached_lengths.split(",")) + if not cached_lengths or any(length <= 0 for length in cached_lengths): + raise ValueError("--cached-lengths must contain positive integers") + dtype = getattr(torch, args.dtype) + results = [ + _benchmark_length( + cached_length=cached_length, + dtype=dtype, + warmup=args.warmup, + blocks=args.blocks, + iterations=args.iterations, + ) + for cached_length in cached_lengths + ] + print( + json.dumps( + { + "device": torch.cuda.get_device_name(), + "torch": torch.__version__, + "hip": torch.version.hip, + "benchmark_scope": "strict_runtime_distinct_vllm_rocm_decoder_layers", + "blocks": args.blocks, + "iterations_per_block": args.iterations, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_rope_inference_dispatch.py b/benchmarks/benchmark_rocm_rope_inference_dispatch.py new file mode 100644 index 00000000..d0c672b5 --- /dev/null +++ b/benchmarks/benchmark_rocm_rope_inference_dispatch.py @@ -0,0 +1,233 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Benchmark ROCm RoPE rollout dispatch and static-frequency preparation. + +The three arms reproduce the old autograd path, direct inference with a freshly +built ``inv_freq``, and direct inference with that position-independent vector +cached. All arms rebuild position-dependent cos/sin and launch the same two +deterministic HIP kernels. Query and key outputs must remain byte-identical. + +Example on one MI300X: + + HIP_VISIBLE_DEVICES=2 CUDA_VISIBLE_DEVICES=2 \ + python benchmarks/benchmark_rocm_rope_inference_dispatch.py \ + --tokens 1,2,32 --blocks 12 --iterations 400 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +from collections.abc import Callable + +import torch + +from rl_engine.kernels.ops.rocm.rotary_embedding.rope import ( + RocmDeterministicRoPEOp, + _forward_rope_pair, + _RocmRoPEPairFunction, +) + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _elapsed_ms(function: Callable[[], object], iterations: int) -> float: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def _benchmark_tokens( + *, + tokens: int, + query_heads: int, + key_heads: int, + head_dim: int, + warmup: int, + blocks: int, + iterations: int, +) -> dict[str, object]: + generator = torch.Generator(device="cpu").manual_seed(390_600 + tokens) + query = torch.randn( + query_heads, + tokens, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + key = torch.randn( + key_heads, + tokens, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + positions = torch.arange(tokens, device="cuda", dtype=torch.int64) + operator = RocmDeterministicRoPEOp() + + @torch.inference_mode() + def autograd_wrapper() -> tuple[torch.Tensor, torch.Tensor]: + # Reproduce the complete public method before the fast path, including + # its validation and independent detach semantics. + operator._validate_input(query) + operator._validate_input(key) + if query.device != key.device or query.dtype != key.dtype: + raise ValueError("paired ROCm RoPE Q/K must share one device and dtype") + if query.shape[-1] != key.shape[-1]: + raise ValueError("paired ROCm RoPE Q/K must share one head dimension") + query_out, key_out = _RocmRoPEPairFunction.apply( + query, + key, + positions, + 1_000_000.0, + ) + if not query.requires_grad: + query_out = query_out.detach() + if not key.requires_grad: + key_out = key_out.detach() + return query_out, key_out + + @torch.inference_mode() + def inference_direct_rebuild() -> tuple[torch.Tensor, torch.Tensor]: + operator._validate_input(query) + operator._validate_input(key) + if query.device != key.device or query.dtype != key.dtype: + raise ValueError("paired ROCm RoPE Q/K must share one device and dtype") + if query.shape[-1] != key.shape[-1]: + raise ValueError("paired ROCm RoPE Q/K must share one head dimension") + query_out, key_out, _cos, _sin = _forward_rope_pair( + query, + key, + positions, + 1_000_000.0, + ) + return query_out, key_out + + @torch.inference_mode() + def inference_cached() -> tuple[torch.Tensor, torch.Tensor]: + return operator.forward_pair(query, key, positions) + + for _ in range(warmup): + autograd_wrapper() + inference_direct_rebuild() + inference_cached() + torch.cuda.synchronize() + + baseline_out = autograd_wrapper() + direct_out = inference_direct_rebuild() + candidate_out = inference_cached() + repeated_out = inference_cached() + torch.cuda.synchronize() + exact = all( + torch.equal(expected.view(torch.uint8), direct.view(torch.uint8)) + and torch.equal(direct.view(torch.uint8), actual.view(torch.uint8)) + and torch.equal(actual.view(torch.uint8), repeated.view(torch.uint8)) + for expected, direct, actual, repeated in zip( + baseline_out, + direct_out, + candidate_out, + repeated_out, + strict=True, + ) + ) + if not exact: + raise AssertionError("ROCm RoPE benchmark arms differ in output bytes") + + samples: dict[str, list[float]] = { + "autograd_wrapper": [], + "inference_direct_rebuild": [], + "inference_cached": [], + } + functions = { + "autograd_wrapper": autograd_wrapper, + "inference_direct_rebuild": inference_direct_rebuild, + "inference_cached": inference_cached, + } + for block in range(blocks): + orders = ( + ("autograd_wrapper", "inference_direct_rebuild", "inference_cached"), + ("autograd_wrapper", "inference_cached", "inference_direct_rebuild"), + ("inference_direct_rebuild", "autograd_wrapper", "inference_cached"), + ("inference_direct_rebuild", "inference_cached", "autograd_wrapper"), + ("inference_cached", "autograd_wrapper", "inference_direct_rebuild"), + ("inference_cached", "inference_direct_rebuild", "autograd_wrapper"), + ) + for name in orders[block % len(orders)]: + samples[name].append(_elapsed_ms(functions[name], iterations)) + + autograd_ms = statistics.median(samples["autograd_wrapper"]) + direct_ms = statistics.median(samples["inference_direct_rebuild"]) + cached_ms = statistics.median(samples["inference_cached"]) + return { + "tokens": tokens, + "query_heads": query_heads, + "key_heads": key_heads, + "head_dim": head_dim, + "autograd_wrapper_ms": autograd_ms, + "inference_direct_rebuild_ms": direct_ms, + "inference_cached_ms": cached_ms, + "dispatch_latency_reduction_percent": 100.0 * (autograd_ms - direct_ms) / autograd_ms, + "inv_freq_cache_latency_reduction_percent": 100.0 * (direct_ms - cached_ms) / direct_ms, + "combined_latency_reduction_percent": 100.0 * (autograd_ms - cached_ms) / autograd_ms, + "inv_freq_cache_speedup": direct_ms / cached_ms, + "combined_speedup": autograd_ms / cached_ms, + "raw_bytes_equal": exact, + "query_sha256": _digest(candidate_out[0]), + "key_sha256": _digest(candidate_out[1]), + "autograd_wrapper_samples_ms": samples["autograd_wrapper"], + "inference_direct_rebuild_samples_ms": samples["inference_direct_rebuild"], + "inference_cached_samples_ms": samples["inference_cached"], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", default="1,2,32") + parser.add_argument("--query-heads", type=int, default=8) + parser.add_argument("--key-heads", type=int, default=2) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--warmup", type=int, default=100) + parser.add_argument("--blocks", type=int, default=12) + parser.add_argument("--iterations", type=int, default=400) + args = parser.parse_args() + + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires ROCm PyTorch and a visible GPU") + results = [ + _benchmark_tokens( + tokens=int(tokens), + query_heads=args.query_heads, + key_heads=args.key_heads, + head_dim=args.head_dim, + warmup=args.warmup, + blocks=args.blocks, + iterations=args.iterations, + ) + for tokens in args.tokens.split(",") + ] + print( + json.dumps( + { + "device": torch.cuda.get_device_name(), + "dtype": "bfloat16", + "blocks": args.blocks, + "iterations_per_block": args.iterations, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_rope_pair.py b/benchmarks/benchmark_rocm_rope_pair.py new file mode 100644 index 00000000..a3c6c1a6 --- /dev/null +++ b/benchmarks/benchmark_rocm_rope_pair.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Benchmark sharing the deterministic ROCm RoPE table across vLLM Q/K. + +The baseline reproduces the pre-optimization strict vLLM adapter: convert the +flattened Q/K tensors to head-major layout and invoke the deterministic RoPE +operator twice. The candidate keeps the same layout transforms and HIP kernel +launches, but builds the FP32 cos/sin table once through ``forward_pair``. + +Example on one MI300X: + + HIP_VISIBLE_DEVICES=2 CUDA_VISIBLE_DEVICES=2 \ + python benchmarks/benchmark_rocm_rope_pair.py \ + --tokens 2,32 --blocks 10 --iterations 400 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +from collections.abc import Callable + +import torch + +from rl_engine.kernels.ops.rocm.rotary_embedding.rope import RocmDeterministicRoPEOp + + +def _parse_dtype(value: str) -> torch.dtype: + dtypes = {"float16": torch.float16, "bfloat16": torch.bfloat16} + try: + return dtypes[value] + except KeyError as error: + raise argparse.ArgumentTypeError(f"unsupported dtype: {value}") from error + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _elapsed_ms(function: Callable[[], object], iterations: int) -> float: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + return start.elapsed_time(end) / iterations + + +def _benchmark_shape( + *, + tokens: int, + query_heads: int, + key_heads: int, + head_dim: int, + dtype: torch.dtype, + warmup: int, + blocks: int, + iterations: int, +) -> dict[str, object]: + generator = torch.Generator(device="cpu").manual_seed(390_000 + tokens) + query = torch.randn( + tokens, + query_heads * head_dim, + generator=generator, + dtype=dtype, + ).cuda() + key = torch.randn( + tokens, + key_heads * head_dim, + generator=generator, + dtype=dtype, + ).cuda() + positions = torch.arange(tokens, device="cuda", dtype=torch.int64) + operator = RocmDeterministicRoPEOp() + + def head_major(value: torch.Tensor, heads: int) -> torch.Tensor: + return value.view(tokens, heads, head_dim).permute(1, 0, 2).contiguous() + + def restore(value: torch.Tensor, reference: torch.Tensor) -> torch.Tensor: + return value.permute(1, 0, 2).reshape_as(reference).contiguous() + + @torch.inference_mode() + def legacy() -> tuple[torch.Tensor, torch.Tensor]: + # Keep the same left-to-right evaluation order as the old adapter's + # ``return apply(query), apply(key)`` implementation. + query_major = head_major(query, query_heads) + query_out = operator(query_major, positions) + query_out = restore(query_out, query) + key_major = head_major(key, key_heads) + key_out = operator(key_major, positions) + return query_out, restore(key_out, key) + + @torch.inference_mode() + def shared_table() -> tuple[torch.Tensor, torch.Tensor]: + query_major = head_major(query, query_heads) + key_major = head_major(key, key_heads) + query_out, key_out = operator.forward_pair(query_major, key_major, positions) + return restore(query_out, query), restore(key_out, key) + + for _ in range(warmup): + legacy() + shared_table() + torch.cuda.synchronize() + + legacy_out = legacy() + shared_out = shared_table() + torch.cuda.synchronize() + exact = all( + torch.equal(expected.view(torch.uint8), actual.view(torch.uint8)) + for expected, actual in zip(legacy_out, shared_out, strict=True) + ) + if not exact: + raise AssertionError("shared-table Q/K output differs from the two-call baseline") + + samples: dict[str, list[float]] = {"legacy_two_tables": [], "shared_table": []} + functions = {"legacy_two_tables": legacy, "shared_table": shared_table} + # Alternating A/B then B/A blocks balances launch-order and clock drift. + for block in range(blocks): + order = ( + ("legacy_two_tables", "shared_table") + if block % 2 == 0 + else ("shared_table", "legacy_two_tables") + ) + for name in order: + samples[name].append(_elapsed_ms(functions[name], iterations)) + + baseline_ms = statistics.median(samples["legacy_two_tables"]) + candidate_ms = statistics.median(samples["shared_table"]) + return { + "tokens": tokens, + "query_heads": query_heads, + "key_heads": key_heads, + "head_dim": head_dim, + "dtype": str(dtype).removeprefix("torch."), + "legacy_two_tables_ms": baseline_ms, + "shared_table_ms": candidate_ms, + "latency_reduction_percent": 100.0 * (baseline_ms - candidate_ms) / baseline_ms, + "speedup": baseline_ms / candidate_ms, + "raw_bytes_equal": exact, + "query_sha256": _digest(shared_out[0]), + "key_sha256": _digest(shared_out[1]), + "legacy_samples_ms": samples["legacy_two_tables"], + "shared_samples_ms": samples["shared_table"], + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", default="2,32") + parser.add_argument("--query-heads", type=int, default=8) + parser.add_argument("--key-heads", type=int, default=2) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--dtype", type=_parse_dtype, default=torch.bfloat16) + parser.add_argument("--warmup", type=int, default=100) + parser.add_argument("--blocks", type=int, default=10) + parser.add_argument("--iterations", type=int, default=400) + args = parser.parse_args() + + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires ROCm PyTorch and a visible GPU") + token_counts = [int(value) for value in args.tokens.split(",")] + results = [ + _benchmark_shape( + tokens=tokens, + query_heads=args.query_heads, + key_heads=args.key_heads, + head_dim=args.head_dim, + dtype=args.dtype, + warmup=args.warmup, + blocks=args.blocks, + iterations=args.iterations, + ) + for tokens in token_counts + ] + print( + json.dumps( + { + "device": torch.cuda.get_device_name(), + "blocks": args.blocks, + "iterations_per_block": args.iterations, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_singleton_lse.py b/benchmarks/benchmark_rocm_singleton_lse.py new file mode 100644 index 00000000..026b7561 --- /dev/null +++ b/benchmarks/benchmark_rocm_singleton_lse.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Benchmark removing singleton LSE concatenations from ROCm paged decode. + +The legacy arm restores exactly the redundant ``torch.cat([lse], dim=0)`` +copies that the production candidate skips. Both arms still execute the same +page gathers and AITER CK launches. This is an isolated strict-runtime timing; +end-to-end benefit must still be measured with the VIME rollout A/B. + +Example on one MI300X: + + HIP_VISIBLE_DEVICES=2 CUDA_VISIBLE_DEVICES=2 \ + python benchmarks/benchmark_rocm_singleton_lse.py \ + --cached-lengths 32,128,512,2048 --layers 36 \ + --blocks 12 --iterations 20 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import time +from collections.abc import Callable + +import torch + +from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore +from rl_engine.kernels.ops.rocm.attention.strict_runtime import StrictRocmAttentionRuntime + + +def _digest(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes() + return hashlib.sha256(raw).hexdigest() + + +def _elapsed_ms(function: Callable[[], object], iterations: int) -> tuple[float, float]: + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + wall_start = time.perf_counter() + start.record() + for _ in range(iterations): + function() + end.record() + end.synchronize() + wall_ms = (time.perf_counter() - wall_start) * 1_000.0 / iterations + return start.elapsed_time(end) / iterations, wall_ms + + +def _benchmark_length( + *, + cached_length: int, + batch: int, + layers: int, + query_heads: int, + key_heads: int, + head_dim: int, + page_size: int, + warmup: int, + blocks: int, + iterations: int, +) -> dict[str, object]: + if batch != 1: + raise ValueError("this exact legacy singleton-cat reconstruction requires batch=1") + pages_per_row = (cached_length + page_size - 1) // page_size + total_pages = batch * pages_per_row + generator = torch.Generator(device="cpu").manual_seed(390_800 + cached_length + batch) + q = torch.randn( + batch, + query_heads, + 1, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + k_cache = torch.randn( + total_pages, + page_size, + key_heads, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + v_cache = torch.randn( + total_pages, + page_size, + key_heads, + head_dim, + dtype=torch.bfloat16, + generator=generator, + ).cuda() + page_table = torch.arange(total_pages, device="cuda", dtype=torch.int32).reshape( + batch, + pages_per_row, + ) + seqused_k = torch.full((batch,), cached_length, device="cuda", dtype=torch.int32) + cached_lengths = (cached_length,) * batch + core = StrictRocmAiterCKAttentionCore() + baseline_runtime = StrictRocmAttentionRuntime(core=core) + candidate_runtime = StrictRocmAttentionRuntime(core=core) + baseline_out = torch.empty_like(q) + candidate_out = torch.empty_like(q) + common = { + "page_table": page_table, + "seqused_k": seqused_k, + "max_seqlen_k": pages_per_row * page_size, + "scale": head_dim**-0.5, + "cached_lengths": cached_lengths, + } + # The old path copied once for the singleton row passed into _run_core and + # once for the singleton outer paged batch. With one KV head, its group + # assembly was a third singleton copy. + legacy_singleton_cats_per_layer = 2 + int(key_heads == 1) + + @torch.inference_mode() + def run_layers(runtime, out, *, restore_legacy_cats): + epoch = runtime.new_page_bounds_epoch() + result = None + returned_lse = None + for _ in range(layers): + result = runtime.forward_paged_with_lse( + q, + k_cache, + v_cache, + out=out, + page_bounds_epoch=epoch, + **common, + ) + returned_lse = result.lse + if restore_legacy_cats: + if key_heads == 1: + returned_lse = torch.cat([returned_lse], dim=1) + returned_lse = torch.cat([returned_lse], dim=0) + returned_lse = torch.cat([returned_lse], dim=0) + if result is None or returned_lse is None: + raise AssertionError("layers must be positive") + return result.out, returned_lse, result.provenance + + def baseline(): + return run_layers(baseline_runtime, baseline_out, restore_legacy_cats=True) + + def candidate(): + return run_layers(candidate_runtime, candidate_out, restore_legacy_cats=False) + + for _ in range(warmup): + baseline() + candidate() + torch.cuda.synchronize() + + baseline_output, baseline_lse, baseline_provenance = baseline() + candidate_output, candidate_lse, candidate_provenance = candidate() + candidate_output_snapshot = candidate_output.clone() + candidate_lse_snapshot = candidate_lse.clone() + replay_output, replay_lse, _ = candidate() + torch.cuda.synchronize() + output_equal = torch.equal( + baseline_output.view(torch.uint8), + candidate_output.view(torch.uint8), + ) and torch.equal( + candidate_output_snapshot.view(torch.uint8), + replay_output.view(torch.uint8), + ) + lse_equal = torch.equal( + baseline_lse.view(torch.uint8), + candidate_lse.view(torch.uint8), + ) and torch.equal( + candidate_lse_snapshot.view(torch.uint8), + replay_lse.view(torch.uint8), + ) + if not output_equal or not lse_equal: + raise AssertionError("singleton LSE fast path changed output or LSE bytes") + if baseline_lse.shape != candidate_lse.shape or baseline_lse.stride() != candidate_lse.stride(): + raise AssertionError("singleton LSE fast path changed LSE layout") + expected_launches = batch * key_heads + if baseline_provenance["core_launch_count"] != expected_launches: + raise AssertionError("baseline AITER launch count changed") + if candidate_provenance["core_launch_count"] != expected_launches: + raise AssertionError("candidate AITER launch count changed") + + functions = {"legacy_singleton_cat": baseline, "singleton_fast_path": candidate} + gpu_samples: dict[str, list[float]] = {name: [] for name in functions} + wall_samples: dict[str, list[float]] = {name: [] for name in functions} + for block in range(blocks): + order = ( + ("legacy_singleton_cat", "singleton_fast_path") + if block % 2 == 0 + else ("singleton_fast_path", "legacy_singleton_cat") + ) + for name in order: + gpu_ms, wall_ms = _elapsed_ms(functions[name], iterations) + gpu_samples[name].append(gpu_ms) + wall_samples[name].append(wall_ms) + + baseline_ms = statistics.median(gpu_samples["legacy_singleton_cat"]) + candidate_ms = statistics.median(gpu_samples["singleton_fast_path"]) + baseline_wall_ms = statistics.median(wall_samples["legacy_singleton_cat"]) + candidate_wall_ms = statistics.median(wall_samples["singleton_fast_path"]) + return { + "cached_length": cached_length, + "batch": batch, + "layers_per_forward": layers, + "query_heads": query_heads, + "key_heads": key_heads, + "head_dim": head_dim, + "benchmark_scope": "strict_runtime_simulated_decoder_layers", + "legacy_singleton_cats_per_forward": legacy_singleton_cats_per_layer * layers, + "aiter_core_launches_per_layer": expected_launches, + "median_ms": { + "gpu": { + "legacy_singleton_cat": baseline_ms, + "singleton_fast_path": candidate_ms, + }, + "wall": { + "legacy_singleton_cat": baseline_wall_ms, + "singleton_fast_path": candidate_wall_ms, + }, + }, + "latency_reduction_percent": 100.0 * (baseline_ms - candidate_ms) / baseline_ms, + "wall_latency_reduction_percent": ( + 100.0 * (baseline_wall_ms - candidate_wall_ms) / baseline_wall_ms + ), + "saved_ms_per_forward": baseline_ms - candidate_ms, + "raw_output_equal": output_equal, + "raw_lse_equal": lse_equal, + "lse_shape": list(candidate_lse.shape), + "lse_stride": list(candidate_lse.stride()), + "output_sha256": _digest(candidate_output), + "lse_sha256": _digest(candidate_lse), + "samples_ms": {"gpu": gpu_samples, "wall": wall_samples}, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cached-lengths", default="32,128,512,2048") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--layers", type=int, default=36) + parser.add_argument("--query-heads", type=int, default=8) + parser.add_argument("--key-heads", type=int, default=2) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--page-size", type=int, default=16) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--blocks", type=int, default=12) + parser.add_argument("--iterations", type=int, default=20) + args = parser.parse_args() + + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("this benchmark requires ROCm PyTorch and a visible GPU") + if ( + min( + args.batch, + args.layers, + args.query_heads, + args.key_heads, + args.head_dim, + args.page_size, + args.blocks, + args.iterations, + ) + <= 0 + ): + raise ValueError("benchmark dimensions, blocks, and iterations must be positive") + if args.query_heads % args.key_heads: + raise ValueError("--query-heads must be divisible by --key-heads") + if args.batch != 1: + raise ValueError("--batch must be 1 for the exact legacy singleton-cat A/B") + results = [ + _benchmark_length( + cached_length=int(cached_length), + batch=args.batch, + layers=args.layers, + query_heads=args.query_heads, + key_heads=args.key_heads, + head_dim=args.head_dim, + page_size=args.page_size, + warmup=args.warmup, + blocks=args.blocks, + iterations=args.iterations, + ) + for cached_length in args.cached_lengths.split(",") + ] + print( + json.dumps( + { + "device": torch.cuda.get_device_name(), + "torch": torch.__version__, + "hip": torch.version.hip, + "dtype": "bfloat16", + "blocks": args.blocks, + "iterations_per_block": args.iterations, + "results": results, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/csrc/cuda/attention/deterministic_attention.cu b/csrc/cuda/attention/deterministic_attention.cu index 2c56d6c7..82ec824d 100644 --- a/csrc/cuda/attention/deterministic_attention.cu +++ b/csrc/cuda/attention/deterministic_attention.cu @@ -718,6 +718,46 @@ __global__ void deterministic_rope_kernel( out[base + pair + half] = static_cast(high * c + low * s); } +template +__global__ void deterministic_rope_token_major_kernel( + const scalar_t* __restrict__ x, + const int64_t* __restrict__ positions, + const float* __restrict__ cos, + const float* __restrict__ sin, + scalar_t* __restrict__ out, + int64_t tokens, + int64_t input_token_stride, + int heads, + int head_dim, + int half, + int table_rows, + float sin_sign) { + const int64_t index = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + const int64_t values_per_token = static_cast(heads) * half; + const int64_t count = tokens * values_per_token; + if (index >= count) { + return; + } + + const int64_t token = index / values_per_token; + const int64_t head_pair = index % values_per_token; + const int head = static_cast(head_pair / half); + const int pair = static_cast(head_pair % half); + const int64_t position = positions[token]; + if (position < 0 || position >= table_rows) { + return; + } + const float c = cos[position * half + pair]; + const float s = sin[position * half + pair] * sin_sign; + const int64_t input_base = token * input_token_stride + head * head_dim; + const int64_t output_base = (token * heads + head) * static_cast(head_dim); + const float low = static_cast(x[input_base + pair]); + const float high = static_cast(x[input_base + pair + half]); + + out[output_base + pair] = static_cast(low * c - high * s); + out[output_base + pair + half] = static_cast(high * c + low * s); +} + } // namespace torch::Tensor deterministic_rope_apply_rocm( @@ -771,4 +811,70 @@ torch::Tensor deterministic_rope_apply_rocm( }); return out; } + +torch::Tensor deterministic_rope_apply_token_major_rocm( + torch::Tensor x, + torch::Tensor positions, + torch::Tensor cos, + torch::Tensor sin, + int64_t head_dim, + double sin_sign) { + TORCH_CHECK(x.is_cuda(), "ROCm token-major RoPE: x must be a GPU tensor"); + TORCH_CHECK(x.dim() == 2 && x.stride(1) == 1, + "ROCm token-major RoPE: x must be [tokens, heads * head_dim] with unit inner stride"); + TORCH_CHECK(head_dim > 0 && head_dim % 2 == 0 && x.size(1) % head_dim == 0, + "ROCm token-major RoPE: invalid head dimension"); + TORCH_CHECK(positions.is_cuda() && positions.scalar_type() == torch::kInt64 && + positions.dim() == 1 && positions.size(0) == x.size(0) && + positions.is_contiguous(), + "ROCm token-major RoPE: positions must be contiguous GPU int64 [tokens]"); + TORCH_CHECK(cos.is_cuda() && sin.is_cuda() && cos.scalar_type() == torch::kFloat32 && + sin.scalar_type() == torch::kFloat32 && cos.is_contiguous() && + sin.is_contiguous() && cos.dim() == 2 && sin.sizes() == cos.sizes(), + "ROCm token-major RoPE: cos/sin must be contiguous GPU FP32 tables"); + TORCH_CHECK(cos.size(0) > 0 && cos.size(1) == head_dim / 2, + "ROCm token-major RoPE: invalid cos/sin table shape"); + TORCH_CHECK(x.get_device() == positions.get_device() && + x.get_device() == cos.get_device() && x.get_device() == sin.get_device(), + "ROCm token-major RoPE: all tensors must share one device"); + + const at::cuda::OptionalCUDAGuard guard(device_of(x)); + auto out = torch::empty({x.size(0), x.size(1)}, x.options()); + const int64_t tokens = x.size(0); + const int heads = static_cast(x.size(1) / head_dim); + const int half = static_cast(head_dim / 2); + const int table_rows = static_cast(cos.size(0)); + const int64_t count = tokens * static_cast(heads) * half; + constexpr int threads = 256; + const int64_t blocks = (count + threads - 1) / threads; + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + x.scalar_type(), + "deterministic_rope_apply_token_major_rocm", + [&] { + hipLaunchKernelGGL( + (deterministic_rope_token_major_kernel), + dim3(blocks), + dim3(threads), + 0, + stream, + x.data_ptr(), + positions.data_ptr(), + cos.data_ptr(), + sin.data_ptr(), + out.data_ptr(), + tokens, + x.stride(0), + heads, + static_cast(head_dim), + half, + table_rows, + static_cast(sin_sign)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} #endif diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 603adceb..026ea23d 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -174,6 +174,13 @@ int64_t deterministic_collective_rocm_ipc_create( void deterministic_collective_rocm_ipc_synchronize(int64_t handle); void deterministic_collective_rocm_ipc_destroy(int64_t handle); void deterministic_collective_rocm_ipc_stage(int64_t handle, torch::Tensor input); +void deterministic_collective_rocm_ipc_prepare_staged( + int64_t handle, + torch::Tensor input); +void deterministic_collective_rocm_ipc_all_reduce_staged( + int64_t handle, + torch::Tensor input, + torch::Tensor output); void deterministic_collective_rocm_ipc_all_reduce( int64_t handle, torch::Tensor output); @@ -408,6 +415,13 @@ torch::Tensor deterministic_rope_apply_rocm( torch::Tensor cos, torch::Tensor sin, double sin_sign); +torch::Tensor deterministic_rope_apply_token_major_rocm( + torch::Tensor x, + torch::Tensor positions, + torch::Tensor cos, + torch::Tensor sin, + int64_t head_dim, + double sin_sign); #endif #if !defined(USE_ROCM) @@ -613,6 +627,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("deterministic_collective_rocm_ipc_stage", &deterministic_collective_rocm_ipc_stage, "Stage an input for ROCm IPC deterministic collectives"); + m.def("deterministic_collective_rocm_ipc_prepare_staged", + &deterministic_collective_rocm_ipc_prepare_staged, + "Reserve the direct-output ROCm IPC payload"); + m.def("deterministic_collective_rocm_ipc_all_reduce_staged", + &deterministic_collective_rocm_ipc_all_reduce_staged, + "Reduce a result already resident in the ROCm IPC payload"); m.def("deterministic_collective_rocm_ipc_all_reduce", &deterministic_collective_rocm_ipc_all_reduce, "Run a direct ROCm IPC fixed-tree all-reduce"); @@ -692,6 +712,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_rope_apply_rocm", &deterministic_rope_apply_rocm, "Deterministic GPT-NeoX RoPE apply for ROCm"); + m.def( + "deterministic_rope_apply_token_major_rocm", + &deterministic_rope_apply_token_major_rocm, + "Deterministic GPT-NeoX token-major RoPE apply for ROCm"); #endif #endif } diff --git a/csrc/rocm/distributed/deterministic_collective.hip b/csrc/rocm/distributed/deterministic_collective.hip index c8448300..0bb882a0 100644 --- a/csrc/rocm/distributed/deterministic_collective.hip +++ b/csrc/rocm/distributed/deterministic_collective.hip @@ -30,6 +30,7 @@ constexpr int64_t kIPCControlBytes = 256; constexpr int64_t kIPCReadyOffset = 0; constexpr int64_t kIPCDoneOffset = 64; constexpr int64_t kIPCCloseOffset = 128; +constexpr int64_t kIPCSequenceOffset = 192; struct PeerPointers { const void* values[kMaxWorldSize]; @@ -39,6 +40,7 @@ struct PeerSignals { uint64_t* ready[kMaxWorldSize]; uint64_t* done[kMaxWorldSize]; uint64_t* closed[kMaxWorldSize]; + uint64_t* sequence[kMaxWorldSize]; }; template @@ -441,6 +443,124 @@ __global__ void ipc_mark_signal_kernel(uint64_t* signal, uint64_t sequence) { } } +__global__ void ipc_reserve_sequence_kernel( + PeerSignals signals, + int64_t rank, + int64_t world_size) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + uint64_t sequence = 0; + if (rank == 0) { + const uint64_t previous = __hip_atomic_load( + signals.sequence[0], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM); + for (int peer = 0; peer < world_size; ++peer) { + while (__hip_atomic_load( + signals.done[peer], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) < previous) { + __builtin_amdgcn_s_sleep(1); + } + } + sequence = previous + 1; + __hip_atomic_store( + signals.sequence[0], + sequence, + __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + } else { + const uint64_t previous = __hip_atomic_load( + signals.sequence[rank], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM); + do { + sequence = __hip_atomic_load( + signals.sequence[0], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM); + if (sequence <= previous) { + __builtin_amdgcn_s_sleep(1); + } + } while (sequence <= previous); + for (int peer = 0; peer < world_size; ++peer) { + while (__hip_atomic_load( + signals.done[peer], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) < sequence - 1) { + __builtin_amdgcn_s_sleep(1); + } + } + __hip_atomic_store( + signals.sequence[rank], + sequence, + __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + } +} + +__global__ void ipc_mark_ready_and_wait_dynamic_kernel( + PeerSignals signals, + int64_t rank, + int64_t world_size) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + const uint64_t sequence = __hip_atomic_load( + signals.sequence[rank], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM); + __hip_atomic_store( + signals.ready[rank], + sequence, + __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + for (int peer = 0; peer < world_size; ++peer) { + while (__hip_atomic_load( + signals.ready[peer], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) != sequence) { + __builtin_amdgcn_s_sleep(1); + } + } +} + +__global__ void ipc_mark_done_dynamic_kernel(PeerSignals signals, int64_t rank) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + const uint64_t sequence = __hip_atomic_load( + signals.sequence[rank], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM); + __hip_atomic_store( + signals.done[rank], + sequence, + __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + } +} + +__global__ void ipc_synchronize_dynamic_kernel( + PeerSignals signals, + int64_t rank, + int64_t world_size) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + const uint64_t sequence = __hip_atomic_load( + signals.sequence[rank], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM); + for (int peer = 0; peer < world_size; ++peer) { + while (__hip_atomic_load( + signals.done[peer], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) < sequence) { + __builtin_amdgcn_s_sleep(1); + } + } +} + __global__ void ipc_mark_ready_and_wait_kernel( PeerSignals signals, int64_t rank, @@ -613,6 +733,72 @@ void launch_mark_ready_and_wait( C10_CUDA_KERNEL_LAUNCH_CHECK(); } +void launch_reserve_sequence( + const PeerSignals& signals, + int64_t rank, + int64_t world_size, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_reserve_sequence_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + rank, + world_size); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_mark_ready_and_wait_dynamic( + const PeerSignals& signals, + int64_t rank, + int64_t world_size, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_mark_ready_and_wait_dynamic_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + rank, + world_size); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_mark_done_dynamic( + const PeerSignals& signals, + int64_t rank, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_mark_done_dynamic_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + rank); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_synchronize_dynamic( + const PeerSignals& signals, + int64_t rank, + int64_t world_size, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_synchronize_dynamic_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + rank, + world_size); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + void launch_close_and_wait( const PeerSignals& signals, int64_t rank, @@ -706,13 +892,7 @@ class ROCmIPCCollectiveState { const int64_t input_bytes = input.numel() * input.element_size(); TORCH_CHECK(input_bytes <= capacity_bytes_, "input exceeds ROCm IPC staging capacity"); - ++sequence_; - launch_wait_signal( - signals_, - sequence_ - 1, - world_size_, - true, - stream); + launch_reserve_sequence(signals_, rank_, world_size_, stream); if (input_bytes > 0) { C10_HIP_CHECK(hipMemcpyAsync( const_cast(peers_.values[rank_]), @@ -721,20 +901,42 @@ class ROCmIPCCollectiveState { hipMemcpyDeviceToDevice, stream)); } - launch_mark_ready_and_wait( - signals_, - rank_, - sequence_, - world_size_, - stream); + launch_mark_ready_and_wait_dynamic(signals_, rank_, world_size_, stream); staged_bytes_ = input_bytes; staged_type_ = input.scalar_type(); } + void prepare_staged(torch::Tensor input, hipStream_t stream) { + check_tensor(input, "direct staging input"); + const int64_t input_bytes = input.numel() * input.element_size(); + TORCH_CHECK(input_bytes <= capacity_bytes_, + "direct staging input exceeds ROCm IPC capacity"); + TORCH_CHECK(input.data_ptr() == peers_.values[rank_], + "direct staging input must start at the local IPC payload"); + TORCH_CHECK(!staging_reserved_, "direct staging is already reserved"); + launch_reserve_sequence(signals_, rank_, world_size_, stream); + staged_bytes_ = input_bytes; + staged_type_ = input.scalar_type(); + staging_reserved_ = true; + } + + void finish_staged(torch::Tensor input, hipStream_t stream) { + check_tensor(input, "direct staging input"); + TORCH_CHECK(staging_reserved_, "direct staging must be reserved before reduction"); + TORCH_CHECK(input.data_ptr() == peers_.values[rank_], + "direct staging input must start at the local IPC payload"); + TORCH_CHECK(input.numel() * input.element_size() == staged_bytes_, + "direct staging input size changed after reserve"); + TORCH_CHECK(input.scalar_type() == staged_type_, + "direct staging input dtype changed after reserve"); + launch_mark_ready_and_wait_dynamic(signals_, rank_, world_size_, stream); + staging_reserved_ = false; + } + void all_reduce(torch::Tensor output, hipStream_t stream) const { check_reduction_output(output, staged_bytes_, "all_reduce"); launch(output, 0, output.numel(), stream); - launch_mark_signal(signals_.done[rank_], sequence_, stream); + launch_mark_done_dynamic(signals_, rank_, stream); } void reduce_scatter(torch::Tensor output, hipStream_t stream) const { @@ -747,7 +949,7 @@ class ROCmIPCCollectiveState { rank_ * output.numel(), output.numel(), stream); - launch_mark_signal(signals_.done[rank_], sequence_, stream); + launch_mark_done_dynamic(signals_, rank_, stream); } void reduce_scatter_many( @@ -775,8 +977,7 @@ class ROCmIPCCollectiveState { total_bytes += input_bytes; } - ++sequence_; - launch_wait_signal(signals_, sequence_ - 1, world_size_, true, stream); + launch_reserve_sequence(signals_, rank_, world_size_, stream); int64_t byte_offset = 0; for (const auto& input : inputs) { const int64_t input_bytes = input.numel() * input.element_size(); @@ -790,12 +991,7 @@ class ROCmIPCCollectiveState { } byte_offset += input_bytes; } - launch_mark_ready_and_wait( - signals_, - rank_, - sequence_, - world_size_, - stream); + launch_mark_ready_and_wait_dynamic(signals_, rank_, world_size_, stream); int64_t element_offset = 0; for (size_t index = 0; index < outputs.size(); ++index) { @@ -808,7 +1004,7 @@ class ROCmIPCCollectiveState { stream); element_offset += input.numel(); } - launch_mark_signal(signals_.done[rank_], sequence_, stream); + launch_mark_done_dynamic(signals_, rank_, stream); } void all_gather(torch::Tensor output, hipStream_t stream) const { @@ -824,11 +1020,11 @@ class ROCmIPCCollectiveState { staged_bytes_, world_size_, stream); - launch_mark_signal(signals_.done[rank_], sequence_, stream); + launch_mark_done_dynamic(signals_, rank_, stream); } void synchronize(hipStream_t stream) const { - launch_wait_signal(signals_, sequence_, world_size_, true, stream); + launch_synchronize_dynamic(signals_, rank_, world_size_, stream); launch_close_and_wait(signals_, rank_, world_size_, stream); } @@ -838,6 +1034,7 @@ class ROCmIPCCollectiveState { signals_.ready[peer] = reinterpret_cast(bytes + kIPCReadyOffset); signals_.done[peer] = reinterpret_cast(bytes + kIPCDoneOffset); signals_.closed[peer] = reinterpret_cast(bytes + kIPCCloseOffset); + signals_.sequence[peer] = reinterpret_cast(bytes + kIPCSequenceOffset); peers_.values[peer] = bytes + kIPCControlBytes; } @@ -923,7 +1120,7 @@ class ROCmIPCCollectiveState { int64_t capacity_bytes_; int64_t staged_bytes_{0}; at::ScalarType staged_type_{at::ScalarType::Undefined}; - uint64_t sequence_{0}; + bool staging_reserved_{false}; PeerPointers peers_{}; PeerSignals signals_{}; std::array imported_bases_{}; @@ -1031,6 +1228,25 @@ void deterministic_collective_rocm_ipc_stage(int64_t handle, torch::Tensor input ipc_state(handle)->stage(input, stream); } +void deterministic_collective_rocm_ipc_prepare_staged( + int64_t handle, + torch::Tensor input) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->prepare_staged(input, stream); +} + +void deterministic_collective_rocm_ipc_all_reduce_staged( + int64_t handle, + torch::Tensor input, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + auto* state = ipc_state(handle); + state->finish_staged(input, stream); + state->all_reduce(output, stream); +} + void deterministic_collective_rocm_ipc_all_reduce( int64_t handle, torch::Tensor output) { diff --git a/examples/vime_rocm_attention_ablation/launch_arm.sh b/examples/vime_rocm_attention_ablation/launch_arm.sh index c040d1e0..7d675741 100644 --- a/examples/vime_rocm_attention_ablation/launch_arm.sh +++ b/examples/vime_rocm_attention_ablation/launch_arm.sh @@ -35,6 +35,7 @@ set -euo pipefail : "${RLK_ABLATION_ROLLOUT_SEED:?}" : "${RLK_ABLATION_RAY_PORT:?}" : "${RLK_ABLATION_RAY_DASHBOARD_PORT:?}" +: "${RLK_ABLATION_RAY_DASHBOARD_AGENT_PORT:?}" : "${RL_KERNEL_READBACK_DIR:?}" : "${RL_KERNEL_MISMATCH_SIDECAR_DIR:?}" : "${RL_KERNEL_VLLM_REAL_VOCAB_SIZE:?}" @@ -135,6 +136,7 @@ export VLLM_ROCM_USE_AITER=1 # the optional shuffled physical layout. export VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=0 export VLLM_ATTENTION_BACKEND=ROCM_AITER_FA +export RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE="${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE:-32}" export MIOPEN_DEBUG_CONV_DIRECT="${MIOPEN_DEBUG_CONV_DIRECT:-0}" mkdir -p \ @@ -188,6 +190,7 @@ names = [ "VLLM_ROCM_USE_AITER", "VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "VLLM_ATTENTION_BACKEND", + "RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE", "MIOPEN_DEBUG_CONV_DIRECT", "RL_KERNEL_ATTENTION_CASE", "RL_KERNEL_FFN_CASE", @@ -221,14 +224,31 @@ ray start --head \ --num-gpus="${RLK_ABLATION_NUM_GPUS}" \ --disable-usage-stats \ --dashboard-host=127.0.0.1 \ - --dashboard-port="${RLK_ABLATION_RAY_DASHBOARD_PORT}" + --dashboard-port="${RLK_ABLATION_RAY_DASHBOARD_PORT}" \ + --dashboard-agent-listen-port="${RLK_ABLATION_RAY_DASHBOARD_AGENT_PORT}" ray_started=1 -# Correctness first: eager mode is frozen across all four arms. The strict -# ROCm projection collective performs Python/lock bookkeeping and is not a -# valid vLLM fullgraph capture target yet. +ray_job_address="http://127.0.0.1:${RLK_ABLATION_RAY_DASHBOARD_PORT}" +ray_job_agent_ready=0 +for _attempt in {1..60}; do + if ray job list --address="${ray_job_address}" >/dev/null 2>&1; then + ray_job_agent_ready=1 + break + fi + sleep 1 +done +if [[ "${ray_job_agent_ready}" != "1" ]]; then + echo "Ray job agent did not become ready within 60 seconds" >&2 + exit 5 +fi +# The dashboard endpoint can answer before the local node's job agent has +# finished registering. Give that registration a bounded stabilization window. +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. ray job submit \ - --address="http://127.0.0.1:${RLK_ABLATION_RAY_DASHBOARD_PORT}" \ + --address="${ray_job_address}" \ --runtime-env-json="${RUNTIME_ENV_JSON}" \ --working-dir="${RLK_ABLATION_VIME_ROOT}" \ -- python3 "${RLK_ABLATION_VIME_ROOT}/train.py" \ @@ -282,8 +302,9 @@ ray job submit \ --router-policy "${RLK_ABLATION_ROUTER_POLICY}" \ --rollout-num-gpus-per-engine "${RLK_ABLATION_ROLLOUT_TP_SIZE}" \ --vllm-gpu-memory-utilization "${VLLM_GPU_MEMORY_UTILIZATION:-0.4}" \ + --vllm-max-cudagraph-capture-size \ + "${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE}" \ --vllm-attention-backend ROCM_AITER_FA \ - --vllm-enforce-eager \ --vllm-disable-custom-all-reduce \ --attention-dropout 0 \ --hidden-dropout 0 \ diff --git a/examples/vime_rocm_attention_ablation/run.py b/examples/vime_rocm_attention_ablation/run.py index 3314b265..863489df 100644 --- a/examples/vime_rocm_attention_ablation/run.py +++ b/examples/vime_rocm_attention_ablation/run.py @@ -211,7 +211,8 @@ def frozen_parameters(self) -> dict[str, Any]: # declares unsupported. The R route owns its deterministic # per-row schedule; the P route remains the native baseline. "vllm_batch_invariant": False, - "enforce_eager": True, + "enforce_eager": False, + "execution_mode": "compiled_hip_graph", "custom_all_reduce": False, "attention_backend": "ROCM_AITER_FA", "shuffle_kv_cache_layout": False, @@ -501,6 +502,7 @@ def build_arm_environment( "VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT": "0", "VLLM_ATTENTION_BACKEND": "ROCM_AITER_FA", + "RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE": "32", "RL_KERNEL_ATTENTION_CASE": case_id, "RL_KERNEL_FFN_CASE": "R/R", "RL_KERNEL_LOGP_CASE": "R/R", @@ -541,6 +543,9 @@ def build_arm_environment( "RLK_ABLATION_RAY_DASHBOARD_PORT": str( config.ray_dashboard_port + arm_index ), + "RLK_ABLATION_RAY_DASHBOARD_AGENT_PORT": str( + config.ray_dashboard_port + arm_index + 10_000 + ), } ) return env @@ -560,6 +565,7 @@ def public_arm_environment(environment: Mapping[str, str]) -> dict[str, str]: "RL_KERNEL_MISMATCH_SIDECAR_DIR", "RL_KERNEL_READBACK_DIR", "RL_KERNEL_VLLM_INTEGRATION", + "RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE", "VLLM_ROCM_USE_AITER", "VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "VLLM_ATTENTION_BACKEND", diff --git a/examples/vime_rocm_attention_ablation/validate_artifacts.py b/examples/vime_rocm_attention_ablation/validate_artifacts.py index e7ae89d9..1fa0b86a 100644 --- a/examples/vime_rocm_attention_ablation/validate_artifacts.py +++ b/examples/vime_rocm_attention_ablation/validate_artifacts.py @@ -42,6 +42,7 @@ STRICT_ROCM_CORE_ID = "rlkernel.attention.rocm.aiter_ck_dense_mha.v1" STRICT_ROCM_SCHEDULE_ID = "single_batch_aiter_ck_dense_mha_no_splitkv" ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID = "rlkernel.rocm.triton_det_gemm" +ROCM_PAGED_ATTENTION_BACKEND_ID = "aiter_mha_batch_prefill_non_split_ck" ROCM_DETERMINISTIC_COLLECTIVE_BACKEND_ID = "rocm_ipc_fixed_tree" ROCM_FFN_BACKEND_ID = "rlkernel.rocm.det_gemm_swiglu" STRICT_FFN_BACKEND_ID = "rlkernel.ffn.qwen3.deterministic.v1" @@ -173,6 +174,7 @@ def validate_launch_manifest(path: Path, case_id: str) -> dict[str, Any]: "VLLM_ATTENTION_BACKEND": "ROCM_AITER_FA", "VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT": "0", + "RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE": "32", } for name, expected in expected_environment.items(): if environment.get(name) != expected: @@ -297,19 +299,30 @@ def _validate_rlkernel_record( f"{ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID!r}" ) - # Triton is permitted only for the declared deterministic QKV/O projection. - # The Attention core itself must remain the native AITER/CK implementation. + # Triton is permitted for deterministic QKV/O projection. Paged Attention + # arithmetic and page-table reads remain inside native AITER/CK. + allowed_triton_backends = { + ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, + } for key, item in _walk_key_values(provenance): if key in _TRITON_KEYS and item is True: continue if isinstance(item, str) and "triton" in item.lower(): - if item != ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID: + if item not in allowed_triton_backends: errors.append(f"{label} reported an unapproved Triton backend {item!r}") if framework == "vllm": layouts = _values_for_keys(provenance, {"framework_layout"}) if "vllm_paged_kv" not in layouts: errors.append(f"{label} did not prove the vLLM paged-KV execution boundary") + if not _has_exact_value( + provenance, + {"paged_kernel"}, + ROCM_PAGED_ATTENTION_BACKEND_ID, + ): + errors.append(f"{label} did not prove direct non-Split-K paged CK execution") + if any(_values_for_keys(provenance, {"dense_kv_materialized"})): + errors.append(f"{label} materialized dense KV during paged decode") tp_values = _values_for_keys(provenance, {"tp_world_size"}) try: tp_world_size = max(int(value) for value in tp_values) @@ -389,8 +402,11 @@ def _validate_strict_dense_record( errors.append(f"{label} did not execute the fixed R/R route") if record.get("backend_id") != expected_backend: errors.append(f"{label} reported backend {record.get('backend_id')!r}") - if record.get("execution_mode", "eager") != "eager": - errors.append(f"{label} did not execute in frozen eager mode") + expected_mode = ( + "compiled_hip_graph" if framework == "vllm" and module == "ffn" else "eager" + ) + if record.get("execution_mode", "eager") != expected_mode: + errors.append(f"{label} did not execute in {expected_mode} mode") if int(record.get("call_count", 0)) <= 0: errors.append(f"{label} had zero executed calls") if _runtime_platform(provenance) != "rocm": @@ -479,8 +495,12 @@ def validate_attention_readbacks( backend_ids.add(str(record.get("backend_id", ""))) if record.get("case_id") != normalized_case: errors.append(f"{record_label} record has the wrong case_id") - if record.get("execution_mode", "eager") != "eager": - errors.append(f"{record_label} did not execute in frozen eager mode") + execution_mode = record.get("execution_mode", "eager") + if framework == "vllm": + if execution_mode not in {"eager", "compiled_hip_graph"}: + errors.append(f"{record_label} has invalid HIP execution mode") + elif execution_mode != "eager": + errors.append(f"{record_label} did not execute in eager mode") if record.get("implementation") != expected_implementation: errors.append( f"{record_label} implementation={record.get('implementation')!r}, " diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 4e60e52b..fb3c4379 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -225,6 +225,20 @@ def deterministic_attention_backward( scale: float, key_padding_mask: torch.Tensor | None, ) -> list[torch.Tensor]: ... +def deterministic_rope_apply_rocm( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + sin_sign: float, +) -> torch.Tensor: ... +def deterministic_rope_apply_token_major_rocm( + x: torch.Tensor, + positions: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + head_dim: int, + sin_sign: float, +) -> torch.Tensor: ... def det_gemm_sm90_compiled() -> bool: ... def det_gemm_fwd(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ... def det_gemm_fwd_rhs_transposed( @@ -270,6 +284,15 @@ def deterministic_collective_rocm_ipc_create( def deterministic_collective_rocm_ipc_synchronize(handle: int) -> None: ... def deterministic_collective_rocm_ipc_destroy(handle: int) -> None: ... def deterministic_collective_rocm_ipc_stage(handle: int, input: torch.Tensor) -> None: ... +def deterministic_collective_rocm_ipc_prepare_staged( + handle: int, + input: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_all_reduce_staged( + handle: int, + input: torch.Tensor, + output: torch.Tensor, +) -> None: ... def deterministic_collective_rocm_ipc_all_reduce( handle: int, output: torch.Tensor, diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index fe0f6eee..0b8ecb9a 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -31,7 +31,12 @@ def _deterministic_all_reduce_(input: torch.Tensor, collective_handle: int) -> N from rl_engine import _C - _C.deterministic_collective_all_reduce_fused(collective_handle, input, input) + if torch.version.hip is not None: + _C.deterministic_collective_rocm_ipc_all_reduce_input( + collective_handle, input, input + ) + else: + _C.deterministic_collective_all_reduce_fused(collective_handle, input, input) @_deterministic_all_reduce_.register_fake @@ -59,7 +64,12 @@ def _deterministic_staging_reserve_(staging: torch.Tensor, collective_handle: in from rl_engine import _C - _C.deterministic_collective_prepare_staged(collective_handle, staging) + if torch.version.hip is not None: + _C.deterministic_collective_rocm_ipc_prepare_staged( + collective_handle, staging + ) + else: + _C.deterministic_collective_prepare_staged(collective_handle, staging) @_deterministic_staging_reserve_.register_fake @@ -80,7 +90,12 @@ def _deterministic_staged_all_reduce( from rl_engine import _C output = torch.empty_like(staging) - _C.deterministic_collective_all_reduce_staged(collective_handle, staging, output) + if torch.version.hip is not None: + _C.deterministic_collective_rocm_ipc_all_reduce_staged( + collective_handle, staging, output + ) + else: + _C.deterministic_collective_all_reduce_staged(collective_handle, staging, output) return output diff --git a/rl_engine/distributed/rocm_collectives.py b/rl_engine/distributed/rocm_collectives.py index ace5dac2..13d2ce3d 100644 --- a/rl_engine/distributed/rocm_collectives.py +++ b/rl_engine/distributed/rocm_collectives.py @@ -5,6 +5,7 @@ import socket import threading +from collections.abc import Iterable from types import TracebackType from typing import Any @@ -17,6 +18,7 @@ _ROCM_IPC_DIRECT_ALL_REDUCE_MAX_BYTES = 768 * 1024 _ROCM_IPC_SHARDED_ALL_REDUCE_MIN_BYTES = 2176 * 1024 _ROCM_IPC_ALL_GATHER_MAX_BYTES = 256 * 1024 +_ROCM_IPC_CONTROL_BYTES = 256 _REDUCTION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) class TorchDistributedDeterministicCollective: @@ -652,6 +654,9 @@ def __init__( ) self._ipc_handle = 0 self._ipc_staging: torch.Tensor | None = None + self._direct_staging_views: dict[ + tuple[tuple[int, ...], torch.dtype], torch.Tensor + ] = {} self._initialize_ipc_transport() @property @@ -674,6 +679,8 @@ def _initialize_ipc_transport(self) -> None: "deterministic_collective_rocm_ipc_synchronize", "deterministic_collective_rocm_ipc_destroy", "deterministic_collective_rocm_ipc_stage", + "deterministic_collective_rocm_ipc_prepare_staged", + "deterministic_collective_rocm_ipc_all_reduce_staged", "deterministic_collective_rocm_ipc_all_reduce", "deterministic_collective_rocm_ipc_all_reduce_input", "deterministic_collective_rocm_ipc_reduce_scatter", @@ -704,6 +711,49 @@ def _initialize_ipc_transport(self) -> None: ) ) self._ipc_staging = staging + self._handle = self._ipc_handle + + def prepare_direct_staging_views( + self, + shapes: Iterable[tuple[int, ...]], + *, + dtype: torch.dtype, + ) -> None: + staging = self._ipc_staging + if not self._ipc_handle or staging is None: + return + if dtype not in _REDUCTION_DTYPES: + raise TypeError(f"unsupported direct-staging dtype {dtype}") + element_size = torch.empty((), dtype=dtype).element_size() + for raw_shape in shapes: + shape = tuple(int(dim) for dim in raw_shape) + numel = 1 + for dim in shape: + if dim < 0: + raise ValueError( + f"direct-staging dimensions must be non-negative, got {shape}" + ) + numel *= dim + size_bytes = numel * element_size + if size_bytes > min( + self.max_size_bytes, + _ROCM_IPC_DIRECT_ALL_REDUCE_MAX_BYTES, + ): + continue + payload = staging.narrow( + 0, + _ROCM_IPC_CONTROL_BYTES, + size_bytes, + ) + self._direct_staging_views[(shape, dtype)] = payload.view(dtype).view(shape) + + def direct_staging_view( + self, + shape: tuple[int, ...], + *, + dtype: torch.dtype, + ) -> torch.Tensor | None: + return self._direct_staging_views.get((tuple(int(dim) for dim in shape), dtype)) def _direct_all_reduce(self, input: torch.Tensor, output: torch.Tensor) -> bool: handle = self._ipc_handle @@ -795,6 +845,8 @@ def close(self) -> None: _C.deterministic_collective_rocm_ipc_synchronize(handle) torch.cuda.synchronize(self.device) self._ipc_handle = 0 + self._handle = 0 _C.deterministic_collective_rocm_ipc_destroy(handle) self._ipc_staging = None + self._direct_staging_views.clear() super().close() diff --git a/rl_engine/integrations/framework_operators.py b/rl_engine/integrations/framework_operators.py index ac981447..52ec4d88 100644 --- a/rl_engine/integrations/framework_operators.py +++ b/rl_engine/integrations/framework_operators.py @@ -378,13 +378,20 @@ def _packed_local_sequence_layout( def _tensor_cache_token(tensor: torch.Tensor) -> tuple[Any, ...]: """Identify one tensor value while it is reused across framework layers.""" + try: + version: int | None = int(tensor._version) + except RuntimeError: + # vLLM warmup runs under inference mode. Inference tensors have no + # version counter, but their storage address and metadata remain valid + # cache identity for the lifetime of that model forward. + version = None return ( tensor.data_ptr(), tuple(tensor.shape), tensor.dtype, tensor.device.type, tensor.device.index, - int(tensor._version), + version, ) @@ -736,9 +743,7 @@ def execute_sequence( "tp_world_size": tp_world, "runtime_platform": runtime_platform, "triton_used": runtime_platform == "rocm", - "deterministic_projection": _strict_attention_projection_provenance( - runtime_platform - ), + "deterministic_projection": _strict_attention_projection_provenance(runtime_platform), "tp_qkv_dgrad_collective": self._tp_collective_backend( module, _MEGATRON_TP_QKV_DGRAD_COLLECTIVE_ATTR, @@ -834,16 +839,22 @@ def __call__( "cp_world_size": cp_world, "tp_world_size": tp_world, "runtime_platform": _device_name(hidden_states), - "actual_backend": "rlkernel.rocm.det_gemm_swiglu" if torch.version.hip is not None else "rlkernel.cuda.det_gemm_swiglu", + "actual_backend": ( + "rlkernel.rocm.det_gemm_swiglu" + if torch.version.hip is not None + else "rlkernel.cuda.det_gemm_swiglu" + ), "gemm_backend": det_gemm_backend_id(), "fallback": False, "gate_up_projection": "separate_strict_launches", "deterministic_all_reduce_backend": ( "none" if tp_world == 1 - else "rocm_ipc_fixed_tree" - if torch.version.hip is not None - else "deterministic_all_reduce.ipc_localized_fixed_tree.v1" + else ( + "rocm_ipc_fixed_tree" + if torch.version.hip is not None + else "deterministic_all_reduce.ipc_localized_fixed_tree.v1" + ) ), "triton_used": torch.version.hip is not None, } @@ -884,13 +895,21 @@ def normalize(plane: torch.Tensor) -> torch.Tensor: if plane.size(0) == num_kv_heads: # [heads, blocks, block, head] (LHBNC after selecting K/V). return plane.permute(1, 2, 0, 3) - raise RuntimeError( - "vLLM K/V cache layout does not expose the declared number of KV heads" - ) + raise RuntimeError("vLLM K/V cache layout does not expose the declared number of KV heads") - # Match vLLM's native AITER layout before the legacy flattened layout. - # With two KV heads both layouts otherwise look like [B, 2, N, 2 * head]. + # RL-Kernel's ROCm backend allocates token-major pages so AITER CK can read + # the paged cache directly. Splitting the packed tail is then a zero-copy + # [blocks, block, heads, head] view with page-major strides. Native AITER + # keeps the head dimension before the block dimension and still needs the + # transpose below. if kv_cache.ndim == 4 and kv_cache.size(-1) == 2 * head_size: + if ( + platform == "rocm" + and num_kv_heads is not None + and kv_cache.size(2) == num_kv_heads + and kv_cache.size(1) != num_kv_heads + ): + return kv_cache.split(head_size, dim=-1) return kv_cache.transpose(1, 2).split(head_size, dim=-1) if ( kv_cache.ndim == 4 @@ -929,9 +948,7 @@ def normalize(plane: torch.Tensor) -> torch.Tensor: key_cache.view(blocks, block_size, num_kv_heads, head_size), value_cache.view(blocks, block_size, num_kv_heads, head_size), ) - raise RuntimeError( - "vLLM Attention KV cache does not match a supported CUDA/ROCm paged layout" - ) + raise RuntimeError("vLLM Attention KV cache does not match a supported CUDA/ROCm paged layout") class VllmAttentionOperator: @@ -954,6 +971,15 @@ def __init__( self._metadata_cache_key: tuple[Any, ...] | None = None self._metadata_cache_owners: set[int] = set() self._metadata_cache_value: tuple[list[dict[str, Any]], dict[str, Any]] | None = None + self._dense_prefill_layout_key: tuple[Any, ...] | None = None + self._dense_prefill_layout_value: tuple[tuple[int, ...], tuple[int, ...]] | None = None + self._dense_prefill_position_ids: dict[tuple[Any, ...], torch.Tensor] = {} + self._rocm_decode_warmup_key: tuple[Any, ...] | None = None + self._rocm_paged_metadata_key: tuple[Any, ...] | None = None + self._rocm_paged_metadata_owners: set[int] = set() + self._rocm_paged_metadata_value: dict[str, Any] | None = None + self._rocm_kv_indptr_cache: dict[tuple[Any, ...], torch.Tensor] = {} + self._phase_provenance: dict[str, dict[str, Any]] = {} def bind_inference(self) -> None: """Resolve the backend after vLLM has selected the worker CUDA device.""" @@ -972,16 +998,110 @@ def bind_inference(self) -> None: ) self._tp_coordinates = (tp_world, tp_rank, tp_group) + def warmup_rocm_decode(self, impl: Any, *, dtype: torch.dtype) -> None: + """Warm the exact strict decode schedule before rollout timing.""" + + if torch.version.hip is None or dtype not in (torch.float16, torch.bfloat16): + return + device = torch.device("cuda", torch.cuda.current_device()) + key = ( + device.index, + dtype, + int(impl.num_heads), + int(impl.num_kv_heads), + int(impl.head_size), + ) + if self._rocm_decode_warmup_key == key: + return + if self._tp_coordinates is None: + self.bind_inference() + assert self._tp_coordinates is not None + tp_world, _tp_rank, _tp_group = self._tp_coordinates + bound = self._handle.get( + torch.empty((1,), device=device, dtype=dtype), + topology={ + "world_size": tp_world, + "tensor_parallel_size": tp_world, + "context_parallel_size": 1, + }, + ) + runtime = bound.bind_accelerator_runtime( + torch.empty((1,), device=device, dtype=dtype) + ) + page_size = 16 + q = torch.empty( + (1, int(impl.num_heads), 1, int(impl.head_size)), + device=device, + dtype=dtype, + ) + cache_shape = ( + 1, + page_size, + int(impl.num_kv_heads), + int(impl.head_size), + ) + k_cache = torch.empty(cache_shape, device=device, dtype=dtype) + v_cache = torch.empty_like(k_cache) + page_table = torch.zeros((1, 1), device=device, dtype=torch.int32) + seqused_k = torch.full((1,), page_size, device=device, dtype=torch.int32) + cu_seqlens_q = torch.tensor((0, 1), device=device, dtype=torch.int32) + kv_indptr = torch.tensor((0, 1), device=device, dtype=torch.int32) + out = torch.empty_like(q) + with torch.inference_mode(): + runtime.forward_paged_with_lse( + q, + k_cache, + v_cache, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_k=page_size, + scale=float(impl.scale), + out=out, + return_lse=False, + page_table_validated=True, + cu_seqlens_q=cu_seqlens_q, + kv_indptr=kv_indptr, + ) + dense_q = torch.empty( + (1, int(impl.num_heads), page_size, int(impl.head_size)), + device=device, + dtype=dtype, + ) + dense_k = k_cache.permute(0, 2, 1, 3).contiguous() + dense_v = v_cache.permute(0, 2, 1, 3).contiguous() + positions = torch.arange(page_size, device=device, dtype=torch.int64).unsqueeze(0) + runtime._core.forward_with_lse( + dense_q, + dense_k, + dense_v, + causal=True, + scale=float(impl.scale), + query_position_ids=positions, + key_position_ids=positions, + ) + torch.cuda.synchronize(device) + self._rocm_decode_warmup_key = key + @property def provenance(self) -> Mapping[str, Any]: + execution = self._phase_provenance.get("decode", self._last_provenance) return { "interface": "vllm.attention.forward", "operator": self.backend_id, "fallback": False, "semantic_instance": self._handle.provenance, - "execution": dict(self._last_provenance), + "execution": { + **dict(execution), + "captured_attention_phases": sorted(self._phase_provenance), + }, } + def _record_phase_provenance( + self, phase: str, provenance: dict[str, Any] + ) -> None: + self._last_provenance = provenance + self._phase_provenance[phase] = provenance + @staticmethod def _metadata_tensor(metadata: Any, *names: str) -> torch.Tensor: for name in names: @@ -990,6 +1110,355 @@ def _metadata_tensor(metadata: Any, *names: str) -> torch.Tensor: return value raise RuntimeError(f"vLLM Attention metadata is missing {'/'.join(names)}") + def _dense_prefill_layout( + self, + starts_cpu: torch.Tensor, + lengths_cpu: torch.Tensor, + *, + num_prefills: int, + num_actual: int, + ) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Cache the host-only request layout across the decoder layers.""" + + key = ( + _tensor_cache_token(starts_cpu), + _tensor_cache_token(lengths_cpu), + num_prefills, + num_actual, + ) + if self._dense_prefill_layout_key == key: + if self._dense_prefill_layout_value is None: + raise RuntimeError("dense prefill layout cache is empty") + return self._dense_prefill_layout_value + starts = tuple(int(v) for v in starts_cpu[: num_prefills + 1].tolist()) + lengths = tuple(int(v) for v in lengths_cpu[:num_prefills].tolist()) + value = (starts, lengths) + self._dense_prefill_layout_key = key + self._dense_prefill_layout_value = value + return value + + def _dense_prefill_positions( + self, + length: int, + *, + device: torch.device, + ) -> torch.Tensor: + key = (int(length), device.type, device.index) + cached = self._dense_prefill_position_ids.get(key) + if cached is not None: + return cached + if len(self._dense_prefill_position_ids) >= 128: + self._dense_prefill_position_ids.pop(next(iter(self._dense_prefill_position_ids))) + value = torch.arange(length, dtype=torch.int64, device=device).unsqueeze(0) + self._dense_prefill_position_ids[key] = value + return value + + def _rocm_dense_prefill( + self, + impl: Any, + query: torch.Tensor, + key: torch.Tensor | None, + value: torch.Tensor | None, + output: torch.Tensor, + attn_metadata: Any, + runtime: Any, + *, + tp_rank: int, + tp_world: int, + num_actual: int, + ) -> torch.Tensor | None: + """Schedule pure prefill as dense per-request AITER launches. + + vLLM hands prefill Q/K/V in logical token order. Reusing those tensors + avoids turning every token into a paged decode row (thousands of + launches for one prompt) while keeping the strict runtime's one KV + group reduction order identical to training. + """ + + if key is None or value is None: + return None + if query.ndim != 3 or key.ndim != 3 or value.ndim != 3: + return None + if int(getattr(attn_metadata, "num_decodes", 0)) != 0: + return None + num_prefills = int(getattr(attn_metadata, "num_prefills", 0)) + if num_prefills <= 0 or int(getattr(attn_metadata, "num_extends", 0)) != 0: + return None + starts_cpu = getattr(attn_metadata, "_rlk_query_start_loc_cpu", None) + lengths_cpu = getattr(attn_metadata, "_rlk_seq_lens_cpu", None) + if not isinstance(starts_cpu, torch.Tensor) or starts_cpu.device.type != "cpu": + return None + if not isinstance(lengths_cpu, torch.Tensor) or lengths_cpu.device.type != "cpu": + return None + if starts_cpu.numel() < num_prefills + 1 or lengths_cpu.numel() < num_prefills: + return None + starts, lengths = self._dense_prefill_layout( + starts_cpu, + lengths_cpu, + num_prefills=num_prefills, + num_actual=num_actual, + ) + if starts[0] != 0 or starts[-1] != num_actual: + return None + if any(end <= start or end - start != length for start, end, length in zip( + starts[:-1], starts[1:], lengths, strict=True + )): + return None + + output_heads = output.view(output.size(0), impl.num_heads, impl.head_size) + for start, end in zip(starts[:-1], starts[1:], strict=True): + query_row = query[start:end].permute(1, 0, 2).unsqueeze(0).contiguous() + key_row = key[start:end].permute(1, 0, 2).unsqueeze(0).contiguous() + value_row = value[start:end].permute(1, 0, 2).unsqueeze(0).contiguous() + position_ids = self._dense_prefill_positions( + end - start, + device=query.device, + ) + result = runtime.forward_with_lse( + query_row, + key_row, + value_row, + contract=_dense_attention_contract( + query_row, + key_row, + role=AttentionRole.INFER, + causal=True, + tp_rank=tp_rank, + tp_world_size=tp_world, + mode=AttentionMode.PREFILL, + global_sequence_length=end - start, + ), + causal=True, + scale=float(impl.scale), + cp_world_size=1, + query_position_ids=position_ids, + key_position_ids=position_ids, + positions_are_sorted=True, + ) + # vLLM uses a rank-3 output buffer at this boundary + # ([tokens, heads, head_dim]); older integrations may provide the + # equivalent flattened rank-2 view. Write through output_heads so + # the adapter preserves the framework-owned layout in both cases. + result_output = result.out.permute(0, 2, 1, 3).squeeze(0) + output_group = output_heads.narrow(0, start, end - start) + if result_output.shape != output_group.shape: + raise RuntimeError( + "strict ROCm dense prefill output shape does not match vLLM: " + f"result={tuple(result_output.shape)}, " + f"output={tuple(output_group.shape)}" + ) + output_group.copy_(result_output) + if num_actual < output.size(0): + output[num_actual:].zero_() + self._record_phase_provenance("prefill", { + "framework_layout": "vllm_dense_qkv_prefill", + "materialization": "direct_dense_qkv_to_aiter_ck", + "tp_world_size": tp_world, + "runtime_platform": "rocm", + "triton_used": True, + "prefill_request_count": num_prefills, + "prefill_token_count": num_actual, + "core_launch_count": num_prefills, + "deterministic_projection": _strict_attention_projection_provenance("rocm"), + "deterministic_all_reduce_backend": "unbound" if tp_world > 1 else "none", + "direct_output_buffer": True, + }) + return output + + def _rocm_direct_paged_metadata( + self, + attn_metadata: Any, + *, + block_table: torch.Tensor, + block_size: int, + num_actual: int, + cache_owner: Any, + ) -> tuple[dict[str, Any], bool] | None: + """Reuse vLLM's sequence-level GPU metadata without token expansion.""" + + num_decodes = int(getattr(attn_metadata, "num_decodes", 0)) + num_prefills = int(getattr(attn_metadata, "num_prefills", 0)) + num_extends = int(getattr(attn_metadata, "num_extends", 0)) + if num_extends or (num_decodes > 0) == (num_prefills > 0): + return None + mode = "prefill" if num_prefills > 0 else "decode" + sequence_count = num_prefills if num_prefills > 0 else num_decodes + if sequence_count <= 0 or num_actual <= 0: + return None + if mode == "prefill": + prefill = getattr(attn_metadata, "prefill_metadata", None) + if prefill is None: + return None + query_start_loc = getattr(prefill, "query_start_loc", None) + max_seqlen_q = int(getattr(prefill, "max_query_len", 0)) + causal = bool(getattr(attn_metadata, "causal", True)) + else: + decode = getattr(attn_metadata, "decode_metadata", None) + if decode is None: + return None + query_start_loc = self._metadata_tensor(attn_metadata, "query_start_loc") + max_seqlen_q = int(getattr(decode, "max_query_len", 0)) + 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") + max_seq_len = int( + getattr(attn_metadata, "max_seq_len", block_table.size(1) * block_size) + ) + 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(block_table), + sequence_count, + num_actual, + max_seqlen_q, + max_seq_len, + page_count, + ) + owner_id = id(cache_owner) + if ( + self._rocm_paged_metadata_key == key + and owner_id not in self._rocm_paged_metadata_owners + and self._rocm_paged_metadata_value is not None + ): + 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( + 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 + ) + if not seq_lens.is_contiguous(): + seq_lens = seq_lens.contiguous() + pages = block_table[:sequence_count, :page_count].to(dtype=torch.int32) + if not pages.is_contiguous(): + pages = pages.contiguous() + indptr_key = ( + block_table.device.type, + block_table.device.index, + sequence_count, + page_count, + ) + kv_indptr = self._rocm_kv_indptr_cache.get(indptr_key) + if kv_indptr is None: + kv_indptr = torch.arange( + sequence_count + 1, + dtype=torch.int32, + device=block_table.device, + ) * page_count + self._rocm_kv_indptr_cache[indptr_key] = kv_indptr + value = { + "mode": mode, + "sequence_count": sequence_count, + "page_count": page_count, + "pages": pages, + "seqused_k": seq_lens, + "cu_seqlens_q": query_start_loc, + "kv_indptr": kv_indptr, + "max_seqlen_q": max_seqlen_q, + "max_seqlen_k": page_count * block_size, + "causal": causal, + } + self._rocm_paged_metadata_key = key + self._rocm_paged_metadata_owners = {owner_id} + self._rocm_paged_metadata_value = value + return value, False + + def _rocm_direct_paged( + self, + impl: Any, + layer: Any, + query: torch.Tensor, + output: torch.Tensor, + attn_metadata: Any, + runtime: Any, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + *, + tp_world: int, + num_actual: int, + ) -> torch.Tensor | None: + direct = getattr(runtime, "forward_paged_varlen_with_lse", None) + if not callable(direct): + return None + metadata_result = self._rocm_direct_paged_metadata( + attn_metadata, + block_table=block_table, + block_size=key_cache.size(1), + num_actual=num_actual, + cache_owner=layer, + ) + if metadata_result is None: + return None + metadata, reused = metadata_result + if not reused: + self._validate_page_bounds_once( + [{"pages": metadata["pages"]}], + num_cache_blocks=key_cache.size(0), + ) + query_ready = query.narrow(0, 0, num_actual) + if not query_ready.is_contiguous(): + query_ready = query_ready.contiguous() + output_heads = output.view(output.size(0), impl.num_heads, impl.head_size) + output_ready = output_heads.narrow(0, 0, num_actual) + result = direct( + query_ready, + key_cache, + value_cache, + page_table=metadata["pages"], + seqused_k=metadata["seqused_k"], + cu_seqlens_q=metadata["cu_seqlens_q"], + kv_indptr=metadata["kv_indptr"], + max_seqlen_q=metadata["max_seqlen_q"], + max_seqlen_k=metadata["max_seqlen_k"], + causal=metadata["causal"], + scale=float(impl.scale), + out=output_ready, + return_lse=False, + page_table_validated=True, + ) + if result.out.data_ptr() != output_ready.data_ptr(): + output_ready.copy_(result.out) + if num_actual < output.size(0): + output[num_actual:].zero_() + operator_provenance = _compact_attention_provenance(result.provenance) + projection_collective_backend = "none" + if tp_world > 1: + projection_collective_backend = "unbound" + if self._projection_collective_backend is not None: + projection_collective_backend = ( + self._projection_collective_backend() or "unbound" + ) + self._record_phase_provenance(metadata["mode"], { + "framework_layout": "vllm_paged_kv", + "materialization": "direct_vllm_paged_kv_to_aiter_batch_prefill_ck", + "dense_kv_materialized": False, + "tp_world_size": tp_world, + "runtime_platform": "rocm", + "triton_used": True, + "attention_phase": metadata["mode"], + "sequence_count": metadata["sequence_count"], + "query_token_count": num_actual, + "launch_group_count": 1, + "metadata_source": "vllm_gpu_sequence_level", + "metadata_reused_across_layers": reused, + "deterministic_projection": _strict_attention_projection_provenance("rocm"), + "deterministic_all_reduce_backend": projection_collective_backend, + "direct_output_buffer": True, + "operator": operator_provenance, + }) + return output + def _materialization_groups( self, attn_metadata: Any, @@ -999,6 +1468,8 @@ def _materialization_groups( block_size: int, num_actual: int, cache_owner: Any | None = None, + include_host_lengths: bool = False, + page_bounds_epoch_factory: Callable[[], object] | None = None, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: """Build one paged launch from vLLM's graph-replayable GPU metadata.""" @@ -1026,6 +1497,12 @@ def _materialization_groups( num_actual, block_size, page_count, + include_host_lengths, + ( + None + if page_bounds_epoch_factory is None + else id(getattr(page_bounds_epoch_factory, "__self__", page_bounds_epoch_factory)) + ), ) owner_id = id(cache_owner) if cache_owner is not None else None # Resolve the cross-layer cache before scheduling any GPU metadata work. @@ -1057,12 +1534,19 @@ def _materialization_groups( seqused_k, torch.ones_like(seqused_k), ) + cached_lengths = ( + tuple(int(value) for value in seqused_k.tolist()) if include_host_lengths else None + ) pages = ( block_table.index_select(0, request_indices)[:, :page_count] .to(dtype=torch.int32) .contiguous() ) + cu_seqlens_q = torch.arange( + num_actual + 1, dtype=torch.int32, device=query.device + ) + kv_indptr = cu_seqlens_q * page_count groups = [ { "page_count": page_count, @@ -1072,6 +1556,12 @@ def _materialization_groups( "query_count": num_actual, "query_contiguous": True, "seqused_k": seqused_k, + "cached_lengths": cached_lengths, + "page_bounds_epoch": ( + None if page_bounds_epoch_factory is None else page_bounds_epoch_factory() + ), + "cu_seqlens_q": cu_seqlens_q, + "kv_indptr": kv_indptr, } ] summary = { @@ -1090,6 +1580,20 @@ def _materialization_groups( self._metadata_cache_value = (groups, summary) return groups, summary + @staticmethod + def _validate_page_bounds_once( + groups: list[dict[str, Any]], + *, + num_cache_blocks: int, + ) -> None: + for group in groups: + pages = group["pages"] + bounds_ok = torch.all((pages >= 0) & (pages < num_cache_blocks)) + if pages.is_cuda: + torch._assert_async(bounds_ok, "page_table entries are outside the KV cache") + elif not bool(bounds_ok.item()): + raise ValueError("page_table entries are outside the KV cache") + def __call__( self, impl: Any, @@ -1103,7 +1607,6 @@ def __call__( output_scale: torch.Tensor | None = None, output_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: - del key, value if output_scale is not None or output_block_scale is not None: raise RuntimeError("strict vLLM Attention does not support quantized output") if attn_metadata is None: @@ -1113,16 +1616,7 @@ def __call__( if query.ndim != 3: raise RuntimeError("vLLM query must use [tokens, heads, head_dim]") runtime_platform = _require_attention_accelerator(query) - key_cache, value_cache = _vllm_kv_cache_views( - kv_cache, - head_size=int(impl.head_size), - num_kv_heads=int(impl.num_kv_heads), - platform=runtime_platform, - ) - if key_cache.dtype != query.dtype or value_cache.dtype != query.dtype: - raise RuntimeError("strict vLLM Attention requires an unquantized KV cache") - block_table = self._metadata_tensor(attn_metadata, "block_table", "block_table_tensor") num_actual = int(getattr(attn_metadata, "num_actual_tokens", query.size(0))) if num_actual < 0 or num_actual > query.size(0): raise RuntimeError("vLLM num_actual_tokens is outside the query buffer") @@ -1135,7 +1629,6 @@ def __call__( if output.size(0) < num_actual: raise RuntimeError("vLLM output buffer is smaller than num_actual_tokens") output_heads = output.view(output.size(0), impl.num_heads, impl.head_size) - block_size = key_cache.size(1) if self._tp_coordinates is None: self._tp_coordinates = _vllm_tp_coordinates() tp_world, tp_rank, tp_group = self._tp_coordinates @@ -1148,6 +1641,53 @@ def __call__( }, ) runtime = operator.bind_accelerator_runtime(query) + page_bounds_epoch_factory: Callable[[], object] | None = None + if runtime_platform == "rocm": + candidate_factory = getattr(runtime, "new_page_bounds_epoch", None) + if callable(candidate_factory): + page_bounds_epoch_factory = candidate_factory + block_table = self._metadata_tensor( + attn_metadata, "block_table", "block_table_tensor" + ) + key_cache, value_cache = _vllm_kv_cache_views( + kv_cache, + head_size=int(impl.head_size), + num_kv_heads=int(impl.num_kv_heads), + platform=runtime_platform, + ) + if key_cache.dtype != query.dtype or value_cache.dtype != query.dtype: + raise RuntimeError("strict vLLM Attention requires an unquantized KV cache") + if runtime_platform == "rocm": + direct_output = self._rocm_direct_paged( + impl, + layer, + query, + output, + attn_metadata, + runtime, + key_cache, + value_cache, + block_table, + tp_world=tp_world, + num_actual=num_actual, + ) + if direct_output is not None: + return direct_output + prefill_output = self._rocm_dense_prefill( + impl, + query, + key, + value, + output, + attn_metadata, + runtime, + tp_rank=tp_rank, + tp_world=tp_world, + num_actual=num_actual, + ) + if prefill_output is not None: + return prefill_output + block_size = key_cache.size(1) groups, metadata_summary = self._materialization_groups( attn_metadata, query=query, @@ -1155,7 +1695,17 @@ def __call__( block_size=block_size, num_actual=num_actual, cache_owner=layer, + include_host_lengths=runtime_platform == "rocm", + page_bounds_epoch_factory=page_bounds_epoch_factory, ) + page_table_validated = False + if runtime_platform == "rocm": + if not metadata_summary["metadata_reused_across_layers"]: + self._validate_page_bounds_once( + groups, + num_cache_blocks=key_cache.size(0), + ) + page_table_validated = True last_operator_provenance: dict[str, Any] = {} next_query_row = 0 for group in groups: @@ -1171,26 +1721,42 @@ def __call__( page_count = int(group["page_count"]) pages = group["pages"] query_indices = group["query_indices"] + runtime_out = None + output_group = None if group["query_contiguous"]: query_start = int(group["query_start"]) query_count = int(group["query_count"]) q_ready = query.narrow(0, query_start, query_count).unsqueeze(2) + output_group = output_heads.narrow(0, query_start, query_count) + runtime_out = output_group.unsqueeze(2) else: q_ready = query.index_select(0, query_indices).unsqueeze(2).contiguous() + runtime_kwargs = { + "page_table": pages, + "seqused_k": group["seqused_k"], + "max_seqlen_k": page_count * block_size, + "scale": float(impl.scale), + "out": runtime_out, + } + if runtime_platform == "rocm": + runtime_kwargs["cached_lengths"] = group["cached_lengths"] + if group["page_bounds_epoch"] is not None: + runtime_kwargs["page_bounds_epoch"] = group["page_bounds_epoch"] + runtime_kwargs["return_lse"] = False + runtime_kwargs["page_table_validated"] = page_table_validated + runtime_kwargs["cu_seqlens_q"] = group["cu_seqlens_q"] + runtime_kwargs["kv_indptr"] = group["kv_indptr"] result = runtime.forward_paged_with_lse( q_ready, key_cache, value_cache, - page_table=pages, - seqused_k=group["seqused_k"], - max_seqlen_k=page_count * block_size, - scale=float(impl.scale), + **runtime_kwargs, ) result_output = result.out.squeeze(2) if group["query_contiguous"]: - output_heads.narrow(0, int(group["query_start"]), int(group["query_count"])).copy_( - result_output - ) + assert output_group is not None + if result_output.data_ptr() != output_group.data_ptr(): + output_group.copy_(result_output) else: output_heads.index_copy_(0, query_indices, result.out.squeeze(2)) last_operator_provenance = _compact_attention_provenance(result.provenance) @@ -1201,25 +1767,29 @@ def __call__( projection_collective_backend = ( self._projection_collective_backend() or "unbound" ) - self._last_provenance = { + phase = ( + "decode" + if int(getattr(attn_metadata, "num_decodes", 0)) > 0 + else "prefill" + ) + self._record_phase_provenance(phase, { "framework_layout": "vllm_paged_kv", "materialization": ( - "logical_paged_kv_to_aiter_ck_dense" + "direct_vllm_paged_kv_to_aiter_batch_prefill_ck" if runtime_platform == "rocm" else "direct_paged_fa4" ), + "dense_kv_materialized": False, "tp_world_size": tp_world, "tp_group_bound": tp_group is not None, "runtime_platform": runtime_platform, "triton_used": runtime_platform == "rocm", - "deterministic_projection": _strict_attention_projection_provenance( - runtime_platform - ), + "deterministic_projection": _strict_attention_projection_provenance(runtime_platform), "deterministic_all_reduce_backend": projection_collective_backend, "direct_output_buffer": direct_output_buffer, **metadata_summary, "operator": last_operator_provenance, - } + }) return output @@ -1657,8 +2227,7 @@ def __call__( if ( strict_provenance.get("deterministic_linear_logp") is not True or strict_provenance.get("actual_backend") != self._linear_logp.backend_id - or strict_provenance.get("strict_entrypoint") - != expected_entrypoint + or strict_provenance.get("strict_entrypoint") != expected_entrypoint ): raise RuntimeError( "strict vLLM rollout linear_logp did not execute the deterministic " diff --git a/rl_engine/integrations/runtime.py b/rl_engine/integrations/runtime.py index 3807ad52..083783e4 100644 --- a/rl_engine/integrations/runtime.py +++ b/rl_engine/integrations/runtime.py @@ -156,7 +156,7 @@ def record_execution( execution_mode: str = "eager", execution_provenance: Mapping[str, Any] | None = None, ) -> None: - """Record one eager or custom-op execution outside Dynamo tracing.""" + """Record one eager or captured custom-op execution outside Dynamo tracing.""" normalized = module.strip().lower() implementation = self.plan.implementation_for(normalized, self.target) @@ -167,7 +167,7 @@ def record_execution( f"{self.framework} {normalized} execution evidence did not use " "the installed RL-Kernel operator" ) - if execution_mode not in {"eager", "compiled_cuda_graph"}: + if execution_mode not in {"eager", "compiled_cuda_graph", "compiled_hip_graph"}: raise ValueError(f"unknown execution mode {execution_mode!r}") backend_id = getattr(selected, "backend_id", None) if not isinstance(backend_id, str) or not backend_id.strip(): diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index 63e43424..e43d9032 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -10,6 +10,7 @@ import pathlib import re import sys +from dataclasses import dataclass from types import MethodType from typing import Any @@ -52,10 +53,19 @@ _STRICT_ROTARY_INIT_MARKER = "__rl_kernel_original_strict_rotary_init__" _STRICT_ROCM_ROPE_PATCH_MARKER = "__rl_kernel_original_strict_rocm_rope_forward__" _STRICT_LM_HEAD_LINEAR_PATCH_MARKER = "__rl_kernel_original_lm_head_linear_apply__" +_STRICT_LM_HEAD_PROCESS_PATCH_MARKER = "__rl_kernel_original_lm_head_process_weights__" +_STRICT_LM_HEAD_TIE_PATCH_MARKER = "__rl_kernel_original_lm_head_tie_weights__" +_STRICT_LM_HEAD_CACHE_BUFFER = "_rl_kernel_lm_head_weight_t" +_STRICT_LM_HEAD_CACHE_STATE = "_rl_kernel_lm_head_weight_cache_state" +_STRICT_LM_HEAD_CACHE_HOOK = "_rl_kernel_lm_head_weight_cache_state_dict_hook" +_STRICT_LM_HEAD_TIED = "_rl_kernel_lm_head_tied_weight" _STRICT_O_PROJ_COLLECTIVE_MARKER = "__rl_kernel_o_proj_collective__" +_STRICT_O_PROJ_FUSED_ALL_REDUCE_MARKER = "__rl_kernel_o_proj_fused_all_reduce__" +_STRICT_O_PROJ_COMPILED_COLLECTIVE_SLOT = "__rl_kernel_o_proj_compiled_collective_slot__" _STRICT_ROW_PARALLEL_PATCH_MARKER = "__rl_kernel_original_row_parallel_forward__" _STRICT_DIRECT_STAGING_MARKER = "__rl_kernel_direct_staging_active__" _STRICT_LAYER_DIAGNOSTIC_PATCH_MARKER = "__rl_kernel_original_layer_diagnostic_forward__" +_STRICT_WEIGHT_CACHE_REFRESH_MARKER = "__rl_kernel_original_finish_weight_update__" _RLK_ATTENTION_BACKEND: type[Any] | None = None _RLK_ATTENTION_IMPL: type[Any] | None = None _RLK_ATTENTION_BUILDER: type[Any] | None = None @@ -63,6 +73,223 @@ _VLLM_LAYER_DIAGNOSTIC_BUFFER: dict[str, Any] | None = None _VLLM_LAYER_DIAGNOSTIC_CALLS = 0 _VLLM_LAYER_DIAGNOSTIC_ACTIVE_LAYER: int | None = None +_ROCM_STATEFUL_GRAPH_SPLITTING_OPS = ( + DETERMINISTIC_ALL_REDUCE_OP, + "rl_kernel::rocm_det_gemm_linear_all_reduce_inference", + "rl_kernel::qwen3_ffn_packed_tp_inference_rocm", +) + + +@dataclass +class _LmHeadWeightCacheState: + source: torch.Tensor + weight_t: torch.Tensor + source_data_ptr: int + source_shape: tuple[int, ...] + source_stride: tuple[int, ...] + source_dtype: torch.dtype + source_device: torch.device + source_version: int | None + cache_data_ptr: int + cache_version: int | None + generation: int + valid: bool + refresh_pending: bool + + +def _tracked_tensor_version(value: torch.Tensor) -> int | None: + try: + return int(value._version) + except RuntimeError: + # Tensors created in inference mode do not carry a version counter. + return None + + +def _invalidate_lm_head_weight_cache(weight: Any) -> None: + state = getattr(weight, _STRICT_LM_HEAD_CACHE_STATE, None) + if isinstance(state, _LmHeadWeightCacheState): + state.valid = False + state.refresh_pending = False + + +def _mark_lm_head_weight_cache_refreshable(weight: Any) -> None: + state = getattr(weight, _STRICT_LM_HEAD_CACHE_STATE, None) + if isinstance(state, _LmHeadWeightCacheState): + state.refresh_pending = True + + +def _keep_lm_head_cache_non_persistent(layer: Any, *_args: Any) -> None: + buffers = getattr(layer, "_buffers", {}) + non_persistent = getattr(layer, "_non_persistent_buffers_set", None) + if _STRICT_LM_HEAD_CACHE_BUFFER in buffers and isinstance(non_persistent, set): + non_persistent.add(_STRICT_LM_HEAD_CACHE_BUFFER) + + +def _record_lm_head_weight_cache_refresh( + state: _LmHeadWeightCacheState, + weight: torch.Tensor, +) -> None: + state.source_data_ptr = int(weight.data_ptr()) + state.source_shape = tuple(int(dim) for dim in weight.shape) + state.source_stride = tuple(int(stride) for stride in weight.stride()) + state.source_dtype = weight.dtype + state.source_device = weight.device + state.source_version = _tracked_tensor_version(weight) + state.cache_data_ptr = int(state.weight_t.data_ptr()) + state.cache_version = _tracked_tensor_version(state.weight_t) + state.generation += 1 + state.valid = True + state.refresh_pending = False + + +def _refresh_lm_head_weight_cache( + layer: Any, + prepare_weight: Any, +) -> _LmHeadWeightCacheState: + weight = getattr(layer, "weight", None) + if not isinstance(weight, torch.Tensor): + raise RuntimeError("strict ROCm LM-head cache requires a tensor weight") + if bool(getattr(layer, _STRICT_LM_HEAD_TIED, False)): + raise RuntimeError("strict ROCm LM-head cache does not support tied embeddings") + + state = getattr(layer, _STRICT_LM_HEAD_CACHE_STATE, None) + if state is None: + weight_t = prepare_weight(weight) + register_buffer = getattr(layer, "register_buffer", None) + if callable(register_buffer): + register_buffer( + _STRICT_LM_HEAD_CACHE_BUFFER, + weight_t, + persistent=False, + ) + else: + # Kept for lightweight integration adapters used outside nn.Module. + setattr(layer, _STRICT_LM_HEAD_CACHE_BUFFER, weight_t) + register_state_dict_pre_hook = getattr(layer, "register_state_dict_pre_hook", None) + if callable(register_state_dict_pre_hook) and not hasattr( + layer, _STRICT_LM_HEAD_CACHE_HOOK + ): + handle = register_state_dict_pre_hook(_keep_lm_head_cache_non_persistent) + setattr(layer, _STRICT_LM_HEAD_CACHE_HOOK, handle) + state = _LmHeadWeightCacheState( + source=weight, + weight_t=weight_t, + source_data_ptr=0, + source_shape=(), + source_stride=(), + source_dtype=weight.dtype, + source_device=weight.device, + source_version=None, + cache_data_ptr=int(weight_t.data_ptr()), + cache_version=None, + generation=0, + valid=False, + refresh_pending=False, + ) + setattr(layer, _STRICT_LM_HEAD_CACHE_STATE, state) + else: + if not isinstance(state, _LmHeadWeightCacheState): + raise RuntimeError("strict ROCm LM-head cache state has an invalid type") + if ( + state.source is not weight + or getattr(layer, _STRICT_LM_HEAD_CACHE_BUFFER, None) is not state.weight_t + ): + # Checkpoint-format layerwise reload temporarily replaces every + # Parameter (and may remove derived buffers) before post-load + # processing, then copies the canonical weight back into the + # original stable storage. Defer the transpose until the first + # forward observes that stable source again. This also avoids + # depending on vLLM copying a non-persistent derived buffer. + state.valid = False + state.refresh_pending = True + return state + state.valid = False + state.refresh_pending = False + cache_data_ptr = int(state.weight_t.data_ptr()) + refreshed = prepare_weight(weight, out=state.weight_t) + if refreshed is not state.weight_t or int(refreshed.data_ptr()) != cache_data_ptr: + raise RuntimeError("strict ROCm LM-head refresh replaced stable cache storage") + _record_lm_head_weight_cache_refresh(state, weight) + + if state.generation == 0: + _record_lm_head_weight_cache_refresh(state, weight) + setattr(weight, _STRICT_LM_HEAD_CACHE_STATE, state) + return state + + +def _validated_lm_head_weight_cache( + layer: Any, + prepare_weight: Any | None = None, +) -> torch.Tensor: + weight = getattr(layer, "weight", None) + if bool(getattr(layer, _STRICT_LM_HEAD_TIED, False)): + raise RuntimeError("strict ROCm LM-head cache does not support tied embeddings") + _keep_lm_head_cache_non_persistent(layer) + if not isinstance(weight, torch.Tensor): + raise RuntimeError("strict ROCm LM-head cache was not prepared after model loading") + state_value = getattr(layer, _STRICT_LM_HEAD_CACHE_STATE, None) + if not isinstance(state_value, _LmHeadWeightCacheState): + raise RuntimeError("strict ROCm LM-head cache was not prepared after model loading") + state: _LmHeadWeightCacheState = state_value + cached_weight = state.weight_t + if not isinstance(cached_weight, torch.Tensor): + raise RuntimeError("strict ROCm LM-head cache state has an invalid weight") + if ( + state.source is not weight + or getattr(weight, _STRICT_LM_HEAD_CACHE_STATE, None) is not state + ): + raise RuntimeError("strict ROCm LM-head cache is not bound to the active weight") + if getattr(layer, _STRICT_LM_HEAD_CACHE_BUFFER, None) is not cached_weight: + raise RuntimeError("strict ROCm LM-head cache buffer was replaced") + + current_source = ( + int(weight.data_ptr()), + tuple(int(dim) for dim in weight.shape), + tuple(int(stride) for stride in weight.stride()), + weight.dtype, + weight.device, + _tracked_tensor_version(weight), + ) + expected_source = ( + state.source_data_ptr, + state.source_shape, + state.source_stride, + state.source_dtype, + state.source_device, + state.source_version, + ) + if current_source[:-1] != expected_source[:-1]: + state.valid = False + state.refresh_pending = False + raise RuntimeError("strict ROCm LM-head weight storage changed without a cache refresh") + if not cached_weight.is_contiguous() or int(cached_weight.data_ptr()) != state.cache_data_ptr: + state.valid = False + state.refresh_pending = False + raise RuntimeError("strict ROCm LM-head cache storage changed after refresh") + if state.valid: + if current_source[-1] != expected_source[-1]: + state.valid = False + state.refresh_pending = False + raise RuntimeError("strict ROCm LM-head weight changed without a cache refresh") + if _tracked_tensor_version(cached_weight) != state.cache_version: + state.valid = False + state.refresh_pending = False + raise RuntimeError("strict ROCm LM-head cache bytes changed after refresh") + return cached_weight + + if not state.refresh_pending or prepare_weight is None: + raise RuntimeError("strict ROCm LM-head cache is invalid during weight update") + cache_data_ptr = int(cached_weight.data_ptr()) + try: + refreshed = prepare_weight(weight, out=cached_weight) + except Exception: + state.refresh_pending = False + raise + if refreshed is not cached_weight or int(refreshed.data_ptr()) != cache_data_ptr: + state.refresh_pending = False + raise RuntimeError("strict ROCm LM-head refresh replaced stable cache storage") + _record_lm_head_weight_cache_refresh(state, weight) + return cached_weight def _alignment_diagnostics_enabled() -> bool: @@ -90,6 +317,7 @@ def _patch_qwen3_layer_alignment_diagnostics() -> None: if not _alignment_diagnostics_enabled(): return from vllm.model_executor.models.qwen3 import Qwen3Attention, Qwen3DecoderLayer + from rl_engine.kernels.ops.rocm.attention.strict_runtime import StrictRocmAttentionRuntime if hasattr(Qwen3DecoderLayer, _STRICT_LAYER_DIAGNOSTIC_PATCH_MARKER): @@ -178,9 +406,7 @@ def forward_wrapped( "outputs": torch.stack(buffer["outputs"]).cpu(), } if "layer0" in buffer: - payload["layer0"] = { - name: value.cpu() for name, value in buffer["layer0"].items() - } + payload["layer0"] = {name: value.cpu() for name, value in buffer["layer0"].items()} _VLLM_LAYER_DIAGNOSTIC_BUFFER = None root = os.getenv("RL_KERNEL_ALIGNMENT_DIAGNOSTICS_DIR", "").strip() if root: @@ -188,11 +414,7 @@ def forward_wrapped( output_dir.mkdir(parents=True, exist_ok=True) torch.save( payload, - output_dir - / ( - f"vllm-pid{os.getpid()}-rank{rank:05d}-" - f"call{call_index:08d}.pt" - ), + output_dir / (f"vllm-pid{os.getpid()}-rank{rank:05d}-" f"call{call_index:08d}.pt"), ) return result @@ -215,15 +437,17 @@ def attention_forward_wrapped( q, k = instance.rotary_emb(positions, q, k) attention_core = instance.attn(q, k, v) output, _ = instance.o_proj(attention_core) - buffer.setdefault("layer0", {}).update({ - "attention_norm": hidden_states.detach(), - "qkv": qkv.detach(), - "query": q.view(*q.shape[:-1], -1, instance.head_dim).detach(), - "key": k.view(*k.shape[:-1], -1, instance.head_dim).detach(), - "value": v.view(*v.shape[:-1], -1, instance.head_dim).detach(), - "attention_core": attention_core.detach(), - "attention_output": output.detach(), - }) + buffer.setdefault("layer0", {}).update( + { + "attention_norm": hidden_states.detach(), + "qkv": qkv.detach(), + "query": q.view(*q.shape[:-1], -1, instance.head_dim).detach(), + "key": k.view(*k.shape[:-1], -1, instance.head_dim).detach(), + "value": v.view(*v.shape[:-1], -1, instance.head_dim).detach(), + "attention_core": attention_core.detach(), + "attention_output": output.detach(), + } + ) return output def gather_paged_row_wrapped( @@ -231,12 +455,15 @@ def gather_paged_row_wrapped( v_cache: torch.Tensor, page_row: torch.Tensor, cached_length: int, + *, + validate_bounds: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: key, value = original_gather_paged_row( k_cache, v_cache, page_row, cached_length, + validate_bounds=validate_bounds, ) buffer = _VLLM_LAYER_DIAGNOSTIC_BUFFER if buffer is not None and _VLLM_LAYER_DIAGNOSTIC_ACTIVE_LAYER == 0: @@ -347,8 +574,14 @@ def _patch_qwen_lm_head_padding() -> None: padding_size = padded_vocab - real_vocab original = ParallelLMHead.__init__ original_weight_loader = VocabParallelEmbedding.weight_loader + original_tie_weights = getattr(ParallelLMHead, "tie_weights", None) def strict_weight_loader(instance: Any, param: Any, loaded_weight: torch.Tensor) -> None: + # vLLM loaders write through ``param.data``, which does not reliably + # advance the Parameter version counter. Explicitly invalidate before + # any write so a failed or incomplete hot update cannot reuse old + # prepared LM-head bytes. + _invalidate_lm_head_weight_cache(param) strict_real_vocab = getattr(instance, "_rl_kernel_real_vocab_size", None) output_dim = getattr(param, "output_dim", None) packed_dim = getattr(param, "packed_dim", None) @@ -368,8 +601,10 @@ def strict_weight_loader(instance: Any, param: Any, loaded_weight: torch.Tensor) param[available:shard_size].data.fill_(0) if shard_size < param.shape[0]: param[shard_size:].data.fill_(0) + _mark_lm_head_weight_cache_refreshable(param) return original_weight_loader(instance, param, loaded_weight) + _mark_lm_head_weight_cache_refreshable(param) def wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: if args: @@ -382,6 +617,7 @@ def wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: kwargs["padding_size"] = padding_size original(instance, *args, **kwargs) instance._rl_kernel_real_vocab_size = real_vocab + setattr(instance, _STRICT_PROJECTION_MARKER, "lm_head") if ( int(getattr(instance, "org_vocab_size", -1)) != padded_vocab or int(getattr(instance, "num_embeddings_padded", -1)) != padded_vocab @@ -392,6 +628,21 @@ def wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: ) setattr(ParallelLMHead, _PATCH_MARKER, original) + if callable(original_tie_weights) and not hasattr( + ParallelLMHead, _STRICT_LM_HEAD_TIE_PATCH_MARKER + ): + + def strict_tie_weights(instance: Any, embed_tokens: Any) -> Any: + result = original_tie_weights(instance, embed_tokens) + setattr(instance, _STRICT_LM_HEAD_TIED, True) + return result + + setattr( + ParallelLMHead, + _STRICT_LM_HEAD_TIE_PATCH_MARKER, + original_tie_weights, + ) + setattr(ParallelLMHead, "tie_weights", strict_tie_weights) if not hasattr(VocabParallelEmbedding, "__rl_kernel_original_weight_loader__"): setattr( VocabParallelEmbedding, @@ -425,6 +676,14 @@ def wrapped( lm_head = getattr(instance, "lm_head", None) if lm_head is None: raise RuntimeError("strict rollout linear_logp requires a Qwen LM-head") + model = getattr(instance, "model", None) + if ( + torch.version.hip is not None + and model is not None + and lm_head is getattr(model, "embed_tokens", None) + ): + setattr(lm_head, _STRICT_LM_HEAD_TIED, True) + raise RuntimeError("strict ROCm LM-head cache does not support tied embeddings") setattr(lm_head, _STRICT_PROJECTION_MARKER, "lm_head") logits = _original(instance, hidden_states) if not get_pp_group().is_last_rank: @@ -459,6 +718,7 @@ def _patch_strict_lm_head_linear( *, linear_method_cls: type[Any] | None = None, det_gemm: Any | None = None, + prepare_weight: Any | None = None, ) -> None: """Route vLLM's existing LM-head projection through RL-Kernel det_gemm.""" @@ -473,6 +733,38 @@ def _patch_strict_lm_head_linear( from rl_engine.kernels.ops.matmul.det_gemm import DetGemmOp det_gemm = DetGemmOp() + if torch.version.hip is not None and prepare_weight is None: + from rl_engine.kernels.ops.rocm.matmul.det_gemm import prepare_det_gemm_linear_weight + + prepare_weight = prepare_det_gemm_linear_weight + + if not hasattr(linear_method_cls, _STRICT_LM_HEAD_PROCESS_PATCH_MARKER): + previous_process_weights = linear_method_cls.process_weights_after_loading + + def process_weights_after_loading( + method: Any, + layer: Any, + ) -> Any: + result = previous_process_weights(method, layer) + if ( + torch.version.hip is not None + and getattr(layer, _STRICT_PROJECTION_MARKER, None) == "lm_head" + ): + if prepare_weight is None: + raise RuntimeError("strict ROCm LM-head weight preparation is unavailable") + _refresh_lm_head_weight_cache(layer, prepare_weight) + return result + + setattr( + linear_method_cls, + _STRICT_LM_HEAD_PROCESS_PATCH_MARKER, + previous_process_weights, + ) + setattr( + linear_method_cls, + "process_weights_after_loading", + process_weights_after_loading, + ) def wrapped( method: Any, @@ -483,7 +775,16 @@ def wrapped( if getattr(layer, _STRICT_PROJECTION_MARKER, None) != "lm_head": return previous_apply(method, layer, x, bias) x_2d = x.reshape(-1, x.shape[-1]) - output_2d = det_gemm.linear(x_2d, layer.weight) + if torch.version.hip is not None: + linear_prepared = getattr(det_gemm, "linear_prepared", None) + if not callable(linear_prepared): + raise RuntimeError("strict ROCm LM-head prepared GEMM is unavailable") + output_2d = linear_prepared( + x_2d, + _validated_lm_head_weight_cache(layer, prepare_weight), + ) + else: + output_2d = det_gemm.linear(x_2d, layer.weight) output = output_2d.reshape(*x.shape[:-1], layer.weight.size(0)) return output if bias is None else output + bias @@ -505,6 +806,41 @@ def _patch_strict_rocm_rotary_embedding(rotary_cls: type[Any]) -> None: operator = RocmDeterministicRoPEOp() original = rotary_cls.forward_cuda + def prepare_tables(instance: Any, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + head_size = int(getattr(instance, "head_size", 0)) + max_positions = int(getattr(instance, "max_position_embeddings", 0)) + theta = float(getattr(instance, "base", 1_000_000.0)) + key = (device.type, device.index, head_size, max_positions, theta) + if getattr(instance, "_rl_kernel_rope_table_key", None) == key: + cos = getattr(instance, "_rl_kernel_rope_cos_fp32", None) + sin = getattr(instance, "_rl_kernel_rope_sin_fp32", None) + if isinstance(cos, torch.Tensor) and isinstance(sin, torch.Tensor): + return cos, sin + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "strict ROCm RoPE FP32 table must be initialized before HIP Graph capture" + ) + cos, sin = operator.build_position_table( + max_positions, + head_size, + device=device, + theta=theta, + ) + def register_table(name: str, table: torch.Tensor) -> None: + if isinstance(instance, torch.nn.Module): + buffers = instance._buffers + if name in buffers: + buffers[name] = table + else: + instance.register_buffer(name, table, persistent=False) + else: + setattr(instance, name, table) + + register_table("_rl_kernel_rope_cos_fp32", cos) + register_table("_rl_kernel_rope_sin_fp32", sin) + instance._rl_kernel_rope_table_key = key + return cos, sin + def strict_forward_cuda( instance: Any, positions: torch.Tensor, @@ -518,12 +854,19 @@ def strict_forward_cuda( "strict ROCm Qwen3 RoPE requires full-dimension rotation: " f"rotary_dim={rotary_dim}, head_size={head_size}" ) - if positions.ndim not in (1, 2) or positions.numel() != query.shape[0]: + if positions.ndim not in (1, 2): raise RuntimeError( "strict ROCm vLLM RoPE positions must align with flattened query rows: " f"positions={tuple(positions.shape)}, query={tuple(query.shape)}" ) + torch._check( + positions.numel() == query.shape[0], + lambda: "strict ROCm vLLM RoPE positions must align with query rows", + ) flat_positions = positions.reshape(-1).to(device=query.device, dtype=torch.int64) + if not flat_positions.is_contiguous(): + flat_positions = flat_positions.contiguous() + cos, sin = prepare_tables(instance, query.device) def apply(value: torch.Tensor | None) -> torch.Tensor | None: if value is None: @@ -532,24 +875,24 @@ def apply(value: torch.Tensor | None) -> torch.Tensor | None: raise RuntimeError( "strict ROCm vLLM RoPE expects flattened [tokens, heads*head_dim] tensors" ) - tokens = value.shape[0] - heads = value.shape[1] // head_size - # The HIP kernel indexes one position table across rows. A - # head-major view avoids duplicating the table for every head and - # keeps the dispatch to one deterministic launch per Q/K tensor. - head_major = value.view(tokens, heads, head_size).permute(1, 0, 2).contiguous() - rotated = operator(head_major, flat_positions) - return rotated.permute(1, 0, 2).reshape_as(value).contiguous() + return operator.forward_token_major( + value, + flat_positions, + cos, + sin, + head_dim=head_size, + ) return apply(query), apply(key) strict_forward_cuda.__name__ = getattr(original, "__name__", "forward_cuda") + setattr(rotary_cls, "_rl_kernel_prepare_strict_rocm_tables", prepare_tables) setattr(rotary_cls, _STRICT_ROCM_ROPE_PATCH_MARKER, original) rotary_cls.forward_cuda = strict_forward_cuda def _configure_strict_ffn_compilation(vllm_config: Any | None = None) -> None: - """Keep the graph-safe TP reduction inside vLLM CUDA graphs.""" + """Keep stateful ROCm TP reductions outside replayed HIP graph segments.""" if vllm_config is None: from vllm.config import get_current_vllm_config_or_none @@ -562,10 +905,47 @@ def _configure_strict_ffn_compilation(vllm_config: Any | None = None) -> None: splitting_ops = compilation.splitting_ops if splitting_ops is None: raise RuntimeError("vLLM splitting operators were not finalized before model init") - # Older strict runtimes split this op because their host-owned sequence - # value was frozen at graph capture. Sequence allocation is now device - # owned, so remove stale entries and preserve the user's graph mode. - splitting_ops[:] = [op for op in splitting_ops if op != DETERMINISTIC_ALL_REDUCE_OP] + if torch.version.hip is None: + # Preserve the CUDA full-graph path introduced by PR 377. + splitting_ops[:] = [ + op for op in splitting_ops if op != DETERMINISTIC_ALL_REDUCE_OP + ] + return + + from vllm.config import CUDAGraphMode + + compilation.cudagraph_mode = CUDAGraphMode.PIECEWISE + # The IPC transport owns peer-visible sequence and staging state that is + # advanced once per invocation. Replaying the transport inside a captured + # graph can reuse capture-time state and stale peer payloads. Keep just + # these opaque reductions eager while vLLM captures the pure compute around + # them as piecewise HIP graphs. + for op in _ROCM_STATEFUL_GRAPH_SPLITTING_OPS: + if op not in splitting_ops: + splitting_ops.append(op) + + +def _patch_rocm_weight_cache_refresh() -> None: + """Refresh stable transpose buffers after vLLM IPC weight transfer.""" + + if torch.version.hip is None: + return + from vllm.v1.worker.gpu_worker import Worker + from rl_engine.kernels.ops.rocm.matmul.det_gemm import ( + refresh_cached_weight_transposes, + ) + + if hasattr(Worker, _STRICT_WEIGHT_CACHE_REFRESH_MARKER): + return + original_finish = Worker.finish_weight_update + + def finish_weight_update_wrapped(instance: Any) -> None: + original_finish(instance) + model = instance.model_runner.get_model() + refresh_cached_weight_transposes(model.parameters()) + + setattr(Worker, _STRICT_WEIGHT_CACHE_REFRESH_MARKER, original_finish) + Worker.finish_weight_update = finish_weight_update_wrapped def _patch_qwen_ffn(integration: VllmIntegration) -> None: @@ -591,7 +971,7 @@ def wrapped_init(instance: Any, *args: Any, **kwargs: Any) -> None: _handle, tp_world_size = operator.bind_packed_inference(instance) if not compiled_evidence_armed: execution_mode = ( - "eager" if getattr(torch.version, "hip", None) is not None + "compiled_hip_graph" if getattr(torch.version, "hip", None) is not None else "compiled_cuda_graph" ) register_packed_inference_observer( @@ -600,7 +980,7 @@ def wrapped_init(instance: Any, *args: Any, **kwargs: Any) -> None: ) ) compiled_evidence_armed = True - if tp_world_size > 1 and getattr(torch.version, "hip", None) is None: + if tp_world_size > 1: _configure_strict_ffn_compilation() setattr(Qwen2MLP, _STRICT_FFN_INIT_MARKER, original_init) @@ -665,6 +1045,16 @@ def _patch_qwen3_strict_model( _patch_strict_rocm_rotary_embedding(rotary_cls) if det_gemm is None: det_gemm = _strict_attention_projection_op() + rocm_linear_all_reduce = None + register_rocm_linear_staging = None + if torch.version.hip is not None: + from rl_engine.kernels.ops.rocm.matmul.det_gemm import ( + det_gemm_linear_all_reduce_inference, + register_det_gemm_all_reduce_staging, + ) + + rocm_linear_all_reduce = det_gemm_linear_all_reduce_inference + register_rocm_linear_staging = register_det_gemm_all_reduce_staging attention_init = attention_cls.__init__ unquantized_apply = linear_method_cls.apply @@ -681,13 +1071,19 @@ def deterministic_linear_apply( x_2d = x.reshape(-1, x.shape[-1]) linear = getattr(det_gemm, "linear", None) collective = getattr(layer, _STRICT_O_PROJ_COLLECTIVE_MARKER, None) + if collective is not None and torch.version.hip is not None: + if bias is not None or rocm_linear_all_reduce is None: + raise RuntimeError("strict ROCm o_proj fusion requires a bias-free linear") + output_2d = rocm_linear_all_reduce( + x_2d, + layer.weight, + collective_handle=int( + getattr(layer, _STRICT_O_PROJ_COMPILED_COLLECTIVE_SLOT) + ), + ) + return output_2d.reshape(*x.shape[:-1], layer.weight.shape[0]) direct_output = None - if ( - torch.version.hip is None - and collective is not None - and bias is None - and linear is not None - ): + if collective is not None and bias is None and linear is not None: direct_output = collective.direct_staging_view( (x_2d.size(0), layer.weight.shape[0]), dtype=x.dtype, @@ -720,28 +1116,23 @@ def strict_attention_rms_norm_forward( raise RuntimeError("strict Attention Q/K RMSNorm requires a weight") return strict_rms_norm(x, weight, eps=instance.variance_epsilon) - def require_rocm_eager_runtime() -> None: + def require_rocm_graph_runtime() -> None: if not production_classes or torch.version.hip is None: return - from vllm.config import ( - CUDAGraphMode, - CompilationMode, - get_current_vllm_config_or_none, - ) + from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config_or_none config = get_current_vllm_config_or_none() model_config = None if config is None else config.model_config compilation_config = None if config is None else config.compilation_config if ( model_config is None - or model_config.enforce_eager is not True + or model_config.enforce_eager is True or compilation_config is None - or compilation_config.mode != CompilationMode.NONE - or compilation_config.cudagraph_mode != CUDAGraphMode.NONE + or compilation_config.mode == CompilationMode.NONE + or compilation_config.cudagraph_mode == CUDAGraphMode.NONE ): raise RuntimeError( - "strict ROCm vLLM Attention requires enforce_eager with compilation " - "and CUDA/HIP graph capture disabled" + "strict ROCm rollout requires compilation and HIP graph capture enabled" ) def bind_attention_rms_norm(attention: Any, name: str) -> None: @@ -794,8 +1185,6 @@ def bind_o_proj_collective(module: Any) -> None: if not isinstance(backend_id, str) or not backend_id.strip(): raise RuntimeError("strict rollout o_proj collective has no backend identity") _RLK_O_PROJ_COLLECTIVE_BACKEND = backend_id.strip() - if torch.version.hip is not None: - return max_capture = int(os.getenv("RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE", "0")) if max_capture <= 0: raise RuntimeError( @@ -805,6 +1194,20 @@ def bind_o_proj_collective(module: Any) -> None: ((batch, int(module.weight.shape[0])) for batch in range(1, max_capture + 1)), dtype=module.weight.dtype, ) + if torch.version.hip is not None and production_classes: + staging = collective.direct_staging_view( + (max_capture, int(module.weight.shape[0])), + dtype=module.weight.dtype, + ) + if staging is None: + raise RuntimeError("strict ROCm o_proj staging allocation failed") + if register_rocm_linear_staging is None: + raise RuntimeError("strict ROCm o_proj staging registry is unavailable") + compiled_slot = register_rocm_linear_staging( + int(collective._handle), staging + ) + setattr(module, _STRICT_O_PROJ_COMPILED_COLLECTIVE_SLOT, compiled_slot) + setattr(module, _STRICT_O_PROJ_FUSED_ALL_REDUCE_MARKER, True) if row_parallel_cls is not None and not hasattr( row_parallel_cls, _STRICT_ROW_PARALLEL_PATCH_MARKER @@ -830,8 +1233,10 @@ def strict_row_parallel_forward(instance: Any, input_: torch.Tensor) -> Any: output_parallel = instance.quant_method.apply(instance, input_parallel, bias_) if instance.reduce_results and instance.tp_size > 1: - if torch.version.hip is not None: - output = collective.all_reduce(output_parallel) + if bool( + getattr(instance, _STRICT_O_PROJ_FUSED_ALL_REDUCE_MARKER, False) + ): + output = output_parallel elif bool(getattr(instance, _STRICT_DIRECT_STAGING_MARKER, False)): output = deterministic_all_reduce_staged( output_parallel, @@ -871,6 +1276,14 @@ def rms_norm_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: def rotary_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: rotary_init(instance, *args, **kwargs) instance._forward_method = instance.forward_cuda + cache = getattr(instance, "cos_sin_cache", None) + prepare = getattr(instance, "_rl_kernel_prepare_strict_rocm_tables", None) + if ( + isinstance(cache, torch.Tensor) + and cache.is_cuda + and callable(prepare) + ): + prepare(cache.device) setattr(rotary_cls, _STRICT_ROTARY_INIT_MARKER, rotary_init) rotary_cls.__init__ = rotary_init_wrapped @@ -879,7 +1292,7 @@ def rotary_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: return def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: - require_rocm_eager_runtime() + require_rocm_graph_runtime() attention_init(instance, *args, **kwargs) setattr(instance.qkv_proj, _STRICT_PROJECTION_MARKER, "qkv") setattr(instance.o_proj, _STRICT_PROJECTION_MARKER, "o_proj") @@ -994,9 +1407,7 @@ def _register_attention_backend(integration: VllmIntegration) -> None: operator: VllmAttentionOperator | None = None if integration.plan.implementation_for("attention", "rollout") is Implementation.RL_KERNEL: - operator = VllmAttentionOperator( - projection_collective_backend=_o_proj_collective_backend - ) + operator = VllmAttentionOperator(projection_collective_backend=_o_proj_collective_backend) integration.install_operator("attention", operator) class RlKernelAttentionImpl(PlatformAttentionImpl): @@ -1004,6 +1415,30 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) if operator is not None: operator.bind_inference() + if torch.version.hip is not None and operator is not None: + from vllm.config import get_current_vllm_config_or_none + + config = get_current_vllm_config_or_none() + model_config = None if config is None else config.model_config + dtype = None if model_config is None else model_config.dtype + if dtype in (torch.float16, torch.bfloat16): + operator.warmup_rocm_decode(self, dtype=dtype) + + def _split_kv_cache( + self, kv_cache: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if torch.version.hip is not None and operator is not None: + if ( + kv_cache.ndim != 4 + or kv_cache.size(2) != int(self.num_kv_heads) + or kv_cache.size(-1) != 2 * int(self.head_size) + ): + raise RuntimeError( + "RL-Kernel ROCm KV cache must use " + "[blocks, block, heads, 2 * head_size]" + ) + return kv_cache.split(int(self.head_size), dim=-1) + return super()._split_kv_cache(kv_cache) def forward(self, *args: Any, **kwargs: Any) -> Any: active = get_active_integration("vllm") @@ -1016,9 +1451,57 @@ def native(_impl: Any, *call_args: Any, **call_kwargs: Any) -> Any: return integration.execute("attention", native, self, *args, **kwargs) class RlKernelAttentionMetadataBuilder(PlatformAttentionMetadataBuilder): - pass + def build( + self, + common_prefix_len: int, + common_attn_metadata: Any, + fast_build: bool = False, + ) -> Any: + metadata = super().build( + common_prefix_len, + common_attn_metadata, + fast_build=fast_build, + ) + # The ROCm RL-Kernel adapter can reuse vLLM's exact CPU snapshot + # for pure decode scheduling. Keep this private metadata on the + # adapter-owned object; native vLLM code remains untouched. + if torch.version.hip is not None: + # attn_utils reconstructs CommonAttentionMetadata without the + # deprecated private CPU cache. Its upper-bound snapshot is + # exact for prefill, which is the only path consuming it here. + host_lengths = getattr(common_attn_metadata, "seq_lens_cpu_upper_bound", None) + if isinstance(host_lengths, torch.Tensor) and host_lengths.device.type == "cpu": + setattr(metadata, "_rlk_seq_lens_cpu", host_lengths) + starts = getattr(common_attn_metadata, "query_start_loc_cpu", None) + if isinstance(starts, torch.Tensor) and starts.device.type == "cpu": + setattr(metadata, "_rlk_query_start_loc_cpu", starts) + return metadata class RlKernelAttentionBackend(PlatformAttentionBackend): + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + if torch.version.hip is not None and operator is not None: + if block_size % 16 != 0: + raise ValueError("Block size must be a multiple of 16.") + # Token-major pages let AITER mha_batch_prefill consume K/V + # directly. The inherited cache-update methods call the + # adapter-owned _split_kv_cache above, so no vLLM source or + # cache-update kernel needs to change. + return (num_blocks, block_size, num_kv_heads, 2 * head_size) + return PlatformAttentionBackend.get_kv_cache_shape( + num_blocks, + block_size, + num_kv_heads, + head_size, + cache_dtype_str, + ) + @staticmethod def get_impl_cls() -> type[Any]: return RlKernelAttentionImpl @@ -1069,6 +1552,7 @@ def install_vllm_integration(plan: IntegrationPlan) -> VllmIntegration: strict_attention = plan.implementation_for("attention", "rollout") is Implementation.RL_KERNEL if strict_attention: _patch_qwen3_strict_model() + _patch_rocm_weight_cache_refresh() _patch_qwen3_layer_alignment_diagnostics() if strict_linear_logp: _patch_qwen_lm_head_padding() @@ -1092,8 +1576,7 @@ def install_vllm_integration(plan: IntegrationPlan) -> VllmIntegration: ) integration.record_installed_hook( "logp", - "vllm.model_executor.models.qwen3.Qwen3ForCausalLM.compute_logits," - f"{sampler_hook}", + "vllm.model_executor.models.qwen3.Qwen3ForCausalLM.compute_logits," f"{sampler_hook}", ) return integration diff --git a/rl_engine/kernels/ops/pytorch/ffn/ffn.py b/rl_engine/kernels/ops/pytorch/ffn/ffn.py index d7c495f5..ed1c761b 100644 --- a/rl_engine/kernels/ops/pytorch/ffn/ffn.py +++ b/rl_engine/kernels/ops/pytorch/ffn/ffn.py @@ -47,6 +47,8 @@ # Backward-compatible test hook; ownership lives in the shared communication layer. _COLLECTIVES = _SHARED_COLLECTIVES _PACKED_INFERENCE_OBSERVERS: list[Callable[[], None]] = [] +_PACKED_INFERENCE_STAGING_BY_HANDLE: dict[int, tuple[int, Tensor, Tensor]] = {} +_PACKED_INFERENCE_SLOT_BY_RUNTIME_HANDLE: dict[int, int] = {} def register_packed_inference_observer(callback: Callable[[], None]) -> None: @@ -137,6 +139,71 @@ def _qwen3_ffn_packed_inference_to_staging_fake( del rmsnorm_output, fused_gate_up_weight, down_weight, output +@torch.library.custom_op( + "rl_kernel::qwen3_ffn_packed_tp_inference_rocm", + mutates_args=(), +) +def _qwen3_ffn_packed_tp_inference_rocm( + rmsnorm_output: Tensor, + fused_gate_up_weight: Tensor, + down_weight: Tensor, + collective_handle: int, +) -> Tensor: + """Keep the ROCm TP FFN behind one eager graph-partition boundary.""" + + binding = _PACKED_INFERENCE_STAGING_BY_HANDLE.get(collective_handle) + if binding is None: + raise RuntimeError("packed ROCm rollout FFN staging handle is not registered") + runtime_handle, staging, stable_output = binding + input_shape = rmsnorm_output.shape + rows = rmsnorm_output.numel() // input_shape[-1] + if rows <= staging.size(0): + # Keep the output address stable across piecewise HIP-graph capture and + # replay so the next captured partition reads the current invocation. + direct_input = staging.narrow(0, 0, rows) + output = stable_output.narrow(0, 0, rows) + _C.deterministic_collective_rocm_ipc_prepare_staged( + runtime_handle, + direct_input, + ) + _qwen3_ffn_packed_inference_to_staging( + rmsnorm_output, + fused_gate_up_weight, + down_weight, + direct_input, + ) + _C.deterministic_collective_rocm_ipc_all_reduce_staged( + runtime_handle, direct_input, output + ) + return output.reshape(*input_shape[:-1], down_weight.shape[0]) + else: + # Profiling and uncaptured prefill can exceed the decode capture bound. + output = _qwen3_ffn_packed_inference( + rmsnorm_output, + fused_gate_up_weight, + down_weight, + ) + _C.deterministic_collective_rocm_ipc_all_reduce_input( + runtime_handle, + output, + output, + ) + return output.reshape(*input_shape[:-1], down_weight.shape[0]) + + +@_qwen3_ffn_packed_tp_inference_rocm.register_fake +def _qwen3_ffn_packed_tp_inference_rocm_fake( + rmsnorm_output: Tensor, + fused_gate_up_weight: Tensor, + down_weight: Tensor, + collective_handle: int, +) -> Tensor: + del fused_gate_up_weight, collective_handle + return rmsnorm_output.new_empty( + (*rmsnorm_output.shape[:-1], down_weight.shape[0]) + ) + + def qwen3_ffn_packed_inference( rmsnorm_output: Tensor, fused_gate_up_weight: Tensor, @@ -154,26 +221,41 @@ def qwen3_ffn_packed_inference( fused_gate_up_weight, down_weight, ) - if getattr(torch.version, "hip", None) is not None: - if collective is None: - raise RuntimeError("packed ROCm rollout FFN requires a bound TP collective") + if collective_handle <= 0: + raise RuntimeError("packed rollout FFN requires a bound TP collective") + if ( + getattr(torch.version, "hip", None) is not None + and collective is not None + and collective_handle not in _PACKED_INFERENCE_STAGING_BY_HANDLE + ): + # Eager ROCm does not reserve graph staging. Keep its established + # in-place fixed-tree path; graph-enabled runs register the handle in + # prepare_packed_inference and stay behind the opaque custom op below. output = _qwen3_ffn_packed_inference( rmsnorm_output, fused_gate_up_weight, down_weight, ) return collective.all_reduce(output, out=output) - if collective_handle <= 0: - raise RuntimeError("packed rollout FFN requires a bound TP collective") + if getattr(torch.version, "hip", None) is not None: + return _qwen3_ffn_packed_tp_inference_rocm( + rmsnorm_output, + fused_gate_up_weight, + down_weight, + collective_handle, + ) input_shape = rmsnorm_output.shape output_shape_2d = ( rmsnorm_output.numel() // input_shape[-1], down_weight.shape[0], ) + direct_staging = None if collective is None else getattr( + collective, "direct_staging_view", None + ) direct_output = ( None - if collective is None - else collective.direct_staging_view(output_shape_2d, dtype=rmsnorm_output.dtype) + if not callable(direct_staging) + else direct_staging(output_shape_2d, dtype=rmsnorm_output.dtype) ) if direct_output is not None: deterministic_staging_reserve( @@ -733,18 +815,45 @@ def prepare_packed_inference( tp_group, min_size_bytes=min_size_bytes, ) - if getattr(torch.version, "hip", None) is not None: - collective_handle = int(collective._handle) - self._packed_inference_collectives[collective_handle] = collective - return collective_handle, tp_world_size max_capture = int(os.getenv("RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE", "0")) if max_capture <= 0: + if getattr(torch.version, "hip", None) is not None: + collective_handle = int(collective._handle) + self._packed_inference_collectives[collective_handle] = collective + return collective_handle, tp_world_size raise RuntimeError("packed rollout FFN requires a positive graph capture size") collective.prepare_direct_staging_views( ((batch, int(down_weight.shape[0])) for batch in range(1, max_capture + 1)), dtype=down_weight.dtype, ) - collective_handle = int(collective._handle) + runtime_handle = int(collective._handle) + collective_handle = runtime_handle + if getattr(torch.version, "hip", None) is not None: + staging = collective.direct_staging_view( + (max_capture, int(down_weight.shape[0])), + dtype=down_weight.dtype, + ) + if staging is None: + raise RuntimeError("packed ROCm rollout FFN staging allocation failed") + collective_handle = _PACKED_INFERENCE_SLOT_BY_RUNTIME_HANDLE.get( + runtime_handle, 0 + ) + if collective_handle == 0: + # Keep the AOT graph identity stable across worker processes; + # resolve its process-local C++ handle inside the custom op. + collective_handle = len(_PACKED_INFERENCE_SLOT_BY_RUNTIME_HANDLE) + 1 + _PACKED_INFERENCE_SLOT_BY_RUNTIME_HANDLE[runtime_handle] = ( + collective_handle + ) + binding = _PACKED_INFERENCE_STAGING_BY_HANDLE.get(collective_handle) + stable_output = ( + torch.empty_like(staging) if binding is None else binding[2] + ) + _PACKED_INFERENCE_STAGING_BY_HANDLE[collective_handle] = ( + runtime_handle, + staging, + stable_output, + ) self._packed_inference_collectives[collective_handle] = collective return collective_handle, tp_world_size diff --git a/rl_engine/kernels/ops/rocm/attention/flash_attn.py b/rl_engine/kernels/ops/rocm/attention/flash_attn.py index 65cd94aa..7a8e7296 100644 --- a/rl_engine/kernels/ops/rocm/attention/flash_attn.py +++ b/rl_engine/kernels/ops/rocm/attention/flash_attn.py @@ -18,6 +18,7 @@ from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, STRICT_ATTENTION_ROCM_SCHEDULE_ID, + SplitKVExecutionPlan, SplitKVMode, SplitKVSpec, ) @@ -67,6 +68,27 @@ ) # Passed by keyword, so only presence matters. _AITER_BWD_REQUIRED_KEYWORDS = frozenset({"rng_state"}) +_AITER_BATCH_PREFILL_POSITIONAL_CONTRACT = ( + "q", + "k", + "v", + "cu_seqlens_q", + "kv_indptr", + "kv_page_indices", + "max_seqlen_q", + "max_seqlen_k", + "dropout_p", + "softmax_scale", + "logits_soft_cap", + "zero_tensors", + "is_causal", + "window_size_left", + "window_size_right", + "sink_size", + "return_softmax_lse", + "return_dropout_randval", +) +_AITER_BATCH_PREFILL_REQUIRED_KEYWORDS = frozenset({"block_table", "seqlen_k"}) # Stable dispatch identity for the strict ROCm attention core. Kept at module # scope so contract-aware dispatch and the Vime adapter name one constant @@ -123,16 +145,19 @@ def _validate_aiter_schema( ) -def _load_aiter_ck_ops() -> tuple[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") except (AttributeError, ImportError, OSError, RuntimeError) as exc: raise StrictRocmAttentionUnavailable( - "strict ROCm Attention requires aiter.ops.mha.mha_fwd and mha_bwd" + "strict ROCm Attention requires AITER dense backward and batch-prefill CK ops" ) from exc - if not callable(mha_fwd) or not callable(mha_bwd): + if not all(callable(op) for op in (mha_fwd, mha_bwd, mha_batch_prefill)): raise StrictRocmAttentionUnavailable("AITER CK MHA entry points are not callable") _validate_aiter_schema("mha_fwd", _AITER_FWD_POSITIONAL_CONTRACT) _validate_aiter_schema( @@ -140,11 +165,16 @@ def _load_aiter_ck_ops() -> tuple[Callable[..., Any], Callable[..., Any], str]: _AITER_BWD_POSITIONAL_CONTRACT, required_keywords=_AITER_BWD_REQUIRED_KEYWORDS, ) + _validate_aiter_schema( + "mha_batch_prefill", + _AITER_BATCH_PREFILL_POSITIONAL_CONTRACT, + required_keywords=_AITER_BATCH_PREFILL_REQUIRED_KEYWORDS, + ) module_file = inspect.getsourcefile(module) if not module_file: raise StrictRocmAttentionUnavailable("cannot fingerprint the AITER MHA source module") source_sha256 = hashlib.sha256(Path(module_file).read_bytes()).hexdigest() - return mha_fwd, mha_bwd, source_sha256 + return mha_fwd, mha_bwd, mha_batch_prefill, source_sha256 class _AiterCKAttentionFn(Function): @@ -223,6 +253,138 @@ def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): ) +class _AiterCKPagedAttentionFn(Function): + """Use one non-Split-K paged CK schedule for train forward and decode.""" + + page_size = 16 + + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + scale: float, + mha_batch_prefill: Callable[..., Any], + mha_bwd: Callable[..., Any], + ) -> tuple[torch.Tensor, torch.Tensor]: + batch, q_heads, q_len, head_dim = q.shape + kv_heads, kv_len = k.size(1), k.size(2) + q_fa = q.transpose(1, 2).contiguous() + k_fa = k.transpose(1, 2).contiguous() + v_fa = v.transpose(1, 2).contiguous() + + page_count = (kv_len + _AiterCKPagedAttentionFn.page_size - 1) // ( + _AiterCKPagedAttentionFn.page_size + ) + padded_kv_len = page_count * _AiterCKPagedAttentionFn.page_size + if padded_kv_len == kv_len: + k_cache = k_fa.reshape( + batch * page_count, + _AiterCKPagedAttentionFn.page_size, + kv_heads, + head_dim, + ) + v_cache = v_fa.reshape_as(k_cache) + else: + padded_shape = (batch, padded_kv_len, kv_heads, head_dim) + k_padded = torch.zeros(padded_shape, dtype=k.dtype, device=k.device) + v_padded = torch.zeros_like(k_padded) + k_padded[:, :kv_len].copy_(k_fa) + v_padded[:, :kv_len].copy_(v_fa) + k_cache = k_padded.reshape( + batch * page_count, + _AiterCKPagedAttentionFn.page_size, + kv_heads, + head_dim, + ) + 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) + 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) + result = mha_batch_prefill( + q_fa.reshape(batch * q_len, q_heads, head_dim), + k_cache, + v_cache, + cu_seqlens_q, + kv_indptr, + block_table.reshape(-1), + q_len, + kv_len, + 0.0, + float(scale), + 0.0, + False, + bool(causal), + -1, + -1, + 0, + True, + False, + block_table=block_table, + seqlen_k=seqlen_k, + ) + if not isinstance(result, (tuple, list)) or len(result) != 4: + raise StrictRocmAttentionUnavailable( + "AITER mha_batch_prefill must return (out, lse, dropout_mask, rng_state)" + ) + out_flat, lse_flat, _dropout_mask, rng_state = result + if not all(isinstance(item, torch.Tensor) for item in (out_flat, lse_flat, rng_state)): + raise StrictRocmAttentionUnavailable("AITER paged CK returned non-tensor state") + expected_out = (batch * q_len, q_heads, head_dim) + if tuple(out_flat.shape) != expected_out: + raise StrictRocmAttentionUnavailable("AITER paged CK output shape changed") + if tuple(lse_flat.shape) != (q_heads, batch * q_len) or lse_flat.dtype != torch.float32: + raise StrictRocmAttentionUnavailable("AITER paged CK must export [H,total_q] FP32 LSE") + + out_fa = out_flat.reshape(batch, q_len, q_heads, head_dim) + lse = lse_flat.reshape(q_heads, batch, q_len).permute(1, 0, 2).contiguous() + ctx.save_for_backward(q_fa, k_fa, v_fa, out_fa, lse, rng_state) + ctx.causal = bool(causal) + ctx.scale = float(scale) + ctx.mha_bwd = mha_bwd + ctx.mark_non_differentiable(lse) + return out_fa.transpose(1, 2).contiguous(), lse + + @staticmethod + @once_differentiable + def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): + q_fa, k_fa, v_fa, out_fa, lse, rng_state = ctx.saved_tensors + grad_out_fa = grad_out.transpose(1, 2).contiguous() + result = ctx.mha_bwd( + grad_out_fa, + q_fa, + k_fa, + v_fa, + out_fa, + lse, + 0.0, + ctx.scale, + ctx.causal, + -1, + -1, + True, + rng_state=rng_state, + ) + if not isinstance(result, (tuple, list)) or len(result) < 3: + raise StrictRocmAttentionUnavailable("AITER mha_bwd must return dQ/dK/dV") + dq, dk, dv = result[:3] + return ( + dq.transpose(1, 2).contiguous(), + dk.transpose(1, 2).contiguous(), + dv.transpose(1, 2).contiguous(), + None, + None, + None, + None, + ) + + class StrictRocmAiterCKAttentionCore: """Shared ROCm production core using the non-Split-K AITER CK dense MHA.""" @@ -247,6 +409,7 @@ def __init__( split_kv: SplitKVSpec | None = None, _mha_fwd: Callable[..., Any] | None = None, _mha_bwd: Callable[..., Any] | None = None, + _mha_batch_prefill: Callable[..., Any] | None = None, _source_sha256: str | None = None, ) -> None: requested = SplitKVSpec.disabled() if split_kv is None else split_kv @@ -257,11 +420,12 @@ def __init__( if (_mha_fwd is None) != (_mha_bwd is None): raise ValueError("test injection requires both AITER forward and backward callables") if _mha_fwd is None: - mha_fwd, mha_bwd, source_sha256 = _load_aiter_ck_ops() + mha_fwd, mha_bwd, mha_batch_prefill, source_sha256 = _load_aiter_ck_ops() else: assert _mha_bwd is not None mha_fwd = _mha_fwd mha_bwd = _mha_bwd + mha_batch_prefill = _mha_batch_prefill source_sha256 = "test-double" if _source_sha256 is None else _source_sha256 if not callable(mha_fwd) or not callable(mha_bwd): raise StrictRocmAttentionUnavailable("AITER CK MHA entry points are not callable") @@ -269,6 +433,12 @@ def __init__( self.source_sha256 = source_sha256 self._mha_fwd = mha_fwd self._mha_bwd = mha_bwd + self._mha_batch_prefill = mha_batch_prefill + self.supports_paged_schedule = callable(mha_batch_prefill) + self._device_description_cache: tuple[torch.device, tuple[str, str]] | None = None + self._split_kv_plan_cache: ( + tuple[tuple[SplitKVSpec, int, str], SplitKVExecutionPlan] | None + ) = None def forward_with_lse( self, @@ -295,21 +465,37 @@ def forward_with_lse( if resolved_dtype != q.dtype: raise ValueError("strict Attention output_dtype must match the Q/K/V input dtype") resolved_scale = 1.0 / math.sqrt(q.size(-1)) if scale is None else float(scale) - out, lse = _AiterCKAttentionFn.apply( - q, - k, - v, - bool(causal), - resolved_scale, - self._mha_fwd, - self._mha_bwd, - ) + if self._mha_batch_prefill is None: + out, lse = _AiterCKAttentionFn.apply( + q, + k, + v, + bool(causal), + resolved_scale, + self._mha_fwd, + self._mha_bwd, + ) + forward_entrypoint = "mha_fwd" + kv_layout = "dense_bshd" + else: + out, lse = _AiterCKPagedAttentionFn.apply( + q, + k, + v, + bool(causal), + resolved_scale, + self._mha_batch_prefill, + self._mha_bwd, + ) + forward_entrypoint = "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: raise StrictRocmAttentionUnavailable("AITER CK output shape/dtype changed") if tuple(lse.shape) != expected_lse_shape or lse.dtype != torch.float32: raise StrictRocmAttentionUnavailable("AITER CK must export [B,H,Sq] FP32 LSE") - device_properties = torch.cuda.get_device_properties(q.device) + gpu_name, gpu_arch = self._device_description(q.device) + split_kv_plan = self._resolve_split_kv_plan(k.size(2)) return DeterministicAttentionCoreResult( out=out, lse=lse, @@ -320,18 +506,344 @@ def forward_with_lse( "platform": "rocm", "torch_version": torch.__version__, "rocm_version": torch.version.hip, - "gpu_name": device_properties.name, - "gpu_arch": getattr(device_properties, "gcnArchName", "unknown"), + "gpu_name": gpu_name, + "gpu_arch": gpu_arch, "aiter_api_source": self.api_source, "aiter_source_sha256": self.source_sha256, "num_splits": self.num_splits, "split_kv_control": self.split_kv_control, "deterministic_backward": self.deterministic_backward, + "forward_entrypoint": forward_entrypoint, + "kv_layout": kv_layout, + "dense_kv_materialized": False, + "dropout_p": 0.0, + "split_kv": split_kv_plan.to_dict(), + "merge_order": self.merge_order, + "accum_dtype": self.accum_dtype, + "downcast_at": self.downcast_at, + "fallback": self.fallback, + "fallback_reason": None, + "native_attention_arithmetic": self.native_attention_arithmetic, + "production_ready": self.production_ready, + "reference_only": self.reference_only, + }, + ) + + def forward_paged_varlen_with_lse( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + page_table: torch.Tensor, + seqused_k: torch.Tensor, + cu_seqlens_q: torch.Tensor, + kv_indptr: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool, + scale: float | None, + out: torch.Tensor | None = None, + return_lse: bool = True, + ) -> DeterministicAttentionCoreResult: + """Run one graph-safe CK launch over packed queries and paged KV.""" + + if self._mha_batch_prefill is None: + raise StrictRocmAttentionUnavailable("AITER paged CK entry point is unavailable") + if q.ndim != 3: + raise ValueError("packed paged Q must use [tokens, heads, head_dim]") + if k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("paged K/V must use [pages, page_size, heads, head_dim]") + if q.size(1) % k_cache.size(2) or q.size(2) != k_cache.size(3): + raise ValueError("paged Q/K head counts or dimensions are incompatible") + if q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict paged Attention supports FP16/BF16 only") + if k_cache.dtype != q.dtype or v_cache.dtype != q.dtype: + raise ValueError("paged Q/K/V must share one dtype") + if not (q.device == k_cache.device == v_cache.device): + raise ValueError("paged Q/K/V must share one ROCm device") + batch = page_table.size(0) + if page_table.ndim != 2 or seqused_k.shape != (batch,): + raise ValueError("paged metadata must carry one row per sequence") + 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) + ): + raise ValueError("paged metadata must be on the Q device") + if not q.is_contiguous() or not page_table.is_contiguous(): + raise ValueError("packed Q and page_table must be contiguous") + if max_seqlen_q <= 0 or max_seqlen_k <= 0: + raise ValueError("paged maximum sequence lengths must be positive") + if max_seqlen_k > page_table.size(1) * k_cache.size(1): + raise ValueError("max_seqlen_k exceeds the page table capacity") + if torch.is_grad_enabled() and any( + tensor.requires_grad for tensor in (q, k_cache, v_cache) + ): + raise RuntimeError("direct paged Attention is inference-only") + if out is not None: + if out.shape != q.shape or out.dtype != q.dtype or out.device != q.device: + raise ValueError("direct paged output must match packed Q") + if not out.is_contiguous(): + raise ValueError("direct paged output must be contiguous") + + resolved_scale = 1.0 / math.sqrt(q.size(-1)) if scale is None else float(scale) + result = self._mha_batch_prefill( + q, + k_cache, + v_cache, + cu_seqlens_q, + kv_indptr, + page_table.reshape(-1), + max_seqlen_q, + max_seqlen_k, + 0.0, + resolved_scale, + 0.0, + False, + bool(causal), + -1, + -1, + 0, + bool(return_lse), + False, + out=out, + block_table=page_table, + seqlen_k=seqused_k, + ) + if not isinstance(result, (tuple, list)) or len(result) != 4: + raise StrictRocmAttentionUnavailable( + "AITER mha_batch_prefill must return (out, lse, dropout_mask, rng_state)" + ) + result_out, lse, _dropout_mask, _rng_state = result + if tuple(result_out.shape) != tuple(q.shape) or result_out.dtype != q.dtype: + raise StrictRocmAttentionUnavailable("AITER paged CK output shape/dtype changed") + if out is not None and result_out.data_ptr() != out.data_ptr(): + raise StrictRocmAttentionUnavailable("AITER paged CK ignored the output buffer") + if return_lse: + if tuple(lse.shape) != (q.size(1), q.size(0)) or lse.dtype != torch.float32: + raise StrictRocmAttentionUnavailable( + "AITER paged CK must export [heads,total_q] FP32 LSE" + ) + else: + lse = torch.empty((0,), dtype=torch.float32, device=q.device) + gpu_name, gpu_arch = self._device_description(q.device) + split_kv_plan = self._resolve_split_kv_plan(max_seqlen_k) + return DeterministicAttentionCoreResult( + out=result_out, + lse=lse, + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "attention_backend": self.backend_id, + "platform": "rocm", + "torch_version": torch.__version__, + "rocm_version": torch.version.hip, + "gpu_name": gpu_name, + "gpu_arch": gpu_arch, + "aiter_api_source": self.api_source, + "aiter_source_sha256": self.source_sha256, + "forward_entrypoint": "mha_batch_prefill", + "kv_layout": "vllm_linear_paged", + "dense_kv_materialized": False, + "num_splits": self.num_splits, + "split_kv_control": "batch_prefill_non_split_ck", + "deterministic_backward": False, + "dropout_p": 0.0, + "split_kv": split_kv_plan.to_dict(), + "merge_order": self.merge_order, + "accum_dtype": self.accum_dtype, + "downcast_at": self.downcast_at, + "direct_output": out is not None, + "packed_query_count": batch, + "packed_query_tokens": q.size(0), + "fallback": self.fallback, + "fallback_reason": None, + "native_attention_arithmetic": self.native_attention_arithmetic, + "production_ready": self.production_ready, + "reference_only": self.reference_only, + }, + ) + + def forward_paged_with_lse( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + page_table: torch.Tensor, + seqused_k: torch.Tensor, + cu_seqlens_q: torch.Tensor, + kv_indptr: torch.Tensor, + max_seqlen_k: int, + causal: bool, + scale: float | None, + out: torch.Tensor | None = None, + return_lse: bool = True, + ) -> DeterministicAttentionCoreResult: + self._validate_paged_inputs(q, k_cache, v_cache, page_table, seqused_k) + if q.size(2) != 1: + raise ValueError("direct paged decode requires one query token per row") + if torch.is_grad_enabled() and any( + tensor.requires_grad for tensor in (q, k_cache, v_cache) + ): + raise RuntimeError("direct paged decode is inference-only") + resolved_scale = 1.0 / math.sqrt(q.size(-1)) if scale is None else float(scale) + q_flat = q.transpose(1, 2).reshape(q.size(0), q.size(1), q.size(3)) + if not q_flat.is_contiguous(): + q_flat = q_flat.contiguous() + out_flat = None + if out is not None: + if out.shape != q.shape or out.dtype != q.dtype or out.device != q.device: + raise ValueError("direct paged output must match Q shape, dtype, and device") + out_flat = out.transpose(1, 2).reshape_as(q_flat) + if not out_flat.is_contiguous(): + raise ValueError("direct paged output must expose a contiguous AITER view") + flat_result = self.forward_paged_varlen_with_lse( + q_flat, + k_cache, + v_cache, + page_table=page_table, + seqused_k=seqused_k, + cu_seqlens_q=cu_seqlens_q, + kv_indptr=kv_indptr, + max_seqlen_q=1, + max_seqlen_k=max_seqlen_k, + causal=causal, + scale=resolved_scale, + out=out_flat, + return_lse=return_lse, + ) + result_flat = flat_result.out + if return_lse: + lse = flat_result.lse.transpose(0, 1).unsqueeze(-1).contiguous() + else: + lse = torch.empty((0,), dtype=torch.float32, device=q.device) + return DeterministicAttentionCoreResult( + out=out if out is not None else result_flat.unsqueeze(1).transpose(1, 2).contiguous(), + lse=lse, + provenance=dict(flat_result.provenance), + ) + + @staticmethod + def _validate_paged_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + page_table: torch.Tensor, + seqused_k: torch.Tensor, + ) -> None: + if torch.version.hip is None: + raise StrictRocmAttentionUnavailable("strict AITER CK core requires ROCm PyTorch") + if q.ndim != 4 or k.ndim != 4 or v.shape != k.shape: + 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: + 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") + if page_table.shape[0] != q.size(0) or seqused_k.shape != (q.size(0),): + raise ValueError("paged metadata must carry one row per query") + if page_table.device != q.device or seqused_k.device != q.device: + raise ValueError("paged metadata must be on the Q device") + + def forward_bshd_with_lse( + self, + q_bshd: torch.Tensor, + k_bshd: torch.Tensor, + v_bshd: torch.Tensor, + *, + causal: bool, + scale: float | None, + out: torch.Tensor | None = None, + ) -> DeterministicAttentionCoreResult: + """Run the identical CK forward arithmetic on already packed BSHD tensors. + + Decode has no backward pass. Keeping paged gather output in AITER's + native layout removes the BHSD round trip, while the same dense + non-Split-K entry point and arguments preserve the strict arithmetic. + ``out`` uses the public BHSD shape; Sq=1 makes its BSHD transpose a + contiguous view suitable for AITER's output argument. + """ + + self._validate_bshd_inputs(q_bshd, k_bshd, v_bshd) + if torch.is_grad_enabled() and any( + tensor.requires_grad for tensor in (q_bshd, k_bshd, v_bshd) + ): + raise RuntimeError("strict BSHD decode entry point is inference-only") + resolved_scale = 1.0 / math.sqrt(q_bshd.size(-1)) if scale is None else float(scale) + out_bshd = None + if out is not None: + expected_out = ( + q_bshd.size(0), + q_bshd.size(2), + q_bshd.size(1), + q_bshd.size(3), + ) + if out.shape != expected_out or out.dtype != q_bshd.dtype: + raise ValueError("strict BSHD decode output buffer has the wrong shape or dtype") + out_bshd = out.transpose(1, 2) + if not out_bshd.is_contiguous(): + raise ValueError("strict BSHD decode output must expose a contiguous AITER view") + + result = self._mha_fwd( + q_bshd, + k_bshd, + v_bshd, + 0.0, + resolved_scale, + bool(causal), + -1, + -1, + 0, + True, + False, + out=out_bshd, + ) + if not isinstance(result, (tuple, list)) or len(result) != 4: + raise StrictRocmAttentionUnavailable( + "AITER mha_fwd must return (out, lse, dropout_mask, rng_state)" + ) + result_bshd, lse, _dropout_mask, _rng_state = result + if not isinstance(result_bshd, torch.Tensor) or not isinstance(lse, torch.Tensor): + raise StrictRocmAttentionUnavailable("AITER mha_fwd returned non-tensor output") + if result_bshd.shape != q_bshd.shape or result_bshd.dtype != q_bshd.dtype: + raise StrictRocmAttentionUnavailable("AITER CK output shape/dtype changed") + expected_lse_shape = (q_bshd.size(0), q_bshd.size(2), q_bshd.size(1)) + if tuple(lse.shape) != expected_lse_shape or lse.dtype != torch.float32: + raise StrictRocmAttentionUnavailable("AITER CK must export [B,H,Sq] FP32 LSE") + if out_bshd is not None and result_bshd.data_ptr() != out_bshd.data_ptr(): + raise StrictRocmAttentionUnavailable("AITER CK ignored the strict decode output buffer") + + gpu_name, gpu_arch = self._device_description(q_bshd.device) + split_kv_plan = self._resolve_split_kv_plan(k_bshd.size(1)) + result_out = out if out is not None else result_bshd.transpose(1, 2).contiguous() + return DeterministicAttentionCoreResult( + out=result_out, + lse=lse.contiguous(), + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "attention_backend": self.backend_id, + "platform": "rocm", + "torch_version": torch.__version__, + "rocm_version": torch.version.hip, + "gpu_name": gpu_name, + "gpu_arch": gpu_arch, + "aiter_api_source": self.api_source, + "aiter_source_sha256": self.source_sha256, + "num_splits": self.num_splits, + "split_kv_control": self.split_kv_control, + "deterministic_backward": False, "dropout_p": 0.0, - "split_kv": self.split_kv.resolve(k.size(2), backend=self.backend_id).to_dict(), + "split_kv": split_kv_plan.to_dict(), "merge_order": self.merge_order, "accum_dtype": self.accum_dtype, "downcast_at": self.downcast_at, + "input_layout": "bshd_direct", + "direct_output": out is not None, "fallback": self.fallback, "fallback_reason": None, "native_attention_arithmetic": self.native_attention_arithmetic, @@ -340,6 +852,53 @@ def forward_with_lse( }, ) + def _device_description(self, device: torch.device) -> tuple[str, str]: + cached = self._device_description_cache + if cached is not None and cached[0] == device: + return cached[1] + properties = torch.cuda.get_device_properties(device) + description = (properties.name, getattr(properties, "gcnArchName", "unknown")) + self._device_description_cache = (device, description) + return description + + def _resolve_split_kv_plan(self, total_kv_tokens: int) -> SplitKVExecutionPlan: + key = (self.split_kv, total_kv_tokens, self.backend_id) + cached = self._split_kv_plan_cache + if cached is not None and cached[0] == key: + return cached[1] + plan = self.split_kv.resolve(total_kv_tokens, backend=self.backend_id) + self._split_kv_plan_cache = (key, plan) + return plan + + @staticmethod + def _validate_bshd_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> None: + if torch.version.hip is None: + raise StrictRocmAttentionUnavailable("strict AITER CK core requires ROCm PyTorch") + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q/k/v must be 4-D [B,S,H,D]") + if q.size(0) != k.size(0) or q.size(0) != v.size(0): + raise ValueError("q/k/v batch sizes must match") + if k.shape != v.shape or q.size(-1) != k.size(-1): + raise ValueError("k/v shapes and q/k/v head dimensions must match") + if q.size(2) % k.size(2) != 0: + raise ValueError("Q heads must be divisible by KV heads for GQA") + if q.size(-1) > 256 or q.size(-1) % 8: + raise ValueError("AITER CK requires head_dim <= 256 and divisible by 8") + if q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict AITER CK core supports FP16/BF16 only") + if k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q/k/v must share one dtype") + if not (q.is_cuda and k.is_cuda and v.is_cuda): + raise ValueError("strict AITER CK core requires ROCm GPU tensors") + if not (q.device == k.device == v.device): + raise ValueError("q/k/v must be on one ROCm device") + if not (q.is_contiguous() and k.is_contiguous() and v.is_contiguous()): + raise ValueError("strict BSHD decode inputs must be contiguous") + @staticmethod def _validate_inputs( q: torch.Tensor, diff --git a/rl_engine/kernels/ops/rocm/attention/paged_gather.py b/rl_engine/kernels/ops/rocm/attention/paged_gather.py new file mode 100644 index 00000000..e43eb9a5 --- /dev/null +++ b/rl_engine/kernels/ops/rocm/attention/paged_gather.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused logical paged-KV gather for the strict ROCm Attention runtime.""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: + _TRITON_AVAILABLE = False + + +if _TRITON_AVAILABLE: + + @triton.jit + def _paged_kv_gather_bhsd_kernel( + k_cache, + v_cache, + page_rows, + k_out, + v_out, + total_elements, + tokens, + heads, + head_dim: tl.constexpr, + page_size, + k_stride_page, + k_stride_token, + k_stride_head, + k_stride_dim, + v_stride_page, + v_stride_token, + v_stride_head, + v_stride_dim, + page_stride_row, + page_stride_col, + BLOCK: tl.constexpr, + ): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < total_elements + dim = offsets % head_dim + rem = offsets // head_dim + token = rem % tokens + rem = rem // tokens + head = rem % heads + row = rem // heads + logical_page = token // page_size + page_offset = token % page_size + physical_page = tl.load( + page_rows + row * page_stride_row + logical_page * page_stride_col, + mask=mask, + other=0, + ).to(tl.int64) + k_offsets = ( + physical_page * k_stride_page + + page_offset * k_stride_token + + head * k_stride_head + + dim * k_stride_dim + ) + v_offsets = ( + physical_page * v_stride_page + + page_offset * v_stride_token + + head * v_stride_head + + dim * v_stride_dim + ) + k_value = tl.load(k_cache + k_offsets, mask=mask) + v_value = tl.load(v_cache + v_offsets, mask=mask) + tl.store(k_out + offsets, k_value, mask=mask) + tl.store(v_out + offsets, v_value, mask=mask) + + +def fused_paged_kv_gather_bhsd( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_rows: torch.Tensor, + page_count: int, + *, + k_out: torch.Tensor, + v_out: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather K and V together into contiguous ``[B, H, S, D]`` buffers. + + Each KV group is sequence-contiguous. Its ``[B, 1, S, D]`` slice can + therefore be transposed to AITER's ``[B, S, 1, D]`` as a view, avoiding + the second layout materialization that a token-major gather would need. + """ + + if not _TRITON_AVAILABLE: + raise RuntimeError("strict ROCm fused paged gather requires Triton") + if page_count <= 0 or page_rows.ndim != 2 or page_rows.size(1) < page_count: + raise ValueError("page_rows must provide a positive page count") + if k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("paged K/V cache must use [pages, page_size, heads, dim]") + rows = page_rows.size(0) + tokens = page_count * k_cache.size(1) + expected = (rows, k_cache.size(2), tokens, k_cache.size(3)) + if k_out.shape != expected or v_out.shape != expected: + raise ValueError("paged gather output buffers have the wrong BHSD shape") + if not k_out.is_contiguous() or not v_out.is_contiguous(): + raise ValueError("paged gather output buffers must be contiguous") + if not ( + k_cache.device == v_cache.device == page_rows.device == k_out.device == v_out.device + ): + raise ValueError("paged gather tensors must share one device") + if k_out.dtype != k_cache.dtype or v_out.dtype != v_cache.dtype: + raise ValueError("paged gather output dtype must match the cache") + + total = k_out.numel() + block = 256 + _paged_kv_gather_bhsd_kernel[(triton.cdiv(total, block),)]( + k_cache, + v_cache, + page_rows, + k_out, + v_out, + total_elements=total, + tokens=tokens, + heads=k_cache.size(2), + head_dim=k_cache.size(3), + page_size=k_cache.size(1), + k_stride_page=k_cache.stride(0), + k_stride_token=k_cache.stride(1), + k_stride_head=k_cache.stride(2), + k_stride_dim=k_cache.stride(3), + v_stride_page=v_cache.stride(0), + v_stride_token=v_cache.stride(1), + v_stride_head=v_cache.stride(2), + v_stride_dim=v_cache.stride(3), + page_stride_row=page_rows.stride(0), + page_stride_col=page_rows.stride(1), + BLOCK=block, + num_warps=4, + ) + return k_out, v_out + + +_WARMED_GATHER_VARIANTS: set[tuple[int, torch.dtype, int]] = set() + + +def warmup_fused_paged_kv_gather( + *, + device: torch.device, + dtype: torch.dtype, + head_dim: int, +) -> None: + """Compile the production gather variant before rollout timing starts.""" + + device = torch.device(device) + if device.type != "cuda" or not _TRITON_AVAILABLE: + return + device_index = torch.cuda.current_device() if device.index is None else device.index + key = (device_index, dtype, head_dim) + if key in _WARMED_GATHER_VARIANTS: + return + + page_size = 16 + cache_shape = (1, page_size, 1, head_dim) + output_shape = (1, 1, page_size, head_dim) + k_cache = torch.empty(cache_shape, dtype=dtype, device=device) + v_cache = torch.empty_like(k_cache) + page_rows = torch.zeros((1, 1), dtype=torch.int32, device=device) + k_out = torch.empty(output_shape, dtype=dtype, device=device) + v_out = torch.empty_like(k_out) + fused_paged_kv_gather_bhsd( + k_cache, + v_cache, + page_rows, + 1, + k_out=k_out, + v_out=v_out, + ) + torch.cuda.synchronize(device) + _WARMED_GATHER_VARIANTS.add(key) + + +__all__ = ["fused_paged_kv_gather_bhsd", "warmup_fused_paged_kv_gather"] diff --git a/rl_engine/kernels/ops/rocm/attention/strict_runtime.py b/rl_engine/kernels/ops/rocm/attention/strict_runtime.py index a1bf816d..53f1de30 100644 --- a/rl_engine/kernels/ops/rocm/attention/strict_runtime.py +++ b/rl_engine/kernels/ops/rocm/attention/strict_runtime.py @@ -27,6 +27,7 @@ from __future__ import annotations +from collections.abc import Sequence from dataclasses import dataclass, replace from typing import Any @@ -39,13 +40,14 @@ ) from rl_engine.kernels.ops.cuda.attention.cp_comm import ( AttentionCPBlockMetadata, - AttentionCPCommunicationUnavailable, AttentionCPCommunicationPlan, + AttentionCPCommunicationUnavailable, AttentionParallelSpec, CUDAAGRSAttentionCPCommunication, ) from rl_engine.kernels.ops.cuda.attention.strict_runtime import StrictCUDAAttentionRuntime from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore +from rl_engine.kernels.ops.rocm.attention.paged_gather import fused_paged_kv_gather_bhsd # The sequence reorder and the position validation are platform-neutral tensor # bookkeeping. They are bound from the CUDA runtime rather than reimplemented @@ -94,6 +96,23 @@ class StrictRocmAttentionResult: provenance: dict[str, Any] +@dataclass(frozen=True) +class _PageBoundsEpoch: + """Opaque proof scope issued by one strict ROCm runtime.""" + + owner: object + + +@dataclass(frozen=True) +class _PageBoundsValidation: + """Keep validated metadata alive so object/address reuse cannot spoof a hit.""" + + epoch: _PageBoundsEpoch + page_table: torch.Tensor + seqused_k: torch.Tensor + signature: tuple[Any, ...] + + class StrictRocmAttentionRuntime: """Run one AITER/CK arithmetic identity at CP=1 or through RCCL AG/RS.""" @@ -126,6 +145,18 @@ def __init__( if getattr(self._core, "strict_schedule", None) != self.strict_schedule: raise RuntimeError("strict ROCm Attention runtime requires the AITER/CK fixed schedule") self.communication_executed = False + self._page_bounds_epoch_owner = object() + self._page_bounds_validation: _PageBoundsValidation | None = None + self._causal_prefill_position_cache: tuple[torch.device, int, torch.Tensor] | None = None + + def new_page_bounds_epoch(self) -> object: + """Issue a proof scope while its page table and lengths remain immutable. + + vLLM creates fresh materialized metadata for the next model forward, + which must receive a fresh epoch as well. + """ + + return _PageBoundsEpoch(self._page_bounds_epoch_owner) def forward_with_lse( self, @@ -177,16 +208,34 @@ def forward_with_lse( v_sorted = _gather_sequence(global_v, k_sort) _validate_global_positions(q_positions_sorted, k_positions_sorted, causal) - out_sorted, lse_sorted, core_provenance, launches = self._run_core( - q_sorted, - k_sorted, - v_sorted, - causal=causal, - scale=scale, - query_position_ids=q_positions_sorted, - key_position_ids=k_positions_sorted, - output_dtype=q.dtype, - ) + paged_schedule = bool(getattr(self._core, "supports_paged_schedule", False)) + if paged_schedule: + core_result = self._core.forward_with_lse( + q_sorted, + k_sorted, + v_sorted, + causal=causal, + scale=scale, + key_padding_mask=None, + query_position_ids=q_positions_sorted, + key_position_ids=k_positions_sorted, + output_dtype=q.dtype, + ) + out_sorted = core_result.out + lse_sorted = core_result.lse + core_provenance = dict(core_result.provenance) + launches = 1 + else: + out_sorted, lse_sorted, core_provenance, launches = self._run_core( + q_sorted, + k_sorted, + v_sorted, + causal=causal, + scale=scale, + query_position_ids=q_positions_sorted, + key_position_ids=k_positions_sorted, + output_dtype=q.dtype, + ) if cp_world_size > 1: if q_sort is None: @@ -224,11 +273,15 @@ def forward_with_lse( "reference_only": False, "split_kv": "disabled", "framework_position_reorder": True, - # Unlike the CUDA runtime's single full-sequence launch, the - # query schedule here is one launch per (batch row, KV group). - "query_schedule": "one_batch_row_one_kv_group", + "query_schedule": ( + "one_local_gqa_batch_paged_ck" + if paged_schedule + else "one_batch_row_one_kv_group" + ), "backward_schedule": "aiter_ck_deterministic_per_kv_group", - "launch_granularity": "one_batch_row_one_kv_group", + "launch_granularity": ( + "one_local_gqa_batch" if paged_schedule else "one_batch_row_one_kv_group" + ), "tp_degree_invariant": True, "invariance_mechanism": "one_kv_group_per_launch", "core_row_count": q_sorted.size(0) * q_sorted.size(2), @@ -240,6 +293,96 @@ def forward_with_lse( }, ) + def forward_paged_varlen_with_lse( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + page_table: torch.Tensor, + seqused_k: torch.Tensor, + cu_seqlens_q: torch.Tensor, + kv_indptr: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + causal: bool, + scale: float | None, + out: torch.Tensor | None = None, + return_lse: bool = False, + page_table_validated: bool = False, + ) -> StrictRocmAttentionResult: + """Run packed prefill/decode directly from vLLM's paged KV cache.""" + + self._require_rocm(q) + direct = getattr(self._core, "forward_paged_varlen_with_lse", None) + if not callable(direct) or not bool( + getattr(self._core, "supports_paged_schedule", False) + ): + raise RuntimeError("strict ROCm direct paged-varlen CK is unavailable") + if q.ndim != 3: + raise ValueError("packed paged Q must use [tokens, heads, head_dim]") + if not page_table_validated: + bounds_ok = torch.all((page_table >= 0) & (page_table < k_cache.size(0))) + torch._assert_async(bounds_ok, "page_table entries are outside the KV cache") + valid_lengths = torch.all((seqused_k > 0) & (seqused_k <= max_seqlen_k)) + torch._assert_async( + valid_lengths, + "seqused_k entries must be positive and within max_seqlen_k", + ) + core_result = direct( + q, + k_cache, + v_cache, + page_table=page_table, + seqused_k=seqused_k, + cu_seqlens_q=cu_seqlens_q, + kv_indptr=kv_indptr, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=causal, + scale=scale, + out=out, + return_lse=return_lse, + ) + core_provenance = dict(core_result.provenance) + backend = ( + core_provenance.get("attention_backend") + or core_provenance.get("actual_backend") + or getattr(self._core, "backend_id", None) + ) + self.communication_executed = False + return StrictRocmAttentionResult( + out=core_result.out, + lse=core_result.lse, + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "communication_backend": "none", + "communication_executed": False, + "native_attention_arithmetic": True, + "production_ready": True, + "fallback": False, + "fallback_reason": None, + "reference_only": False, + "split_kv": "disabled", + "query_schedule": "paged_varlen_batch", + "paged_execution": "direct_vllm_pages_to_aiter_batch_prefill_ck", + "paged_kernel": "aiter_mha_batch_prefill_non_split_ck", + "dense_kv_materialized": False, + "lse_returned": bool(return_lse), + "launch_granularity": "one_local_gqa_batch", + "tp_degree_invariant": True, + "invariance_mechanism": "matched_train_and_rollout_paged_ck_schedule", + "core_row_count": q.size(0), + "core_launch_count": 1, + "core_batch_size": page_table.size(0), + "core_query_length": max_seqlen_q, + "core_actual_backends": [] if backend is None else [str(backend)], + "core": core_provenance, + }, + ) + def forward_paged_with_lse( self, q: torch.Tensor, @@ -251,6 +394,12 @@ def forward_paged_with_lse( max_seqlen_k: int, scale: float | None, out: torch.Tensor | None = None, + cached_lengths: Sequence[int] | None = None, + page_bounds_epoch: object | None = None, + return_lse: bool = True, + page_table_validated: bool = False, + cu_seqlens_q: torch.Tensor | None = None, + kv_indptr: torch.Tensor | None = None, ) -> StrictRocmAttentionResult: """Run strict decode Attention over a paged KV cache. @@ -287,58 +436,251 @@ def forward_paged_with_lse( raise ValueError("paged Attention out must match the Q dtype and device") if not out.is_contiguous(): raise ValueError("paged Attention out must be contiguous") + direct_core_out = ( + out is not None + and not torch.is_grad_enabled() + and not any(tensor.requires_grad for tensor in (q, k_cache, v_cache, out)) + and q.size(2) == 1 + and callable(getattr(self._core, "forward_decode_with_lse_into", None)) + and self._storage_is_disjoint( + out, + q, + k_cache, + v_cache, + page_table, + seqused_k, + ) + ) - row_outs: list[torch.Tensor] = [] - row_lses: list[torch.Tensor] = [] - core_provenance: dict[str, Any] | None = None - launches = 0 - for row in range(q.size(0)): - cached_length = int(seqused_k[row].item()) + direct_paged = callable(getattr(self._core, "forward_paged_with_lse", None)) and bool( + getattr(self._core, "supports_paged_schedule", False) + ) + if direct_paged: + if not page_table_validated: + valid_lengths = torch.all((seqused_k > 0) & (seqused_k <= max_seqlen_k)) + torch._assert_async( + valid_lengths, + "seqused_k entries must be positive and within max_seqlen_k", + ) + bounds_ok = torch.all((page_table >= 0) & (page_table < k_cache.size(0))) + torch._assert_async(bounds_ok, "page_table entries are outside the KV cache") + if cu_seqlens_q is None: + cu_seqlens_q = torch.arange( + q.size(0) + 1, dtype=torch.int32, device=q.device + ) + if kv_indptr is None: + kv_indptr = torch.arange( + q.size(0) + 1, dtype=torch.int32, device=q.device + ) * page_table.size(1) + core_result = self._core.forward_paged_with_lse( + q, + k_cache, + v_cache, + page_table=page_table, + seqused_k=seqused_k, + cu_seqlens_q=cu_seqlens_q, + kv_indptr=kv_indptr, + max_seqlen_k=max_seqlen_k, + causal=False, + scale=scale, + out=out, + return_lse=return_lse, + ) + core_provenance = dict(core_result.provenance) + backend = ( + core_provenance.get("attention_backend") + or core_provenance.get("actual_backend") + or getattr(self._core, "backend_id", None) + ) + self.communication_executed = False + return StrictRocmAttentionResult( + out=core_result.out, + lse=core_result.lse, + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "communication_backend": "none", + "communication_executed": False, + "native_attention_arithmetic": True, + "production_ready": True, + "fallback": False, + "fallback_reason": None, + "reference_only": False, + "split_kv": "disabled", + "query_schedule": "paged_single_query_batch", + "paged_execution": "direct_vllm_pages_to_aiter_batch_prefill_ck", + "paged_kernel": "aiter_mha_batch_prefill_non_split_ck", + "dense_kv_materialized": False, + "lse_returned": bool(return_lse), + "launch_granularity": "one_local_gqa_batch", + "tp_degree_invariant": True, + "invariance_mechanism": "matched_train_and_rollout_paged_ck_schedule", + "core_row_count": q.size(0) * q.size(2), + "core_launch_count": 1, + "core_batch_size": q.size(0), + "core_query_length": q.size(2), + "core_actual_backends": [] if backend is None else [str(backend)], + "core": core_provenance, + }, + ) + + if cached_lengths is None: + cached_lengths = tuple(int(value) for value in seqused_k.tolist()) + else: + cached_lengths = tuple(int(value) for value in cached_lengths) + if len(cached_lengths) != q.size(0): + raise ValueError("cached_lengths must carry one length per query") + for row, cached_length in enumerate(cached_lengths): if cached_length <= 0 or cached_length > max_seqlen_k: raise ValueError( "seqused_k entries must be positive and within max_seqlen_k; " f"row {row} requested {cached_length}" ) + causal_prefill = ( + q.size(0) > 1 + and q.size(2) == 1 + and not torch.is_grad_enabled() + and not any(tensor.requires_grad for tensor in (q, k_cache, v_cache)) + and cached_lengths == tuple(range(1, q.size(0) + 1)) + ) + use_fused_gather = False + + resolved_epoch: _PageBoundsEpoch | None + if page_bounds_epoch is None: + resolved_epoch = None + elif ( + not isinstance(page_bounds_epoch, _PageBoundsEpoch) + or page_bounds_epoch.owner is not self._page_bounds_epoch_owner + ): + raise ValueError("page_bounds_epoch was not issued by this ROCm runtime") + else: + resolved_epoch = page_bounds_epoch + bounds_signature = ( + None + if resolved_epoch is None + else self._page_bounds_signature( + page_table, + seqused_k, + cached_lengths=cached_lengths, + cache_pages=k_cache.size(0), + page_size=k_cache.size(1), + max_seqlen_k=max_seqlen_k, + ) + ) + cached_validation = self._page_bounds_validation + bounds_reused = bool( + resolved_epoch is not None + and cached_validation is not None + and cached_validation.epoch is resolved_epoch + and cached_validation.page_table is page_table + and cached_validation.seqused_k is seqused_k + and cached_validation.signature == bounds_signature + ) + + if causal_prefill: + required_pages = (cached_lengths[-1] + k_cache.size(1) - 1) // k_cache.size(1) + if not (page_table_validated or bounds_reused): + page_rows_equal = torch.all( + page_table[:, :required_pages] == page_table[-1:, :required_pages] + ) + if page_table.is_cuda: + torch._assert_async( + page_rows_equal, + "causal prefill rows must share one logical page table", + ) + elif not bool(page_rows_equal.item()): + raise ValueError("causal prefill rows must share one logical page table") k_row, v_row = self._gather_paged_row( k_cache, v_cache, - page_table[row], - cached_length, + page_table[-1], + cached_lengths[-1], + validate_bounds=not (page_table_validated or bounds_reused), ) - # Decode attends over the whole cached prefix, so the mask is not - # causal within this launch. The logical positions are still passed - # for provenance-grade auditing of what each launch consumed. - key_positions = torch.arange( - cached_length, - dtype=torch.int64, - device=q.device, - ).unsqueeze(0) - query_positions = key_positions[:, -q.size(2) :] - row_out, row_lse, row_provenance, row_launches = self._run_core( - q[row : row + 1], + q_sequence = q.squeeze(2).permute(1, 0, 2).unsqueeze(0) + positions = self._causal_prefill_positions(q.size(0), q.device) + sequence_out, sequence_lse, core_provenance, launches = self._run_core( + q_sequence, k_row, v_row, - causal=False, + causal=True, scale=scale, - query_position_ids=query_positions, - key_position_ids=key_positions, + query_position_ids=positions, + key_position_ids=positions, output_dtype=q.dtype, + collect_lse=return_lse, ) - row_outs.append(row_out) - row_lses.append(row_lse) - launches += row_launches - if core_provenance is None: - core_provenance = row_provenance + reordered_out = sequence_out.permute(2, 1, 0, 3) + if out is None: + result_out = reordered_out.contiguous() + else: + out.copy_(reordered_out) + result_out = out + result_lse = ( + sequence_lse.permute(2, 1, 0).contiguous() + if return_lse + else torch.empty((0,), dtype=torch.float32, device=q.device) + ) + else: + row_outs: list[torch.Tensor] = [] + row_lses: list[torch.Tensor] = [] + core_provenance = None + launches = 0 + for row, cached_length in enumerate(cached_lengths): + k_row, v_row = self._gather_paged_row( + k_cache, + v_cache, + page_table[row], + cached_length, + validate_bounds=not (page_table_validated or bounds_reused), + ) + # Decode attends over the whole cached prefix, so neither the mask + # nor the AITER call consumes position IDs. Avoid allocating two + # dead tensors for every row of every decoder layer. + row_out, row_lse, row_provenance, row_launches = self._run_core( + q[row : row + 1], + k_row, + v_row, + causal=False, + scale=scale, + query_position_ids=None, + key_position_ids=None, + output_dtype=q.dtype, + out=None if out is None else out[row : row + 1], + direct_core_out=direct_core_out, + collect_lse=return_lse, + ) + if out is None: + row_outs.append(row_out) + if return_lse: + row_lses.append(row_lse) + launches += row_launches + if core_provenance is None: + core_provenance = row_provenance - if core_provenance is None: - raise RuntimeError("strict ROCm paged Attention executed no core launch") + if core_provenance is None: + raise RuntimeError("strict ROCm paged Attention executed no core launch") - result_out = torch.cat(row_outs, dim=0) - result_lse = torch.cat(row_lses, dim=0) - if out is not None: - out.copy_(result_out) - result_out = out + result_out = out if out is not None else torch.cat(row_outs, dim=0) + result_lse = ( + row_lses[0] + if return_lse and direct_core_out and len(row_lses) == 1 + else ( + torch.cat(row_lses, dim=0) + if return_lse + else torch.empty((0,), dtype=torch.float32, device=q.device) + ) + ) self.communication_executed = False + if resolved_epoch is not None and not bounds_reused: + assert bounds_signature is not None + self._page_bounds_validation = _PageBoundsValidation( + epoch=resolved_epoch, + page_table=page_table, + seqused_k=seqused_k, + signature=bounds_signature, + ) backend = ( core_provenance.get("attention_backend") @@ -360,29 +702,63 @@ def forward_paged_with_lse( "fallback_reason": None, "reference_only": False, "split_kv": "disabled", - "query_schedule": "paged_single_query_batch", + "query_schedule": ( + "paged_causal_prefill_batch" if causal_prefill else "paged_single_query_batch" + ), # The dense core runs; the pages are gathered first. Recorded so # a reader never mistakes this for a native paged kernel. - "paged_execution": "logical_kv_gather_then_dense_core", - "paged_kernel": "none", + "paged_execution": ( + "fused_paged_kv_gather_to_aiter_ck_bshd" + if use_fused_gather + else "logical_kv_gather_then_dense_core" + ), + "paged_kernel": ( + "triton_fused_kv_gather_bhsd" if use_fused_gather else "none" + ), + "lse_returned": bool(return_lse), "launch_granularity": "one_batch_row_one_kv_group", "tp_degree_invariant": True, "invariance_mechanism": "one_kv_group_per_launch", "core_row_count": q.size(0) * q.size(2), "core_launch_count": launches, - "core_batch_size": q.size(0), - "core_query_length": q.size(2), + "core_batch_size": 1 if causal_prefill else q.size(0), + "core_query_length": q.size(0) if causal_prefill else q.size(2), "core_actual_backends": [] if backend is None else [str(backend)], + "core_output_staging": ( + "runtime_causal_prefill" + if causal_prefill + else ("aiter_direct_caller_group" if direct_core_out else "runtime_group_cat") + ), + "causal_prefill_collapsed": causal_prefill, + "page_bounds_validation_reused": bounds_reused, "core": core_provenance, }, ) + def _causal_prefill_positions( + self, + token_count: int, + device: torch.device, + ) -> torch.Tensor: + cached = self._causal_prefill_position_cache + if cached is not None and cached[0] == device and cached[1] == token_count: + return cached[2] + positions = torch.arange( + token_count, + dtype=torch.int64, + device=device, + ).unsqueeze(0) + self._causal_prefill_position_cache = (device, token_count, positions) + return positions + @staticmethod def _gather_paged_row( k_cache: torch.Tensor, v_cache: torch.Tensor, page_row: torch.Tensor, cached_length: int, + *, + validate_bounds: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: """Materialize one row's cached KV in logical order as ``[1, H, S, D]``. @@ -395,17 +771,164 @@ def _gather_paged_row( page_count = (cached_length + page_size - 1) // page_size if page_count > page_row.numel(): raise ValueError("page_table row is shorter than the cached length requires") - pages = page_row[:page_count].to(dtype=torch.int64) - if int(pages.min().item()) < 0 or int(pages.max().item()) >= k_cache.size(0): - raise ValueError("page_table entries are outside the KV cache") + pages = page_row[:page_count] + if validate_bounds: + bounds_ok = torch.all((pages >= 0) & (pages < k_cache.size(0))) + if pages.is_cuda: + # Keep malformed metadata fail-closed without synchronizing the host + # twice per row. vLLM page tables are already int32, which + # index_select accepts directly on ROCm. + torch._assert_async(bounds_ok, "page_table entries are outside the KV cache") + elif not bool(bounds_ok.item()): + raise ValueError("page_table entries are outside the KV cache") + + use_head_major_gather = ( + not torch.is_grad_enabled() + and not k_cache.requires_grad + and not v_cache.requires_grad + # With either singleton dimension, the legacy transpose is already + # contiguous and does not pay the second materialization copy. + and k_cache.size(2) > 1 + and cached_length > 1 + ) def _gather(cache: torch.Tensor) -> torch.Tensor: + if use_head_major_gather: + # Index the non-contiguous [H, pages, page, D] view directly. + # index_select writes head-major storage, leaving only views + # before the runtime slices one contiguous KV group at a time. + selected = cache.permute(2, 0, 1, 3).index_select(1, pages) + return selected.flatten(1, 2)[:, :cached_length].unsqueeze(0) selected = cache.index_select(0, pages) flat = selected.reshape(page_count * page_size, cache.size(2), cache.size(3)) return flat[:cached_length].permute(1, 0, 2).unsqueeze(0).contiguous() return _gather(k_cache), _gather(v_cache) + def _gather_paged_rows_fused_bhsd( + self, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_rows: torch.Tensor, + page_count: int, + *, + validate_bounds: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + pages = page_rows[:, :page_count] + flat_pages = pages.reshape(-1) + if validate_bounds: + bounds_ok = torch.all((flat_pages >= 0) & (flat_pages < k_cache.size(0))) + torch._assert_async(bounds_ok, "page_table entries are outside the KV cache") + rows = pages.size(0) + shape = ( + rows, + k_cache.size(2), + page_count * k_cache.size(1), + k_cache.size(3), + ) + key = ( + k_cache.device.type, + k_cache.device.index, + k_cache.dtype, + *shape, + ) + buffers = self._paged_bhsd_workspaces.get(key) + if buffers is None: + if len(self._paged_bhsd_workspaces) >= 128: + self._paged_bhsd_workspaces.pop(next(iter(self._paged_bhsd_workspaces))) + buffers = ( + torch.empty(shape, dtype=k_cache.dtype, device=k_cache.device), + torch.empty(shape, dtype=v_cache.dtype, device=v_cache.device), + ) + self._paged_bhsd_workspaces[key] = buffers + return fused_paged_kv_gather_bhsd( + k_cache, + v_cache, + pages, + page_count, + k_out=buffers[0], + v_out=buffers[1], + ) + + @staticmethod + def _gather_paged_rows_by_page_count( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_rows: torch.Tensor, + page_count: int, + *, + validate_bounds: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Gather one page-count group with shared index-selects. + + if page_count <= 0 or page_rows.ndim != 2: + raise ValueError("page_rows must be 2-D with a positive page count") + if page_rows.size(1) < page_count: + raise ValueError("page_table rows are shorter than the requested page count") + pages = page_rows[:, :page_count] + flat_pages = pages.reshape(-1) + if validate_bounds: + bounds_ok = torch.all((flat_pages >= 0) & (flat_pages < k_cache.size(0))) + if flat_pages.is_cuda: + torch._assert_async(bounds_ok, "page_table entries are outside the KV cache") + elif not bool(bounds_ok.item()): + raise ValueError("page_table entries are outside the KV cache") + + rows = pages.size(0) + page_size = k_cache.size(1) + def _gather(cache: torch.Tensor) -> torch.Tensor: + selected = cache.index_select(0, flat_pages) + flat = selected.reshape( + rows, + page_count * page_size, + cache.size(2), + cache.size(3), + ) + return flat.permute(0, 2, 1, 3).contiguous() + + return _gather(k_cache), _gather(v_cache) + + @staticmethod + def _page_bounds_signature( + page_table: torch.Tensor, + seqused_k: torch.Tensor, + *, + cached_lengths: Sequence[int], + cache_pages: int, + page_size: int, + max_seqlen_k: int, + ) -> tuple[Any, ...]: + """Fingerprint metadata that can affect which physical pages are read.""" + + def tensor_signature(tensor: torch.Tensor) -> tuple[Any, ...]: + try: + version: int | None = tensor._version + except RuntimeError: + # Inference tensors have no version counter. Their lifetime is + # still fenced by the adapter-issued materialization epoch. + version = None + return ( + id(tensor), + tensor.untyped_storage().data_ptr(), + tensor.storage_offset(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + version, + ) + + stream = torch.cuda.current_stream(page_table.device) if page_table.is_cuda else None + return ( + tensor_signature(page_table), + tensor_signature(seqused_k), + tuple(cached_lengths), + int(cache_pages), + int(page_size), + int(max_seqlen_k), + stream, + ) + @staticmethod def _validate_paged_inputs( q: torch.Tensor, @@ -441,6 +964,66 @@ def _validate_paged_inputs( if max_seqlen_k <= 0 or max_seqlen_k > page_table.size(1) * k_cache.size(1): raise ValueError("max_seqlen_k exceeds the page table capacity") + def _run_paged_core_bshd( + self, + q: torch.Tensor, + k_bhsd: torch.Tensor, + v_bhsd: torch.Tensor, + *, + scale: float | None, + output_dtype: torch.dtype, + out: torch.Tensor | None, + collect_lse: bool, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any], int]: + """Launch strict CK from fused group-major KV without layout copies.""" + + if output_dtype != q.dtype: + raise ValueError("strict paged Attention output dtype must match Q") + local_kv_heads = k_bhsd.size(1) + if local_kv_heads <= 0 or q.size(1) % local_kv_heads: + raise RuntimeError("paged Q heads must be divisible by local KV heads") + group_size = q.size(1) // local_kv_heads + group_outs: list[torch.Tensor] = [] + group_lses: list[torch.Tensor] = [] + core_provenance: dict[str, Any] | None = None + direct = getattr(self._core, "forward_bshd_with_lse") + for group in range(local_kv_heads): + q_lo, q_hi = group * group_size, (group + 1) * group_size + q_bshd = q[:, q_lo:q_hi].transpose(1, 2) + k_group_bshd = k_bhsd[:, group : group + 1].transpose(1, 2) + v_group_bshd = v_bhsd[:, group : group + 1].transpose(1, 2) + if not q_bshd.is_contiguous(): + q_bshd = q_bshd.contiguous() + if not k_group_bshd.is_contiguous() or not v_group_bshd.is_contiguous(): + raise RuntimeError("fused paged gather did not produce group-contiguous KV") + group_out = None if out is None else out[:, q_lo:q_hi] + result = direct( + q_bshd, + k_group_bshd, + v_group_bshd, + causal=False, + scale=scale, + out=group_out, + ) + if out is None: + group_outs.append(result.out) + if collect_lse: + group_lses.append(result.lse) + if core_provenance is None: + core_provenance = dict(result.provenance) + if core_provenance is None: + raise RuntimeError("strict ROCm paged Attention executed no CK launch") + return ( + out if out is not None else torch.cat(group_outs, dim=1), + ( + torch.cat(group_lses, dim=1) + if collect_lse + else torch.empty((0,), dtype=torch.float32, device=q.device) + ), + core_provenance, + local_kv_heads, + ) + def _run_core( self, q: torch.Tensor, @@ -449,9 +1032,12 @@ def _run_core( *, causal: bool, scale: float | None, - query_position_ids: torch.Tensor, - key_position_ids: torch.Tensor, + query_position_ids: torch.Tensor | None, + key_position_ids: torch.Tensor | None, output_dtype: torch.dtype, + out: torch.Tensor | None = None, + direct_core_out: bool = False, + collect_lse: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any], int]: """Launch the core once per ``(batch row, KV group)`` and concatenate. @@ -465,46 +1051,152 @@ def _run_core( f"local Q heads={q.size(1)} must be divisible by local KV heads={local_kv_heads}" ) group_size = q.size(1) // local_kv_heads + direct_output = q + grouped_decode = False + if direct_core_out: + if out is None or causal or q.size(2) != 1: + raise RuntimeError("direct ROCm core output is only valid for paged decode") + direct_output = out + if torch.is_grad_enabled() or any( + tensor.requires_grad for tensor in (q, k, v, direct_output) + ): + raise RuntimeError("direct ROCm core output requires disabled gradient mode") + if not callable(getattr(self._core, "forward_decode_with_lse_into", None)): + raise RuntimeError("strict ROCm core has no direct decode output entry point") + grouped_decode = local_kv_heads > 1 and callable( + getattr(self._core, "forward_grouped_decode_with_lse_into", None) + ) + + if grouped_decode: + batch_size, _query_heads, query_length, head_dim = q.shape + key_length = k.size(2) + grouped_q = q.reshape( + batch_size, + local_kv_heads, + group_size, + query_length, + head_dim, + ).reshape(batch_size * local_kv_heads, group_size, query_length, head_dim) + grouped_k = k.reshape(batch_size * local_kv_heads, 1, key_length, head_dim) + grouped_v = v.reshape(batch_size * local_kv_heads, 1, key_length, head_dim) + grouped_out = direct_output.reshape( + batch_size, + local_kv_heads, + group_size, + query_length, + head_dim, + ).reshape(batch_size * local_kv_heads, group_size, query_length, head_dim) + result = self._core.forward_grouped_decode_with_lse_into( + grouped_q, + grouped_k, + grouped_v, + out=grouped_out, + scale=scale, + output_dtype=output_dtype, + ) + if result.out.data_ptr() != grouped_out.data_ptr(): + raise RuntimeError("strict ROCm grouped core did not write to its output view") + grouped_lse = ( + result.lse.reshape( + batch_size, + local_kv_heads, + group_size, + query_length, + ).reshape( + batch_size, + local_kv_heads * group_size, + query_length, + ) + if collect_lse + else torch.empty((0,), dtype=torch.float32, device=q.device) + ) + return ( + direct_output, + grouped_lse, + dict(result.provenance), + 1, + ) row_outs: list[torch.Tensor] = [] row_lses: list[torch.Tensor] = [] core_provenance: dict[str, Any] | None = None launches = 0 for row in range(q.size(0)): - row_query_positions = query_position_ids[row : row + 1] - row_key_positions = key_position_ids[row : row + 1] + row_query_positions = ( + None if query_position_ids is None else query_position_ids[row : row + 1] + ) + row_key_positions = ( + None if key_position_ids is None else key_position_ids[row : row + 1] + ) group_outs: list[torch.Tensor] = [] group_lses: list[torch.Tensor] = [] for group in range(local_kv_heads): q_lo, q_hi = group * group_size, (group + 1) * group_size - result = self._core.forward_with_lse( - q[row : row + 1, q_lo:q_hi], - k[row : row + 1, group : group + 1], - v[row : row + 1, group : group + 1], - causal=causal, - scale=scale, - key_padding_mask=None, - query_position_ids=row_query_positions if causal else None, - key_position_ids=row_key_positions if causal else None, - output_dtype=output_dtype, - ) - group_outs.append(result.out) - group_lses.append(result.lse) + group_out = None + if direct_core_out: + group_out = direct_output[row : row + 1, q_lo:q_hi] + result = self._core.forward_decode_with_lse_into( + q[row : row + 1, q_lo:q_hi], + k[row : row + 1, group : group + 1], + v[row : row + 1, group : group + 1], + out=group_out, + scale=scale, + output_dtype=output_dtype, + ) + else: + result = self._core.forward_with_lse( + q[row : row + 1, q_lo:q_hi], + k[row : row + 1, group : group + 1], + v[row : row + 1, group : group + 1], + causal=causal, + scale=scale, + key_padding_mask=None, + query_position_ids=row_query_positions if causal else None, + key_position_ids=row_key_positions if causal else None, + output_dtype=output_dtype, + ) + if group_out is None: + group_outs.append(result.out) + elif result.out.data_ptr() != group_out.data_ptr(): + raise RuntimeError("strict ROCm core did not write to its output slice") + if collect_lse: + group_lses.append(result.lse) launches += 1 if core_provenance is None: core_provenance = dict(result.provenance) - row_outs.append(torch.cat(group_outs, dim=1)) - row_lses.append(torch.cat(group_lses, dim=1)) + if out is None: + row_outs.append(torch.cat(group_outs, dim=1)) + elif not direct_core_out: + torch.cat(group_outs, dim=1, out=out[row : row + 1]) + if collect_lse: + row_lses.append( + group_lses[0] + if direct_core_out and len(group_lses) == 1 + else torch.cat(group_lses, dim=1) + ) if core_provenance is None: raise RuntimeError("strict ROCm Attention runtime executed no core launch") return ( - torch.cat(row_outs, dim=0), - torch.cat(row_lses, dim=0), + out if out is not None else torch.cat(row_outs, dim=0), + ( + row_lses[0] + if collect_lse and direct_core_out and len(row_lses) == 1 + else ( + torch.cat(row_lses, dim=0) + if collect_lse + else torch.empty((0,), dtype=torch.float32, device=q.device) + ) + ), core_provenance, launches, ) + @staticmethod + def _storage_is_disjoint(output: torch.Tensor, *inputs: torch.Tensor) -> bool: + output_storage = output.untyped_storage().data_ptr() + return all(tensor.untyped_storage().data_ptr() != output_storage for tensor in inputs) + @staticmethod def _require_rocm(tensor: torch.Tensor) -> None: if tensor.device.type != "cuda" or torch.version.hip is None: diff --git a/rl_engine/kernels/ops/rocm/matmul/det_gemm.py b/rl_engine/kernels/ops/rocm/matmul/det_gemm.py index 6fed72b0..06febba6 100644 --- a/rl_engine/kernels/ops/rocm/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/rocm/matmul/det_gemm.py @@ -24,6 +24,12 @@ _TRITON_BACKEND = "triton" _ROUTE_REPORTED = False _ROUTE_REPORT_LOCK = Lock() +_WEIGHT_TRANSPOSE_CACHE: dict[ + tuple[int, str, tuple[int, ...], torch.dtype], tuple[int | None, torch.Tensor] +] = {} +_WEIGHT_TRANSPOSE_CACHE_LIMIT = 256 +_DIRECT_STAGING_BY_SLOT: dict[int, tuple[int, torch.Tensor, torch.Tensor]] = {} +_DIRECT_STAGING_SLOT_BY_HANDLE: dict[int, int] = {} def _requested_det_gemm_backend() -> str: @@ -54,6 +60,60 @@ def det_gemm_backend_id() -> str: return "rlkernel.det_gemm.triton_tree_rocm.v1" +def _tensor_version(tensor: torch.Tensor) -> int | None: + """Return the mutation version when the tensor tracks one.""" + + try: + return int(tensor._version) + except RuntimeError: + # Inference tensors intentionally omit version counters. Their cached + # transposes are refreshed explicitly after every IPC weight update. + return None + + +def _cached_weight_transpose(weight: torch.Tensor) -> torch.Tensor: + """Reuse inference-only transposed weights until the parameter is updated.""" + + # Training must retain the original autograd graph and therefore cannot + # retain a detached transpose in a process-global cache. + if torch.is_grad_enabled(): + return weight.t().contiguous() + key = (weight.data_ptr(), str(weight.device), tuple(weight.shape), weight.dtype) + version = _tensor_version(weight) + cached = _WEIGHT_TRANSPOSE_CACHE.get(key) + if cached is not None: + if cached[0] != version: + # Preserve the allocation captured by HIP Graph while refreshing + # its contents after an in-place framework weight update. + cached[1].copy_(weight.t()) + _WEIGHT_TRANSPOSE_CACHE[key] = (version, cached[1]) + return cached[1] + transposed = weight.t().contiguous() + if len(_WEIGHT_TRANSPOSE_CACHE) >= _WEIGHT_TRANSPOSE_CACHE_LIMIT: + _WEIGHT_TRANSPOSE_CACHE.pop(next(iter(_WEIGHT_TRANSPOSE_CACHE))) + _WEIGHT_TRANSPOSE_CACHE[key] = (version, transposed) + return transposed + + +@torch.inference_mode() +def refresh_cached_weight_transposes(weights: object) -> int: + """Refresh graph-captured transpose buffers after an IPC weight update.""" + + refreshed = 0 + for weight in weights: + if not isinstance(weight, torch.Tensor) or weight.ndim != 2: + continue + key = (weight.data_ptr(), str(weight.device), tuple(weight.shape), weight.dtype) + cached = _WEIGHT_TRANSPOSE_CACHE.get(key) + if cached is None: + continue + cached[1].copy_(weight.t()) + version = _tensor_version(weight) + _WEIGHT_TRANSPOSE_CACHE[key] = (version, cached[1]) + refreshed += 1 + return refreshed + + def _report_strict_route_once() -> None: global _ROUTE_REPORTED if torch._dynamo.is_compiling() or not route_report_enabled(): @@ -80,7 +140,201 @@ def det_gemm_linear( """Apply a native [N,K] weight through the strict ROCm backend.""" del native_op - return _triton_tree_gemm(a, weight.t().contiguous(), out=out) + return _triton_tree_gemm(a, _cached_weight_transpose(weight), out=out) + + +@torch.library.custom_op("rl_kernel::rocm_det_gemm_linear_inference", mutates_args=()) +def _det_gemm_linear_inference( + a: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor: + return det_gemm_linear(a, weight) + + +@_det_gemm_linear_inference.register_fake +def _det_gemm_linear_inference_fake( + a: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor: + return a.new_empty((a.shape[0], weight.shape[0])) + + +@torch.library.custom_op( + "rl_kernel::rocm_det_gemm_linear_inference_out", + mutates_args={"out"}, +) +def _det_gemm_linear_inference_out( + a: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, +) -> None: + det_gemm_linear(a, weight, out=out) + + +@_det_gemm_linear_inference_out.register_fake +def _det_gemm_linear_inference_out_fake( + a: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, +) -> None: + del a, weight, out + + +@torch.library.custom_op( + "rl_kernel::rocm_det_gemm_linear_all_reduce_inference", + mutates_args=(), +) +def _det_gemm_linear_all_reduce_inference( + a: torch.Tensor, + weight: torch.Tensor, + collective_handle: int, +) -> torch.Tensor: + """Run row-parallel inference without exposing shape dispatch to Dynamo.""" + + from rl_engine import _C + + binding = _DIRECT_STAGING_BY_SLOT.get(collective_handle) + if binding is None: + raise RuntimeError("strict ROCm row-parallel staging handle is not registered") + runtime_handle, staging, stable_output = binding + if a.size(0) <= staging.size(0): + # A piecewise HIP graph replays against capture-time addresses. Return + # the registered allocation for graph-eligible token counts so that + # the following captured partition always sees the same input address. + direct_input = staging.narrow(0, 0, a.size(0)) + output = stable_output.narrow(0, 0, a.size(0)) + _C.deterministic_collective_rocm_ipc_prepare_staged( + runtime_handle, + direct_input, + ) + det_gemm_linear(a, weight, out=direct_input) + _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) + _C.deterministic_collective_rocm_ipc_all_reduce_input( + runtime_handle, + output, + output, + ) + return output + + +@_det_gemm_linear_all_reduce_inference.register_fake +def _det_gemm_linear_all_reduce_inference_fake( + a: torch.Tensor, + weight: torch.Tensor, + collective_handle: int, +) -> torch.Tensor: + del collective_handle + return a.new_empty((a.shape[0], weight.shape[0])) + + +def register_det_gemm_all_reduce_staging( + collective_handle: int, + staging: torch.Tensor, +) -> int: + if collective_handle <= 0: + raise ValueError("collective_handle must be positive") + slot = _DIRECT_STAGING_SLOT_BY_HANDLE.get(collective_handle) + if slot is None: + # Dynamo persists scalar custom-op arguments in its cross-process AOT + # cache. Use registration order as the stable graph identity and resolve + # the process-local C++ handle only when the custom op executes. + slot = len(_DIRECT_STAGING_SLOT_BY_HANDLE) + 1 + _DIRECT_STAGING_SLOT_BY_HANDLE[collective_handle] = slot + binding = _DIRECT_STAGING_BY_SLOT.get(slot) + stable_output = torch.empty_like(staging) if binding is None else binding[2] + _DIRECT_STAGING_BY_SLOT[slot] = (collective_handle, staging, stable_output) + return slot + + +def det_gemm_linear_all_reduce_inference( + a: torch.Tensor, + weight: torch.Tensor, + *, + collective_handle: int, +) -> torch.Tensor: + return _det_gemm_linear_all_reduce_inference( + a, + weight, + collective_handle, + ) + + +def prepare_det_gemm_linear_weight( + weight: torch.Tensor, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Prepare a stable contiguous ``[K,N]`` inference weight. + + ROCm's strict tree kernel consumes its RHS in ``[K,N]`` layout. vLLM + stores an unquantized linear weight as ``[N,K]`` instead, so materializing + the transpose in every decode call is especially expensive for the + LM-head. The caller owns freshness: create or refresh this tensor from a + verified model load/update lifecycle, never from a version-counter guess. + """ + + if weight.dim() != 2: + raise ValueError("prepared deterministic linear weights must be 2-D") + if weight.dtype != torch.bfloat16: + raise TypeError("prepared deterministic linear weights must be BF16") + if not weight.is_cuda: + raise RuntimeError("prepared deterministic linear weights must be on ROCm") + if not weight.is_contiguous(): + raise ValueError("source deterministic linear weights must be contiguous") + + expected_shape = (weight.size(1), weight.size(0)) + if out is None: + # vLLM invokes post-load hooks from inference mode. Allocate an + # ordinary tensor so later hot-weight refreshes may update it in place. + with torch.inference_mode(False), torch.no_grad(): + out = torch.empty( + expected_shape, + dtype=weight.dtype, + device=weight.device, + ) + else: + if tuple(out.shape) != expected_shape: + raise ValueError( + f"prepared deterministic linear weight must have shape " + f"{expected_shape}, got {tuple(out.shape)}" + ) + if out.dtype != weight.dtype: + raise TypeError("prepared deterministic linear weight dtype must match its source") + if out.device != weight.device: + raise RuntimeError("prepared deterministic linear weight device must match its source") + if not out.is_contiguous(): + raise ValueError("prepared deterministic linear weight must be contiguous") + if out.requires_grad: + raise ValueError("prepared deterministic linear weight must not require gradients") + if torch._C._overlaps(out, weight): + raise ValueError("prepared deterministic linear weight must not alias its source") + + with torch.no_grad(): + out.copy_(weight.t()) + return out + + +def det_gemm_linear_prepared( + a: torch.Tensor, + weight_t: torch.Tensor, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Apply a lifecycle-managed contiguous ``[K,N]`` inference weight.""" + + if torch.is_grad_enabled() and (a.requires_grad or weight_t.requires_grad): + raise RuntimeError("prepared deterministic linear GEMM is inference-only") + if not weight_t.is_contiguous(): + 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) def det_gemm_linear_input_gradient( @@ -124,16 +378,8 @@ def backward(ctx, grad_out): grad_out = grad_out.contiguous() if grad_out.dtype != torch.bfloat16: grad_out = grad_out.to(torch.bfloat16) - da = ( - det_gemm_linear_input_gradient(grad_out, weight) - if ctx.needs_input_grad[0] - else None - ) - dweight = ( - det_gemm_linear_weight_gradient(a, grad_out) - if ctx.needs_input_grad[1] - else None - ) + da = det_gemm_linear_input_gradient(grad_out, weight) if ctx.needs_input_grad[0] else None + dweight = det_gemm_linear_weight_gradient(a, grad_out) if ctx.needs_input_grad[1] else None record_backward( "det_gemm", kernel_id=det_gemm_backend_id(), @@ -174,9 +420,25 @@ def linear( if out is not None: if torch.is_grad_enabled() and (a.requires_grad or weight.requires_grad): raise RuntimeError("direct-output deterministic GEMM is inference-only") - return det_gemm_linear(a, weight, out=out) + _det_gemm_linear_inference_out(a, weight, out) + return out + if not torch.is_grad_enabled() or not (a.requires_grad or weight.requires_grad): + return _det_gemm_linear_inference(a, weight) return _DetLinearFn.apply(a, weight) + def linear_prepared( + self, + a: torch.Tensor, + weight_t: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Apply a post-load prepared RHS without a hot-path transpose.""" + + assert a.dtype == torch.bfloat16 and weight_t.dtype == torch.bfloat16, "BF16 only" + assert a.is_cuda and weight_t.is_cuda, "Inputs must be on ROCm device" + return det_gemm_linear_prepared(a.contiguous(), weight_t, out=out) + def forward_fp32(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16, "BF16 only" assert a.is_cuda and b.is_cuda, "Inputs must be on ROCm device" @@ -214,7 +476,12 @@ def deterministic_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: "det_gemm_backend_id", "det_gemm_fallback_reason", "det_gemm_linear", + "det_gemm_linear_prepared", "det_gemm_linear_input_gradient", "det_gemm_linear_weight_gradient", "deterministic_gemm", + "prepare_det_gemm_linear_weight", + "det_gemm_linear_all_reduce_inference", + "register_det_gemm_all_reduce_staging", + "refresh_cached_weight_transposes", ] diff --git a/rl_engine/kernels/ops/rocm/rotary_embedding/rope.py b/rl_engine/kernels/ops/rocm/rotary_embedding/rope.py index 6164523f..360b4c5b 100644 --- a/rl_engine/kernels/ops/rocm/rotary_embedding/rope.py +++ b/rl_engine/kernels/ops/rocm/rotary_embedding/rope.py @@ -4,11 +4,117 @@ from __future__ import annotations +from threading import Lock + import torch from torch import Tensor from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE -from rl_engine.kernels.ops.cuda.rotary_embedding.rope import _restore_rope, _rope_table +from rl_engine.kernels.ops.cuda.rotary_embedding.rope import ( + _build_cos_sin, + _restore_rope, + _rope_table, +) + + +@torch.library.custom_op("rl_kernel::deterministic_rope_apply_rocm", mutates_args=()) +def _deterministic_rope_apply_rocm( + x: Tensor, + cos: Tensor, + sin: Tensor, + direction: float, +) -> Tensor: + return _C.deterministic_rope_apply_rocm(x, cos, sin, direction) + + +@_deterministic_rope_apply_rocm.register_fake +def _deterministic_rope_apply_rocm_fake( + x: Tensor, + cos: Tensor, + sin: Tensor, + direction: float, +) -> Tensor: + del cos, sin, direction + return torch.empty_like(x) + + +@torch.library.custom_op( + "rl_kernel::deterministic_rope_apply_token_major_rocm", mutates_args=() +) +def _deterministic_rope_apply_token_major_rocm( + x: Tensor, + positions: Tensor, + cos: Tensor, + sin: Tensor, + head_dim: int, + direction: float, +) -> Tensor: + return _C.deterministic_rope_apply_token_major_rocm( + x, + positions, + cos, + sin, + head_dim, + direction, + ) + + +@_deterministic_rope_apply_token_major_rocm.register_fake +def _deterministic_rope_apply_token_major_rocm_fake( + x: Tensor, + positions: Tensor, + cos: Tensor, + sin: Tensor, + head_dim: int, + direction: float, +) -> Tensor: + del positions, cos, sin, head_dim, direction + return torch.empty(x.shape, dtype=x.dtype, device=x.device) + + +def _forward_rope_pair( + query: Tensor, + key: Tensor, + positions: Tensor, + theta: float, + *, + inv_freq: Tensor | None = None, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Run the paired forward arithmetic and return its shared table.""" + + if positions.dim() != 1: + raise ValueError("paired ROCm RoPE requires a flat position tensor") + if inv_freq is None: + query_2d, cos, sin = _rope_table(query, positions, theta) + else: + head_dim = query.shape[-1] + if head_dim % 2 != 0: + raise ValueError(f"RoPE head_dim must be even, got {head_dim}") + table_len = int(positions.shape[0]) + query_2d = query.contiguous().reshape(-1, head_dim) + if query_2d.shape[0] % table_len != 0: + raise ValueError( + f"row count {query_2d.shape[0]} not divisible by seq length {table_len}; " + "expected a [..., S, D] contiguous layout." + ) + if ( + inv_freq.shape != (head_dim // 2,) + or inv_freq.dtype != torch.float32 + or inv_freq.device != query.device + ): + raise RuntimeError("cached ROCm RoPE inv_freq does not match the input") + position_rows = positions.to(device=query.device, dtype=torch.float32).reshape(-1, 1) + frequencies = position_rows * inv_freq + cos = frequencies.cos().contiguous() + sin = frequencies.sin().contiguous() + key_2d = key.contiguous().reshape(-1, key.shape[-1]) + if key_2d.size(0) % cos.size(0): + raise ValueError( + f"key row count {key_2d.size(0)} is not divisible by " f"position count {cos.size(0)}" + ) + query_out = _C.deterministic_rope_apply_rocm(query_2d, cos, sin, 1.0) + key_out = _C.deterministic_rope_apply_rocm(key_2d, cos, sin, 1.0) + return query_out.reshape(query.shape), key_out.reshape(key.shape), cos, sin class _RocmRoPEFunction(torch.autograd.Function): @@ -18,7 +124,7 @@ def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: ctx.save_for_backward(cos, sin) ctx.x_shape = tuple(x.shape) ctx.pos_dim = positions.dim() - out_2d = _C.deterministic_rope_apply_rocm(x_2d, cos, sin, 1.0) + out_2d = _deterministic_rope_apply_rocm(x_2d, cos, sin, 1.0) return _restore_rope(out_2d, x, positions) @staticmethod @@ -28,26 +134,57 @@ def backward(ctx, grad_out: Tensor): if ctx.needs_input_grad[0]: if ctx.pos_dim == 2 and len(ctx.x_shape) == 4: g_2d = grad_out.permute(1, 0, 2, 3).contiguous().reshape(-1, ctx.x_shape[-1]) - out_2d = _C.deterministic_rope_apply_rocm(g_2d, cos, sin, -1.0) + out_2d = _deterministic_rope_apply_rocm(g_2d, cos, sin, -1.0) heads, batch, seq, dim = ( ctx.x_shape[1], ctx.x_shape[0], ctx.x_shape[2], ctx.x_shape[3], ) - grad_x = ( - out_2d.reshape(heads, batch, seq, dim) - .permute(1, 0, 2, 3) - .contiguous() - ) + grad_x = out_2d.reshape(heads, batch, seq, dim).permute(1, 0, 2, 3).contiguous() else: g_2d = grad_out.contiguous().reshape(-1, grad_out.shape[-1]) - grad_x = _C.deterministic_rope_apply_rocm(g_2d, cos, sin, -1.0).reshape( + grad_x = _deterministic_rope_apply_rocm(g_2d, cos, sin, -1.0).reshape( grad_out.shape ) return grad_x, None, None +class _RocmRoPEPairFunction(torch.autograd.Function): + """Rotate Q/K with one shared position table and independent HIP launches.""" + + @staticmethod + def forward( + ctx, + query: Tensor, + key: Tensor, + positions: Tensor, + theta: float, + ) -> tuple[Tensor, Tensor]: + ctx.set_materialize_grads(False) + query_out, key_out, cos, sin = _forward_rope_pair(query, key, positions, theta) + ctx.save_for_backward(cos, sin) + ctx.query_shape = tuple(query.shape) + ctx.key_shape = tuple(key.shape) + return query_out, key_out + + @staticmethod + def backward(ctx, grad_query: Tensor, grad_key: Tensor): + cos, sin = ctx.saved_tensors + + def rotate_gradient(grad: Tensor | None, shape: tuple[int, ...]) -> Tensor | None: + if grad is None: + return None + grad_2d = grad.contiguous().reshape(-1, shape[-1]) + return _C.deterministic_rope_apply_rocm(grad_2d, cos, sin, -1.0).reshape(shape) + + query_grad = ( + rotate_gradient(grad_query, ctx.query_shape) if ctx.needs_input_grad[0] else None + ) + key_grad = rotate_gradient(grad_key, ctx.key_shape) if ctx.needs_input_grad[1] else None + return query_grad, key_grad, None, None + + class RocmDeterministicRoPEOp: """Precompiled HIP RoPE path shared by ROCm training and rollout.""" @@ -58,20 +195,145 @@ class RocmDeterministicRoPEOp: def __init__(self) -> None: if torch.version.hip is None: raise RuntimeError("RocmDeterministicRoPEOp requires a ROCm PyTorch build") - if not _EXT_AVAILABLE or not hasattr(_C, "deterministic_rope_apply_rocm"): + if not _EXT_AVAILABLE or not all( + hasattr(_C, symbol) + for symbol in ( + "deterministic_rope_apply_rocm", + "deterministic_rope_apply_token_major_rocm", + ) + ): raise RuntimeError( "ROCm deterministic RoPE is unavailable; rebuild rl_engine._C for ROCm" ) + self._inference_inv_freq_cache: dict[tuple[torch.device, object, int, str], Tensor] = {} + self._inference_inv_freq_lock = Lock() def __call__(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: return self.forward(x, positions, theta=theta) def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: + self._validate_input(x) + return _RocmRoPEFunction.apply(x, positions, theta) + + def forward_pair( + self, + query: Tensor, + key: Tensor, + positions: Tensor, + *, + theta: float = 1_000_000.0, + ) -> tuple[Tensor, Tensor]: + """Rotate Q/K with one FP32 cos/sin table and unchanged HIP arithmetic.""" + + self._validate_input(query) + self._validate_input(key) + if query.device != key.device or query.dtype != key.dtype: + raise ValueError("paired ROCm RoPE Q/K must share one device and dtype") + if query.shape[-1] != key.shape[-1]: + raise ValueError("paired ROCm RoPE Q/K must share one head dimension") + if positions.dim() != 1: + raise ValueError("paired ROCm RoPE requires a flat position tensor") + if not torch.is_grad_enabled(): + # vLLM executes rollout under inference/no-grad mode. Calling the + # custom autograd Function there cannot contribute a backward but + # still pays its dispatcher/context cost in every decoder layer. + # Keep the exact same table construction and HIP launches while + # bypassing only that unused autograd wrapper. + query_out, key_out, _cos, _sin = _forward_rope_pair( + query, + key, + positions, + theta, + inv_freq=self._cached_inference_inv_freq(query, theta), + ) + return query_out, key_out + query_out, key_out = _RocmRoPEPairFunction.apply(query, key, positions, theta) + # A multi-output autograd Function marks both outputs differentiable if + # either input requires grad. Preserve the two independent-call API. + if not query.requires_grad: + query_out = query_out.detach() + if not key.requires_grad: + key_out = key_out.detach() + return query_out, key_out + + def _cached_inference_inv_freq(self, query: Tensor, theta: float) -> Tensor: + """Reuse only the position-independent FP32 part of the RoPE table.""" + + head_dim = query.shape[-1] + if head_dim % 2 != 0: + raise ValueError(f"RoPE head_dim must be even, got {head_dim}") + resolved_theta = float(theta) + # ``float.hex`` keeps +0.0 and -0.0 distinct, unlike float equality. + stream = torch.cuda.current_stream(query.device) + key = (query.device, stream, head_dim, resolved_theta.hex()) + cached = self._inference_inv_freq_cache.get(key) + if cached is not None: + return cached + half = head_dim // 2 + inv_freq = 1.0 / ( + resolved_theta + ** ( + torch.arange( + 0, + half, + dtype=torch.float32, + device=query.device, + ) + / half + ) + ) + # Build outside the lock, then publish once. The dict never evicts: + # keys retain their streams and values remain alive for any HIP graph + # that captured them. A same-key race returns the first published + # tensor, whose producer and consumers all use that key's stream. + with self._inference_inv_freq_lock: + return self._inference_inv_freq_cache.setdefault(key, inv_freq) + + @staticmethod + def _validate_input(x: Tensor) -> None: if not x.is_cuda: raise RuntimeError("ROCm deterministic RoPE requires a GPU tensor") if x.dtype not in (torch.float16, torch.bfloat16): raise ValueError("ROCm deterministic RoPE requires FP16 or BF16") - return _RocmRoPEFunction.apply(x, positions, theta) + + @staticmethod + def build_position_table( + max_positions: int, + head_dim: int, + *, + device: torch.device, + theta: float, + ) -> tuple[Tensor, Tensor]: + """Build the training-identical FP32 table once before graph capture.""" + + if max_positions <= 0 or head_dim <= 0 or head_dim % 2: + raise ValueError("ROCm RoPE table dimensions must be positive and even") + positions = torch.arange(max_positions, dtype=torch.int64, device=device) + return _build_cos_sin(positions, head_dim // 2, float(theta), device) + + @staticmethod + def forward_token_major( + x: Tensor, + positions: Tensor, + cos: Tensor, + sin: Tensor, + *, + head_dim: int, + ) -> Tensor: + """Rotate vLLM's strided token-major Q/K without layout copies.""" + + if x.ndim != 2 or x.stride(1) != 1: + raise ValueError("token-major ROCm RoPE requires unit inner stride") + if positions.ndim != 1 or positions.numel() != x.size(0): + raise ValueError("token-major ROCm RoPE positions must match token rows") + return _deterministic_rope_apply_token_major_rocm( + x, + positions, + cos, + sin, + head_dim, + 1.0, + ) __all__ = ["RocmDeterministicRoPEOp"] diff --git a/rl_engine/kernels/ops/triton/matmul/det_gemm.py b/rl_engine/kernels/ops/triton/matmul/det_gemm.py index 68db1379..a4fc1529 100644 --- a/rl_engine/kernels/ops/triton/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/triton/matmul/det_gemm.py @@ -252,7 +252,7 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: if _TRITON_AVAILABLE: - @triton.jit + @triton.jit(do_not_specialize=["M"]) def _det_gemm_tree_leaf_kernel( a_ptr, b_ptr, @@ -260,7 +260,7 @@ def _det_gemm_tree_leaf_kernel( leaf_starts_ptr, leaf_lengths_ptr, leaf_nodes_ptr, - M: tl.constexpr, + M, N: tl.constexpr, K: tl.constexpr, stride_am: tl.constexpr, @@ -322,13 +322,13 @@ 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 + @triton.jit(do_not_specialize=["M"]) def _det_gemm_tree_reduce_kernel( workspace_ptr, lower_nodes_ptr, upper_nodes_ptr, output_nodes_ptr, - M: tl.constexpr, + M, N: tl.constexpr, BLOCK: tl.constexpr, ): @@ -357,6 +357,40 @@ def _det_gemm_tree_reduce_kernel( mask=mask, ) + @triton.jit + def _det_gemm_tree_reduce_to_output_kernel( + workspace_ptr, + output_ptr, + lower_nodes_ptr, + upper_nodes_ptr, + M: tl.constexpr, + N: tl.constexpr, + BLOCK: tl.constexpr, + ): + offsets = (tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)).to(tl.int64) + elements = M * N + mask = offsets < elements + lower_node = tl.load(lower_nodes_ptr).to(tl.int64) + upper_node = tl.load(upper_nodes_ptr).to(tl.int64) + lower = tl.load( + workspace_ptr + lower_node * elements + offsets, + mask=mask, + other=0.0, + ).to(tl.float32) + upper = tl.load( + workspace_ptr + upper_node * elements + offsets, + mask=mask, + other=0.0, + ).to(tl.float32) + result = lower + upper + # Preserve the canonical root's FP32 add and BF16 store boundary; only + # its destination changes, avoiding a device-to-device copy launch. + tl.store( + output_ptr + offsets, + result.to(output_ptr.dtype.element_ty), + mask=mask, + ) + @triton.jit def _copy_tree_root_kernel( workspace_ptr, @@ -594,21 +628,42 @@ def _triton_tree_gemm( num_warps=leaf_config.num_warps, ) reduction_block = 256 - for operations, (lower, upper, output) in zip( - plan.host.reduction_levels, - plan.reduction_levels, - strict=True, + 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, + ) ): - 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, + 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),) + _det_gemm_tree_reduce_to_output_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: @@ -626,7 +681,7 @@ def _triton_tree_gemm( BLOCK_M=transpose_block, BLOCK_N=transpose_block, ) - else: + elif not direct_root_output or not plan.host.reduction_levels: _copy_tree_root_kernel[(triton.cdiv(result.numel(), copy_block),)]( workspace, result, diff --git a/tests/test_attention_correctness.py b/tests/test_attention_correctness.py index a77e195c..97ad0569 100644 --- a/tests/test_attention_correctness.py +++ b/tests/test_attention_correctness.py @@ -558,6 +558,139 @@ def fake_bwd( assert q.grad is not None and k.grad is not None and v.grad is not None +def test_strict_rocm_aiter_ck_direct_decode_uses_callers_output(monkeypatch): + seen_out = None + + def fake_fwd(q, k, v, *_args, out=None): + nonlocal seen_out + seen_out = out + out.copy_(q) + return ( + out, + torch.zeros(q.size(0), q.size(2), q.size(1), dtype=torch.float32), + torch.empty(0), + torch.zeros(2, dtype=torch.int64), + ) + + core = StrictRocmAiterCKAttentionCore( + _mha_fwd=fake_fwd, + _mha_bwd=lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr(core, "_validate_inputs", lambda *_args: None) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: type("Props", (), {"name": "test-gpu", "gcnArchName": "gfx-test"})(), + ) + q = torch.randn(1, 4, 1, 8, dtype=torch.bfloat16) + k = torch.randn(1, 1, 7, 8, dtype=torch.bfloat16) + v = torch.randn_like(k) + out = torch.empty_like(q) + + with torch.no_grad(): + result = core.forward_decode_with_lse_into(q, k, v, out=out, scale=0.125) + + assert seen_out is not None and seen_out.data_ptr() == out.data_ptr() + assert result.out is out + assert torch.equal(out, q) + assert result.provenance["core_output_staging"] == "aiter_direct_caller_group" + + +def test_strict_rocm_aiter_ck_reuses_immutable_provenance_inputs(monkeypatch): + def fake_fwd(q, k, v, *_args, out=None): + out.copy_(q) + return ( + out, + torch.zeros(q.size(0), q.size(2), q.size(1), dtype=torch.float32), + torch.empty(0), + torch.zeros(2, dtype=torch.int64), + ) + + core = StrictRocmAiterCKAttentionCore( + _mha_fwd=fake_fwd, + _mha_bwd=lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr(core, "_validate_inputs", lambda *_args: None) + device_lookups = 0 + + def fake_device_properties(_device): + nonlocal device_lookups + device_lookups += 1 + return type("Props", (), {"name": "test-gpu", "gcnArchName": "gfx-test"})() + + monkeypatch.setattr(torch.cuda, "get_device_properties", fake_device_properties) + split_resolves = 0 + split_type = type(core.split_kv) + original_resolve = split_type.resolve + + def counted_resolve(self, total_kv_tokens, *, backend): + nonlocal split_resolves + split_resolves += 1 + return original_resolve(self, total_kv_tokens, backend=backend) + + monkeypatch.setattr(split_type, "resolve", counted_resolve) + q = torch.randn(1, 4, 1, 8, dtype=torch.bfloat16) + + def run(kv_tokens): + k = torch.randn(1, 1, kv_tokens, 8, dtype=torch.bfloat16) + with torch.no_grad(): + return core.forward_decode_with_lse_into(q, k, k.clone(), out=torch.empty_like(q)) + + first = run(7) + repeated = run(7) + changed_length = run(9) + restored_length = run(7) + + assert device_lookups == 1 + assert split_resolves == 3 + assert first.provenance == repeated.provenance + assert first.provenance is not repeated.provenance + assert first.provenance["split_kv"] is not repeated.provenance["split_kv"] + first_boundaries = first.provenance["split_kv"]["actual_split_boundaries"] + repeated_boundaries = repeated.provenance["split_kv"]["actual_split_boundaries"] + assert first_boundaries is not repeated_boundaries + assert first_boundaries[0] is not repeated_boundaries[0] + assert changed_length.provenance["split_kv"]["actual_split_boundaries"] == [[0, 9]] + assert restored_length.provenance["split_kv"]["actual_split_boundaries"] == [[0, 7]] + first_boundaries[0][0] = 3 + assert repeated_boundaries == [[0, 7]] + after_mutation = run(7) + assert after_mutation.provenance["split_kv"]["actual_split_boundaries"] == [[0, 7]] + assert split_resolves == 3 + + assert core._device_description(torch.device("cuda:1")) == ("test-gpu", "gfx-test") + assert core._device_description(torch.device("cuda:1")) == ("test-gpu", "gfx-test") + assert device_lookups == 2 + + +def test_strict_rocm_aiter_ck_direct_decode_rejects_ignored_output(monkeypatch): + def fake_fwd(q, k, v, *_args, out=None): + return ( + q.clone(), + torch.zeros(q.size(0), q.size(2), q.size(1), dtype=torch.float32), + torch.empty(0), + torch.zeros(2, dtype=torch.int64), + ) + + core = StrictRocmAiterCKAttentionCore( + _mha_fwd=fake_fwd, + _mha_bwd=lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr(core, "_validate_inputs", lambda *_args: None) + q = torch.randn(1, 4, 1, 8, dtype=torch.bfloat16) + k = torch.randn(1, 1, 7, 8, dtype=torch.bfloat16) + out = torch.empty_like(q) + + with ( + torch.no_grad(), + pytest.raises( + StrictRocmAttentionUnavailable, + match="requested output buffer", + ), + ): + core.forward_decode_with_lse_into(q, k, k, out=out) + + def test_strict_rocm_aiter_ck_core_rejects_non_fp32_lse(monkeypatch): def fake_fwd(q, k, v, *_args): return ( diff --git a/tests/test_det_gemm.py b/tests/test_det_gemm.py index c1d0992f..0c4408c7 100644 --- a/tests/test_det_gemm.py +++ b/tests/test_det_gemm.py @@ -25,7 +25,10 @@ _GFX942_QWEN_TP_SHARD_FORWARD_LEAF_CONFIGS, _GFX942_QWEN_TP_SHARD_WGRAD_LEAF_CONFIGS, _GFX942_QWEN_WGRAD_LEAF_CONFIGS, + _copy_tree_root_kernel, _det_gemm_tree_leaf_kernel, + _det_gemm_tree_reduce_kernel, + _det_gemm_tree_reduce_to_output_kernel, _device_tree_plan, _gfx942_qwen_tree_leaf_config, _triton_tree_gemm, @@ -538,6 +541,68 @@ def test_triton_tree_matches_python_reference_backward_bitwise(): assert torch.equal(expected_grad, triton_input.grad) +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") +def test_triton_final_tree_level_direct_output_matches_legacy_copy_raw_bytes(): + """Changing the root destination must not change its BF16 bit pattern.""" + + m_size, n_size = 3, 97 + workspace = _special_bf16((3, m_size, n_size), offset=5) + legacy_workspace = workspace.clone() + lower = torch.tensor([0], device=DEV, dtype=torch.int64) + upper = torch.tensor([1], device=DEV, dtype=torch.int64) + root = torch.tensor([2], device=DEV, dtype=torch.int64) + block = 256 + grid = (1, triton.cdiv(m_size * n_size, block)) + _det_gemm_tree_reduce_kernel[grid]( + legacy_workspace, + lower, + upper, + root, + M=m_size, + N=n_size, + BLOCK=block, + ) + expected = torch.empty((m_size, n_size), device=DEV, dtype=torch.bfloat16) + _copy_tree_root_kernel[(triton.cdiv(expected.numel(), block),)]( + legacy_workspace, + expected, + 2, + expected.numel(), + BLOCK=block, + ) + + actual = torch.empty_like(expected) + _det_gemm_tree_reduce_to_output_kernel[(triton.cdiv(actual.numel(), block),)]( + workspace, + actual, + lower, + upper, + M=m_size, + N=n_size, + BLOCK=block, + ) + _assert_same_raw_bytes(actual, expected) + + +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") +@pytest.mark.parametrize("k_size", (31, 65)) +def test_triton_forward_out_buffer_preserves_identity_version_and_raw_bytes(k_size): + """Caller storage works for both the leaf-root and reduced-root paths.""" + + torch.manual_seed(51 + k_size) + a = _rand(3, k_size) + b = _rand(k_size, 17) + expected = _triton_tree_gemm(a, b) + output_buffer = torch.empty_like(expected) + version_before = output_buffer._version + + actual = _triton_tree_gemm(a, b, out=output_buffer) + + assert actual is output_buffer + assert output_buffer._version == version_before + 1 + _assert_same_raw_bytes(actual, expected) + + @pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") @pytest.mark.parametrize( "shape", diff --git a/tests/test_deterministic_attention_cuda.py b/tests/test_deterministic_attention_cuda.py index f981e840..32dc0682 100644 --- a/tests/test_deterministic_attention_cuda.py +++ b/tests/test_deterministic_attention_cuda.py @@ -47,7 +47,7 @@ _OP_AVAILABLE = False if IS_ROCM: - from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RocmDeterministicRoPEOp + from rl_engine.kernels.ops.rocm.rotary_embedding.rope import RocmDeterministicRoPEOp pytestmark = [ pytestmark, @@ -785,6 +785,305 @@ def test_rocm_rope_matches_fp32_rotate_half_reference(): torch.testing.assert_close(actual, reference, atol=0, rtol=0) +@ROCM_ONLY +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("tokens", (1, 2, 7, 32)) +def test_rocm_rope_pair_matches_two_single_calls_raw_bytes(dtype, tokens): + generator = torch.Generator(device="cpu").manual_seed(810 + tokens) + query = torch.randn(8, tokens, D, dtype=dtype, generator=generator).to(DEVICE) + key = torch.randn(2, tokens, D, dtype=dtype, generator=generator).to(DEVICE) + position_storage = torch.arange(tokens * 2, device=DEVICE, dtype=torch.int64) + positions = position_storage[::2] + assert not positions.is_contiguous() or tokens == 1 + rope = RocmDeterministicRoPEOp() + + expected = (rope(query, positions), rope(key, positions)) + actual = rope.forward_pair(query, key, positions) + repeated = rope.forward_pair(query, key, positions) + for expected_tensor, actual_tensor, repeated_tensor in zip( + expected, + actual, + repeated, + strict=True, + ): + assert torch.equal( + expected_tensor.contiguous().view(torch.uint8), + actual_tensor.contiguous().view(torch.uint8), + ) + assert torch.equal( + actual_tensor.contiguous().view(torch.uint8), + repeated_tensor.contiguous().view(torch.uint8), + ) + + +@ROCM_ONLY +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +@pytest.mark.parametrize("tokens", (1, 2, 32)) +def test_rocm_rope_pair_inference_bypasses_autograd_raw_bytes(monkeypatch, dtype, tokens): + from rl_engine.kernels.ops.rocm.rotary_embedding.rope import _RocmRoPEPairFunction + + generator = torch.Generator(device="cpu").manual_seed(813 + tokens) + position_storage = torch.arange(tokens * 2, device=DEVICE, dtype=torch.int64) + positions = position_storage[::2] + query = torch.randn(8, tokens, D, dtype=dtype, generator=generator).to(DEVICE) + key = torch.randn(2, tokens, D, dtype=dtype, generator=generator).to(DEVICE) + query.requires_grad_() + key.requires_grad_() + rope = RocmDeterministicRoPEOp() + + with torch.inference_mode(): + expected = _RocmRoPEPairFunction.apply(query, key, positions, 1_000_000.0) + + def unexpected_apply(*args, **kwargs): + raise AssertionError("no-grad paired RoPE entered its autograd Function") + + monkeypatch.setattr(_RocmRoPEPairFunction, "apply", unexpected_apply) + for mode in (torch.no_grad, torch.inference_mode): + with mode(): + actual = rope.forward_pair(query, key, positions) + repeated = rope.forward_pair(query, key, positions) + + for expected_tensor, actual_tensor, repeated_tensor in zip( + expected, + actual, + repeated, + strict=True, + ): + assert actual_tensor.requires_grad is False + assert actual_tensor.grad_fn is None + assert torch.equal(expected_tensor.view(torch.uint8), actual_tensor.view(torch.uint8)) + assert torch.equal(actual_tensor.view(torch.uint8), repeated_tensor.view(torch.uint8)) + + +@ROCM_ONLY +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_rocm_rope_pair_inference_reuses_only_position_independent_frequency(dtype): + from rl_engine.kernels.ops.rocm.rotary_embedding.rope import _forward_rope_pair + + generator = torch.Generator(device="cpu").manual_seed(814) + query = torch.randn(8, 3, D, dtype=dtype, generator=generator).to(DEVICE) + key = torch.randn(2, 3, D, dtype=dtype, generator=generator).to(DEVICE) + first_positions = torch.tensor([0, 17, 2048], device=DEVICE) + second_positions = torch.tensor([4095, 3, 29], device=DEVICE) + rope = RocmDeterministicRoPEOp() + + with torch.inference_mode(): + first = rope.forward_pair(query, key, first_positions) + assert len(rope._inference_inv_freq_cache) == 1 + cached_inv_freq = next(iter(rope._inference_inv_freq_cache.values())) + expected_second = _forward_rope_pair( + query, + key, + second_positions, + 1_000_000.0, + ) + cached_second = _forward_rope_pair( + query, + key, + second_positions, + 1_000_000.0, + inv_freq=cached_inv_freq, + ) + with torch.no_grad(): + second = rope.forward_pair(query, key, second_positions) + + assert len(rope._inference_inv_freq_cache) == 1 + assert next(iter(rope._inference_inv_freq_cache.values())) is cached_inv_freq + assert not torch.equal(first[0].view(torch.uint8), second[0].view(torch.uint8)) + for expected_tensor, cached_tensor in zip(expected_second, cached_second, strict=True): + assert torch.equal(expected_tensor.view(torch.uint8), cached_tensor.view(torch.uint8)) + for expected_tensor, actual_tensor in zip(expected_second[:2], second, strict=True): + assert torch.equal(expected_tensor.view(torch.uint8), actual_tensor.view(torch.uint8)) + + +@ROCM_ONLY +def test_rocm_rope_pair_inference_frequency_cache_keys_shape_and_theta(): + from rl_engine.kernels.ops.rocm.rotary_embedding.rope import _forward_rope_pair + + rope = RocmDeterministicRoPEOp() + positions = torch.tensor([1, 2], device=DEVICE) + + def run(head_dim, theta): + query = torch.randn(2, 2, head_dim, dtype=torch.bfloat16, device=DEVICE) + key = torch.randn(1, 2, head_dim, dtype=torch.bfloat16, device=DEVICE) + with torch.inference_mode(): + expected = _forward_rope_pair(query, key, positions, theta)[:2] + actual = rope.forward_pair(query, key, positions, theta=theta) + for expected_tensor, actual_tensor in zip(expected, actual, strict=True): + assert torch.equal(expected_tensor.view(torch.uint8), actual_tensor.view(torch.uint8)) + return rope._cached_inference_inv_freq(query, theta) + + grad_query = torch.randn(2, 2, D, dtype=torch.bfloat16, device=DEVICE).requires_grad_() + grad_key = torch.randn(1, 2, D, dtype=torch.bfloat16, device=DEVICE).requires_grad_() + rope.forward_pair(grad_query, grad_key, positions) + assert not rope._inference_inv_freq_cache + + initial = run(D, 1_000_000.0) + assert run(D, 1_000_000.0) is initial + changed_theta = run(D, 10_000.0) + assert changed_theta is not initial + changed_shape = run(D // 2, 10_000.0) + assert changed_shape is not changed_theta + assert changed_shape.numel() == D // 4 + + cache_entries = len(rope._inference_inv_freq_cache) + trained_query = grad_query.detach().clone().requires_grad_() + trained_key = grad_key.detach().clone().requires_grad_() + trained_out = rope.forward_pair(trained_query, trained_key, positions) + (trained_out[0].float().sum() + trained_out[1].float().sum()).backward() + assert trained_query.grad is not None + assert trained_key.grad is not None + assert len(rope._inference_inv_freq_cache) == cache_entries + + +@ROCM_ONLY +def test_rocm_rope_pair_inference_frequency_cache_is_stream_local(): + from rl_engine.kernels.ops.rocm.rotary_embedding.rope import _forward_rope_pair + + rope = RocmDeterministicRoPEOp() + query = torch.randn(2, 2, D, dtype=torch.bfloat16, device=DEVICE) + key = torch.randn(1, 2, D, dtype=torch.bfloat16, device=DEVICE) + positions = torch.tensor([7, 31], device=DEVICE) + streams = (torch.cuda.Stream(device=DEVICE), torch.cuda.Stream(device=DEVICE)) + outputs = [] + repeated_outputs = [] + expected = [] + torch.cuda.synchronize(DEVICE) + + for stream in streams: + theta = 1_000_000.0 + with torch.cuda.stream(stream), torch.inference_mode(): + expected.append(_forward_rope_pair(query, key, positions, theta)[:2]) + outputs.append(rope.forward_pair(query, key, positions, theta=theta)) + repeated_outputs.append(rope.forward_pair(query, key, positions, theta=theta)) + torch.cuda.synchronize(DEVICE) + + assert len(rope._inference_inv_freq_cache) == 2 + assert len({id(value) for value in rope._inference_inv_freq_cache.values()}) == 2 + for expected_pair, actual_pair, repeated_pair in zip( + expected, + outputs, + repeated_outputs, + strict=True, + ): + for expected_tensor, actual_tensor, repeated_tensor in zip( + expected_pair, + actual_pair, + repeated_pair, + strict=True, + ): + assert torch.equal(expected_tensor.view(torch.uint8), actual_tensor.view(torch.uint8)) + assert torch.equal(actual_tensor.view(torch.uint8), repeated_tensor.view(torch.uint8)) + + +@ROCM_ONLY +@pytest.mark.parametrize("tokens", (2, 32)) +def test_vllm_rocm_rope_pair_adapter_matches_legacy_raw_bytes(tokens): + from rl_engine.integrations.vllm_runtime import _patch_strict_rocm_rotary_embedding + + generator = torch.Generator(device="cpu").manual_seed(812 + tokens) + query = torch.randn(tokens, 8 * D, dtype=torch.bfloat16, generator=generator).to(DEVICE) + key = torch.randn(tokens, 2 * D, dtype=torch.bfloat16, generator=generator).to(DEVICE) + positions = torch.arange(tokens, device=DEVICE, dtype=torch.int64) + rope = RocmDeterministicRoPEOp() + + def legacy(value): + heads = value.shape[1] // D + head_major = value.view(tokens, heads, D).permute(1, 0, 2).contiguous() + return rope(head_major, positions).permute(1, 0, 2).reshape_as(value).contiguous() + + class Rotary: + head_size = D + rotary_dim = D + + def forward_cuda(self, positions, query, key=None): + return query, key + + expected = legacy(query), legacy(key) + _patch_strict_rocm_rotary_embedding(Rotary) + actual = Rotary().forward_cuda(positions, query, key) + for expected_tensor, actual_tensor in zip(expected, actual, strict=True): + assert torch.equal(expected_tensor.view(torch.uint8), actual_tensor.view(torch.uint8)) + + +@ROCM_ONLY +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_rocm_rope_pair_backward_matches_two_single_calls_raw_bytes(dtype): + generator = torch.Generator(device="cpu").manual_seed(811) + positions = torch.tensor([9, 2, 41, 2, 77, 5, 103], device=DEVICE) + query = torch.randn(8, positions.numel(), D, dtype=dtype, generator=generator) + key = torch.randn(2, positions.numel(), D, dtype=dtype, generator=generator) + query = query.to(DEVICE) + key = key.to(DEVICE) + query_grad = torch.randn(query.shape, dtype=query.dtype, generator=generator).to(DEVICE) + key_grad = torch.randn(key.shape, dtype=key.dtype, generator=generator).to(DEVICE) + rope = RocmDeterministicRoPEOp() + + expected_inputs = tuple(tensor.detach().clone().requires_grad_() for tensor in (query, key)) + expected_outputs = tuple(rope(tensor, positions) for tensor in expected_inputs) + (expected_outputs[0].float() * query_grad.float()).sum().backward() + (expected_outputs[1].float() * key_grad.float()).sum().backward() + + actual_inputs = tuple(tensor.detach().clone().requires_grad_() for tensor in (query, key)) + actual_outputs = rope.forward_pair(actual_inputs[0], actual_inputs[1], positions) + loss = (actual_outputs[0].float() * query_grad.float()).sum() + loss = loss + (actual_outputs[1].float() * key_grad.float()).sum() + loss.backward() + + for expected_input, actual_input in zip(expected_inputs, actual_inputs, strict=True): + assert torch.equal( + expected_input.grad.contiguous().view(torch.uint8), + actual_input.grad.contiguous().view(torch.uint8), + ) + + +@ROCM_ONLY +@pytest.mark.parametrize("query_requires_grad", (False, True)) +def test_rocm_rope_pair_preserves_independent_autograd_semantics(query_requires_grad): + positions = torch.tensor([1, 3], device=DEVICE) + query = torch.randn(8, 2, D, device=DEVICE, dtype=torch.bfloat16) + key = torch.randn(2, 2, D, device=DEVICE, dtype=torch.bfloat16) + query.requires_grad_(query_requires_grad) + key.requires_grad_(not query_requires_grad) + + query_out, key_out = RocmDeterministicRoPEOp().forward_pair(query, key, positions) + + assert query_out.requires_grad is query_requires_grad + assert key_out.requires_grad is (not query_requires_grad) + + +@ROCM_ONLY +@pytest.mark.parametrize("cached_length", (32, 127, 512)) +def test_rocm_aiter_decode_direct_output_matches_staged_raw_bytes(cached_length): + from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore + + generator = torch.Generator(device="cpu").manual_seed(820 + cached_length) + q = torch.randn(1, 4, 1, D, dtype=torch.bfloat16, generator=generator).to(DEVICE) + k = torch.randn(1, 1, cached_length, D, dtype=torch.bfloat16, generator=generator).to(DEVICE) + v = torch.randn(1, 1, cached_length, D, dtype=torch.bfloat16, generator=generator).to(DEVICE) + core = StrictRocmAiterCKAttentionCore() + + with torch.inference_mode(): + staged = core.forward_with_lse(q, k, v, causal=False) + caller_out = torch.empty_like(q) + direct = core.forward_decode_with_lse_into(q, k, v, out=caller_out) + repeated_out = torch.empty_like(q) + repeated = core.forward_decode_with_lse_into(q, k, v, out=repeated_out) + + assert direct.out is caller_out + assert direct.out.data_ptr() == caller_out.data_ptr() + assert direct.provenance["core_output_staging"] == "aiter_direct_caller_group" + for expected, actual, replay in ( + (staged.out, direct.out, repeated.out), + (staged.lse, direct.lse, repeated.lse), + ): + expected_bytes = expected.contiguous().view(torch.uint8) + actual_bytes = actual.contiguous().view(torch.uint8) + replay_bytes = replay.contiguous().view(torch.uint8) + assert torch.equal(expected_bytes, actual_bytes) + assert torch.equal(actual_bytes, replay_bytes) + + def test_strict_core_rejects_split_k(): with pytest.raises(ValueError, match="Split-KV"): RLKernelDeterministicAttentionCore(split_kv=SplitKVSpec.fixed(32)) diff --git a/tests/test_framework_runtime_adapters.py b/tests/test_framework_runtime_adapters.py index 34d92e48..2fa2900a 100644 --- a/tests/test_framework_runtime_adapters.py +++ b/tests/test_framework_runtime_adapters.py @@ -37,9 +37,10 @@ ) from rl_engine.integrations.runtime import FrameworkOperatorIntegration from rl_engine.integrations.state import clear_active_integration -from rl_engine.integrations.vllm_runtime import ( - _patch_qwen3_strict_model, - _register_attention_backend, +from rl_engine.integrations.vllm_runtime import ( + _patch_qwen3_strict_model, + _patch_strict_rocm_rotary_embedding, + _register_attention_backend, configure_vllm_environment, ) from rl_engine.kernels.attention_contract import ( @@ -419,7 +420,7 @@ def get(self, tensor, *, topology): assert adapter.provenance["execution"]["runtime_platform"] == "rocm" -def test_vllm_attention_routes_paged_cache_to_rocm_runtime(monkeypatch): +def test_vllm_attention_routes_paged_cache_to_rocm_runtime(monkeypatch): runtime_calls = [] class Runtime: @@ -465,16 +466,65 @@ def get(self, tensor, *, topology): output = adapter(impl, object(), query, query, query, kv_cache, metadata) - assert output.shape == (1, 16) - assert len(runtime_calls) == 1 - assert runtime_calls[0][0].shape == (1, 2, 1, 8) - assert adapter.provenance["execution"]["runtime_platform"] == "rocm" - assert ( - adapter.provenance["execution"]["materialization"] - == "logical_paged_kv_to_aiter_ck_dense" - ) - - + assert output.shape == (1, 16) + assert len(runtime_calls) == 1 + assert runtime_calls[0][0].shape == (1, 2, 1, 8) + assert torch.equal(runtime_calls[0][3]["cu_seqlens_q"], torch.tensor([0, 1])) + assert torch.equal(runtime_calls[0][3]["kv_indptr"], torch.tensor([0, 1])) + assert adapter.provenance["execution"]["runtime_platform"] == "rocm" + assert ( + adapter.provenance["execution"]["materialization"] + == "direct_vllm_paged_kv_to_aiter_batch_prefill_ck" + ) + assert adapter.provenance["execution"]["dense_kv_materialized"] is False + + +def test_vllm_rocm_metadata_stays_on_device_and_is_reused_across_layers(monkeypatch): + original_tolist = torch.Tensor.tolist + tolist_calls = 0 + + def counting_tolist(tensor): + nonlocal tolist_calls + tolist_calls += 1 + return original_tolist(tensor) + + monkeypatch.setattr(torch.Tensor, "tolist", counting_tolist) + query = torch.zeros(2, 2, 8, dtype=torch.bfloat16) + block_table = torch.tensor([[0], [1]], dtype=torch.int32) + metadata = SimpleNamespace( + query_start_loc=torch.tensor([0, 1, 2], dtype=torch.int32), + seq_lens=torch.tensor([3, 5], dtype=torch.int32), + max_seq_len=5, + ) + adapter = VllmAttentionOperator() + first_layer = object() + second_layer = object() + + first, _ = adapter._materialization_groups( + metadata, + query=query, + block_table=block_table, + block_size=8, + num_actual=2, + cache_owner=first_layer, + ) + second, summary = adapter._materialization_groups( + metadata, + query=query, + block_table=block_table, + block_size=8, + num_actual=2, + cache_owner=second_layer, + ) + + assert first is second + assert torch.equal(first[0]["seqused_k"], torch.tensor([3, 5], dtype=torch.int32)) + assert torch.equal(first[0]["cu_seqlens_q"], torch.tensor([0, 1, 2], dtype=torch.int32)) + assert torch.equal(first[0]["kv_indptr"], torch.tensor([0, 1, 2], dtype=torch.int32)) + assert tolist_calls == 0 + assert summary["metadata_reused_across_layers"] is True + + def test_vllm_current_flash_attention_kv_cache_layout_is_materialized(): cache = torch.arange(2 * 3 * 4 * 10).reshape(2, 3, 4, 10) @@ -502,6 +552,25 @@ def test_vllm_rocm_native_kv_cache_with_two_heads_is_not_treated_as_pair_axis(): assert torch.equal(value, cache.transpose(1, 2)[..., 5:]) +def test_vllm_rocm_rlkernel_token_major_kv_cache_is_zero_copy(): + cache = torch.arange(3 * 4 * 2 * 10).reshape(3, 4, 2, 10) + + key, value = _vllm_kv_cache_views( + cache, + head_size=5, + num_kv_heads=2, + platform="rocm", + ) + + assert key.shape == (3, 4, 2, 5) + assert value.shape == (3, 4, 2, 5) + assert key.data_ptr() == cache.data_ptr() + assert value.untyped_storage().data_ptr() == cache.untyped_storage().data_ptr() + assert torch.equal(key, cache[..., :5]) + assert torch.equal(value, cache[..., 5:]) + assert key.stride(1) >= key.size(2) * key.stride(2) + + def test_vllm_rocm_kv_cache_pair_axis_is_materialized(): cache = torch.arange(3 * 2 * 4 * 1 * 5).reshape(3, 2, 4, 1, 5) @@ -654,7 +723,7 @@ def all_reduce(self, value): assert torch.equal(value.grad, torch.ones_like(value)) -def test_vllm_qwen3_strict_model_installs_without_debug_environment(monkeypatch): +def test_vllm_qwen3_strict_model_installs_without_debug_environment(monkeypatch): monkeypatch.delenv("RL_KERNEL_MODEL_DEBUG_DIR", raising=False) class RMSNorm: @@ -706,10 +775,52 @@ def __init__(self): method.apply(LinearLayer(), value), torch.full((1, 3), -1.0), ) - assert torch.equal( - norm.forward_cuda(value), - torch.nn.functional.rms_norm(value, (2,), norm.weight, 1e-6), - ) + assert torch.equal( + norm.forward_cuda(value), + torch.nn.functional.rms_norm(value, (2,), norm.weight, 1e-6), + ) + + +def test_vllm_rocm_rotary_reuses_one_table_for_query_and_key(monkeypatch): + calls = [] + + class FakeOperator: + def __call__(self, value, positions): + calls.append(("single", tuple(value.shape), positions.clone())) + return value + 1 + + def forward_pair(self, query, key, positions): + calls.append(("pair", tuple(query.shape), tuple(key.shape), positions.clone())) + return query + 2, key + 3 + + class Rotary: + head_size = 4 + rotary_dim = 4 + + def forward_cuda(self, positions, query, key=None): + return query, key + + monkeypatch.setattr(torch.version, "hip", "test") + monkeypatch.setattr( + "rl_engine.kernels.ops.rocm.rotary_embedding.rope.RocmDeterministicRoPEOp", + FakeOperator, + ) + _patch_strict_rocm_rotary_embedding(Rotary) + rotary = Rotary() + positions = torch.tensor([7, 2]) + query = torch.arange(24, dtype=torch.float32).reshape(2, 12) + key = torch.arange(8, dtype=torch.float32).reshape(2, 4) + + query_out, key_out = rotary.forward_cuda(positions, query, key) + assert torch.equal(query_out, query + 2) + assert torch.equal(key_out, key + 3) + assert calls[0][0] == "pair" + assert calls[0][1:3] == ((3, 2, 4), (1, 2, 4)) + + query_only, absent_key = rotary.forward_cuda(positions, query) + assert torch.equal(query_only, query + 1) + assert absent_key is None + assert calls[1][0] == "single" def test_vllm_logp_replaces_every_duplicate_sampled_token_column(): diff --git a/tests/test_rocm_aiter_api_contract.py b/tests/test_rocm_aiter_api_contract.py index f7ef6dce..a75d3781 100644 --- a/tests/test_rocm_aiter_api_contract.py +++ b/tests/test_rocm_aiter_api_contract.py @@ -18,7 +18,10 @@ from rl_engine.kernels.ops.rocm.attention.flash_attn import ( _AITER_BWD_POSITIONAL_CONTRACT, _AITER_BWD_REQUIRED_KEYWORDS, + _AITER_BATCH_PREFILL_POSITIONAL_CONTRACT, + _AITER_BATCH_PREFILL_REQUIRED_KEYWORDS, _AITER_FWD_POSITIONAL_CONTRACT, + _AITER_FWD_REQUIRED_KEYWORDS, StrictRocmAttentionUnavailable, _validate_aiter_schema, ) @@ -59,6 +62,7 @@ def test_positional_contract_matches_the_strict_call_sites() -> None: "return_softmax_lse", "return_dropout_randval", ) + assert "out" in _AITER_FWD_REQUIRED_KEYWORDS # The backward pins determinism positionally, so its slot must not move. assert _AITER_BWD_POSITIONAL_CONTRACT[-1] == "deterministic" assert _AITER_BWD_POSITIONAL_CONTRACT.index("softmax_lse") == 5 @@ -67,12 +71,21 @@ def test_positional_contract_matches_the_strict_call_sites() -> None: @requires_aiter def test_installed_aiter_satisfies_the_strict_contract() -> None: - _validate_aiter_schema("mha_fwd", _AITER_FWD_POSITIONAL_CONTRACT) + _validate_aiter_schema( + "mha_fwd", + _AITER_FWD_POSITIONAL_CONTRACT, + required_keywords=_AITER_FWD_REQUIRED_KEYWORDS, + ) _validate_aiter_schema( "mha_bwd", _AITER_BWD_POSITIONAL_CONTRACT, required_keywords=_AITER_BWD_REQUIRED_KEYWORDS, ) + _validate_aiter_schema( + "mha_batch_prefill", + _AITER_BATCH_PREFILL_POSITIONAL_CONTRACT, + required_keywords=_AITER_BATCH_PREFILL_REQUIRED_KEYWORDS, + ) @requires_aiter diff --git a/tests/test_rocm_lm_head_weight_cache.py b/tests/test_rocm_lm_head_weight_cache.py new file mode 100644 index 00000000..d0a6281c --- /dev/null +++ b/tests/test_rocm_lm_head_weight_cache.py @@ -0,0 +1,641 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +import rl_engine.integrations.vllm_runtime as vllm_runtime +from rl_engine.kernels.ops.rocm.matmul.det_gemm import ( + RocmDetGemmOp, + det_gemm_linear_prepared, + prepare_det_gemm_linear_weight, +) + + +class _CpuRocmTensor(torch.Tensor): + """CPU tensor that reaches the allocation-free ROCm helper validations.""" + + @staticmethod + def __new__(cls, value: torch.Tensor) -> _CpuRocmTensor: + return torch.Tensor._make_subclass(cls, value, value.requires_grad) + + @property + def is_cuda(self) -> bool: + return True + + +def _cpu_rocm_tensor(value: torch.Tensor) -> _CpuRocmTensor: + return _CpuRocmTensor(value) + + +def _prepare_on_cpu( + weight: torch.Tensor, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + transposed = weight.detach().transpose(0, 1).contiguous() + if out is None: + return transposed + with torch.no_grad(): + out.copy_(transposed) + return out + + +class _FakeLmHead(torch.nn.Module): + def __init__(self, weight: torch.Tensor, *, strict: bool = True) -> None: + super().__init__() + self.weight = torch.nn.Parameter(weight, requires_grad=False) + if strict: + setattr( + self, + vllm_runtime._STRICT_PROJECTION_MARKER, + "lm_head", + ) + + +def _linear_method_class(): + class FakeUnquantizedEmbeddingMethod: + def __init__(self) -> None: + self.original_apply_calls = [] + self.process_calls = [] + + def process_weights_after_loading(self, layer): + self.process_calls.append(layer) + return "processed" + + def apply(self, layer, x, bias=None): + self.original_apply_calls.append((layer, x, bias)) + return x + + return FakeUnquantizedEmbeddingMethod + + +class _FakeDetGemm: + def __init__(self) -> None: + self.prepared_calls = [] + self.linear_calls = [] + + def linear_prepared(self, x: torch.Tensor, weight_t: torch.Tensor) -> torch.Tensor: + self.prepared_calls.append((x, weight_t)) + return x.float().matmul(weight_t.float()) + + def linear(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + self.linear_calls.append((x, weight)) + return x.float().matmul(weight.float().transpose(0, 1)) + + +def test_prepare_det_gemm_linear_weight_creates_and_refreshes_stable_storage(): + source = _cpu_rocm_tensor(torch.arange(12, dtype=torch.bfloat16).reshape(3, 4)) + + prepared = prepare_det_gemm_linear_weight(source) + original_ptr = prepared.data_ptr() + assert prepared.shape == (4, 3) + assert prepared.is_contiguous() + assert not prepared.requires_grad + assert torch.equal(prepared, source.as_subclass(torch.Tensor).transpose(0, 1)) + + with torch.no_grad(): + source.add_(16) + refreshed = prepare_det_gemm_linear_weight(source, out=prepared) + + assert refreshed is prepared + assert refreshed.data_ptr() == original_ptr + assert torch.equal(refreshed, source.as_subclass(torch.Tensor).transpose(0, 1)) + + +@pytest.mark.parametrize( + ("source", "error", "message"), + ( + (torch.ones(4, dtype=torch.bfloat16), ValueError, "must be 2-D"), + (torch.ones(2, 3, dtype=torch.float16), TypeError, "must be BF16"), + ( + torch.ones(2, 3, dtype=torch.bfloat16), + RuntimeError, + "must be on ROCm", + ), + ( + _cpu_rocm_tensor(torch.arange(12, dtype=torch.bfloat16).reshape(4, 3).transpose(0, 1)), + ValueError, + "source deterministic linear weights must be contiguous", + ), + ), +) +def test_prepare_det_gemm_linear_weight_rejects_invalid_sources( + source: torch.Tensor, + error: type[Exception], + message: str, +): + with pytest.raises(error, match=message): + prepare_det_gemm_linear_weight(source) + + +@pytest.mark.parametrize( + ("out", "error", "message"), + ( + ( + torch.empty(3, 4, dtype=torch.bfloat16), + ValueError, + "must have shape", + ), + ( + torch.empty(4, 3, dtype=torch.float16), + TypeError, + "dtype must match", + ), + ( + torch.empty(3, 4, dtype=torch.bfloat16).transpose(0, 1), + ValueError, + "must be contiguous", + ), + ( + torch.empty(4, 3, dtype=torch.bfloat16, requires_grad=True), + ValueError, + "must not require gradients", + ), + ), +) +def test_prepare_det_gemm_linear_weight_rejects_invalid_refresh_buffers( + out: torch.Tensor, + error: type[Exception], + message: str, +): + source = _cpu_rocm_tensor(torch.ones(3, 4, dtype=torch.bfloat16)) + + with pytest.raises(error, match=message): + prepare_det_gemm_linear_weight(source, out=out) + + +def test_prepare_det_gemm_linear_weight_rejects_source_storage_overlap(): + storage = torch.arange(18, dtype=torch.bfloat16) + source = _cpu_rocm_tensor(storage[:9].reshape(3, 3)) + overlapping_out = storage[1:10].reshape(3, 3) + + with pytest.raises(ValueError, match="must not alias its source"): + prepare_det_gemm_linear_weight(source, out=overlapping_out) + + +@pytest.mark.parametrize("alias", ("activation", "weight")) +def test_det_gemm_linear_prepared_rejects_output_alias(alias): + # A square projection makes both aliases otherwise-valid output buffers. + activation = torch.ones(3, 3, dtype=torch.bfloat16) + weight_t = torch.ones(3, 3, dtype=torch.bfloat16) + out = activation if alias == "activation" else weight_t + + with pytest.raises(ValueError, match="output must not alias its inputs"): + det_gemm_linear_prepared(activation, weight_t, out=out) + + +def test_vllm_rocm_lm_head_post_load_refreshes_cache_in_place(monkeypatch): + monkeypatch.setattr(torch.version, "hip", "test-rocm") + method_cls = _linear_method_class() + det_gemm = _FakeDetGemm() + prepare_calls = [] + + def prepare(weight, *, out=None): + prepare_calls.append((weight, out)) + return _prepare_on_cpu(weight, out=out) + + vllm_runtime._patch_strict_lm_head_linear( + linear_method_cls=method_cls, + det_gemm=det_gemm, + prepare_weight=prepare, + ) + method = method_cls() + layer = _FakeLmHead(torch.arange(20, dtype=torch.bfloat16).reshape(5, 4)) + + assert method.process_weights_after_loading(layer) == "processed" + state = getattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_STATE) + cached = getattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER) + cache_ptr = cached.data_ptr() + assert method.process_calls == [layer] + assert prepare_calls == [(layer.weight, None)] + assert state.generation == 1 + assert state.valid + assert state.weight_t is cached + assert cached.data_ptr() == cache_ptr + assert vllm_runtime._validated_lm_head_weight_cache(layer) is cached + assert torch.equal(cached, layer.weight.detach().transpose(0, 1)) + assert vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER in layer._non_persistent_buffers_set + + x = torch.arange(24, dtype=torch.bfloat16).reshape(2, 3, 4) + bias = torch.arange(5, dtype=torch.float32) + output = method.apply(layer, x, bias) + expected = x.reshape(-1, 4).float().matmul(cached.float()) + expected = expected.reshape(2, 3, 5) + bias + torch.testing.assert_close(output, expected) + assert len(det_gemm.prepared_calls) == 1 + called_x, called_weight_t = det_gemm.prepared_calls[0] + assert torch.equal(called_x, x.reshape(-1, 4)) + assert called_weight_t is cached + assert det_gemm.linear_calls == [] + + replacement = torch.arange(20, 40, dtype=torch.bfloat16).reshape(5, 4) + layer.weight.data.copy_(replacement) + vllm_runtime._invalidate_lm_head_weight_cache(layer.weight) + with pytest.raises(RuntimeError, match="invalid during weight update"): + method.apply(layer, x) + + assert method.process_weights_after_loading(layer) == "processed" + refreshed_state = getattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_STATE) + assert refreshed_state is state + assert refreshed_state.generation == 2 + assert refreshed_state.valid + assert refreshed_state.weight_t is cached + assert cached.data_ptr() == cache_ptr + assert prepare_calls[-1] == (layer.weight, cached) + assert torch.equal(cached, replacement.transpose(0, 1)) + + +def test_vllm_layerwise_reload_defers_refresh_until_stable_weight_is_restored( + monkeypatch, +): + monkeypatch.setattr(torch.version, "hip", "test-rocm") + method_cls = _linear_method_class() + det_gemm = _FakeDetGemm() + prepare_calls = [] + + def prepare(weight, *, out=None): + prepare_calls.append((weight, out)) + return _prepare_on_cpu(weight, out=out) + + vllm_runtime._patch_strict_lm_head_linear( + linear_method_cls=method_cls, + det_gemm=det_gemm, + prepare_weight=prepare, + ) + method = method_cls() + layer = _FakeLmHead(torch.zeros(5, 4, dtype=torch.bfloat16)) + method.process_weights_after_loading(layer) + state = getattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_STATE) + stable_weight = layer.weight + stable_cache = state.weight_t + stable_cache_ptr = stable_cache.data_ptr() + + # vLLM checkpoint-format layerwise reload removes the stable tensors, + # materializes temporary Parameters, runs post-load processing, then copies + # back and restores the original objects. + delattr(layer, "weight") + delattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER) + replacement = torch.arange(20, dtype=torch.bfloat16).reshape(5, 4) + temporary_weight = torch.nn.Parameter(replacement.clone(), requires_grad=False) + layer.register_parameter("weight", temporary_weight) + + method.process_weights_after_loading(layer) + assert not state.valid + assert state.refresh_pending + assert len(prepare_calls) == 1 + + stable_weight.data.copy_(temporary_weight) + delattr(layer, "weight") + layer.register_parameter("weight", stable_weight) + # vLLM 43914dd74 restores saved buffers with register_buffer's default + # persistence. The cache hook must repair that metadata before state_dict + # serialization, and forward validation must preserve it thereafter. + layer.register_buffer(vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER, stable_cache) + assert vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER not in layer._non_persistent_buffers_set + assert vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER not in layer.state_dict() + assert vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER in layer._non_persistent_buffers_set + + x = torch.arange(4, dtype=torch.bfloat16).reshape(1, 4) + output = method.apply(layer, x) + expected = x.float().matmul(replacement.transpose(0, 1).float()) + torch.testing.assert_close(output, expected) + assert state.valid + assert not state.refresh_pending + assert state.generation == 2 + assert state.weight_t.data_ptr() == stable_cache_ptr + assert prepare_calls[-1] == (stable_weight, stable_cache) + assert torch.equal(stable_cache, replacement.transpose(0, 1)) + + +def test_vllm_rocm_non_lm_head_uses_original_apply(monkeypatch): + monkeypatch.setattr(torch.version, "hip", "test-rocm") + method_cls = _linear_method_class() + det_gemm = _FakeDetGemm() + vllm_runtime._patch_strict_lm_head_linear( + linear_method_cls=method_cls, + det_gemm=det_gemm, + prepare_weight=_prepare_on_cpu, + ) + method = method_cls() + layer = _FakeLmHead(torch.ones(3, 2, dtype=torch.bfloat16), strict=False) + x = torch.ones(4, 2, dtype=torch.bfloat16) + bias = torch.ones(2, dtype=torch.bfloat16) + + assert method.process_weights_after_loading(layer) == "processed" + assert not hasattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_STATE) + assert method.apply(layer, x, bias) is x + assert method.original_apply_calls == [(layer, x, bias)] + assert det_gemm.prepared_calls == [] + assert det_gemm.linear_calls == [] + + +def test_vllm_rocm_lm_head_fails_closed_for_missing_or_mutated_cache(monkeypatch): + monkeypatch.setattr(torch.version, "hip", "test-rocm") + method_cls = _linear_method_class() + det_gemm = _FakeDetGemm() + vllm_runtime._patch_strict_lm_head_linear( + linear_method_cls=method_cls, + det_gemm=det_gemm, + prepare_weight=_prepare_on_cpu, + ) + method = method_cls() + layer = _FakeLmHead(torch.ones(3, 2, dtype=torch.bfloat16)) + x = torch.ones(1, 2, dtype=torch.bfloat16) + + with pytest.raises(RuntimeError, match="was not prepared after model loading"): + method.apply(layer, x) + assert det_gemm.prepared_calls == [] + + method.process_weights_after_loading(layer) + with torch.no_grad(): + layer.weight.add_(1) + with pytest.raises(RuntimeError, match="changed without a cache refresh"): + method.apply(layer, x) + state = getattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_STATE) + assert not state.valid + assert det_gemm.prepared_calls == [] + + +def test_vllm_rocm_lm_head_fails_closed_for_mutated_or_replaced_cache(monkeypatch): + monkeypatch.setattr(torch.version, "hip", "test-rocm") + method_cls = _linear_method_class() + det_gemm = _FakeDetGemm() + vllm_runtime._patch_strict_lm_head_linear( + linear_method_cls=method_cls, + det_gemm=det_gemm, + prepare_weight=_prepare_on_cpu, + ) + method = method_cls() + layer = _FakeLmHead(torch.ones(3, 2, dtype=torch.bfloat16)) + method.process_weights_after_loading(layer) + state = getattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_STATE) + + with torch.no_grad(): + state.weight_t.add_(1) + with pytest.raises(RuntimeError, match="cache bytes changed after refresh"): + method.apply(layer, torch.ones(1, 2, dtype=torch.bfloat16)) + + state.valid = True + replacement = state.weight_t.clone() + setattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER, replacement) + with pytest.raises(RuntimeError, match="cache buffer was replaced"): + method.apply(layer, torch.ones(1, 2, dtype=torch.bfloat16)) + + +def test_vllm_rocm_lm_head_failed_lazy_refresh_cannot_reuse_old_cache(): + layer = _FakeLmHead(torch.ones(3, 2, dtype=torch.bfloat16)) + state = vllm_runtime._refresh_lm_head_weight_cache(layer, _prepare_on_cpu) + vllm_runtime._invalidate_lm_head_weight_cache(layer.weight) + vllm_runtime._mark_lm_head_weight_cache_refreshable(layer.weight) + + def fail_prepare(_weight, *, out=None): + del out + raise RuntimeError("injected refresh failure") + + with pytest.raises(RuntimeError, match="injected refresh failure"): + vllm_runtime._validated_lm_head_weight_cache(layer, fail_prepare) + assert not state.valid + assert not state.refresh_pending + with pytest.raises(RuntimeError, match="invalid during weight update"): + vllm_runtime._validated_lm_head_weight_cache(layer, _prepare_on_cpu) + + +def test_vllm_rocm_lm_head_cache_survives_level_two_buffer_restore(): + layer = _FakeLmHead(torch.arange(6, dtype=torch.bfloat16).reshape(3, 2)) + state = vllm_runtime._refresh_lm_head_weight_cache(layer, _prepare_on_cpu) + saved_buffers = {name: buffer.cpu().clone() for name, buffer in layer.named_buffers()} + + # vLLM's sleep allocator remaps storage, then wake_up restores with + # ``buffer.data.copy_``; neither operation advances the Tensor's counter. + state.weight_t.data.zero_() + for name, buffer in layer.named_buffers(): + buffer.data.copy_(saved_buffers[name].data) + + assert torch.equal(state.weight_t, layer.weight.detach().transpose(0, 1)) + assert vllm_runtime._validated_lm_head_weight_cache(layer) is state.weight_t + + +def test_qwen_weight_loader_invalidates_cache_before_data_write( + monkeypatch, +): + monkeypatch.setenv("RL_KERNEL_VLLM_REAL_VOCAB_SIZE", "6") + monkeypatch.setenv("RL_KERNEL_VLLM_PADDED_VOCAB_SIZE", "8") + + class FakeVocabParallelEmbedding: + def weight_loader(self, param, loaded_weight): + self.original_loader_calls += 1 + param.data.copy_(loaded_weight) + + class FakeParallelLMHead(FakeVocabParallelEmbedding, torch.nn.Module): + def __init__( + self, + num_embeddings, + embedding_dim=3, + *, + org_num_embeddings=None, + padding_size=None, + ): + super().__init__() + del padding_size + self.original_loader_calls = 0 + self.org_vocab_size = int(org_num_embeddings or num_embeddings) + self.num_embeddings_padded = int(num_embeddings) + self.shard_indices = SimpleNamespace( + org_vocab_start_index=0, + org_vocab_end_index=int(num_embeddings), + ) + self.weight = torch.nn.Parameter( + torch.zeros( + int(num_embeddings), + int(embedding_dim), + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + self.weight.output_dim = 0 + self.weight.packed_dim = None + + def tie_weights(self, embed_tokens): + self.weight = embed_tokens.weight + return self + + module_name = "vllm.model_executor.layers.vocab_parallel_embedding" + fake_module = ModuleType(module_name) + fake_module.ParallelLMHead = FakeParallelLMHead + fake_module.VocabParallelEmbedding = FakeVocabParallelEmbedding + monkeypatch.setitem(sys.modules, module_name, fake_module) + + vllm_runtime._patch_qwen_lm_head_padding() + layer = FakeParallelLMHead(6, 3) + state = vllm_runtime._refresh_lm_head_weight_cache(layer, _prepare_on_cpu) + cache_ptr = state.weight_t.data_ptr() + loaded = torch.arange(18, dtype=torch.bfloat16).reshape(6, 3) + + layer.weight_loader(layer.weight, loaded) + + assert not state.valid + assert layer.original_loader_calls == 0 + assert torch.equal(layer.weight[:6], loaded) + assert torch.count_nonzero(layer.weight[6:]) == 0 + refreshed_weight = vllm_runtime._validated_lm_head_weight_cache(layer, _prepare_on_cpu) + assert refreshed_weight is state.weight_t + assert state.generation == 2 + assert state.weight_t.data_ptr() == cache_ptr + assert torch.equal(state.weight_t, layer.weight.detach().transpose(0, 1)) + + embedding = torch.nn.Embedding(8, 3, dtype=torch.bfloat16) + assert layer.tie_weights(embedding) is layer + with pytest.raises(RuntimeError, match="does not support tied embeddings"): + vllm_runtime._refresh_lm_head_weight_cache(layer, _prepare_on_cpu) + with pytest.raises(RuntimeError, match="does not support tied embeddings"): + vllm_runtime._validated_lm_head_weight_cache(layer, _prepare_on_cpu) + + +def test_qwen_original_weight_loader_refreshes_only_after_success(monkeypatch): + monkeypatch.setenv("RL_KERNEL_VLLM_REAL_VOCAB_SIZE", "6") + monkeypatch.setenv("RL_KERNEL_VLLM_PADDED_VOCAB_SIZE", "8") + + class FakeVocabParallelEmbedding: + def weight_loader(self, param, loaded_weight): + self.original_loader_calls += 1 + if self.raise_during_load: + raise RuntimeError("injected loader failure") + param.data.copy_(loaded_weight) + + class FakeParallelLMHead(FakeVocabParallelEmbedding, torch.nn.Module): + def __init__(self, num_embeddings, embedding_dim=3, **_kwargs): + super().__init__() + self.original_loader_calls = 0 + self.raise_during_load = False + self.org_vocab_size = int(num_embeddings) + self.num_embeddings_padded = int(num_embeddings) + self.weight = torch.nn.Parameter( + torch.zeros(num_embeddings, embedding_dim, dtype=torch.bfloat16), + requires_grad=False, + ) + self.weight.output_dim = 0 + # A packed destination selects vLLM's original-loader branch. + self.weight.packed_dim = 0 + + module_name = "vllm.model_executor.layers.vocab_parallel_embedding" + fake_module = ModuleType(module_name) + fake_module.ParallelLMHead = FakeParallelLMHead + fake_module.VocabParallelEmbedding = FakeVocabParallelEmbedding + monkeypatch.setitem(sys.modules, module_name, fake_module) + + vllm_runtime._patch_qwen_lm_head_padding() + layer = FakeParallelLMHead(6, 3) + state = vllm_runtime._refresh_lm_head_weight_cache(layer, _prepare_on_cpu) + loaded = torch.arange(24, dtype=torch.bfloat16).reshape(8, 3) + + layer.weight_loader(layer.weight, loaded) + assert layer.original_loader_calls == 1 + assert not state.valid + assert state.refresh_pending + refreshed = vllm_runtime._validated_lm_head_weight_cache(layer, _prepare_on_cpu) + assert torch.equal(refreshed, loaded.transpose(0, 1)) + + layer.raise_during_load = True + with pytest.raises(RuntimeError, match="injected loader failure"): + layer.weight_loader(layer.weight, loaded + 1) + assert not state.valid + assert not state.refresh_pending + with pytest.raises(RuntimeError, match="invalid during weight update"): + vllm_runtime._validated_lm_head_weight_cache(layer, _prepare_on_cpu) + + +def test_qwen_compute_logits_rejects_directly_shared_embedding(monkeypatch): + monkeypatch.setattr(torch.version, "hip", "test-rocm") + + class FakeQwenForCausalLM: + def __init__(self): + shared = torch.nn.Embedding(8, 3, dtype=torch.bfloat16) + self.model = SimpleNamespace(embed_tokens=shared) + # vLLM 43914dd74 assigns this directly for Qwen tied weights; it + # does not call ParallelLMHead.tie_weights. + self.lm_head = shared + + def compute_logits(self, hidden_states): + return hidden_states + + distributed_module = ModuleType("vllm.distributed") + distributed_module.get_pp_group = lambda: SimpleNamespace(is_last_rank=True) + distributed_module.get_tp_group = lambda: SimpleNamespace(world_size=1) + qwen2_module = ModuleType("vllm.model_executor.models.qwen2") + qwen2_module.Qwen2ForCausalLM = FakeQwenForCausalLM + qwen3_module = ModuleType("vllm.model_executor.models.qwen3") + qwen3_module.Qwen3ForCausalLM = FakeQwenForCausalLM + monkeypatch.setitem(sys.modules, "vllm.distributed", distributed_module) + monkeypatch.setitem(sys.modules, "vllm.model_executor.models.qwen2", qwen2_module) + monkeypatch.setitem(sys.modules, "vllm.model_executor.models.qwen3", qwen3_module) + integration = SimpleNamespace(record_installed_hook=lambda *_args: None) + + vllm_runtime._patch_qwen_compute_logits(integration) + model = FakeQwenForCausalLM() + with pytest.raises(RuntimeError, match="does not support tied embeddings"): + model.compute_logits(torch.ones(1, 3, dtype=torch.bfloat16)) + assert getattr(model.lm_head, vllm_runtime._STRICT_LM_HEAD_TIED) + + +@pytest.mark.skipif( + torch.version.hip is None or not torch.cuda.is_available(), + reason="requires ROCm", +) +def test_rocm_layerwise_reload_uses_new_weight_with_raw_byte_identity(): + method_cls = _linear_method_class() + operator = RocmDetGemmOp() + vllm_runtime._patch_strict_lm_head_linear( + linear_method_cls=method_cls, + det_gemm=operator, + prepare_weight=prepare_det_gemm_linear_weight, + ) + generator = torch.Generator(device="cpu").manual_seed(20260906) + initial = torch.randn(1024, 256, dtype=torch.bfloat16, generator=generator) + replacement = torch.randn(1024, 256, dtype=torch.bfloat16, generator=generator) + x = torch.randn(3, 256, dtype=torch.bfloat16, generator=generator).cuda() + layer = _FakeLmHead(initial.cuda()) + method = method_cls() + + method.process_weights_after_loading(layer) + state = getattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_STATE) + stable_weight = layer.weight + stable_cache = state.weight_t + stable_cache_ptr = stable_cache.data_ptr() + with torch.inference_mode(): + initial_output = method.apply(layer, x) + initial_reference = operator.linear(x, stable_weight) + torch.cuda.synchronize() + assert torch.equal( + initial_output.view(torch.uint8), + initial_reference.view(torch.uint8), + ) + + delattr(layer, "weight") + delattr(layer, vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER) + temporary_weight = torch.nn.Parameter(replacement.cuda(), requires_grad=False) + layer.register_parameter("weight", temporary_weight) + method.process_weights_after_loading(layer) + stable_weight.data.copy_(temporary_weight) + delattr(layer, "weight") + layer.register_parameter("weight", stable_weight) + layer.register_buffer(vllm_runtime._STRICT_LM_HEAD_CACHE_BUFFER, stable_cache) + + with torch.inference_mode(): + updated_output = method.apply(layer, x) + updated_reference = operator.linear(x, stable_weight) + torch.cuda.synchronize() + assert stable_cache.data_ptr() == stable_cache_ptr + assert torch.equal( + updated_output.view(torch.uint8), + updated_reference.view(torch.uint8), + ) + assert not torch.equal( + initial_output.view(torch.uint8), + updated_output.view(torch.uint8), + ) diff --git a/tests/test_rocm_strict_paged_attention.py b/tests/test_rocm_strict_paged_attention.py index 53ef1463..a836ff96 100644 --- a/tests/test_rocm_strict_paged_attention.py +++ b/tests/test_rocm_strict_paged_attention.py @@ -37,23 +37,40 @@ def __init__(self) -> None: self.calls: list[dict[str, Any]] = [] def forward_with_lse(self, q, k, v, **kwargs) -> Any: + requested_out = kwargs.get("out") self.calls.append( { "q": q, "k": k.clone(), "v": v.clone(), + "out": requested_out, "causal": kwargs.get("causal"), "query_position_ids": kwargs.get("query_position_ids"), "key_position_ids": kwargs.get("key_position_ids"), } ) + result_out = torch.zeros( + q.size(0), + q.size(1), + q.size(2), + _HEAD_DIM, + dtype=q.dtype, + ) + if requested_out is not None: + requested_out.copy_(result_out) + result_out = requested_out + class _Result: - out = torch.zeros(q.size(0), q.size(1), q.size(2), _HEAD_DIM, dtype=q.dtype) lse = torch.zeros(q.size(0), q.size(1), q.size(2), dtype=torch.float32) provenance = {"attention_backend": "aiter.rocm.ck_dense_mha"} - return _Result() + result = _Result() + result.out = result_out + return result + + def forward_decode_with_lse_into(self, q, k, v, *, out, **kwargs) -> Any: + return self.forward_with_lse(q, k, v, out=out, causal=False, **kwargs) def _runtime() -> StrictRocmAttentionRuntime: @@ -69,7 +86,39 @@ def _cache(pages: int, kv_heads: int = 1) -> torch.Tensor: ) -def _paged_call(runtime, *, page_table, seqused_k, q_heads=1, kv_heads=1, pages=4): +def _packed_cache_views( + pages: int, + kv_heads: int = 2, + dtype: torch.dtype = torch.bfloat16, +) -> tuple[torch.Tensor, torch.Tensor]: + total = pages * kv_heads * _PAGE_SIZE * 2 * _HEAD_DIM + packed = ( + torch.arange(total, dtype=torch.float32) + .reshape(pages, kv_heads, _PAGE_SIZE, 2 * _HEAD_DIM) + .to(dtype) + ) + return packed.transpose(1, 2).split(_HEAD_DIM, dim=-1) + + +def _legacy_gather(cache, page_row, cached_length): + page_size = cache.size(1) + page_count = (cached_length + page_size - 1) // page_size + selected = cache.index_select(0, page_row[:page_count]) + flat = selected.reshape(page_count * page_size, cache.size(2), cache.size(3)) + return flat[:cached_length].permute(1, 0, 2).unsqueeze(0).contiguous() + + +def _paged_call( + runtime, + *, + page_table, + seqused_k, + q_heads=1, + kv_heads=1, + pages=4, + cached_lengths=None, + page_bounds_epoch=None, +): k_cache = _cache(pages, kv_heads) v_cache = _cache(pages, kv_heads) + 1 q = torch.zeros(page_table.size(0), q_heads, 1, _HEAD_DIM, dtype=torch.bfloat16) @@ -82,6 +131,8 @@ def _paged_call(runtime, *, page_table, seqused_k, q_heads=1, kv_heads=1, pages= seqused_k=seqused_k, max_seqlen_k=page_table.size(1) * _PAGE_SIZE, scale=None, + cached_lengths=cached_lengths, + page_bounds_epoch=page_bounds_epoch, ), k_cache, v_cache, @@ -138,6 +189,107 @@ def test_paged_decode_truncates_to_the_cached_length() -> None: assert core.calls[0]["key_position_ids"] is None +@pytest.mark.parametrize("page_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize("cache_dtype", [torch.float16, torch.bfloat16]) +def test_paged_decode_head_major_gather_matches_legacy_bytes(page_dtype, cache_dtype) -> None: + k_cache, v_cache = _packed_cache_views(4, dtype=cache_dtype) + page_row = torch.tensor([2, 0], dtype=page_dtype) + cached_length = 5 + + with torch.no_grad(): + actual_k, actual_v = StrictRocmAttentionRuntime._gather_paged_row( + k_cache, + v_cache, + page_row, + cached_length, + ) + expected_k = _legacy_gather(k_cache, page_row, cached_length) + expected_v = _legacy_gather(v_cache, page_row, cached_length) + + assert torch.equal(actual_k.contiguous().view(torch.uint8), expected_k.view(torch.uint8)) + assert torch.equal(actual_v.contiguous().view(torch.uint8), expected_v.view(torch.uint8)) + assert not actual_k.is_contiguous() + assert not actual_v.is_contiguous() + for group in range(k_cache.size(2)): + assert actual_k[:, group : group + 1].transpose(1, 2).is_contiguous() + assert actual_v[:, group : group + 1].transpose(1, 2).is_contiguous() + + +@pytest.mark.parametrize( + ("kv_heads", "cached_length"), + [(1, 5), (2, 1)], +) +def test_paged_decode_head_major_gather_keeps_small_legacy_cases_contiguous( + kv_heads, + cached_length, +) -> None: + k_cache, v_cache = _packed_cache_views(2, kv_heads=kv_heads) + page_row = torch.tensor([1, 0], dtype=torch.int32) + + with torch.no_grad(): + actual_k, actual_v = StrictRocmAttentionRuntime._gather_paged_row( + k_cache, + v_cache, + page_row, + cached_length, + ) + + assert actual_k.is_contiguous() + assert actual_v.is_contiguous() + assert torch.equal(actual_k, _legacy_gather(k_cache, page_row, cached_length)) + assert torch.equal(actual_v, _legacy_gather(v_cache, page_row, cached_length)) + + +def test_paged_decode_head_major_gather_keeps_grad_enabled_layout() -> None: + k_cache, v_cache = _packed_cache_views(2) + page_row = torch.tensor([1, 0], dtype=torch.int32) + + assert torch.is_grad_enabled() + actual_k, actual_v = StrictRocmAttentionRuntime._gather_paged_row( + k_cache, + v_cache, + page_row, + 5, + ) + + assert actual_k.is_contiguous() + assert actual_v.is_contiguous() + assert torch.equal(actual_k, _legacy_gather(k_cache, page_row, 5)) + assert torch.equal(actual_v, _legacy_gather(v_cache, page_row, 5)) + + +def test_paged_decode_head_major_gather_preserves_gradient_path_bytes() -> None: + k_source, v_source = _packed_cache_views(4) + actual_k_cache = k_source.detach().clone().requires_grad_() + actual_v_cache = v_source.detach().clone().requires_grad_() + expected_k_cache = k_source.detach().clone().requires_grad_() + expected_v_cache = v_source.detach().clone().requires_grad_() + page_row = torch.tensor([2, 0, 2], dtype=torch.int64) + cached_length = 10 + + actual_k, actual_v = StrictRocmAttentionRuntime._gather_paged_row( + actual_k_cache, + actual_v_cache, + page_row, + cached_length, + ) + expected_k = _legacy_gather(expected_k_cache, page_row, cached_length) + expected_v = _legacy_gather(expected_v_cache, page_row, cached_length) + (actual_k.float().sum() + actual_v.float().sum()).backward() + (expected_k.float().sum() + expected_v.float().sum()).backward() + + assert actual_k.is_contiguous() + assert actual_v.is_contiguous() + assert torch.equal(actual_k.view(torch.uint8), expected_k.view(torch.uint8)) + assert torch.equal(actual_v.view(torch.uint8), expected_v.view(torch.uint8)) + assert torch.equal( + actual_k_cache.grad.view(torch.uint8), expected_k_cache.grad.view(torch.uint8) + ) + assert torch.equal( + actual_v_cache.grad.view(torch.uint8), expected_v_cache.grad.view(torch.uint8) + ) + + def test_paged_decode_is_not_causal_within_a_launch() -> None: """Decode attends over the whole cached prefix, so the launch is not causal.""" @@ -177,6 +329,54 @@ def test_paged_decode_keeps_one_kv_group_per_launch() -> None: assert result.provenance["tp_degree_invariant"] is True +def test_paged_prefill_collapses_one_fresh_request_to_causal_sequence() -> None: + core = _RecordingCore() + runtime = StrictRocmAttentionRuntime(core=core) + page_table = torch.tensor([[1], [1], [1], [1]], dtype=torch.int32) + k_cache = _cache(2, kv_heads=2) + q = torch.zeros(4, 4, 1, _HEAD_DIM, dtype=torch.bfloat16) + out = torch.full_like(q, 7) + + with torch.no_grad(): + result = runtime.forward_paged_with_lse( + q, + k_cache, + k_cache + 1, + page_table=page_table, + seqused_k=torch.tensor([1, 2, 3, 4], dtype=torch.int32), + max_seqlen_k=_PAGE_SIZE, + scale=None, + out=out, + cached_lengths=(1, 2, 3, 4), + ) + + assert result.out is out + assert torch.equal(out, torch.zeros_like(out)) + assert result.lse.shape == (4, 4, 1) + assert result.lse.is_contiguous() + assert len(core.calls) == 2 + assert all(call["causal"] is True for call in core.calls) + assert all(call["q"].shape == (1, 2, 4, _HEAD_DIM) for call in core.calls) + assert all(call["k"].shape == (1, 1, 4, _HEAD_DIM) for call in core.calls) + assert torch.equal(core.calls[0]["query_position_ids"], torch.arange(4).unsqueeze(0)) + assert result.provenance["core_launch_count"] == 2 + assert result.provenance["query_schedule"] == "paged_causal_prefill_batch" + assert result.provenance["core_output_staging"] == "runtime_causal_prefill" + + +def test_paged_prefill_fails_closed_when_rows_do_not_share_pages() -> None: + runtime = _runtime() + with torch.no_grad(), pytest.raises(ValueError, match="share one logical page table"): + _paged_call( + runtime, + page_table=torch.tensor([[0], [1]], dtype=torch.int32), + seqused_k=torch.tensor([1, 2], dtype=torch.int32), + q_heads=2, + kv_heads=1, + cached_lengths=(1, 2), + ) + + def test_paged_decode_provenance_does_not_claim_a_paged_kernel() -> None: """The gather is the implementation; the provenance must say so.""" @@ -213,6 +413,221 @@ def test_paged_decode_fails_closed_on_bad_metadata(page_table, seqused_k, match) _paged_call(runtime, page_table=page_table, seqused_k=seqused_k) +def test_paged_decode_reuses_only_runtime_scoped_page_bounds_validation(monkeypatch) -> None: + runtime = _runtime() + page_table = torch.tensor([[0]], dtype=torch.int32) + seqused_k = torch.tensor([4], dtype=torch.int32) + validation_flags = [] + original = StrictRocmAttentionRuntime._gather_paged_row + + def recording_gather(*args, validate_bounds=True, **kwargs): + validation_flags.append(validate_bounds) + return original(*args, validate_bounds=validate_bounds, **kwargs) + + monkeypatch.setattr( + StrictRocmAttentionRuntime, + "_gather_paged_row", + staticmethod(recording_gather), + ) + epoch = runtime.new_page_bounds_epoch() + first, _k, _v = _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=epoch, + ) + second, _k, _v = _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=epoch, + ) + next_epoch, _k, _v = _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=runtime.new_page_bounds_epoch(), + ) + unscoped, _k, _v = _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + ) + + assert validation_flags == [True, False, True, True] + assert first.provenance["page_bounds_validation_reused"] is False + assert second.provenance["page_bounds_validation_reused"] is True + assert next_epoch.provenance["page_bounds_validation_reused"] is False + assert unscoped.provenance["page_bounds_validation_reused"] is False + + +def test_paged_decode_page_bounds_proof_fails_closed_on_metadata_mutation() -> None: + runtime = _runtime() + page_table = torch.tensor([[0]], dtype=torch.int32) + seqused_k = torch.tensor([4], dtype=torch.int32) + epoch = runtime.new_page_bounds_epoch() + _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=epoch, + ) + + page_table.fill_(9) + with pytest.raises(ValueError, match="outside"): + _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=epoch, + ) + + +def test_paged_decode_revalidates_an_equivalent_metadata_tensor(monkeypatch) -> None: + runtime = _runtime() + page_table = torch.tensor([[0]], dtype=torch.int32) + seqused_k = torch.tensor([4], dtype=torch.int32) + validation_flags = [] + original = StrictRocmAttentionRuntime._gather_paged_row + + def recording_gather(*args, validate_bounds=True, **kwargs): + validation_flags.append(validate_bounds) + return original(*args, validate_bounds=validate_bounds, **kwargs) + + monkeypatch.setattr( + StrictRocmAttentionRuntime, + "_gather_paged_row", + staticmethod(recording_gather), + ) + epoch = runtime.new_page_bounds_epoch() + _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=epoch, + ) + _paged_call( + runtime, + page_table=page_table.clone(), + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=epoch, + ) + + assert validation_flags == [True, True] + + +def test_paged_decode_validates_every_batch_row_only_once_per_epoch(monkeypatch) -> None: + runtime = _runtime() + page_table = torch.tensor([[0], [1]], dtype=torch.int32) + seqused_k = torch.tensor([4, 4], dtype=torch.int32) + validation_flags = [] + original = StrictRocmAttentionRuntime._gather_paged_row + + def recording_gather(*args, validate_bounds=True, **kwargs): + validation_flags.append(validate_bounds) + return original(*args, validate_bounds=validate_bounds, **kwargs) + + monkeypatch.setattr( + StrictRocmAttentionRuntime, + "_gather_paged_row", + staticmethod(recording_gather), + ) + epoch = runtime.new_page_bounds_epoch() + for _ in range(2): + _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4, 4), + page_bounds_epoch=epoch, + ) + + assert validation_flags == [True, True, False, False] + + +def test_paged_decode_does_not_cache_a_failed_page_bounds_validation() -> None: + runtime = _runtime() + page_table = torch.tensor([[9]], dtype=torch.int32) + seqused_k = torch.tensor([4], dtype=torch.int32) + epoch = runtime.new_page_bounds_epoch() + + with pytest.raises(ValueError, match="outside"): + _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=epoch, + ) + + page_table.zero_() + first_valid, _k, _v = _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=epoch, + ) + reused, _k, _v = _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=epoch, + ) + + assert first_valid.provenance["page_bounds_validation_reused"] is False + assert reused.provenance["page_bounds_validation_reused"] is True + + +def test_paged_decode_revalidates_inference_metadata_in_a_new_epoch() -> None: + runtime = _runtime() + with torch.inference_mode(): + page_table = torch.tensor([[0]], dtype=torch.int32) + seqused_k = torch.tensor([4], dtype=torch.int32) + _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=runtime.new_page_bounds_epoch(), + ) + + # Inference tensors have no version counter. The adapter's fresh epoch + # is therefore the fail-closed boundary between model forwards. + page_table.fill_(9) + with pytest.raises(ValueError, match="outside"): + _paged_call( + runtime, + page_table=page_table, + seqused_k=seqused_k, + cached_lengths=(4,), + page_bounds_epoch=runtime.new_page_bounds_epoch(), + ) + + +def test_paged_decode_rejects_page_bounds_epoch_from_another_runtime() -> None: + runtime = _runtime() + foreign_epoch = _runtime().new_page_bounds_epoch() + + with pytest.raises(ValueError, match="not issued by this ROCm runtime"): + _paged_call( + runtime, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + cached_lengths=(4,), + page_bounds_epoch=foreign_epoch, + ) + + def test_paged_decode_rejects_a_mismatched_out_buffer() -> None: runtime = _runtime() k_cache = _cache(2) @@ -231,6 +646,201 @@ def test_paged_decode_rejects_a_mismatched_out_buffer() -> None: ) +def test_paged_decode_writes_into_the_callers_output_buffer() -> None: + core = _RecordingCore() + runtime = StrictRocmAttentionRuntime(core=core) + k_cache = _cache(2) + q = torch.zeros(1, 1, 1, _HEAD_DIM, dtype=torch.bfloat16) + out = torch.full_like(q, 7) + + with torch.no_grad(): + result = runtime.forward_paged_with_lse( + q, + k_cache, + k_cache + 1, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + max_seqlen_k=_PAGE_SIZE, + scale=None, + out=out, + cached_lengths=(4,), + ) + + assert result.out is out + assert torch.equal(out, torch.zeros_like(out)) + assert core.calls[0]["out"].data_ptr() == out.data_ptr() + assert result.provenance["core_output_staging"] == "aiter_direct_caller_group" + + +@pytest.mark.parametrize( + ("q_heads", "kv_heads", "expected_lse_cat_parts"), + [(1, 1, []), (4, 2, [2])], +) +def test_paged_direct_decode_skips_only_singleton_lse_cats( + monkeypatch, + q_heads, + kv_heads, + expected_lse_cat_parts, +) -> None: + original_cat = torch.cat + lse_cat_parts = [] + + def recording_cat(tensors, *args, **kwargs): + tensors = tuple(tensors) + if tensors and tensors[0].dtype == torch.float32: + lse_cat_parts.append(len(tensors)) + return original_cat(tensors, *args, **kwargs) + + monkeypatch.setattr(torch, "cat", recording_cat) + runtime = _runtime() + k_cache = _cache(2, kv_heads=kv_heads) + q = torch.zeros(1, q_heads, 1, _HEAD_DIM, dtype=torch.bfloat16) + + with torch.no_grad(): + result = runtime.forward_paged_with_lse( + q, + k_cache, + k_cache + 1, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + max_seqlen_k=_PAGE_SIZE, + scale=None, + out=torch.empty_like(q), + cached_lengths=(4,), + ) + + assert lse_cat_parts == expected_lse_cat_parts + assert result.lse.shape == (1, q_heads, 1) + assert result.lse.is_contiguous() + assert torch.equal(result.lse, torch.zeros_like(result.lse)) + + +def test_paged_decode_writes_each_kv_group_directly_to_its_output_slice() -> None: + core = _RecordingCore() + runtime = StrictRocmAttentionRuntime(core=core) + k_cache = _cache(2, kv_heads=2) + q = torch.zeros(1, 4, 1, _HEAD_DIM, dtype=torch.bfloat16) + out = torch.full_like(q, 7) + + with torch.no_grad(): + result = runtime.forward_paged_with_lse( + q, + k_cache, + k_cache + 1, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + max_seqlen_k=_PAGE_SIZE, + scale=None, + out=out, + cached_lengths=(4,), + ) + + assert result.out is out + assert len(core.calls) == 2 + assert core.calls[0]["out"].data_ptr() == out[:, :2].data_ptr() + assert core.calls[1]["out"].data_ptr() == out[:, 2:].data_ptr() + assert torch.equal(out, torch.zeros_like(out)) + assert result.provenance["core_output_staging"] == "aiter_direct_caller_group" + + +def test_paged_decode_keeps_staging_when_gradient_mode_is_enabled() -> None: + core = _RecordingCore() + runtime = StrictRocmAttentionRuntime(core=core) + k_cache = _cache(2) + q = torch.zeros(1, 1, 1, _HEAD_DIM, dtype=torch.bfloat16) + out = torch.full_like(q, 7) + + result = runtime.forward_paged_with_lse( + q, + k_cache, + k_cache + 1, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + max_seqlen_k=_PAGE_SIZE, + scale=None, + out=out, + cached_lengths=(4,), + ) + + assert result.out is out + assert core.calls[0]["out"] is None + assert result.provenance["core_output_staging"] == "runtime_group_cat" + + +def test_paged_staged_decode_keeps_singleton_lse_cats(monkeypatch) -> None: + original_cat = torch.cat + lse_cat_parts = [] + + def recording_cat(tensors, *args, **kwargs): + tensors = tuple(tensors) + if tensors and tensors[0].dtype == torch.float32: + lse_cat_parts.append(len(tensors)) + return original_cat(tensors, *args, **kwargs) + + monkeypatch.setattr(torch, "cat", recording_cat) + runtime = _runtime() + k_cache = _cache(2) + q = torch.zeros(1, 1, 1, _HEAD_DIM, dtype=torch.bfloat16) + result = runtime.forward_paged_with_lse( + q, + k_cache, + k_cache + 1, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + max_seqlen_k=_PAGE_SIZE, + scale=None, + out=torch.empty_like(q), + cached_lengths=(4,), + ) + + assert lse_cat_parts == [1, 1, 1] + assert result.lse.shape == (1, 1, 1) + + +def test_paged_decode_keeps_staging_when_output_aliases_an_input() -> None: + core = _RecordingCore() + runtime = StrictRocmAttentionRuntime(core=core) + k_cache = _cache(2) + q = torch.zeros(1, 1, 1, _HEAD_DIM, dtype=torch.bfloat16) + + with torch.no_grad(): + result = runtime.forward_paged_with_lse( + q, + k_cache, + k_cache + 1, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + max_seqlen_k=_PAGE_SIZE, + scale=None, + out=q, + cached_lengths=(4,), + ) + + assert result.out is q + assert core.calls[0]["out"] is None + assert result.provenance["core_output_staging"] == "runtime_group_cat" + + +def test_paged_decode_can_skip_unused_lse_assembly() -> None: + runtime = _runtime() + k_cache = _cache(2) + q = torch.zeros(1, 1, 1, _HEAD_DIM, dtype=torch.bfloat16) + + result = runtime.forward_paged_with_lse( + q, + k_cache, + k_cache + 1, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([_PAGE_SIZE], dtype=torch.int32), + max_seqlen_k=_PAGE_SIZE, + scale=None, + return_lse=False, + ) + + assert result.lse.shape == (0,) + assert result.provenance["lse_returned"] is False + + def test_rocm_registry_does_not_claim_decode_before_a_caller_routes_to_it() -> None: """The paged entry point exists, but nothing dispatches to it yet. diff --git a/tests/test_vllm_page_bounds_epoch.py b/tests/test_vllm_page_bounds_epoch.py new file mode 100644 index 00000000..1a11d8f7 --- /dev/null +++ b/tests/test_vllm_page_bounds_epoch.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""vLLM adapter coverage for runtime-scoped ROCm page validation.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from rl_engine.integrations import framework_operators +from rl_engine.integrations.framework_operators import VllmAttentionOperator + + +def test_vllm_attention_routes_page_bounds_epoch_to_rocm_runtime(monkeypatch): + runtime_calls = [] + page_bounds_epoch = object() + + class Runtime: + def new_page_bounds_epoch(self): + return page_bounds_epoch + + def forward_paged_with_lse(self, q, k, v, **kwargs): + runtime_calls.append((q, k, v, kwargs)) + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:-1], dtype=torch.float32), + provenance={ + "actual_backend": "rlkernel.rocm.attention.aiter_ck_ag_rs.v1", + "fallback": False, + }, + ) + + runtime = Runtime() + + class Operator: + def bind_accelerator_runtime(self, tensor, *, process_group=None): + assert process_group is None + return runtime + + class Handle: + provenance = {} + + def get(self, tensor, *, topology): + assert topology["context_parallel_size"] == 1 + return Operator() + + monkeypatch.setattr( + framework_operators, + "_require_attention_accelerator", + lambda tensor: "rocm", + ) + query = torch.zeros(1, 2, 8, dtype=torch.bfloat16) + kv_cache = torch.zeros(2, 1, 4, 16, dtype=torch.bfloat16) + metadata = SimpleNamespace( + block_table=torch.tensor([[0]], dtype=torch.int32), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + seq_lens=torch.tensor([1], dtype=torch.int32), + num_actual_tokens=1, + max_seq_len=1, + ) + impl = SimpleNamespace(head_size=8, num_heads=2, num_kv_heads=1, scale=8**-0.5) + + output = VllmAttentionOperator(handle=Handle())( + impl, + object(), + query, + query, + query, + kv_cache, + metadata, + ) + + assert output.shape == (1, 16) + assert len(runtime_calls) == 1 + assert runtime_calls[0][3]["cached_lengths"] == (1,) + assert runtime_calls[0][3]["page_bounds_epoch"] is page_bounds_epoch + + +def test_vllm_page_bounds_epoch_is_scoped_to_one_model_forward(): + query = torch.zeros(2, 2, 8, dtype=torch.bfloat16) + block_table = torch.tensor([[0], [1]], dtype=torch.int32) + metadata = SimpleNamespace( + query_start_loc=torch.tensor([0, 1, 2], dtype=torch.int32), + seq_lens=torch.tensor([3, 5], dtype=torch.int32), + max_seq_len=5, + ) + adapter = VllmAttentionOperator() + layers = [object() for _ in range(36)] + issued_epochs = [] + + def issue_epoch(): + epoch = object() + issued_epochs.append(epoch) + return epoch + + common = { + "query": query, + "block_table": block_table, + "block_size": 8, + "num_actual": 2, + "include_host_lengths": True, + "page_bounds_epoch_factory": issue_epoch, + } + groups_by_layer = [] + summaries = [] + for layer in layers: + groups, summary = adapter._materialization_groups( + metadata, + cache_owner=layer, + **common, + ) + groups_by_layer.append(groups) + summaries.append(summary) + next_forward, next_summary = adapter._materialization_groups( + metadata, + cache_owner=layers[0], + **common, + ) + + first = groups_by_layer[0] + assert all(groups is first for groups in groups_by_layer[1:]) + assert summaries[0]["metadata_reused_across_layers"] is False + assert all(summary["metadata_reused_across_layers"] is True for summary in summaries[1:]) + assert next_forward is not first + assert next_forward[0]["page_bounds_epoch"] is not first[0]["page_bounds_epoch"] + assert next_summary["metadata_reused_across_layers"] is False + assert issued_epochs == [ + first[0]["page_bounds_epoch"], + next_forward[0]["page_bounds_epoch"], + ]