diff --git a/benchmarks/profile_rocm_rollout_decode.py b/benchmarks/profile_rocm_rollout_decode.py new file mode 100644 index 00000000..523c0aed --- /dev/null +++ b/benchmarks/profile_rocm_rollout_decode.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Profile one strict (or native) vLLM decode workload on ROCm with torch.profiler. + +Starts an in-process ``vllm.LLM`` with the same engine settings the Vime +launcher uses for the PR377 workload (TP4, HIP Graph ``FULL_AND_PIECEWISE``, +capture size 32, AITER FA backend), warms the strict route, then records a +profiler trace of a decode-heavy generation through ``VLLM_TORCH_PROFILER_DIR``. +Summarize the per-rank traces afterwards with ``summarize_rollout_trace.py``. + +Example:: + + RL_KERNEL_CASE=R/R RL_KERNEL_ROCM_ATTENTION_BACKEND=triton \ + HIP_VISIBLE_DEVICES=0,1,2,3 python benchmarks/profile_rocm_rollout_decode.py \ + --trace-dir /tmp/rollout-trace --prompt-tokens 4000 --decode-tokens 64 --batch 4 +""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from pathlib import Path + + +def configure_environment(case: str, trace_dir: Path, capture_size: int) -> None: + # vLLM workers are separate processes: make them import this checkout's + # rl_engine (and its vLLM plugin) rather than whatever editable install + # is registered in site-packages. + root = str(Path(__file__).resolve().parents[1]) + existing = os.environ.get("PYTHONPATH", "") + if root not in existing.split(os.pathsep): + os.environ["PYTHONPATH"] = root + (os.pathsep + existing if existing else "") + # AITER JIT builds resolve GPU_ARCHS=native to nothing inside workers. + if os.environ.get("GPU_ARCHS", "native") == "native": + os.environ["GPU_ARCHS"] = os.environ.get("PYTORCH_ROCM_ARCH", "gfx942") + os.environ.setdefault("VLLM_ROCM_USE_AITER", "1") + os.environ.setdefault("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "0") + os.environ.setdefault("VLLM_ATTENTION_BACKEND", "ROCM_AITER_FA") + os.environ.setdefault("RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE", str(capture_size)) + os.environ.setdefault("RL_KERNEL_VLLM_REAL_VOCAB_SIZE", "151936") + os.environ.setdefault("RL_KERNEL_VLLM_PADDED_VOCAB_SIZE", "152064") + os.environ.setdefault("RL_KERNEL_ROCM_FIXED_PAGED_TILE", "128") + os.environ.setdefault("RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS", "8192") + os.environ["RL_KERNEL_ATTENTION_CASE"] = case + os.environ["RL_KERNEL_FFN_CASE"] = case + os.environ["RL_KERNEL_LOGP_CASE"] = case + os.environ["RL_KERNEL_VLLM_INTEGRATION"] = "1" + os.environ.setdefault("RL_KERNEL_READBACK_DIR", str(trace_dir / "readbacks")) + os.environ.setdefault("RL_KERNEL_MISMATCH_SIDECAR_DIR", str(trace_dir / "sidecars")) + os.environ["VLLM_TORCH_PROFILER_DIR"] = str(trace_dir) + os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + for name in ( + "http_proxy", + "https_proxy", + "all_proxy", + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + ): + os.environ.pop(name, None) + Path(os.environ["RL_KERNEL_READBACK_DIR"]).mkdir(parents=True, exist_ok=True) + Path(os.environ["RL_KERNEL_MISMATCH_SIDECAR_DIR"]).mkdir(parents=True, exist_ok=True) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", default="/app/model/Qwen3-8B") + parser.add_argument("--case", default=os.environ.get("RL_KERNEL_CASE", "R/R")) + parser.add_argument("--trace-dir", type=Path, required=True) + parser.add_argument("--tensor-parallel-size", type=int, default=4) + parser.add_argument("--gpu-memory-utilization", type=float, default=0.38) + parser.add_argument("--max-model-len", type=int, default=40960) + parser.add_argument("--capture-size", type=int, default=32) + parser.add_argument("--batch", type=int, default=4) + parser.add_argument("--prompt-tokens", type=int, default=4000) + parser.add_argument("--decode-tokens", type=int, default=64) + parser.add_argument("--warmup-decode-tokens", type=int, default=16) + parser.add_argument("--max-num-seqs", type=int, default=None) + parser.add_argument("--max-num-batched-tokens", type=int, default=None) + args = parser.parse_args() + + args.trace_dir.mkdir(parents=True, exist_ok=True) + configure_environment(args.case, args.trace_dir, args.capture_size) + + import torch + from vllm import LLM, SamplingParams + + engine_limits = {} + if args.max_num_seqs is not None: + engine_limits["max_num_seqs"] = args.max_num_seqs + if args.max_num_batched_tokens is not None: + engine_limits["max_num_batched_tokens"] = args.max_num_batched_tokens + llm = LLM( + model=args.model, + **engine_limits, + tensor_parallel_size=args.tensor_parallel_size, + gpu_memory_utilization=args.gpu_memory_utilization, + max_model_len=args.max_model_len, + disable_custom_all_reduce=True, + enable_prefix_caching=True, + seed=1234, + trust_remote_code=True, + compilation_config={ + "cudagraph_mode": "FULL_AND_PIECEWISE", + "max_cudagraph_capture_size": args.capture_size, + }, + ) + generator = torch.Generator().manual_seed(7) + vocab = 150000 + prompts = [ + { + "prompt_token_ids": torch.randint( + 1000, vocab, (args.prompt_tokens,), generator=generator + ).tolist() + } + for _ in range(args.batch) + ] + warm = SamplingParams( + max_tokens=args.warmup_decode_tokens, temperature=1.0, ignore_eos=True, seed=1 + ) + t0 = time.time() + llm.generate(prompts, warm) + print(f"warmup generate: {time.time() - t0:.1f}s", flush=True) + + params = SamplingParams(max_tokens=args.decode_tokens, temperature=1.0, ignore_eos=True, seed=2) + llm.start_profile() + t0 = time.time() + outputs = llm.generate(prompts, params) + elapsed = time.time() - t0 + llm.stop_profile() + tokens = sum(len(o.outputs[0].token_ids) for o in outputs) + summary = { + "case": args.case, + "batch": args.batch, + "prompt_tokens": args.prompt_tokens, + "decode_tokens": args.decode_tokens, + "generated_tokens": tokens, + "elapsed_s": elapsed, + "step_ms_estimate": 1000.0 * elapsed / max(args.decode_tokens, 1), + "attention_backend": os.environ.get("RL_KERNEL_ROCM_ATTENTION_BACKEND", "ck"), + "det_gemm_backend": os.environ.get("RL_KERNEL_DET_GEMM_BACKEND", "auto"), + } + (args.trace_dir / "summary.json").write_text(json.dumps(summary, indent=2)) + print(json.dumps(summary, indent=2), flush=True) + time.sleep(5) # let worker profiler exports finish + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/results/pr396_rocm_mi300x/report.md b/benchmarks/results/pr396_rocm_mi300x/report.md new file mode 100644 index 00000000..2820cf3c --- /dev/null +++ b/benchmarks/results/pr396_rocm_mi300x/report.md @@ -0,0 +1,184 @@ +# PR 396: ROCm strict R/R with MFMA GEMM and chunked Triton attention + +Measurements behind https://github.com/RL-Align/RL-Kernel/pull/396 on one node +of 8 × AMD Instinct MI300X VF (gfx942), torch 2.12.0+rocm7.14, HIP 7.14.60850, +Triton 3.7.0, vLLM 0.26.1rc1, AITER with CK headers, Vime `c80200e`. +Vime rounds were produced with `examples/vime_rocm_attention_ablation/run_pr377_workload.py` +and tabulated with `summarize_pr377_runs.py`; kernel numbers are medians of +`triton.testing.do_bench` over the shapes listed in each table. + +## 1. Where the strict R/R time went + +PR 394's own trace showed the strict route ~3.9x slower than native end to end +(step 295 s vs 77 s). Two structural costs explain most of it: + +1. The gfx942 deterministic GEMM (`triton/matmul/det_gemm.py`) evaluates + scalar-FMA leaves of 32 K values and a BF16 midpoint tree over them. It + never touches the MFMA units, and every K/32 leaf writes a full `M x N` + BF16 partial to HBM. On Qwen3-8B TP4 shapes it is 3-30x slower than + hipBLASLt (table below); with 4 GEMMs per layer this alone put the decode + step near 12 ms of GEMM work and the 4096-token training microbatch near + 1 s per forward. +2. Attention decode ran the AITER/CK fixed-M128 prefill template one program + per (request, KV head) over the whole cache: ~316 us for four 7168-token + requests, 36 layers -> ~11 ms per decode step. + +## 2. MFMA batch-invariant GEMM (`rlkernel.det_gemm.triton_mfma_rocm.v1`) + +Contract: `v_mfma_f32_16x16x16_bf16` with `matrix_instr_nonkdim=16` and +`kpack=2` pinned, `BLOCK_K=64` tiles inside `CHUNK_K=1024` chunks, FP32 chunk +partials combined ascending, one BF16 rounding. Experimentally, `kpack` +changes the result bits (it changes the in-tile K order); `BLOCK_K` (32/64/128), +`BLOCK_M`, `BLOCK_N`, `num_warps`, `num_stages`, `waves_per_eu`, weight +layout, row count and the split schedule do not. `tests/test_rocm_mfma_gemm.py` +pins all of this. + +Median latency in us (k-tree = previous contract with its inference schedule): + +| shape (K x N) | M | hipBLASLt | k-tree | MFMA | +|---|---:|---:|---:|---:| +| qkv 4096x1536 | 8 | 10.2 | 81.7 | 26.0 | +| qkv 4096x1536 | 4096 | 103.3 | 3271.1 | 125.0 | +| o_proj 1024x4096 | 8 | 8.1 | 86.0 | 10.9 | +| o_proj 1024x4096 | 4096 | 72.7 | 2287.8 | 90.4 | +| gate_up 4096x6144 | 8 | 18.8 | 81.7 | 22.9 | +| gate_up 4096x6144 | 4096 | 340.6 | 15494.2 | 482.4 | +| down 3072x4096 | 8 | 11.6 | 83.4 | 27.4 | +| down 3072x4096 | 4096 | 189.1 | 6806.1 | 242.5 | +| lm_head 4096x37984 | 8 | 130.3 | 285.8 | 123.0 | +| lm_head 4096x37984 | 4096 | 2251.2 | 99358.0 | 3056.9 | + +Per decode layer (TP4, M=8) the four projections drop from ~333 us to ~87 us; +per 4096-token training forward layer from ~28 ms to ~0.94 ms. Large-M +forward stays 1.2-1.4x behind hipBLASLt; the weight gradient uses a +reduction-major copy of `dY` because a column-major A operand loads 3-4x +slower on gfx942 (offline sweep: gate_up wgrad 1463 us as a strided view vs +~480 us after the copy). + +Decode chunk sweep (M=8, sum over the four projections): `CHUNK_K=1024/BLOCK_K=64` +79.3 us; 512/64 84.5 us; 256/64 84.4 us; 1024/32 87.7 us. + +## 3. Chunked Triton attention (`rlkernel.rocm.triton_chunked_flash_attention.v1`) + +Contract: 64-key blocks inside 512-token KV chunks, each chunk's online +softmax from an empty state, chunks merged ascending with the exact FA2 +rescale, fully masked blocks/chunks leave a row untouched, `P` rounded to the +input dtype before `P.V`, hardware `exp2`, FP contraction off. The +monolithic schedule (query tiles) and the split schedule (one program per +sequence/KV head/chunk plus an ascending merge) are bit-identical, so a +Megatron full-sequence forward, a vLLM prefill, a prefix-cached extend and a +single-token decode agree exactly (`tests/test_rocm_triton_chunked_attention.py`). + +Eager single-call latency in us (HQ=8, HKV=2, D=128, 16-token pages; +Triton with the unmasked fast path for fully visible key blocks): + +| case | Triton chunked | CK fixed M128 | +|---|---:|---:| +| prefill 1 x 4096 | 270 | 179 | +| prefill 1 x 1024 | 56 | 53 | +| prefill 4 x 1024 | 93 | 59 | +| decode 4 x 7168 | 68 | 316 | +| decode 4 x 2048 | 68 | 95 | +| decode 8 x 4096 | 69 | 184 | + +Training core (`StrictRocmAiterCKAttentionCore.forward_with_lse`) forward / +forward+backward in ms, B=1, AITER deterministic backward in both cases: + +| S | CK fwd | CK fwd+bwd | Triton fwd | Triton fwd+bwd | peak | +|---:|---:|---:|---:|---:|---:| +| 4096 | 0.67 | 3.15 | 0.84 | 2.51 | 4.1 GiB | +| 8192 | 0.98 | 7.62 | 1.51 | 6.46 | 16.2 GiB | + +The backward dominates training attention and is unchanged; the deterministic +AITER backward is O(S^2) in workspace (16 GiB at S=8192). + +## 4. End-to-end, PR377 workload, one round + +Qwen3-8B, actor TP4/CP2/PP1, two TP4 vLLM engines, 8 samples, 7168-token +response limit, 4096 training tokens/GPU, seeds 1234, HIP Graph +FULL_AND_PIECEWISE (capture 32), `RL_KERNEL_ROCM_FIXED_PAGED_TILE=128`. + +### 4a. Healthy node after host GPU reset (vLLM memory utilization 0.38) + +Primary numbers. Health check before the pair: 8-GPU copy 3.82-3.90 TB/s, +decode GEMM 23-28 us, prefill GEMM 335-360 us, TP4 all-reduce 1 MB 43-53 us. +Runs `mxs-pair-v6-p-p` and `mxs-pair-v6-r-r-triton`. + +| Config | Mismatch Count | Max \|dlogp\| | torch.equal | +|---|---:|---:|:---:| +| P/P native | 23294 / 42042 | 2.342051 | false | +| R/R strict (MFMA GEMM + Triton attention) | **0 / 28652** | **0** | **true** | + +| Metric | P/P native | R/R (Triton attn) | R/R vs P/P | +|---|---:|---:|---:| +| rollout time | **56.31 s** | 67.17 s | 19.3% slower | +| effective tokens/GPU/s | **93.32** | 53.32 | 42.9% lower | +| update weights | 2.59 s | **1.18 s** | **54.6% faster** | +| log probs | 8.71 s | **7.19 s** | **17.4% faster** | +| actor train | 14.92 s | **10.67 s** | **28.5% faster** | +| train time | 24.31 s | **18.36 s** | **24.5% faster** | +| actor train tok/s | 2886.6 | 2779.7 | 3.7% lower | +| end-to-end step | **83.69 s** | 87.71 s | 4.8% slower | + +Mean sampled response length was 5255 tokens (P/P) vs 3581 (R/R), so the per-token rollout throughput (93 vs 53 tokens/GPU/s, 1.75x) is the honest measure of the remaining gap, not the 19% rollout-time difference. Against the #394 R/R baseline on a healthy node (round 0: 39 tokens/GPU/s, log probs 25.5 s, actor train 50.1 s, step 183.0 s) this branch is 1.36x faster in rollout throughput and 4.7x faster in actor train. + +### 4b. Same degraded node, same conditions (vLLM memory utilization 0.30) + +| Config | Mismatch Count | Max \|dlogp\| | torch.equal | +|---|---:|---:|:---:| +| P/P native | 20305 / 37034 | 2.840120 | false | +| R/R strict, MFMA GEMM + CK attention | 0 / 33320 | 0 | true | +| R/R strict, MFMA GEMM + Triton attention | 0 / 28652 | 0 | true | + +| Metric | P/P native | R/R (CK attn) | R/R (Triton attn) | Triton R/R vs P/P | +|---|---:|---:|---:|---:| +| rollout time | 172.13 s | 336.20 s | 312.16 s | 81.4% slower | +| effective tokens/GPU/s | 26.89 | 12.39 | 11.47 | 57.3% lower | +| update weights | 7.39 s | 3.53 s | 8.35 s | 13.0% slower | +| log probs | 31.32 s | 44.33 s | 37.63 s | 20.1% slower | +| actor train | 63.25 s | 54.56 s | 48.15 s | 23.9% faster | +| train time | 95.58 s | 99.45 s | 86.38 s | 9.6% faster | +| actor train tok/s | 601.6 | 629.3 | 616.2 | 2.4% higher | +| end-to-end step | 276.24 s | 441.30 s | 409.79 s | 48.3% slower | + +### 4c. Healthy-node points measured before the degradation + +| Run | rollout | log probs | actor train | step | consistency | +|---|---:|---:|---:|---:|---| +| P/P native (v167 round 0, morning) | 42.14 s | 9.44 s | 14.07 s | 69.02 s | 94330/133216 over 3 rounds, false | +| R/R PR 394 baseline (v165 round 0) | 104.87 s | 25.46 s | 50.05 s | 182.98 s | 0/130954, true | +| R/R this branch, MFMA + CK attn | 132.01 s | 21.00 s | 11.59 s | 169.08 s | 0/33320, true | + +Response lengths differ between arms (the sampled tokens differ because the +arithmetic differs), which is why rollout time is best read together with +tokens/GPU/s: 98.4 (P/P), 39.3 (PR 394 R/R), 31.6 (this branch, CK attn). + +## 5. Node degradation and the tools added for it + +During this work the node degraded: a native P/P round went from 69 s to +276 s with every phase ~4x slower, while single-GPU copy bandwidth +(3.85 TB/s) and GEMM latencies stayed unchanged. The +container's PID 1 is `sleep infinity`, so killed Ray/vLLM workers become +permanent zombies; reaping them (a ptrace-injected `wait4` on PID 1) removed +2285 zombies but the driver still lists the GPU contexts of eight killed vLLM +workers (73 GB and 11 GB per GPU, hardware queues still mapped) under +`/sys/class/kfd/kfd/proc/`. A host-side GPU reset restored the node (section 4a was measured after +it); the absolute numbers in 4b are inflated for every arm, the relative +ones are same-condition. + +Two harness pitfalls fixed on the way: operator-shell `http_proxy` variables +inherited by Ray turned every generation longer than 30 s into a 502 retry +loop, and AITER's JIT resolved `GPU_ARCHS=native` to an empty offload list +inside vLLM workers (compiled for gfx906 and failed); `launch_arm.sh` now +strips the proxy variables and pins `GPU_ARCHS`. + +## 6. What is still missing for a net win + +- Rollout decode is ~1.75x native per token even though the GEMMs are + within 1.3x of hipBLASLt and the attention core is faster than CK. The + decode step runs inside a full HIP graph, so the gap has to be attributed + kernel by kernel. `profile_rocm_rollout_decode.py` + `summarize_rollout_trace.py` + capture per-rank kernel traces of one decode workload for this analysis. +- Training attention at CP2 still all-gathers Q/K/V and computes the whole + sequence on every CP rank; the per-row-invariant kernel allows computing + only the local zigzag rows against the gathered K/V. diff --git a/benchmarks/summarize_rollout_trace.py b/benchmarks/summarize_rollout_trace.py new file mode 100644 index 00000000..0950a543 --- /dev/null +++ b/benchmarks/summarize_rollout_trace.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Aggregate GPU kernel time from torch.profiler traces exported by vLLM workers. + +Example:: + + python benchmarks/summarize_rollout_trace.py /tmp/rollout-trace --top 40 +""" + +from __future__ import annotations + +import argparse +import collections +import gzip +import json +import re +from pathlib import Path + + +def load_trace(path: Path) -> dict: + opener = gzip.open if path.suffix == ".gz" else open + with opener(path, "rt") as handle: + return json.load(handle) + + +def summarize(path: Path, top: int, name_width: int = 110) -> None: + trace = load_trace(path) + events = trace["traceEvents"] + agg: dict[str, list[float]] = collections.defaultdict(lambda: [0.0, 0]) + gpu_busy = 0.0 + first = float("inf") + last = 0.0 + kernel_events = [] + for event in events: + if event.get("ph") != "X": + continue + cat = event.get("cat", "") + if cat in ("kernel", "gpu_memcpy", "gpu_memset", "Kernel"): + name = re.sub(r"\(.*$", "", event["name"])[:name_width] + agg[name][0] += event["dur"] + agg[name][1] += 1 + gpu_busy += event["dur"] + first = min(first, event["ts"]) + last = max(last, event["ts"] + event["dur"]) + kernel_events.append((event["ts"], event["dur"])) + if not kernel_events: + print(f"{path.name}: no GPU kernel events") + return + kernel_events.sort() + idle = 0.0 + cursor = kernel_events[0][0] + for ts, dur in kernel_events: + if ts > cursor: + idle += ts - cursor + cursor = max(cursor, ts + dur) + wall = last - first + print( + f"{path.name}: wall={wall / 1e3:.1f}ms gpu_busy={gpu_busy / 1e3:.1f}ms " + f"({100 * gpu_busy / wall:.1f}%) idle_gaps={idle / 1e3:.1f}ms kernels={len(kernel_events)}" + ) + for name, (dur, count) in sorted(agg.items(), key=lambda kv: -kv[1][0])[:top]: + print(f" {dur / 1e3:9.2f}ms {100 * dur / gpu_busy:5.1f}% {count:7d} {name}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("trace_dir", type=Path) + parser.add_argument("--top", type=int, default=30) + parser.add_argument("--rank", type=int, default=None, help="only summarize this rank") + args = parser.parse_args() + paths = sorted(args.trace_dir.rglob("*.pt.trace.json*")) + if not paths: + raise SystemExit(f"no traces under {args.trace_dir}") + for path in paths: + if ( + args.rank is not None + and f"rank{args.rank}" not in path.name + and f"_{args.rank}" not in path.name + ): + continue + summarize(path, args.top) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/vime_rocm_attention_ablation/launch_arm.sh b/examples/vime_rocm_attention_ablation/launch_arm.sh index c51cc630..b04a2294 100644 --- a/examples/vime_rocm_attention_ablation/launch_arm.sh +++ b/examples/vime_rocm_attention_ablation/launch_arm.sh @@ -118,6 +118,10 @@ do done unset CUBLASLT_WORKSPACE_SIZE CUBLAS_WORKSPACE_CONFIG NCCL_ALGO +# Vime's rollout client (httpx) and the router honor proxy variables. A host +# proxy inherited from the operator's shell turns every generation longer than +# the proxy's upstream timeout into a 502 retry loop, so never forward them. +unset http_proxy https_proxy all_proxy HTTP_PROXY HTTPS_PROXY ALL_PROXY no_proxy NO_PROXY unset RL_KERNEL_CUDA_ONLY RL_KERNEL_DET_GEMM_SM90_ONLY RL_KERNEL_PRECOMPILE_FA4 unset VLLM_BATCH_INVARIANT NVTE_FUSED_ATTN NVTE_FLASH_ATTN NVTE_UNFUSED_ATTN @@ -128,6 +132,10 @@ export CUDA_VISIBLE_DEVICES="${HIP_VISIBLE_DEVICES}" export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1 export PYTORCH_ROCM_ARCH="${PYTORCH_ROCM_ARCH:-gfx942}" +# AITER's JIT resolves GPU_ARCHS=native to an empty offload list inside vLLM +# workers and then compiles CK for the compiler default (gfx906), which fails. +# Pin the build target to the same architecture PyTorch targets. +export GPU_ARCHS="${PYTORCH_ROCM_ARCH}" export PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" export CUDA_DEVICE_MAX_CONNECTIONS=1 export NCCL_NVLS_ENABLE=0 @@ -184,6 +192,7 @@ names = [ "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES", "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", "PYTORCH_ROCM_ARCH", + "GPU_ARCHS", "PYTORCH_ALLOC_CONF", "CUDA_DEVICE_MAX_CONNECTIONS", "NCCL_NVLS_ENABLE", @@ -203,7 +212,12 @@ names = [ "RL_KERNEL_MISMATCH_SIDECAR_DIR", ] env_vars = {name: os.environ[name] for name in names} -for name in ("RL_KERNEL_ROCM_FIXED_PAGED_TILE", "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS"): +for name in ( + "RL_KERNEL_ROCM_FIXED_PAGED_TILE", + "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS", + "RL_KERNEL_DET_GEMM_BACKEND", + "RL_KERNEL_ROCM_ATTENTION_BACKEND", +): if name in os.environ: env_vars[name] = os.environ[name] print(json.dumps({"env_vars": env_vars})) diff --git a/examples/vime_rocm_attention_ablation/run.py b/examples/vime_rocm_attention_ablation/run.py index b3c31524..9a757541 100644 --- a/examples/vime_rocm_attention_ablation/run.py +++ b/examples/vime_rocm_attention_ablation/run.py @@ -42,6 +42,18 @@ CASE_ORDER = ("P/P", "P/R", "R/P", "R/R") RL_KERNEL_PLUGIN_ENTRY_POINT = "rl_engine.integrations.vllm_runtime:register_vllm_plugin" +# Proxy settings inherited from an operator shell break long rollout requests +# (httpx and the router honor them; a proxy's upstream timeout yields 502s). +_PROXY_ENVIRONMENT = ( + "http_proxy", + "https_proxy", + "all_proxy", + "no_proxy", + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + "NO_PROXY", +) _CUDA_ONLY_ENVIRONMENT = ( "CUBLASLT_WORKSPACE_SIZE", "CUBLAS_WORKSPACE_CONFIG", @@ -482,6 +494,8 @@ def build_arm_environment( env = dict(os.environ if base_environment is None else base_environment) for name in _CUDA_ONLY_ENVIRONMENT: env.pop(name, None) + for name in _PROXY_ENVIRONMENT: + env.pop(name, None) existing_pythonpath = env.get("PYTHONPATH", "") python_paths = [ str((config.rl_kernel_root / "examples").resolve()), @@ -561,6 +575,8 @@ def public_arm_environment(environment: Mapping[str, str]) -> dict[str, str]: "RL_KERNEL_ATTENTION_CASE", "RL_KERNEL_ROCM_FIXED_PAGED_TILE", "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS", + "RL_KERNEL_DET_GEMM_BACKEND", + "RL_KERNEL_ROCM_ATTENTION_BACKEND", "RL_KERNEL_FFN_CASE", "RL_KERNEL_LOGP_CASE", "RL_KERNEL_VLLM_REAL_VOCAB_SIZE", diff --git a/examples/vime_rocm_attention_ablation/run_pr377_workload.py b/examples/vime_rocm_attention_ablation/run_pr377_workload.py new file mode 100644 index 00000000..3ba83875 --- /dev/null +++ b/examples/vime_rocm_attention_ablation/run_pr377_workload.py @@ -0,0 +1,289 @@ +"""Run one P/P or R/R arm of the PR377 Qwen3-8B TP4/CP2 workload on ROCm. + +The P/P arm selects production attention, FFN and logp on both frameworks and +uses Vime's rollout-logprob consistency mode. The R/R arm selects the strict +RL-Kernel route on both frameworks and validates bitwise train/rollout +agreement. Paths and the round count are CLI arguments so the same script +serves every machine layout; nothing is hard-coded. + +Example:: + + python -m examples.vime_rocm_attention_ablation.run_pr377_workload \ + --case R/R --num-rollout 3 \ + --run-dir /app/model/vime-runs/mfma-rr-3round \ + --rl-kernel-root /work/RL-Kernel --vime-root /work/vime \ + --megatron-root /work/Megatron-LM-vime +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +from examples.vime_rocm_attention_ablation.run import ( + MatrixConfig, + _canonical_fingerprint, + _prepare_run_dir, + build_arm_environment, + frozen_input_manifest, + public_arm_environment, +) +from examples.vime_rocm_attention_ablation.validate_artifacts import ( + CASE_IMPLEMENTATIONS, + compare_train_rollout_logps, + load_readbacks, + load_rollout_identity, + validate_arm, + write_report, +) + +# Strict-path knobs forwarded into the Ray runtime environment when set. +FORWARDED_STRICT_ENVIRONMENT = ( + "RL_KERNEL_DET_GEMM_BACKEND", + "RL_KERNEL_ROCM_FIXED_PAGED_TILE", + "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS", + "RL_KERNEL_ROCM_ATTENTION_BACKEND", +) + + +class WorkloadConfig(MatrixConfig): + case_id: str = "R/R" + + def frozen_parameters(self): + value = super().frozen_parameters() + value["ffn_case"] = self.case_id + value["logp_case"] = self.case_id + if self.case_id == "P/P": + value["framework_consistency"] = { + "use_rollout_logprobs": True, + "get_mismatch_metrics": True, + "custom_tis_function": ( + "vime_rocm_attention_ablation.tis_metrics.metrics_only_tis" + ), + } + return value + + +def sealed_manifest(config: MatrixConfig): + value = frozen_input_manifest(config) + value["fingerprint"] = _canonical_fingerprint( + {key: item for key, item in value.items() if key != "fingerprint"} + ) + return value + + +def validate_native_readbacks(readback_dir: Path): + errors = [] + frameworks = set() + paths = [] + for record in load_readbacks(readback_dir): + paths.append(record["_path"]) + framework = record.get("framework") + if framework in {"megatron", "vllm"}: + frameworks.add(framework) + if record.get("fallbacks"): + errors.append(f"{record['_path']}: unexpected adapter fallback") + operators = record.get("operators", {}) + for module in ("attention", "ffn", "logp"): + operator = operators.get(module) + if not isinstance(operator, dict): + errors.append(f"{record['_path']}: missing {module} readback") + continue + if operator.get("case_id") != "P/P": + errors.append(f"{record['_path']}: {module} case is {operator.get('case_id')!r}") + if operator.get("implementation") != "production": + errors.append( + f"{record['_path']}: {module} implementation is " + f"{operator.get('implementation')!r}" + ) + if frameworks != {"megatron", "vllm"}: + errors.append(f"readback frameworks are {sorted(frameworks)!r}") + return {"passed": not errors, "errors": errors, "paths": paths} + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--case", choices=["P/P", "R/R"], required=True) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--num-rollout", type=int, default=3) + parser.add_argument("--rl-kernel-root", type=Path, default=Path(__file__).resolve().parents[2]) + parser.add_argument( + "--vime-root", type=Path, default=Path(os.environ.get("VIME_ROOT", "/work/vime")) + ) + parser.add_argument( + "--megatron-root", + type=Path, + default=Path(os.environ.get("MEGATRON_ROOT", "/work/Megatron-LM-vime")), + ) + parser.add_argument("--model-root", type=Path, default=Path("/app/model/Qwen3-8B")) + parser.add_argument( + "--reference-checkpoint", type=Path, default=Path("/app/model/Qwen3-8B_torch_dist") + ) + parser.add_argument( + "--prompt-data", type=Path, default=Path("/app/model/dapo-math-17k/dapo-math-17k.jsonl") + ) + parser.add_argument("--samples-per-prompt", type=int, default=8) + parser.add_argument("--global-batch-size", type=int, default=8) + parser.add_argument("--max-response-length", type=int, default=7168) + parser.add_argument("--max-tokens-per-gpu", type=int, default=4096) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--rollout-seed", type=int, default=1234) + parser.add_argument("--fixed-paged-tile", default="128") + parser.add_argument("--paged-kv-max-tokens", default="8192") + parser.add_argument("--vllm-gpu-memory-utilization", default="0.38") + parser.add_argument("--ray-port", type=int, default=6385) + parser.add_argument("--ray-dashboard-port", type=int, default=28265) + return parser.parse_args(argv) + + +def main(argv=None) -> int: + args = parse_args(argv) + case_id = args.case + root = args.rl_kernel_root.resolve() + config = WorkloadConfig( + vime_root=args.vime_root, + rl_kernel_root=root, + megatron_root=args.megatron_root, + model_root=args.model_root, + reference_checkpoint=args.reference_checkpoint, + prompt_data=args.prompt_data, + run_dir=args.run_dir, + launcher=root / "examples/vime_rocm_attention_ablation/launch_arm.sh", + num_rollout=args.num_rollout, + rollout_batch_size=1, + samples_per_prompt=args.samples_per_prompt, + global_batch_size=args.global_batch_size, + max_response_length=args.max_response_length, + max_tokens_per_gpu=args.max_tokens_per_gpu, + seed=args.seed, + rollout_seed=args.rollout_seed, + ray_port=args.ray_port, + ray_dashboard_port=args.ray_dashboard_port, + ) + config.case_id = case_id + config.validate(require_paths=True) + _prepare_run_dir(args.run_dir) + frozen_before = sealed_manifest(config) + write_report(args.run_dir / "frozen-inputs.before.json", frozen_before) + + arm_slug = case_id.lower().replace("/", "-") + arm_dir = args.run_dir / "arms" / arm_slug + for directory in ( + arm_dir / "readbacks", + arm_dir / "dump", + arm_dir / "checkpoint", + arm_dir / "mismatch_sidecars", + ): + directory.mkdir(parents=True, exist_ok=False) + + environment = build_arm_environment( + config, case_id, arm_dir, arm_index=0 if case_id == "P/P" else 3 + ) + environment.update( + { + "RL_KERNEL_ATTENTION_CASE": case_id, + "RL_KERNEL_FFN_CASE": case_id, + "RL_KERNEL_LOGP_CASE": case_id, + "RL_KERNEL_ROCM_FIXED_PAGED_TILE": args.fixed_paged_tile, + "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS": args.paged_kv_max_tokens, + "VLLM_GPU_MEMORY_UTILIZATION": args.vllm_gpu_memory_utilization, + } + ) + if case_id == "P/P": + environment["RLK_ABLATION_USE_ROLLOUT_LOGPROBS"] = "1" + launch = { + "schema_version": "rlkernel.vime_rocm_attention_arm_launch.v1", + "case_id": case_id, + "expected_implementations": CASE_IMPLEMENTATIONS[case_id], + "framework_consistency": ( + {"use_rollout_logprobs": True, "mismatch_metrics_recompute": True} + if case_id == "P/P" + else {"use_rollout_logprobs": False, "strict_linear_logp": True} + ), + "strict_environment": { + name: environment[name] for name in FORWARDED_STRICT_ENVIRONMENT if name in environment + }, + "frozen_input_fingerprint": frozen_before["fingerprint"], + "command": ["bash", str(config.launcher.resolve())], + "environment": public_arm_environment(environment), + "started_at": datetime.now(timezone.utc).isoformat(), + } + write_report(arm_dir / "launch.json", launch) + + with (arm_dir / "launcher.log").open("w", encoding="utf-8") as log_handle: + process = subprocess.run( + ["bash", str(config.launcher.resolve())], + cwd=config.rl_kernel_root, + env=environment, + stdout=log_handle, + stderr=subprocess.STDOUT, + check=False, + ) + + if case_id == "R/R": + report = validate_arm(arm_dir, case_id, launcher_returncode=process.returncode) + else: + errors = [] + if process.returncode: + errors.append(f"Vime launcher exited with status {process.returncode}") + try: + readbacks = validate_native_readbacks(arm_dir / "readbacks") + except Exception as exc: # pragma: no cover - runtime evidence failure + readbacks = {"passed": False, "errors": [str(exc)], "paths": []} + errors.extend(readbacks["errors"]) + rollout_identity = load_rollout_identity(arm_dir / "dump" / "rollout_data") + errors.extend(rollout_identity["errors"]) + try: + metrics = compare_train_rollout_logps( + arm_dir / "mismatch_sidecars", + require_exact=False, + tensor_parallel_size=config.tensor_parallel_size, + context_parallel_size=config.context_parallel_size, + ) + except Exception as exc: # pragma: no cover - runtime evidence failure + metrics = {"passed": False, "errors": [str(exc)]} + errors.extend(metrics["errors"]) + report = { + "case_id": case_id, + "launcher_returncode": process.returncode, + "passed": not errors, + "errors": errors, + "readbacks": readbacks, + "rollout_identity": rollout_identity, + "metrics": metrics, + } + frozen_after = sealed_manifest(config) + write_report(args.run_dir / "frozen-inputs.after.json", frozen_after) + frozen_match = frozen_before["fingerprint"] == frozen_after["fingerprint"] + if not frozen_match: + report["errors"] = list(report.get("errors", [])) + [ + "frozen source fingerprint changed during the run" + ] + report["passed"] = False + report["frozen_sources_match"] = frozen_match + report["run_dir"] = str(args.run_dir) + report["num_rollout"] = args.num_rollout + write_report(arm_dir / "validation.json", report) + summary = { + "run_dir": str(args.run_dir), + "case_id": case_id, + "num_rollout": args.num_rollout, + "launcher_returncode": process.returncode, + "passed": report["passed"], + "errors": report["errors"], + "metrics": report.get("metrics"), + "strict_environment": launch["strict_environment"], + "frozen_sources_match": frozen_match, + } + write_report(args.run_dir / "single-arm-summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True, default=str), flush=True) + return 0 if report["passed"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/vime_rocm_attention_ablation/summarize_pr377_runs.py b/examples/vime_rocm_attention_ablation/summarize_pr377_runs.py new file mode 100644 index 00000000..8a1715db --- /dev/null +++ b/examples/vime_rocm_attention_ablation/summarize_pr377_runs.py @@ -0,0 +1,181 @@ +"""Summarize P/P and R/R PR377-workload runs into consistency and performance tables. + +Reads the ``perf N:`` dictionaries Vime prints into ``launcher.log`` for each +arm, averages them across rounds and renders the same two tables used for the +CUDA PR377 report: an exactness table (mismatch count, max |dlogp|, +``torch.equal``) and a per-phase mean-time table with the R/R-relative-to-P/P +column. + +Example:: + + python -m examples.vime_rocm_attention_ablation.summarize_pr377_runs \ + --pp-run /app/model/vime-runs/pp-30round --rr-run /app/model/vime-runs/rr-30round +""" + +from __future__ import annotations + +import argparse +import ast +import json +import re +import sys +from pathlib import Path + +PERF_LINE = re.compile(r"perf (\d+): (\{.*\})\s*$") + +ROWS = [ + ("rollout time", "perf/rollout_time", "s", "lower"), + ("effective tokens/GPU/s", "perf/effective_tokens_per_gpu_per_sec", "", "higher"), + ("update weights", "perf/update_weights_time", "s", "lower"), + ("reference log probs", "perf/ref_log_probs_time", "s", "lower"), + ("log probs", "perf/log_probs_time", "s", "lower"), + ("actor train", "perf/actor_train_time", "s", "lower"), + ("train time", "perf/train_time", "s", "lower"), + ("actor train tok/s", "perf/actor_train_tok_per_s", "", "higher"), + ("end-to-end step", "perf/step_time", "s", "lower"), +] + + +def load_rounds(run_dir: Path) -> dict[int, dict[str, float]]: + arms = sorted((run_dir / "arms").glob("*")) + if not arms: + raise FileNotFoundError(f"no arms under {run_dir}") + log = arms[0] / "launcher.log" + rounds: dict[int, dict[str, float]] = {} + for line in log.read_text(encoding="utf-8", errors="replace").splitlines(): + match = PERF_LINE.search(line) + if not match: + continue + payload = ast.literal_eval(match.group(2)) + bucket = rounds.setdefault(int(match.group(1)), {}) + for key, value in payload.items(): + if key.startswith("perf/") and isinstance(value, (int, float)): + bucket[key] = float(value) + return rounds + + +def load_metrics(run_dir: Path) -> dict: + arms = sorted((run_dir / "arms").glob("*")) + validation = json.loads((arms[0] / "validation.json").read_text(encoding="utf-8")) + metrics = validation.get("metrics", {}) or {} + return { + "passed": validation.get("passed"), + "errors": validation.get("errors", []), + "metrics": metrics, + } + + +def _metric(metrics: dict, *names, default=None): + for name in names: + if name in metrics: + return metrics[name] + aggregate = metrics.get("aggregate") or metrics.get("summary") or {} + for name in names: + if name in aggregate: + return aggregate[name] + return default + + +def mean_rows(rounds: dict[int, dict[str, float]], keys: list[str]) -> dict[str, float | None]: + out: dict[str, float | None] = {} + for key in keys: + values = [r[key] for r in rounds.values() if key in r] + out[key] = sum(values) / len(values) if values else None + return out + + +def fmt(value, unit): + if value is None: + return "n/a" + return f"{value:.6f}{(' ' + unit) if unit else ''}" + + +def relative(pp, rr, direction): + if pp is None or rr is None or pp == 0: + return "n/a" + if direction == "lower": + delta = (rr - pp) / pp + return f"快 {-delta * 100:.2f}%" if delta < 0 else f"慢 {delta * 100:.2f}%" + delta = (rr - pp) / pp + return f"高 {delta * 100:.2f}%" if delta > 0 else f"低 {-delta * 100:.2f}%" + + +def render(pp_dir: Path, rr_dir: Path) -> str: + pp_rounds = load_rounds(pp_dir) + rr_rounds = load_rounds(rr_dir) + pp_info = load_metrics(pp_dir) + rr_info = load_metrics(rr_dir) + n_pp, n_rr = len(pp_rounds), len(rr_rounds) + lines = [] + lines.append(f"## {min(n_pp, n_rr)} 轮一致性结果\n") + lines.append("| 配置 | Mismatch Count | Max \\|Δlogp\\| | torch.equal |") + lines.append("|---|---:|---:|:---:|") + for label, info in (("P/P 原生", pp_info), ("R/R 严格", rr_info)): + m = info["metrics"] + mismatch = _metric(m, "mismatch_count", "mismatched_tokens", "mismatch_tokens") + total = _metric(m, "element_count", "token_count", "compared_tokens", "total_tokens") + max_abs = _metric(m, "max_abs_diff", "max_abs_drift", "max_abs_delta") + equal = _metric(m, "torch_equal", "bitwise_equal", "exact") + count = f"{mismatch} / {total}" if total is not None else f"{mismatch}" + lines.append( + f"| {label} | {count} | {max_abs if max_abs is None else f'{max_abs:g}'} | " + f"{str(equal).lower()} |" + ) + lines.append("") + lines.append(f"## {min(n_pp, n_rr)} 轮平均性能结果\n") + lines.append(f"(P/P rounds={n_pp}, R/R rounds={n_rr})\n") + lines.append("| 指标 | P/P 原生 | R/R 严格 | R/R 相对 P/P |") + lines.append("|---|---:|---:|---:|") + keys = [row[1] for row in ROWS] + pp_mean = mean_rows(pp_rounds, keys) + rr_mean = mean_rows(rr_rounds, keys) + for label, key, unit, direction in ROWS: + pp_value, rr_value = pp_mean[key], rr_mean[key] + if pp_value is None and rr_value is None: + continue + lines.append( + f"| {label} | {fmt(pp_value, unit)} | {fmt(rr_value, unit)} | " + f"{relative(pp_value, rr_value, direction)} |" + ) + lines.append("") + lines.append("### 每轮明细\n") + lines.append( + "| round | P/P rollout | R/R rollout | P/P actor train | R/R actor train " + "| P/P step | R/R step |" + ) + lines.append("|---:|---:|---:|---:|---:|---:|---:|") + for index in sorted(set(pp_rounds) | set(rr_rounds)): + pp = pp_rounds.get(index, {}) + rr = rr_rounds.get(index, {}) + + def cell(bucket, key): + return f"{bucket[key]:.3f}" if key in bucket else "-" + + lines.append( + f"| {index} | {cell(pp, 'perf/rollout_time')} | {cell(rr, 'perf/rollout_time')} | " + f"{cell(pp, 'perf/actor_train_time')} | {cell(rr, 'perf/actor_train_time')} | " + f"{cell(pp, 'perf/step_time')} | {cell(rr, 'perf/step_time')} |" + ) + lines.append("") + lines.append( + f"validation: P/P passed={pp_info['passed']} errors={len(pp_info['errors'])}; " + f"R/R passed={rr_info['passed']} errors={len(rr_info['errors'])}" + ) + return "\n".join(lines) + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pp-run", type=Path, required=True) + parser.add_argument("--rr-run", type=Path, required=True) + parser.add_argument("--output", type=Path) + args = parser.parse_args(argv) + text = render(args.pp_run, args.rr_run) + if args.output: + args.output.write_text(text + "\n", encoding="utf-8") + print(text) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/vime_rocm_attention_ablation/validate_artifacts.py b/examples/vime_rocm_attention_ablation/validate_artifacts.py index 1fa0b86a..a9e491b8 100644 --- a/examples/vime_rocm_attention_ablation/validate_artifacts.py +++ b/examples/vime_rocm_attention_ablation/validate_artifacts.py @@ -43,6 +43,8 @@ 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_TRITON_PAGED_ATTENTION_BACKEND_ID = "rlkernel.rocm.triton_chunked_flash_attention.v1" +ROCM_TRITON_DET_GEMM_BACKEND_ID = "rlkernel.det_gemm.triton_mfma_rocm.v1" 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" @@ -303,6 +305,8 @@ def _validate_rlkernel_record( # arithmetic and page-table reads remain inside native AITER/CK. allowed_triton_backends = { ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, + ROCM_TRITON_PAGED_ATTENTION_BACKEND_ID, + ROCM_TRITON_DET_GEMM_BACKEND_ID, } for key, item in _walk_key_values(provenance): if key in _TRITON_KEYS and item is True: @@ -315,12 +319,16 @@ def _validate_rlkernel_record( 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, + if not ( + _has_exact_value(provenance, {"paged_kernel"}, ROCM_PAGED_ATTENTION_BACKEND_ID) + or _has_exact_value( + provenance, {"paged_kernel"}, ROCM_TRITON_PAGED_ATTENTION_BACKEND_ID + ) ): - errors.append(f"{label} did not prove direct non-Split-K paged CK execution") + errors.append( + f"{label} did not prove direct fixed-schedule paged execution " + "(CK non-Split-K or the Triton chunked contract)" + ) 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"}) diff --git a/rl_engine/integrations/framework_operators.py b/rl_engine/integrations/framework_operators.py index 425263c3..08b361c6 100644 --- a/rl_engine/integrations/framework_operators.py +++ b/rl_engine/integrations/framework_operators.py @@ -1047,6 +1047,19 @@ def warmup_rocm_decode(self, impl: Any, *, dtype: torch.dtype) -> None: torch.empty((1,), device=device, dtype=dtype) ) page_size = 16 + core = getattr(runtime, "_core", None) + if getattr(core, "attention_backend", "ck") == "triton": + from rl_engine.kernels.ops.triton.attention import chunked_flash_attn + + chunked_flash_attn.warmup( + device, + num_q_heads=int(impl.num_heads), + num_kv_heads=int(impl.num_kv_heads), + dtype=dtype, + ) + from rl_engine.kernels.ops.triton.matmul import mfma_gemm + + mfma_gemm.warmup(device) q = torch.empty( (1, int(impl.num_heads), 1, int(impl.head_size)), device=device, @@ -1291,19 +1304,51 @@ def _rocm_direct_paged_metadata( block_size: int, num_actual: int, cache_owner: Any, + runtime_core: Any = None, ) -> tuple[dict[str, Any], bool] | None: """Prepare graph-safe paged metadata once and share it across layers.""" num_decodes = int(getattr(attn_metadata, "num_decodes", 0)) num_prefills = int(getattr(attn_metadata, "num_prefills", 0)) num_extends = int(getattr(attn_metadata, "num_extends", 0)) - if num_extends or (num_decodes > 0) == (num_prefills > 0): + unified = bool(getattr(runtime_core, "supports_mixed_paged_batches", False)) + if unified: + # The chunked Triton contract serves every request kind from one + # causal launch, so mixed decode/extend/prefill batches need no + # per-kind expansion. + if num_decodes + num_extends + num_prefills <= 0 or num_actual <= 0: + return None + if num_extends or (num_decodes > 0 and num_prefills > 0): + mode = "mixed" + elif num_prefills > 0: + mode = "prefill" + else: + mode = "decode" + sequence_count = num_decodes + num_extends + num_prefills + query_start_loc = self._metadata_tensor(attn_metadata, "query_start_loc") + max_seqlen_q = int(getattr(attn_metadata, "max_query_len", 0) or 0) + if max_seqlen_q <= 0: + lengths = [ + int(getattr(meta, "max_query_len", 0) or 0) + for meta in ( + getattr(attn_metadata, "decode_metadata", None), + getattr(attn_metadata, "extend_metadata", None), + getattr(attn_metadata, "prefill_metadata", None), + ) + if meta is not None + ] + max_seqlen_q = max(lengths) if lengths else 0 + causal = True + elif 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 + else: + 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": + if unified: + pass + elif mode == "prefill": prefill = getattr(attn_metadata, "prefill_metadata", None) if prefill is None: return None @@ -1336,6 +1381,7 @@ def _rocm_direct_paged_metadata( ) key = ( mode, + unified, _tensor_cache_token(query_starts_source), _tensor_cache_token(seq_lens_source), _tensor_cache_token(block_table), @@ -1365,7 +1411,23 @@ def _rocm_direct_paged_metadata( seq_lens = seq_lens.contiguous() # Graph metadata is request-level while packed Q is query-token-level. # Expand decode page rows to query rows without materializing KV. - if mode == "decode": + seq_of_token = None + if unified: + query_start_loc = query_start_loc[: sequence_count + 1] + seq_lens = seq_lens[:sequence_count] + active_rows = seq_lens > 0 + seqused_k = seq_lens + pages = block_table[:sequence_count, :page_count] + if mode == "decode" and num_actual == sequence_count: + seq_of_token = torch.arange( + num_actual, dtype=torch.int32, device=block_table.device + ) + else: + tokens = torch.arange(num_actual, dtype=torch.int32, device=block_table.device) + seq_of_token = torch.searchsorted( + query_start_loc[1:], tokens, right=True + ).to(torch.int32) + elif mode == "decode": query_starts = query_start_loc[: sequence_count + 1] query_ends = query_starts[1:] query_indices = torch.arange( @@ -1399,7 +1461,7 @@ def _rocm_direct_paged_metadata( pages = block_table[:sequence_count, :page_count] if not pages.is_contiguous(): pages = pages.contiguous() - if mode != "decode": + if mode != "decode" and not unified: active_rows = seqused_k > 0 if configured_kv_limit is not None: # Zero-length rows are legal vLLM graph padding. They must not @@ -1455,6 +1517,7 @@ def _rocm_direct_paged_metadata( "max_seqlen_k": kernel_max_seqlen_k, "configured_kv_limit": configured_kv_limit, "causal": causal, + "seq_of_token": seq_of_token, } self._rocm_paged_metadata_key = key self._rocm_paged_metadata_owners = {owner_id} @@ -1485,6 +1548,7 @@ def _rocm_direct_paged( block_size=key_cache.size(1), num_actual=num_actual, cache_owner=layer, + runtime_core=getattr(runtime, "_core", None), ) if metadata_result is None: return None diff --git a/rl_engine/kernels/ops/rocm/attention/flash_attn.py b/rl_engine/kernels/ops/rocm/attention/flash_attn.py index 3462a367..5d06ecce 100644 --- a/rl_engine/kernels/ops/rocm/attention/flash_attn.py +++ b/rl_engine/kernels/ops/rocm/attention/flash_attn.py @@ -94,6 +94,19 @@ # scope so contract-aware dispatch and the Vime adapter name one constant # instead of duplicating the string. BACKEND_ID = "aiter.rocm.ck_dense_mha" +AITER_PAGED_KERNEL_ID = "aiter_mha_batch_prefill_non_split_ck" +TRITON_CHUNKED_FLASH_ATTENTION_ID = "rlkernel.rocm.triton_chunked_flash_attention.v1" +_ATTENTION_BACKEND_ENV = "RL_KERNEL_ROCM_ATTENTION_BACKEND" + + +def _requested_attention_backend() -> str: + """Return ``ck`` (AITER/CK entry points) or ``triton`` (chunked flash contract).""" + + value = os.environ.get(_ATTENTION_BACKEND_ENV, "ck").strip().lower() + value = {"aiter": "ck", "chunked_flash": "triton"}.get(value, value) + if value not in {"ck", "triton"}: + raise RuntimeError(f"{_ATTENTION_BACKEND_ENV} must be 'ck' or 'triton', got {value!r}") + return value class StrictRocmAttentionUnavailable(RuntimeError): @@ -431,8 +444,19 @@ def __init__( mha_batch_prefill = _mha_batch_prefill source_sha256 = "test-double" if _source_sha256 is None else _source_sha256 self._fixed_paged_tile = 0 + self._attention_backend = "ck" fixed_tile = os.environ.get("RL_KERNEL_ROCM_FIXED_PAGED_TILE", "0") - if _mha_fwd is None and fixed_tile != "0": + requested_backend = _requested_attention_backend() + if _mha_fwd is None and requested_backend == "triton": + from rl_engine.kernels.ops.triton.attention import chunked_flash_attn + + self._attention_backend = "triton" + mha_batch_prefill = chunked_flash_attn.triton_paged_prefill + source_digest = hashlib.sha256(source_sha256.encode()) + source_digest.update(Path(inspect.getsourcefile(chunked_flash_attn)).read_bytes()) + source_digest.update(chunked_flash_attn.CHUNKED_FLASH_ATTENTION_CONTRACT_ID.encode()) + source_sha256 = source_digest.hexdigest() + elif _mha_fwd is None and fixed_tile != "0": if fixed_tile not in ("64", "128"): raise ValueError("RL_KERNEL_ROCM_FIXED_PAGED_TILE must be 0, 64 or 128") from .fixed_paged_ck import fixed_paged_prefill @@ -456,11 +480,32 @@ def __init__( self._mha_bwd = mha_bwd self._mha_batch_prefill = mha_batch_prefill self.supports_paged_schedule = callable(mha_batch_prefill) + # The Triton contract serves decode, extend and prefill rows from one + # launch; the CK entry point needs one query row per decode request. + self.supports_mixed_paged_batches = self._attention_backend == "triton" 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 + @property + def attention_backend(self) -> str: + return self._attention_backend + + @property + def paged_entrypoint_id(self) -> str: + if self._attention_backend == "triton": + return TRITON_CHUNKED_FLASH_ATTENTION_ID + if self._fixed_paged_tile: + return f"rl_kernel_fixed_paged_ck_m{self._fixed_paged_tile}" + return "mha_batch_prefill" + + @property + def paged_kernel_id(self) -> str: + if self._attention_backend == "triton": + return TRITON_CHUNKED_FLASH_ATTENTION_ID + return AITER_PAGED_KERNEL_ID + def forward_with_lse( self, q: torch.Tensor, @@ -508,11 +553,7 @@ def forward_with_lse( self._mha_batch_prefill, self._mha_bwd, ) - forward_entrypoint = ( - f"rl_kernel_fixed_paged_ck_m{self._fixed_paged_tile}" - if self._fixed_paged_tile - else "mha_batch_prefill" - ) + forward_entrypoint = self.paged_entrypoint_id 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: @@ -668,11 +709,8 @@ def forward_paged_varlen_with_lse( "gpu_arch": gpu_arch, "aiter_api_source": self.api_source, "aiter_source_sha256": self.source_sha256, - "forward_entrypoint": ( - f"rl_kernel_fixed_paged_ck_m{self._fixed_paged_tile}" - if self._fixed_paged_tile - else "mha_batch_prefill" - ), + "forward_entrypoint": self.paged_entrypoint_id, + "paged_kernel": self.paged_kernel_id, "kv_layout": "vllm_linear_paged", "dense_kv_materialized": False, "num_splits": self.num_splits, @@ -820,7 +858,7 @@ def forward_bshd_with_lse( if not out_bshd.is_contiguous(): raise ValueError("strict BSHD decode output must expose a contiguous AITER view") - if self._fixed_paged_tile: + if self._fixed_paged_tile or self._attention_backend == "triton": # Keep the same arithmetic if a caller already materialized BSHD # inputs (e.g. a mixed prefill/decode fallback). fixed_out, fixed_lse = _AiterCKPagedAttentionFn.apply( @@ -840,7 +878,7 @@ def forward_bshd_with_lse( lse=fixed_lse, provenance={ "actual_backend": self.backend_id, - "forward_entrypoint": f"rl_kernel_fixed_paged_ck_m{self._fixed_paged_tile}", + "forward_entrypoint": self.paged_entrypoint_id, "aiter_source_sha256": self.source_sha256, "dense_kv_materialized": True, "fallback": False, diff --git a/rl_engine/kernels/ops/rocm/attention/strict_runtime.py b/rl_engine/kernels/ops/rocm/attention/strict_runtime.py index e6693b0e..14cad16a 100644 --- a/rl_engine/kernels/ops/rocm/attention/strict_runtime.py +++ b/rl_engine/kernels/ops/rocm/attention/strict_runtime.py @@ -452,7 +452,7 @@ def forward_paged_varlen_with_lse( "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", + "paged_kernel": self._paged_kernel_id(), "dense_kv_materialized": False, "lse_returned": bool(return_lse), "launch_granularity": "one_local_gqa_batch", @@ -594,7 +594,7 @@ def forward_paged_with_lse( "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", + "paged_kernel": self._paged_kernel_id(), "dense_kv_materialized": False, "lse_returned": bool(return_lse), "launch_granularity": "one_local_gqa_batch", @@ -1276,6 +1276,9 @@ def _run_core( launches, ) + def _paged_kernel_id(self) -> str: + return str(getattr(self._core, "paged_kernel_id", "aiter_mha_batch_prefill_non_split_ck")) + @staticmethod def _storage_is_disjoint(output: torch.Tensor, *inputs: torch.Tensor) -> bool: output_storage = output.untyped_storage().data_ptr() diff --git a/rl_engine/kernels/ops/rocm/matmul/det_gemm.py b/rl_engine/kernels/ops/rocm/matmul/det_gemm.py index aed6c8f1..14871417 100644 --- a/rl_engine/kernels/ops/rocm/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/rocm/matmul/det_gemm.py @@ -13,15 +13,38 @@ from rl_engine.kernels.ops.backward_runtime import record_backward from rl_engine.kernels.ops.triton.matmul.det_gemm import ( TritonDetGemmOp, + _device_arch, _triton_gemm_fp32, _triton_tree_gemm, deterministic_gemm_triton, ) +from rl_engine.kernels.ops.triton.matmul.mfma_gemm import ( + MFMA_GEMM_CONTRACT_ID, + MfmaGemmFn, + MfmaLinearFn, + mfma_gemm, + mfma_linear, + mfma_linear_input_gradient, + mfma_linear_weight_gradient, +) from rl_engine.runtime_mode import rl_kernel_mode, route_report_enabled _BACKEND_ENV = "RL_KERNEL_DET_GEMM_BACKEND" _AUTO_BACKEND = "auto" -_TRITON_BACKEND = "triton" +# ``triton_tree``: scalar-FMA leaves plus the canonical BF16 midpoint K-tree +# (TP-degree invariant, matches the native reference kernel bit for bit). +# ``triton_mfma``: pinned-order MFMA accumulation with FP32 chunk partials +# (batch invariant and train/rollout bitwise on one fixed TP sharding, several +# times faster). ``auto`` selects MFMA on gfx942 and the tree elsewhere. +_TREE_BACKEND = "triton_tree" +_MFMA_BACKEND = "triton_mfma" +_TRITON_BACKEND = _TREE_BACKEND +_BACKEND_ALIASES = { + "rocm": _TREE_BACKEND, + "triton": _TREE_BACKEND, + "tree": _TREE_BACKEND, + "mfma": _MFMA_BACKEND, +} _ROUTE_REPORTED = False _ROUTE_REPORT_LOCK = Lock() _WEIGHT_TRANSPOSE_CACHE: dict[ @@ -34,22 +57,42 @@ def _requested_det_gemm_backend() -> str: value = os.getenv(_BACKEND_ENV, _AUTO_BACKEND).strip().lower() - value = {"rocm": _TRITON_BACKEND}.get(value, value) - if value not in {_AUTO_BACKEND, _TRITON_BACKEND}: + value = _BACKEND_ALIASES.get(value, value) + if value not in {_AUTO_BACKEND, _TREE_BACKEND, _MFMA_BACKEND}: raise RuntimeError( - f"{_BACKEND_ENV} must be '{_AUTO_BACKEND}' or " - f"'{_TRITON_BACKEND}' on ROCm, got {value!r}" + f"{_BACKEND_ENV} must be '{_AUTO_BACKEND}', '{_TREE_BACKEND}' or " + f"'{_MFMA_BACKEND}' on ROCm, got {value!r}" ) return value _REQUESTED_BACKEND = _requested_det_gemm_backend() +_RESOLVED_BACKEND: str | None = None + + +def _mfma_supported() -> bool: + if not torch.cuda.is_available(): + return False + try: + return _device_arch(torch.cuda.current_device()) == "gfx942" + except Exception: # pragma: no cover - device query failure + return False def det_gemm_backend() -> str: """Return the strict ROCm GEMM implementation.""" - return _TRITON_BACKEND + global _RESOLVED_BACKEND + if _RESOLVED_BACKEND is None: + if _REQUESTED_BACKEND == _AUTO_BACKEND: + _RESOLVED_BACKEND = _MFMA_BACKEND if _mfma_supported() else _TREE_BACKEND + else: + _RESOLVED_BACKEND = _REQUESTED_BACKEND + return _RESOLVED_BACKEND + + +def _use_mfma() -> bool: + return det_gemm_backend() == _MFMA_BACKEND def det_gemm_fallback_reason() -> str | None: @@ -57,6 +100,8 @@ def det_gemm_fallback_reason() -> str | None: def det_gemm_backend_id() -> str: + if _use_mfma(): + return MFMA_GEMM_CONTRACT_ID return "rlkernel.det_gemm.triton_tree_rocm.v1" @@ -141,6 +186,9 @@ def det_gemm_linear( """Apply a native [N,K] weight through the strict ROCm backend.""" del native_op + if _use_mfma(): + del inference_schedule + return mfma_linear(a, weight, out=out) return _triton_tree_gemm( a, _cached_weight_transpose(weight), @@ -351,6 +399,8 @@ def det_gemm_linear_prepared( raise ValueError("prepared deterministic linear weight must be contiguous") if out is not None and (torch._C._overlaps(out, a) or torch._C._overlaps(out, weight_t)): raise ValueError("prepared deterministic linear output must not alias its inputs") + if _use_mfma(): + return mfma_gemm(a, weight_t, out=out) return _triton_tree_gemm( a, weight_t, @@ -368,6 +418,8 @@ def det_gemm_linear_input_gradient( """Compute ``dX = dY @ weight`` through the strict ROCm backend.""" del native_op + if _use_mfma(): + return mfma_linear_input_gradient(grad_output, weight) return deterministic_gemm_triton(grad_output, weight) @@ -380,6 +432,8 @@ def det_gemm_linear_weight_gradient( """Compute ``dWeight = dY.T @ X`` through the strict ROCm backend.""" del native_op + if _use_mfma(): + return mfma_linear_weight_gradient(a, grad_output) return _triton_tree_gemm( a.t(), grad_output, @@ -412,7 +466,7 @@ def backward(ctx, grad_out): class RocmDetGemmOp: - """Batch-invariant deterministic GEMM backed by the ROCm Triton tree.""" + """Batch-invariant deterministic GEMM backed by the selected ROCm Triton path.""" def __init__(self): det_gemm_backend() @@ -424,7 +478,7 @@ def __init__(self): def __call__(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" - return deterministic_gemm_triton(a.contiguous(), b.contiguous()) + return deterministic_gemm(a, b) def linear( self, @@ -446,6 +500,8 @@ def linear( return out if not torch.is_grad_enabled() or not (a.requires_grad or weight.requires_grad): return _det_gemm_linear_inference(a, weight) + if _use_mfma(): + return MfmaLinearFn.apply(a, weight) return _DetLinearFn.apply(a, weight) def linear_prepared( @@ -486,9 +542,13 @@ def parameter_vjp_contributions_fp32(self, *, a, b, grad_output): def deterministic_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """Functional strict ROCm GEMM entry.""" + """Functional strict ROCm GEMM entry for ``[M, K] @ [K, N]``.""" - return deterministic_gemm_triton(a, b) + if _use_mfma(): + if torch.is_grad_enabled() and (a.requires_grad or b.requires_grad): + return MfmaGemmFn.apply(a, b) + return mfma_gemm(a, b) + return deterministic_gemm_triton(a.contiguous(), b.contiguous()) __all__ = [ diff --git a/rl_engine/kernels/ops/triton/attention/chunked_flash_attn.py b/rl_engine/kernels/ops/triton/attention/chunked_flash_attn.py new file mode 100644 index 00000000..6b2d9183 --- /dev/null +++ b/rl_engine/kernels/ops/triton/attention/chunked_flash_attn.py @@ -0,0 +1,810 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Chunked-KV flash attention for ROCm with one arithmetic contract for every role. + +Contract ``rlkernel.rocm.triton_chunked_flash_attention.v1``: + +* BF16/FP16 Q, K, V with ``head_dim=128`` read from vLLM-style token-major + pages ``[pages, 16, kv_heads, head_dim]`` through a 2-D page table. +* Scores ``Q.K^T`` on ``v_mfma_f32_16x16x16`` (``matrix_instr_nonkdim=16``, + ``kpack=2``), scaled once in FP32 by ``scale * log2(e)``. +* The key axis is consumed in ascending ``BLOCK_N=64`` blocks inside fixed + ``CHUNK_KV=512``-token chunks. Every chunk runs the online softmax from an + empty state (``m=-inf, l=0, acc=0``); chunk states are merged in ascending + order with the exact FA2 rescale. ``P`` is rounded to the input dtype before + ``P.V``; ``exp2`` is the hardware instruction; FP contraction is disabled. +* Blocks and chunks that are entirely masked for a row leave that row's state + bit-for-bit untouched, so a row's result never depends on which other rows + share its tile. + +Every query row's output is therefore a pure function of its own Q, the +visible K/V prefix and the contract. The monolithic schedule (one program per +query tile, chunks evaluated sequentially) and the split schedule (one program +per ``(sequence, kv head, chunk)`` plus an ascending merge) produce identical +bits, which is what lets a Megatron full-sequence forward, a vLLM prefill, a +prefix-cached extend and a single-token paged decode agree exactly. +""" + +from __future__ import annotations + +import math + +import torch + +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: # pragma: no cover - exercised only without Triton + _TRITON_AVAILABLE = False + +CHUNKED_FLASH_ATTENTION_CONTRACT_ID = "rlkernel.rocm.triton_chunked_flash_attention.v1" +BLOCK_M = 64 +BLOCK_N = 64 +CHUNK_KV = 512 +HEAD_DIM = 128 +PAGE_SIZE = 16 +NUM_WARPS = 4 +WAVES_PER_EU = 2 +MATRIX_INSTR_NONKDIM = 16 +KPACK = 2 +_LOG2E = 1.4426950408889634 +_LN2 = 0.6931471805599453 +_COMPILE_OPTIONS = dict( + num_warps=NUM_WARPS, + waves_per_eu=WAVES_PER_EU, + matrix_instr_nonkdim=MATRIX_INSTR_NONKDIM, + kpack=KPACK, + enable_fp_fusion=False, +) + + +if _TRITON_AVAILABLE: + + @triton.jit + def _kv_tile_ptrs( + base_ptr, + block_table_ptr, + bt_stride, + seq, + start_n, + stride_page, + stride_tok, + stride_head, + kv_head, + BLOCK_N: tl.constexpr, + PAGE: tl.constexpr, + D: tl.constexpr, + ): + offs_n = start_n + tl.arange(0, BLOCK_N) + page = tl.load(block_table_ptr + seq * bt_stride + offs_n // PAGE) + rows = ( + page.to(tl.int64) * stride_page + (offs_n % PAGE) * stride_tok + kv_head * stride_head + ) + offs_d = tl.arange(0, D) + return base_ptr + rows[:, None] + offs_d[None, :] + + @triton.jit + def _full_block_update( + q, + k, + v, + m_i, + l_i, + acc, + scale_log2, + ): + """One key block that every row sees completely (no masking needed). + + Identical operation sequence to the masked update with an all-true + mask, so the two paths are bit-for-bit interchangeable. + """ + + x = tl.dot(q, tl.trans(k)) * scale_log2 + row_max = tl.max(x, 1) + m_new = tl.maximum(m_i, row_max) + m_safe = tl.where(m_new == float("-inf"), 0.0, m_new) + alpha = tl.exp2(m_i - m_safe) + p = tl.exp2(x - m_safe[:, None]) + l_ij = tl.sum(p, 1) + pv = tl.dot(p.to(v.dtype), v) + acc = acc * alpha[:, None] + pv + l_i = l_i * alpha + l_ij + return m_new, l_i, acc + + @triton.jit + def _chunk_state( + q, + k_ptr, + v_ptr, + block_table_ptr, + bt_stride, + seq, + kv_head, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + chunk_start, + chunk_end, + kv_len, + q_pos, + full_limit, + scale_log2, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + PAGE: tl.constexpr, + D: tl.constexpr, + FAST_PATH: tl.constexpr, + ): + """Online softmax over keys ``[chunk_start, chunk_end)`` from an empty state. + + ``full_limit`` is a BLOCK_N multiple below which every row of the tile + sees every key; blocks under it skip the masking work when + ``FAST_PATH`` is set. Block boundaries are the same in both paths. + """ + + m_i = tl.full((BLOCK_M,), float("-inf"), dtype=tl.float32) + l_i = tl.zeros((BLOCK_M,), dtype=tl.float32) + acc = tl.zeros((BLOCK_M, D), dtype=tl.float32) + masked_start = chunk_start + if FAST_PATH: + full_end = tl.minimum(chunk_end, full_limit) + for start_n in range(chunk_start, full_end, BLOCK_N): + k_ptrs = _kv_tile_ptrs( + k_ptr, + block_table_ptr, + bt_stride, + seq, + start_n, + stride_kp, + stride_kt, + stride_kh, + kv_head, + BLOCK_N, + PAGE, + D, + ) + v_ptrs = _kv_tile_ptrs( + v_ptr, + block_table_ptr, + bt_stride, + seq, + start_n, + stride_vp, + stride_vt, + stride_vh, + kv_head, + BLOCK_N, + PAGE, + D, + ) + k = tl.load(k_ptrs) + v = tl.load(v_ptrs) + m_i, l_i, acc = _full_block_update(q, k, v, m_i, l_i, acc, scale_log2) + masked_start = tl.maximum(chunk_start, full_end) + for start_n in range(masked_start, chunk_end, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + k_ptrs = _kv_tile_ptrs( + k_ptr, + block_table_ptr, + bt_stride, + seq, + start_n, + stride_kp, + stride_kt, + stride_kh, + kv_head, + BLOCK_N, + PAGE, + D, + ) + v_ptrs = _kv_tile_ptrs( + v_ptr, + block_table_ptr, + bt_stride, + seq, + start_n, + stride_vp, + stride_vt, + stride_vh, + kv_head, + BLOCK_N, + PAGE, + D, + ) + kv_valid = offs_n < kv_len + k = tl.load(k_ptrs, mask=kv_valid[:, None], other=0.0) + v = tl.load(v_ptrs, mask=kv_valid[:, None], other=0.0) + x = tl.dot(q, tl.trans(k)) * scale_log2 + visible = kv_valid[None, :] & (offs_n[None, :] <= q_pos[:, None]) + x = tl.where(visible, x, float("-inf")) + row_max = tl.max(x, 1) + block_empty = row_max == float("-inf") + m_new = tl.maximum(m_i, row_max) + m_safe = tl.where(m_new == float("-inf"), 0.0, m_new) + alpha = tl.exp2(m_i - m_safe) + p = tl.exp2(x - m_safe[:, None]) + l_ij = tl.sum(p, 1) + pv = tl.dot(p.to(v.dtype), v) + acc_new = acc * alpha[:, None] + pv + l_new = l_i * alpha + l_ij + acc = tl.where(block_empty[:, None], acc, acc_new) + l_i = tl.where(block_empty, l_i, l_new) + m_i = m_new + return m_i, l_i, acc + + @triton.jit + def _merge_state(m, row_sum, acc, mc, lc, accc): + """Merge chunk state ``(mc, lc, accc)`` into ``(m, l, acc)``; empty chunks are no-ops.""" + + chunk_empty = mc == float("-inf") + prev_empty = m == float("-inf") + m_new = tl.maximum(m, mc) + m_safe = tl.where(m_new == float("-inf"), 0.0, m_new) + a = tl.exp2(m - m_safe) + b = tl.exp2(mc - m_safe) + acc_merged = acc * a[:, None] + accc * b[:, None] + l_merged = row_sum * a + lc * b + acc_out = tl.where( + chunk_empty[:, None], acc, tl.where(prev_empty[:, None], accc, acc_merged) + ) + l_out = tl.where(chunk_empty, row_sum, tl.where(prev_empty, lc, l_merged)) + return m_new, l_out, acc_out + + @triton.jit + def _attn_fwd_monolithic_kernel( + q_ptr, + k_ptr, + v_ptr, + o_ptr, + lse_ptr, + cu_seqlens_q_ptr, + seqlen_k_ptr, + block_table_ptr, + bt_stride, + stride_qt, + stride_qh, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + stride_ot, + stride_oh, + stride_lh, + stride_lt, + scale_log2, + group_size, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + CHUNK: tl.constexpr, + PAGE: tl.constexpr, + D: tl.constexpr, + WRITE_LSE: tl.constexpr, + FAST_PATH: tl.constexpr, + ): + tile = tl.program_id(0) + head = tl.program_id(1) + seq = tl.program_id(2) + q_start = tl.load(cu_seqlens_q_ptr + seq) + q_len = tl.load(cu_seqlens_q_ptr + seq + 1) - q_start + start_m = tile * BLOCK_M + if start_m >= q_len: + return + kv_len = tl.load(seqlen_k_ptr + seq) + kv_head = head // group_size + offs_m = start_m + tl.arange(0, BLOCK_M) + row_valid = offs_m < q_len + offs_d = tl.arange(0, D) + q_rows = (q_start + offs_m).to(tl.int64) + q = tl.load( + q_ptr + q_rows[:, None] * stride_qt + head * stride_qh + offs_d[None, :], + mask=row_valid[:, None], + other=0.0, + ) + q_pos = kv_len - q_len + offs_m + tile_kv_limit = tl.minimum(kv_len, kv_len - q_len + start_m + BLOCK_M) + # Every row of this tile sees all keys below the first row's causal + # bound; round down to a block boundary so both paths share blocks. + first_row_visible = kv_len - q_len + start_m + 1 + full_limit = tl.maximum((tl.minimum(first_row_visible, kv_len) // BLOCK_N) * BLOCK_N, 0) + m_i = tl.full((BLOCK_M,), float("-inf"), dtype=tl.float32) + l_i = tl.zeros((BLOCK_M,), dtype=tl.float32) + acc = tl.zeros((BLOCK_M, D), dtype=tl.float32) + num_chunks = tl.cdiv(tile_kv_limit, CHUNK) + for chunk in range(0, num_chunks): + chunk_start = chunk * CHUNK + chunk_end = tl.minimum(chunk_start + CHUNK, tile_kv_limit) + mc, lc, accc = _chunk_state( + q, + k_ptr, + v_ptr, + block_table_ptr, + bt_stride, + seq, + kv_head, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + chunk_start, + chunk_end, + kv_len, + q_pos, + full_limit, + scale_log2, + BLOCK_M, + BLOCK_N, + PAGE, + D, + FAST_PATH, + ) + m_i, l_i, acc = _merge_state(m_i, l_i, acc, mc, lc, accc) + out = acc / l_i[:, None] + tl.store( + o_ptr + q_rows[:, None] * stride_ot + head * stride_oh + offs_d[None, :], + out.to(o_ptr.dtype.element_ty), + mask=row_valid[:, None], + ) + if WRITE_LSE: + lse = m_i * 0.6931471805599453 + tl.log(l_i) + tl.store(lse_ptr + head * stride_lh + q_rows * stride_lt, lse, mask=row_valid) + + @triton.jit + def _attn_fwd_split_partial_kernel( + q_ptr, + k_ptr, + v_ptr, + pm_ptr, + pl_ptr, + pacc_ptr, + cu_seqlens_q_ptr, + seqlen_k_ptr, + block_table_ptr, + bt_stride, + stride_qt, + stride_qh, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + stride_pc, + stride_pt, + stride_ph, + stride_pacc_c, + stride_pacc_t, + stride_pacc_h, + scale_log2, + group_size, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + CHUNK: tl.constexpr, + PAGE: tl.constexpr, + D: tl.constexpr, + FAST_PATH: tl.constexpr, + ): + """One program per ``(sequence, kv head, chunk)``; rows pack the GQA group x queries.""" + + chunk = tl.program_id(0) + kv_head = tl.program_id(1) + seq = tl.program_id(2) + q_start = tl.load(cu_seqlens_q_ptr + seq) + q_len = tl.load(cu_seqlens_q_ptr + seq + 1) - q_start + kv_len = tl.load(seqlen_k_ptr + seq) + chunk_start = chunk * CHUNK + if chunk_start >= kv_len: + return + chunk_end = tl.minimum(chunk_start + CHUNK, kv_len) + rows = tl.arange(0, BLOCK_M) + row_q = rows // group_size + row_h = rows % group_size + row_valid = row_q < q_len + head = kv_head * group_size + row_h + offs_d = tl.arange(0, D) + tok = (q_start + row_q).to(tl.int64) + q = tl.load( + q_ptr + tok[:, None] * stride_qt + head[:, None] * stride_qh + offs_d[None, :], + mask=row_valid[:, None], + other=0.0, + ) + q_pos = tl.where(row_valid, kv_len - q_len + row_q, -1) + first_row_visible = kv_len - q_len + 1 + full_limit = tl.maximum((tl.minimum(first_row_visible, kv_len) // BLOCK_N) * BLOCK_N, 0) + mc, lc, accc = _chunk_state( + q, + k_ptr, + v_ptr, + block_table_ptr, + bt_stride, + seq, + kv_head, + stride_kp, + stride_kt, + stride_kh, + stride_vp, + stride_vt, + stride_vh, + chunk_start, + chunk_end, + kv_len, + q_pos, + full_limit, + scale_log2, + BLOCK_M, + BLOCK_N, + PAGE, + D, + FAST_PATH, + ) + state = chunk * stride_pc + tok * stride_pt + head * stride_ph + tl.store(pm_ptr + state, mc, mask=row_valid) + tl.store(pl_ptr + state, lc, mask=row_valid) + tl.store( + pacc_ptr + + chunk * stride_pacc_c + + tok[:, None] * stride_pacc_t + + head[:, None] * stride_pacc_h + + offs_d[None, :], + accc, + mask=row_valid[:, None], + ) + + @triton.jit + def _attn_fwd_split_merge_kernel( + pm_ptr, + pl_ptr, + pacc_ptr, + o_ptr, + lse_ptr, + seq_of_token_ptr, + seqlen_k_ptr, + stride_pc, + stride_pt, + stride_ph, + stride_pacc_c, + stride_pacc_t, + stride_pacc_h, + stride_ot, + stride_oh, + stride_lh, + stride_lt, + CHUNK: tl.constexpr, + D: tl.constexpr, + WRITE_LSE: tl.constexpr, + ): + tok = tl.program_id(0).to(tl.int64) + head = tl.program_id(1) + seq = tl.load(seq_of_token_ptr + tok) + kv_len = tl.load(seqlen_k_ptr + seq) + num_chunks = tl.cdiv(kv_len, CHUNK) + offs_d = tl.arange(0, D) + state = tok * stride_pt + head * stride_ph + acc_base = pacc_ptr + tok * stride_pacc_t + head * stride_pacc_h + offs_d + m = tl.load(pm_ptr + state) + row_sum = tl.load(pl_ptr + state) + acc = tl.load(acc_base) + for chunk in range(1, num_chunks): + mc = tl.load(pm_ptr + chunk * stride_pc + state) + lc = tl.load(pl_ptr + chunk * stride_pc + state) + accc = tl.load(acc_base + chunk * stride_pacc_c) + chunk_empty = mc == float("-inf") + prev_empty = m == float("-inf") + m_new = tl.maximum(m, mc) + m_safe = tl.where(m_new == float("-inf"), 0.0, m_new) + a = tl.exp2(m - m_safe) + b = tl.exp2(mc - m_safe) + acc_merged = acc * a + accc * b + l_merged = row_sum * a + lc * b + acc = tl.where(chunk_empty, acc, tl.where(prev_empty, accc, acc_merged)) + row_sum = tl.where(chunk_empty, row_sum, tl.where(prev_empty, lc, l_merged)) + m = m_new + out = acc / row_sum + tl.store( + o_ptr + tok * stride_ot + head * stride_oh + offs_d, out.to(o_ptr.dtype.element_ty) + ) + if WRITE_LSE: + lse = m * 0.6931471805599453 + tl.log(row_sum) + tl.store(lse_ptr + head * stride_lh + tok * stride_lt, lse) + + +def _validate(q, k_cache, v_cache, cu_seqlens_q, block_table, seqlen_k, max_seqlen_q): + if not _TRITON_AVAILABLE: + raise RuntimeError("Triton is unavailable") + if q.ndim != 3 or q.size(-1) != HEAD_DIM: + raise ValueError(f"packed Q must be [tokens, heads, {HEAD_DIM}]") + if k_cache.ndim != 4 or v_cache.shape != k_cache.shape or k_cache.size(1) != PAGE_SIZE: + raise ValueError(f"paged K/V must be [pages, {PAGE_SIZE}, kv_heads, {HEAD_DIM}]") + if k_cache.size(-1) != HEAD_DIM or q.size(1) % k_cache.size(2): + raise ValueError("paged Q/K head counts or dimensions are incompatible") + if q.dtype not in (torch.float16, torch.bfloat16) or k_cache.dtype != q.dtype: + raise ValueError("chunked flash attention supports one BF16/FP16 dtype for Q/K/V") + if v_cache.dtype != q.dtype: + raise ValueError("chunked flash attention supports one BF16/FP16 dtype for Q/K/V") + if q.stride(-1) != 1 or k_cache.stride(-1) != 1 or v_cache.stride(-1) != 1: + raise ValueError("Q/K/V head dimension must be contiguous") + batch = block_table.size(0) + if block_table.ndim != 2 or block_table.stride(-1) != 1: + raise ValueError("block_table must be a 2-D row-major page table") + if tuple(cu_seqlens_q.shape) != (batch + 1,) or tuple(seqlen_k.shape) != (batch,): + raise ValueError("cu_seqlens_q and seqlen_k must carry batch + 1 and batch entries") + for name, tensor in ( + ("cu_seqlens_q", cu_seqlens_q), + ("seqlen_k", seqlen_k), + ("block_table", block_table), + ): + if tensor.dtype != torch.int32: + raise ValueError(f"{name} must be int32") + if tensor.device != q.device: + raise ValueError(f"{name} must be on the Q device") + if not cu_seqlens_q.is_contiguous() or not seqlen_k.is_contiguous(): + raise ValueError("cu_seqlens_q and seqlen_k must be contiguous") + if max_seqlen_q <= 0: + raise ValueError("max_seqlen_q must be positive") + + +def paged_attention_forward( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + block_table: torch.Tensor, + seqlen_k: torch.Tensor, + max_seqlen_q: int, + scale: float, + out: torch.Tensor | None = None, + return_lse: bool = True, + schedule: str = "auto", + seq_of_token: torch.Tensor | None = None, + fast_path: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """Causal paged attention; returns ``(out [tokens, heads, D], lse [heads, tokens])``. + + Query ``i`` of a sequence sits at key position ``seqlen_k - q_len + i``, + which covers full prefill, prefix-cached extend and single-token decode. + """ + + _validate(q, k_cache, v_cache, cu_seqlens_q, block_table, seqlen_k, max_seqlen_q) + total_q, num_q_heads, _ = q.shape + num_kv_heads = k_cache.size(2) + group = num_q_heads // num_kv_heads + batch = block_table.size(0) + if out is None: + out = torch.empty_like(q) + elif out.shape != q.shape or out.dtype != q.dtype or out.stride(-1) != 1: + raise ValueError("out must match Q shape and dtype with a contiguous head dimension") + lse = torch.empty( + (num_q_heads, total_q) if return_lse else (0,), dtype=torch.float32, device=q.device + ) + if total_q == 0: + return out, lse + scale_log2 = float(scale) * _LOG2E + if schedule == "auto": + schedule = "split" if max_seqlen_q * group <= BLOCK_M else "monolithic" + if schedule not in ("split", "monolithic"): + raise ValueError("schedule must be 'auto', 'split' or 'monolithic'") + common = dict( + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + CHUNK=CHUNK_KV, + PAGE=PAGE_SIZE, + D=HEAD_DIM, + FAST_PATH=bool(fast_path), + ) + if schedule == "monolithic": + grid = (triton.cdiv(max_seqlen_q, BLOCK_M), num_q_heads, batch) + _attn_fwd_monolithic_kernel[grid]( + q, + k_cache, + v_cache, + out, + lse, + cu_seqlens_q, + seqlen_k, + block_table, + block_table.stride(0), + q.stride(0), + q.stride(1), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + out.stride(0), + out.stride(1), + lse.stride(0) if return_lse else 0, + 1, + scale_log2, + group, + WRITE_LSE=return_lse, + **common, + **_COMPILE_OPTIONS, + ) + return out, lse + if max_seqlen_q * group > BLOCK_M: + raise ValueError("split schedule requires max_seqlen_q * gqa_group <= BLOCK_M") + num_chunks = triton.cdiv(int(block_table.size(1)) * PAGE_SIZE, CHUNK_KV) + pm = torch.empty((num_chunks, total_q, num_q_heads), dtype=torch.float32, device=q.device) + pl = torch.empty_like(pm) + pacc = torch.empty( + (num_chunks, total_q, num_q_heads, HEAD_DIM), dtype=torch.float32, device=q.device + ) + grid = (num_chunks, num_kv_heads, batch) + _attn_fwd_split_partial_kernel[grid]( + q, + k_cache, + v_cache, + pm, + pl, + pacc, + cu_seqlens_q, + seqlen_k, + block_table, + block_table.stride(0), + q.stride(0), + q.stride(1), + k_cache.stride(0), + k_cache.stride(1), + k_cache.stride(2), + v_cache.stride(0), + v_cache.stride(1), + v_cache.stride(2), + pm.stride(0), + pm.stride(1), + pm.stride(2), + pacc.stride(0), + pacc.stride(1), + pacc.stride(2), + scale_log2, + group, + **common, + **_COMPILE_OPTIONS, + ) + if seq_of_token is None: + if total_q == batch: + seq_of_token = torch.arange(total_q, dtype=torch.int32, device=q.device) + else: + tokens = torch.arange(total_q, dtype=torch.int32, device=q.device) + seq_of_token = torch.searchsorted(cu_seqlens_q[1:], tokens, right=True).to(torch.int32) + _attn_fwd_split_merge_kernel[(total_q, num_q_heads)]( + pm, + pl, + pacc, + out, + lse, + seq_of_token, + seqlen_k, + pm.stride(0), + pm.stride(1), + pm.stride(2), + pacc.stride(0), + pacc.stride(1), + pacc.stride(2), + out.stride(0), + out.stride(1), + lse.stride(0) if return_lse else 0, + 1, + CHUNK=CHUNK_KV, + D=HEAD_DIM, + WRITE_LSE=return_lse, + num_warps=1, + enable_fp_fusion=False, + ) + return out, lse + + +def triton_paged_prefill( + q, + k, + v, + cuq, + indptr, + flat_pages, + maxq, + maxk, + dropout, + scale, + softcap, + zero_tensors, + causal, + window_left, + window_right, + sink, + return_lse, + return_dropout, + *, + block_table, + seqlen_k, + out=None, +): + """AITER ``mha_batch_prefill``-shaped entry point over the chunked contract. + + Mirrors the fixed CK entry point so the strict core can bind either one. + Single-token rows see their whole prefix under causal masking, so the + ``causal=False`` decode request is served by the same causal kernel. + """ + + del indptr, flat_pages, maxk + if dropout or softcap or sink or return_dropout or zero_tensors: + raise ValueError( + "chunked flash attention requires no dropout, softcap, sink or dropout mask" + ) + if window_left != -1 or window_right != -1: + raise ValueError("chunked flash attention does not support sliding windows") + if not causal and int(maxq) != 1: + raise ValueError("non-causal attention is only served for single-token decode rows") + resolved_scale = 1.0 / math.sqrt(q.size(-1)) if scale is None else float(scale) + out, lse = paged_attention_forward( + q, + k, + v, + cu_seqlens_q=cuq, + block_table=block_table, + seqlen_k=seqlen_k, + max_seqlen_q=int(maxq), + scale=resolved_scale, + out=out, + return_lse=bool(return_lse), + ) + rng_state = torch.empty((2,), dtype=torch.int64, device=q.device) + return out, lse, torch.empty((0,), dtype=q.dtype, device=q.device), rng_state + + +def warmup( + device: torch.device, *, num_q_heads: int, num_kv_heads: int, dtype: torch.dtype +) -> None: + """Compile both schedules so no Triton JIT happens inside HIP Graph capture.""" + + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("warm chunked flash attention before HIP Graph capture") + group = num_q_heads // num_kv_heads + pages = 3 + k_cache = torch.zeros((pages, PAGE_SIZE, num_kv_heads, HEAD_DIM), device=device, dtype=dtype) + v_cache = torch.zeros_like(k_cache) + block_table = torch.arange(pages, device=device, dtype=torch.int32).reshape(1, pages) + scale = 1.0 / math.sqrt(HEAD_DIM) + with torch.inference_mode(): + for q_len in (1, BLOCK_M // group, BLOCK_M + 1): + q = torch.zeros((q_len, num_q_heads, HEAD_DIM), device=device, dtype=dtype) + cu = torch.tensor((0, q_len), device=device, dtype=torch.int32) + seqlen = torch.tensor((pages * PAGE_SIZE,), device=device, dtype=torch.int32) + for schedule in ("split", "monolithic"): + if schedule == "split" and q_len * group > BLOCK_M: + continue + for return_lse in (False, True): + paged_attention_forward( + q, + k_cache, + v_cache, + cu_seqlens_q=cu, + block_table=block_table, + seqlen_k=seqlen, + max_seqlen_q=q_len, + scale=scale, + return_lse=return_lse, + schedule=schedule, + ) + torch.cuda.synchronize(device) + + +__all__ = [ + "BLOCK_M", + "BLOCK_N", + "CHUNKED_FLASH_ATTENTION_CONTRACT_ID", + "CHUNK_KV", + "HEAD_DIM", + "PAGE_SIZE", + "paged_attention_forward", + "triton_paged_prefill", + "warmup", +] diff --git a/rl_engine/kernels/ops/triton/matmul/mfma_gemm.py b/rl_engine/kernels/ops/triton/matmul/mfma_gemm.py new file mode 100644 index 00000000..e9bcd3fe --- /dev/null +++ b/rl_engine/kernels/ops/triton/matmul/mfma_gemm.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Batch-invariant MFMA GEMM for ROCm gfx942 with a fixed chunked-K contract. + +Arithmetic contract ``rlkernel.det_gemm.triton_mfma_rocm.v1``: + +* BF16 operands, FP32 accumulation on ``v_mfma_f32_16x16x16_bf16`` + (``matrix_instr_nonkdim=16``, ``kpack=2`` pinned; both change the in-tile + K order and are therefore part of the contract). +* K is consumed in ascending ``BLOCK_K=64`` tiles inside fixed + ``CHUNK_K``-wide chunks. Every chunk accumulates from zero; chunk partials + are combined in ascending order in FP32 and rounded to BF16 exactly once. +* No Split-K other than the chunk decomposition above, no autotuning. + +Every output element is a pure function of its own A row, its own B column and +the contract. Tile shape, warp count, pipelining depth, grid order, operand +strides, the row count ``M`` and whether the chunks are evaluated by one +program (monolithic schedule) or by ``num_chunks`` programs plus a fixed-order +reduction (split schedule) do not change a single bit. The split schedule +exists only to give small-M decode GEMMs enough programs to saturate HBM. + +The kernels accept arbitrary positive strides, so native ``[N, K]`` weights, +``[K, N]`` prepared weights and transposed activation views are consumed +without materializing a copy. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: # pragma: no cover - exercised only without Triton + _TRITON_AVAILABLE = False + +MFMA_GEMM_CONTRACT_ID = "rlkernel.det_gemm.triton_mfma_rocm.v1" +BLOCK_K = 64 +CHUNK_K = 1024 +MATRIX_INSTR_NONKDIM = 16 +KPACK = 2 +# Rows at or below this count use the split schedule when K spans several +# chunks. Purely a performance threshold; both schedules are bit-identical. +SPLIT_SCHEDULE_MAX_ROWS = 64 + + +@dataclass(frozen=True) +class MfmaGemmConfig: + block_m: int + block_n: int + num_warps: int + waves_per_eu: int = 2 + num_stages: int = 2 + group_m: int = 8 + + +_DECODE_CONFIG = MfmaGemmConfig(16, 32, 2, waves_per_eu=0, num_stages=2, group_m=1) +_SMALL_CONFIG = MfmaGemmConfig(64, 128, 4, waves_per_eu=2, num_stages=2, group_m=8) +_LARGE_CONFIG = MfmaGemmConfig(128, 128, 4, waves_per_eu=2, num_stages=2, group_m=8) + + +def select_config(m_size: int, n_size: int, k_size: int) -> MfmaGemmConfig: + """Pick a performance configuration. Never affects the result bits.""" + + del n_size, k_size + if m_size <= SPLIT_SCHEDULE_MAX_ROWS: + return _DECODE_CONFIG + if m_size <= 1024: + return _SMALL_CONFIG + return _LARGE_CONFIG + + +if _TRITON_AVAILABLE: + + @triton.jit + def _chunk_dot( + a_ptrs, + b_ptrs, + offs_k, + K, + chunk, + stride_ak, + stride_bk, + BLOCK_K: tl.constexpr, + CHUNK_K: tl.constexpr, + EVEN_K: tl.constexpr, + ): + """Accumulate one K chunk from zero in the pinned MFMA order.""" + + TILES_PER_CHUNK: tl.constexpr = CHUNK_K // BLOCK_K + # A trailing chunk shorter than CHUNK_K evaluates only its live tiles; + # tiles past K are never issued, so no exact-zero products enter the + # accumulator. + chunk_start = chunk * CHUNK_K + num_tiles = min(TILES_PER_CHUNK, tl.cdiv(K - chunk_start, BLOCK_K)) + acc = tl.zeros((a_ptrs.shape[0], b_ptrs.shape[1]), dtype=tl.float32) + for tile in range(0, num_tiles): + k0 = chunk_start + tile * BLOCK_K + if EVEN_K: + a = tl.load(a_ptrs + k0 * stride_ak) + b = tl.load(b_ptrs + k0 * stride_bk) + else: + kmask = offs_k < K - k0 + a = tl.load(a_ptrs + k0 * stride_ak, mask=kmask[None, :], other=0.0) + b = tl.load(b_ptrs + k0 * stride_bk, mask=kmask[:, None], other=0.0) + acc = tl.dot(a, b, acc) + return acc + + @triton.jit + def _mfma_gemm_kernel( + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + CHUNK_K: tl.constexpr, + GROUP_M: tl.constexpr, + EVEN_K: tl.constexpr, + ): + pid = tl.program_id(0) + num_pid_m = tl.cdiv(M, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + num_pid_in_group = GROUP_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_M) + pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + offs_k = tl.arange(0, BLOCK_K) + a_ptrs = a_ptr + offs_am[:, None].to(tl.int64) * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :].to(tl.int64) * stride_bn + + num_chunks = tl.cdiv(K, CHUNK_K) + total = _chunk_dot( + a_ptrs, b_ptrs, offs_k, K, 0, stride_ak, stride_bk, BLOCK_K, CHUNK_K, EVEN_K + ) + for chunk in range(1, num_chunks): + partial = _chunk_dot( + a_ptrs, b_ptrs, offs_k, K, chunk, stride_ak, stride_bk, BLOCK_K, CHUNK_K, EVEN_K + ) + total = total + partial + + offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + c_ptrs = ( + c_ptr + + offs_cm[:, None].to(tl.int64) * stride_cm + + offs_cn[None, :].to(tl.int64) * stride_cn + ) + mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, total.to(c_ptr.dtype.element_ty), mask=mask) + + @triton.jit + def _mfma_gemm_split_partial_kernel( + a_ptr, + b_ptr, + p_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_pc, + stride_pm, + stride_pn, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + CHUNK_K: tl.constexpr, + EVEN_K: tl.constexpr, + ): + pid_n = tl.program_id(0) + pid_m = tl.program_id(1) + chunk = tl.program_id(2) + offs_am = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_am = tl.where(offs_am < M, offs_am, 0) + offs_bn = tl.where(offs_bn < N, offs_bn, 0) + offs_k = tl.arange(0, BLOCK_K) + a_ptrs = a_ptr + offs_am[:, None].to(tl.int64) * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :].to(tl.int64) * stride_bn + acc = _chunk_dot( + a_ptrs, b_ptrs, offs_k, K, chunk, stride_ak, stride_bk, BLOCK_K, CHUNK_K, EVEN_K + ) + offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + p_ptrs = ( + p_ptr + + chunk * stride_pc + + offs_cm[:, None].to(tl.int64) * stride_pm + + offs_cn[None, :].to(tl.int64) * stride_pn + ) + mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(p_ptrs, acc, mask=mask) + + @triton.jit + def _mfma_gemm_split_reduce_kernel( + p_ptr, + c_ptr, + M, + N, + num_chunks, + stride_pc, + stride_pm, + stride_pn, + stride_cm, + stride_cn, + BLOCK: tl.constexpr, + ): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < M * N + m = offs // N + n = offs % N + src = p_ptr + m.to(tl.int64) * stride_pm + n.to(tl.int64) * stride_pn + total = tl.load(src, mask=mask, other=0.0) + for chunk in range(1, num_chunks): + total = total + tl.load(src + chunk * stride_pc, mask=mask, other=0.0) + dst = c_ptr + m.to(tl.int64) * stride_cm + n.to(tl.int64) * stride_cn + tl.store(dst, total.to(c_ptr.dtype.element_ty), mask=mask) + + +def _validate_operands(a: torch.Tensor, b: torch.Tensor) -> None: + if a.dim() != 2 or b.dim() != 2: + raise ValueError("MFMA GEMM expects two 2-D tensors") + if a.size(1) != b.size(0): + raise ValueError(f"MFMA GEMM K mismatch: {a.size(1)} and {b.size(0)}") + if a.dtype != torch.bfloat16 or b.dtype != torch.bfloat16: + raise TypeError("MFMA GEMM requires BF16 inputs") + if not a.is_cuda or not b.is_cuda or a.device != b.device: + raise RuntimeError("MFMA GEMM inputs must share one ROCm device") + for name, tensor in (("a", a), ("b", b)): + if min(tensor.stride()) < 0: + raise ValueError(f"MFMA GEMM operand {name} must not use negative strides") + if tensor.stride(0) != 1 and tensor.stride(1) != 1 and tensor.numel() > 1: + raise ValueError(f"MFMA GEMM operand {name} needs one unit-stride dimension") + + +def mfma_gemm( + a: torch.Tensor, + b: torch.Tensor, + *, + out: torch.Tensor | None = None, + config: MfmaGemmConfig | None = None, + force_split: bool | None = None, +) -> torch.Tensor: + """Return ``a @ b`` in BF16 under the pinned MFMA contract. + + ``a`` is ``[M, K]`` and ``b`` is ``[K, N]``; both may be strided views + (for example ``weight.t()`` for a native ``[N, K]`` weight). ``out`` may be + a preallocated BF16 ``[M, N]`` buffer, which is also allowed to be a + narrowed view of a larger staging allocation. + """ + + if not _TRITON_AVAILABLE: + raise RuntimeError("Triton is unavailable") + _validate_operands(a, b) + m_size, k_size = a.shape + n_size = b.size(1) + if out is None: + out = torch.empty((m_size, n_size), dtype=torch.bfloat16, device=a.device) + else: + if tuple(out.shape) != (m_size, n_size): + raise ValueError( + f"MFMA GEMM output must have shape {(m_size, n_size)}, got {tuple(out.shape)}" + ) + if out.dtype != torch.bfloat16: + raise TypeError(f"MFMA GEMM output must be BF16, got {out.dtype}") + if out.device != a.device: + raise RuntimeError("MFMA GEMM output must live on the input device") + if out.stride(1) != 1: + raise ValueError("MFMA GEMM output rows must be contiguous") + if out.requires_grad: + raise ValueError("MFMA GEMM output buffer must not require gradients") + if m_size == 0 or n_size == 0: + return out + if k_size == 0: + return out.zero_() + if config is None: + config = select_config(m_size, n_size, k_size) + even_k = k_size % BLOCK_K == 0 + num_chunks = triton.cdiv(k_size, CHUNK_K) + use_split = ( + num_chunks > 1 and m_size <= SPLIT_SCHEDULE_MAX_ROWS + if force_split is None + else bool(force_split) and num_chunks > 1 + ) + common = dict( + BLOCK_M=config.block_m, + BLOCK_N=config.block_n, + BLOCK_K=BLOCK_K, + CHUNK_K=CHUNK_K, + EVEN_K=even_k, + num_warps=config.num_warps, + waves_per_eu=config.waves_per_eu, + num_stages=config.num_stages, + matrix_instr_nonkdim=MATRIX_INSTR_NONKDIM, + kpack=KPACK, + ) + if not use_split: + grid = (triton.cdiv(m_size, config.block_m) * triton.cdiv(n_size, config.block_n),) + _mfma_gemm_kernel[grid]( + a, + b, + out, + m_size, + n_size, + k_size, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + out.stride(0), + out.stride(1), + GROUP_M=config.group_m, + **common, + ) + return out + partial = torch.empty((num_chunks, m_size, n_size), dtype=torch.float32, device=a.device) + grid = ( + triton.cdiv(n_size, config.block_n), + triton.cdiv(m_size, config.block_m), + num_chunks, + ) + _mfma_gemm_split_partial_kernel[grid]( + a, + b, + partial, + m_size, + n_size, + k_size, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + partial.stride(0), + partial.stride(1), + partial.stride(2), + **common, + ) + reduce_block = 1024 + _mfma_gemm_split_reduce_kernel[(triton.cdiv(m_size * n_size, reduce_block),)]( + partial, + out, + m_size, + n_size, + num_chunks, + partial.stride(0), + partial.stride(1), + partial.stride(2), + out.stride(0), + out.stride(1), + BLOCK=reduce_block, + num_warps=4, + ) + return out + + +def mfma_linear( + a: torch.Tensor, + weight: torch.Tensor, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """``a @ weight.T`` for a native ``[N, K]`` weight, no transpose copy.""" + + return mfma_gemm(a, weight.t(), out=out) + + +def mfma_linear_input_gradient(grad_output: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``dX = dY @ W`` with ``W`` in native ``[N, K]`` layout.""" + + return mfma_gemm(grad_output, weight) + + +def _reduction_major_copy(tensor: torch.Tensor) -> torch.Tensor: + """Return ``tensor.t()`` with the reduction axis contiguous. + + A column-major A operand (``dY.t()`` for the weight gradient) loads several + times slower than a row-major one on gfx942. The transpose copy is a + memory-bound pass that costs a few percent of the GEMM and does not change + any loaded value, so the result stays bitwise identical. + """ + + transposed = tensor.t() + if transposed.stride(1) == 1: + return transposed + return transposed.contiguous() + + +def mfma_linear_weight_gradient(a: torch.Tensor, grad_output: torch.Tensor) -> torch.Tensor: + """``dW = dY.T @ X`` returned in native ``[N, K]`` layout.""" + + return mfma_gemm(_reduction_major_copy(grad_output), a) + + +def warmup(device: torch.device | None = None) -> None: + """Compile every schedule/config variant so no JIT runs inside graph capture.""" + + if device is None: + device = torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("warm the MFMA GEMM before HIP Graph capture") + with torch.inference_mode(): + for k_size in (CHUNK_K, 2 * CHUNK_K + BLOCK_K, CHUNK_K + 8): + b = torch.zeros((k_size, 64), dtype=torch.bfloat16, device=device) + for rows in (1, SPLIT_SCHEDULE_MAX_ROWS + 1, 1025): + a = torch.zeros((rows, k_size), dtype=torch.bfloat16, device=device) + mfma_gemm(a, b) + mfma_gemm(a, b.t().contiguous().t()) + torch.cuda.synchronize(device) + + +class MfmaLinearFn(torch.autograd.Function): + @staticmethod + def forward(ctx, a: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(a, weight) + return mfma_linear(a, weight) + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + a, weight = ctx.saved_tensors + if grad_out.dtype != torch.bfloat16: + grad_out = grad_out.to(torch.bfloat16) + if grad_out.stride(1) != 1 and grad_out.stride(0) != 1: + grad_out = grad_out.contiguous() + da = mfma_linear_input_gradient(grad_out, weight) if ctx.needs_input_grad[0] else None + dw = mfma_linear_weight_gradient(a, grad_out) if ctx.needs_input_grad[1] else None + return da, dw + + +class MfmaGemmFn(torch.autograd.Function): + @staticmethod + def forward(ctx, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(a, b) + return mfma_gemm(a, b) + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + a, b = ctx.saved_tensors + if grad_out.dtype != torch.bfloat16: + grad_out = grad_out.to(torch.bfloat16) + if grad_out.stride(1) != 1 and grad_out.stride(0) != 1: + grad_out = grad_out.contiguous() + da = mfma_gemm(grad_out, b.t()) if ctx.needs_input_grad[0] else None + db = mfma_gemm(_reduction_major_copy(a), grad_out) if ctx.needs_input_grad[1] else None + return da, db + + +__all__ = [ + "BLOCK_K", + "CHUNK_K", + "KPACK", + "MATRIX_INSTR_NONKDIM", + "MFMA_GEMM_CONTRACT_ID", + "MfmaGemmConfig", + "MfmaGemmFn", + "MfmaLinearFn", + "mfma_gemm", + "mfma_linear", + "mfma_linear_input_gradient", + "mfma_linear_weight_gradient", + "select_config", + "warmup", +] diff --git a/tests/test_rocm_mfma_gemm.py b/tests/test_rocm_mfma_gemm.py new file mode 100644 index 00000000..1f12c201 --- /dev/null +++ b/tests/test_rocm_mfma_gemm.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Bitwise contract tests for the ROCm MFMA deterministic GEMM.""" + +from __future__ import annotations + +import importlib +import os + +import pytest +import torch + +IS_ROCM = getattr(torch.version, "hip", None) is not None +IS_GFX942 = ( + IS_ROCM + and torch.cuda.is_available() + and str(getattr(torch.cuda.get_device_properties(0), "gcnArchName", "")).startswith("gfx942") +) + +pytestmark = pytest.mark.skipif(not IS_GFX942, reason="MFMA det GEMM targets ROCm gfx942") + +if IS_GFX942: + from rl_engine.kernels.ops.triton.matmul import mfma_gemm as M + +DEV = "cuda" +# Qwen3-8B TP4 local shapes as (K, N). +QWEN_TP4_SHAPES = [(4096, 1536), (1024, 4096), (4096, 6144), (3072, 4096)] + + +def _rand(*shape, scale=1.0): + return (torch.randn(*shape, device=DEV) * scale).to(torch.bfloat16) + + +def _configs(): + return [ + M.MfmaGemmConfig(16, 16, 1, waves_per_eu=0, num_stages=1, group_m=1), + M.MfmaGemmConfig(16, 32, 2, waves_per_eu=0, num_stages=2, group_m=1), + M.MfmaGemmConfig(32, 128, 4, waves_per_eu=2, num_stages=2, group_m=8), + M.MfmaGemmConfig(64, 64, 4, waves_per_eu=2, num_stages=2, group_m=4), + M.MfmaGemmConfig(64, 128, 4, waves_per_eu=2, num_stages=3, group_m=8), + M.MfmaGemmConfig(128, 128, 4, waves_per_eu=2, num_stages=2, group_m=8), + M.MfmaGemmConfig(128, 128, 8, waves_per_eu=1, num_stages=2, group_m=16), + M.MfmaGemmConfig(128, 256, 8, waves_per_eu=2, num_stages=2, group_m=8), + M.MfmaGemmConfig(256, 128, 8, waves_per_eu=1, num_stages=2, group_m=8), + ] + + +@pytest.mark.parametrize("k_size,n_size", QWEN_TP4_SHAPES) +def test_forward_matches_fp32_reference(k_size, n_size): + a = _rand(300, k_size) + w = _rand(n_size, k_size, scale=0.02) + out = M.mfma_linear(a, w) + ref = a.float() @ w.float().t() + assert out.shape == (300, n_size) + tol = 8e-3 * ref.abs().max().item() + assert (out.float() - ref).abs().max().item() <= tol + + +@pytest.mark.parametrize("k_size,n_size", QWEN_TP4_SHAPES) +def test_every_schedule_and_layout_is_bitwise_identical(k_size, n_size): + a = _rand(4096 + 37, k_size) + w = _rand(n_size, k_size, scale=0.02) + reference = M.mfma_gemm(a, w.t()) + layouts = {"nk_view": w.t(), "kn_contiguous": w.t().contiguous()} + for config in _configs(): + for name, b in layouts.items(): + for force_split in (False, True): + out = M.mfma_gemm(a, b, config=config, force_split=force_split) + assert torch.equal(out, reference), (config, name, force_split) + + +@pytest.mark.parametrize("k_size,n_size", QWEN_TP4_SHAPES) +def test_rows_are_batch_invariant(k_size, n_size): + a = _rand(4096 + 37, k_size) + w = _rand(n_size, k_size, scale=0.02) + reference = M.mfma_linear(a, w) + configs = _configs() + for rows, start in ((1, 5), (7, 100), (8, 0), (32, 1000), (33, 4000), (129, 77), (1024, 3000)): + sub = a[start : start + rows] + for config in (configs[0], configs[2], configs[8]): + for force_split in (False, True): + out = M.mfma_gemm(sub, w.t(), config=config, force_split=force_split) + assert torch.equal(out, reference[start : start + rows]), (rows, start, config) + strided = a[::3][:64] + assert torch.equal(M.mfma_linear(strided, w), reference[::3][:64]) + + +def test_column_major_a_operand_is_bitwise_identical(): + grad = _rand(4096, 6144) + a = _rand(4096, 4096) + reference = M.mfma_gemm(grad.t().contiguous(), a) + assert torch.equal(M.mfma_gemm(grad.t(), a), reference) + assert torch.equal(M.mfma_linear_weight_gradient(a, grad), reference) + ref = grad.float().t() @ a.float() + assert (reference.float() - ref).abs().max().item() <= 8e-3 * ref.abs().max().item() + + +def test_output_buffer_and_narrowed_staging_views(): + a = _rand(24, 4096) + w = _rand(1536, 4096, scale=0.02) + reference = M.mfma_linear(a, w) + staging = torch.zeros((32, 1536), dtype=torch.bfloat16, device=DEV) + out = M.mfma_linear(a, w, out=staging.narrow(0, 0, 24)) + assert out.data_ptr() == staging.data_ptr() + assert torch.equal(staging[:24], reference) + assert torch.equal(staging[24:], torch.zeros_like(staging[24:])) + + +def test_ragged_k_and_ragged_mn_match_padded_computation(): + a = _rand(45, 1000) + b = _rand(1000, 77) + out = M.mfma_gemm(a, b) + a_pad = torch.zeros((45, 1024), dtype=torch.bfloat16, device=DEV) + a_pad[:, :1000] = a + b_pad = torch.zeros((1024, 77), dtype=torch.bfloat16, device=DEV) + b_pad[:1000] = b + assert torch.equal(out, M.mfma_gemm(a_pad, b_pad)) + ref = a.float() @ b.float() + assert (out.float() - ref).abs().max().item() <= 8e-3 * ref.abs().max().item() + + +def test_backward_is_deterministic_and_close_to_reference(): + a = _rand(512, 4096).requires_grad_(True) + w = _rand(1536, 4096, scale=0.02).requires_grad_(True) + grad = _rand(512, 1536) + outputs = [] + for _ in range(2): + out = M.MfmaLinearFn.apply(a, w) + out.backward(grad) + outputs.append((out.detach().clone(), a.grad.clone(), w.grad.clone())) + a.grad = None + w.grad = None + assert torch.equal(outputs[0][0], outputs[1][0]) + assert torch.equal(outputs[0][1], outputs[1][1]) + assert torch.equal(outputs[0][2], outputs[1][2]) + ref_da = grad.float() @ w.detach().float() + ref_dw = grad.float().t() @ a.detach().float() + assert (outputs[0][1].float() - ref_da).abs().max().item() <= 8e-3 * ref_da.abs().max().item() + assert (outputs[0][2].float() - ref_dw).abs().max().item() <= 8e-3 * ref_dw.abs().max().item() + + +def test_generic_gemm_backward_matches_linear_backward_bitwise(): + a = _rand(256, 3072).requires_grad_(True) + w = _rand(4096, 3072, scale=0.02).requires_grad_(True) + grad = _rand(256, 4096) + linear_out = M.MfmaLinearFn.apply(a, w) + linear_out.backward(grad) + linear_grads = (a.grad.clone(), w.grad.clone()) + a.grad = None + w.grad = None + generic_out = M.MfmaGemmFn.apply(a, w.t()) + generic_out.backward(grad) + assert torch.equal(linear_out, generic_out) + assert torch.equal(linear_grads[0], a.grad) + assert torch.equal(linear_grads[1], w.grad) + + +def test_facade_routes_to_mfma_by_default(monkeypatch): + monkeypatch.delenv("RL_KERNEL_DET_GEMM_BACKEND", raising=False) + from rl_engine.kernels.ops.rocm.matmul import det_gemm as facade + + facade = importlib.reload(facade) + assert facade.det_gemm_backend() == "triton_mfma" + assert facade.det_gemm_backend_id() == M.MFMA_GEMM_CONTRACT_ID + a = _rand(40, 4096) + w = _rand(1536, 4096, scale=0.02) + assert torch.equal(facade.det_gemm_linear(a, w), M.mfma_linear(a, w)) + assert torch.equal(facade.det_gemm_linear_prepared(a, w.t().contiguous()), M.mfma_linear(a, w)) + grad = _rand(40, 1536) + assert torch.equal( + facade.det_gemm_linear_input_gradient(grad, w), M.mfma_linear_input_gradient(grad, w) + ) + assert torch.equal( + facade.det_gemm_linear_weight_gradient(a, grad), M.mfma_linear_weight_gradient(a, grad) + ) + op = facade.RocmDetGemmOp() + assert torch.equal(op(a, w.t().contiguous()), M.mfma_linear(a, w)) + assert torch.equal(op.linear(a, w), M.mfma_linear(a, w)) + + +def test_facade_tree_backend_remains_selectable(monkeypatch): + monkeypatch.setenv("RL_KERNEL_DET_GEMM_BACKEND", "triton_tree") + from rl_engine.kernels.ops.rocm.matmul import det_gemm as facade + + facade = importlib.reload(facade) + try: + assert facade.det_gemm_backend() == "triton_tree" + assert facade.det_gemm_backend_id() == "rlkernel.det_gemm.triton_tree_rocm.v1" + from rl_engine.kernels.ops.triton.matmul.det_gemm import _triton_tree_gemm + + a = _rand(8, 1024) + w = _rand(256, 1024, scale=0.02) + assert torch.equal(facade.det_gemm_linear(a, w), _triton_tree_gemm(a, w.t().contiguous())) + finally: + monkeypatch.delenv("RL_KERNEL_DET_GEMM_BACKEND", raising=False) + importlib.reload(facade) + + +def test_facade_rejects_unknown_backend(monkeypatch): + monkeypatch.setenv("RL_KERNEL_DET_GEMM_BACKEND", "cublas") + from rl_engine.kernels.ops.rocm.matmul import det_gemm as facade + + with pytest.raises(RuntimeError, match="RL_KERNEL_DET_GEMM_BACKEND"): + importlib.reload(facade) + monkeypatch.delenv("RL_KERNEL_DET_GEMM_BACKEND", raising=False) + importlib.reload(facade) + assert os.getenv("RL_KERNEL_DET_GEMM_BACKEND") is None diff --git a/tests/test_rocm_triton_chunked_attention.py b/tests/test_rocm_triton_chunked_attention.py new file mode 100644 index 00000000..ebacbc71 --- /dev/null +++ b/tests/test_rocm_triton_chunked_attention.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Bitwise contract tests for the ROCm Triton chunked flash attention.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +IS_ROCM = getattr(torch.version, "hip", None) is not None +IS_GFX942 = ( + IS_ROCM + and torch.cuda.is_available() + and str(getattr(torch.cuda.get_device_properties(0), "gcnArchName", "")).startswith("gfx942") +) + +pytestmark = pytest.mark.skipif(not IS_GFX942, reason="chunked flash attention targets ROCm gfx942") + +if IS_GFX942: + from rl_engine.kernels.ops.triton.attention import chunked_flash_attn as A + +DEV = "cuda" +HQ, HKV, D = 8, 2, 128 # Qwen3-8B TP4 local heads +PAGE = 16 +SCALE = 1.0 / math.sqrt(D) + + +def _paged(seqs, *, extra_pages=7): + total_pages = sum((L + PAGE - 1) // PAGE for L in seqs) + extra_pages + perm = torch.randperm(total_pages, device=DEV) + k_cache = (torch.randn(total_pages, PAGE, HKV, D, device=DEV) * 0.5).bfloat16() + v_cache = (torch.randn(total_pages, PAGE, HKV, D, device=DEV) * 0.5).bfloat16() + max_pages = max((L + PAGE - 1) // PAGE for L in seqs) + table = torch.zeros(len(seqs), max_pages, dtype=torch.int32, device=DEV) + cursor = 0 + dense = [] + for row, length in enumerate(seqs): + count = (length + PAGE - 1) // PAGE + pages = perm[cursor : cursor + count] + cursor += count + table[row, :count] = pages.to(torch.int32) + table[row, count:] = pages[0] + dense.append( + ( + k_cache[pages].reshape(-1, HKV, D)[:length], + v_cache[pages].reshape(-1, HKV, D)[:length], + ) + ) + return k_cache, v_cache, table, dense + + +def _reference(q, k, v, q_pos): + group = HQ // HKV + kf = k.float().repeat_interleave(group, dim=1) + vf = v.float().repeat_interleave(group, dim=1) + scores = torch.einsum("qhd,khd->hqk", q.float(), kf) * SCALE + key_pos = torch.arange(k.size(0), device=DEV) + scores = scores.masked_fill((key_pos[None, :] > q_pos[:, None])[None], float("-inf")) + return torch.einsum("hqk,khd->qhd", torch.softmax(scores, -1), vf) + + +def _run(q, k_cache, v_cache, table, cu, seqlen_k, max_q, *, schedule="auto"): + return A.paged_attention_forward( + q, + k_cache, + v_cache, + cu_seqlens_q=cu, + block_table=table, + seqlen_k=seqlen_k, + max_seqlen_q=max_q, + scale=SCALE, + schedule=schedule, + ) + + +@pytest.mark.parametrize("length", [1, 17, 64, 65, 500, 1024, 1500, 4096]) +def test_prefill_decode_and_extend_rows_are_bitwise_identical(length): + torch.manual_seed(length) + k_cache, v_cache, table, dense = _paged([length]) + q = (torch.randn(length, HQ, D, device=DEV) * 0.5).bfloat16() + cu = torch.tensor([0, length], dtype=torch.int32, device=DEV) + seqlen_k = torch.tensor([length], dtype=torch.int32, device=DEV) + out_prefill, lse_prefill = _run( + q, k_cache, v_cache, table, cu, seqlen_k, length, schedule="monolithic" + ) + reference = _reference(q, dense[0][0], dense[0][1], torch.arange(length, device=DEV)) + assert (out_prefill.float() - reference).abs().max().item() <= 1e-2 + + positions = sorted( + { + p + for p in ( + 0, + 1, + 15, + 16, + 63, + 64, + 65, + 127, + 128, + 511, + 512, + 513, + 1023, + 1024, + length - 2, + length - 1, + ) + if 0 <= p < length + } + ) + q_decode = q[positions] + cu_decode = torch.arange(len(positions) + 1, dtype=torch.int32, device=DEV) + seqlen_decode = torch.tensor([p + 1 for p in positions], dtype=torch.int32, device=DEV) + table_decode = table.expand(len(positions), -1).contiguous() + for schedule in ("split", "monolithic"): + out_decode, lse_decode = _run( + q_decode, k_cache, v_cache, table_decode, cu_decode, seqlen_decode, 1, schedule=schedule + ) + assert torch.equal(out_decode, out_prefill[positions]), schedule + assert torch.equal(lse_decode, lse_prefill[:, positions]), schedule + + if length > 40: + start = length - 37 + cu_extend = torch.tensor([0, length - start], dtype=torch.int32, device=DEV) + for schedule in ("monolithic", "split"): + if schedule == "split" and (length - start) * (HQ // HKV) > A.BLOCK_M: + continue + out_extend, _ = _run( + q[start:], + k_cache, + v_cache, + table, + cu_extend, + seqlen_k, + length - start, + schedule=schedule, + ) + assert torch.equal(out_extend, out_prefill[start:]), schedule + + +def test_batch_composition_is_invariant(): + torch.manual_seed(7) + lengths = [700, 1, 4096, 33] + k_cache, v_cache, table, dense = _paged(lengths) + queries = [(torch.randn(L, HQ, D, device=DEV) * 0.5).bfloat16() for L in lengths] + q_all = torch.cat(queries) + cu = torch.tensor( + [0, *torch.cumsum(torch.tensor(lengths), 0).tolist()], dtype=torch.int32, device=DEV + ) + seqlen_k = torch.tensor(lengths, dtype=torch.int32, device=DEV) + out_all, _ = _run( + q_all, k_cache, v_cache, table, cu, seqlen_k, max(lengths), schedule="monolithic" + ) + offset = 0 + for index, length in enumerate(lengths): + out_one, _ = _run( + queries[index], + k_cache, + v_cache, + table[index : index + 1], + torch.tensor([0, length], dtype=torch.int32, device=DEV), + seqlen_k[index : index + 1], + length, + schedule="monolithic", + ) + assert torch.equal(out_one, out_all[offset : offset + length]) + reference = _reference( + queries[index], dense[index][0], dense[index][1], torch.arange(length, device=DEV) + ) + assert (out_one.float() - reference).abs().max().item() <= 1e-2 + offset += length + q_last = torch.stack([queries[index][-1] for index in range(4)]) + out_decode, _ = _run( + q_last, + k_cache, + v_cache, + table, + torch.arange(5, dtype=torch.int32, device=DEV), + seqlen_k, + 1, + schedule="split", + ) + for index in range(4): + assert torch.equal(out_decode[index], out_all[cu[index + 1] - 1]) + + +def test_mixed_decode_extend_prefill_batch_matches_isolated_rows(): + torch.manual_seed(11) + lengths = [2048, 300, 900] + query_lengths = [1, 17, 900] # decode, extend, fresh prefill + k_cache, v_cache, table, _dense = _paged(lengths) + full = [(torch.randn(L, HQ, D, device=DEV) * 0.5).bfloat16() for L in lengths] + isolated = [] + for index, length in enumerate(lengths): + out, _ = _run( + full[index], + k_cache, + v_cache, + table[index : index + 1], + torch.tensor([0, length], dtype=torch.int32, device=DEV), + torch.tensor([length], dtype=torch.int32, device=DEV), + length, + schedule="monolithic", + ) + isolated.append(out[length - query_lengths[index] :]) + q_mixed = torch.cat([full[i][lengths[i] - query_lengths[i] :] for i in range(3)]) + cu = torch.tensor( + [0, *torch.cumsum(torch.tensor(query_lengths), 0).tolist()], dtype=torch.int32, device=DEV + ) + seqlen_k = torch.tensor(lengths, dtype=torch.int32, device=DEV) + out_mixed, _ = _run(q_mixed, k_cache, v_cache, table, cu, seqlen_k, max(query_lengths)) + assert torch.equal(out_mixed, torch.cat(isolated)) + + +@pytest.mark.parametrize("length", [1, 64, 65, 700, 1500, 4096]) +def test_unmasked_fast_path_matches_masked_path_bitwise(length): + torch.manual_seed(length + 100) + k_cache, v_cache, table, _dense = _paged([length, 33]) + q = (torch.randn(length + 33, HQ, D, device=DEV) * 0.5).bfloat16() + cu = torch.tensor([0, length, length + 33], dtype=torch.int32, device=DEV) + seqlen_k = torch.tensor([length, 33], dtype=torch.int32, device=DEV) + max_q = max(length, 33) + for schedule in ("monolithic", "split"): + if schedule == "split" and max_q * (HQ // HKV) > A.BLOCK_M: + continue + slow, slow_lse = A.paged_attention_forward( + q, + k_cache, + v_cache, + cu_seqlens_q=cu, + block_table=table, + seqlen_k=seqlen_k, + max_seqlen_q=max_q, + scale=SCALE, + schedule=schedule, + fast_path=False, + ) + fast, fast_lse = A.paged_attention_forward( + q, + k_cache, + v_cache, + cu_seqlens_q=cu, + block_table=table, + seqlen_k=seqlen_k, + max_seqlen_q=max_q, + scale=SCALE, + schedule=schedule, + fast_path=True, + ) + assert torch.equal(slow, fast), schedule + assert torch.equal(slow_lse, fast_lse), schedule + # decode rows through the split schedule against the masked prefill rows + positions = [p for p in (0, 63, 64, 511, 512, 1000, length - 1) if 0 <= p < length] + q_decode = q[positions] + cu_decode = torch.arange(len(positions) + 1, dtype=torch.int32, device=DEV) + seqlen_decode = torch.tensor([p + 1 for p in positions], dtype=torch.int32, device=DEV) + table_decode = table[:1].expand(len(positions), -1).contiguous() + prefill, _ = A.paged_attention_forward( + q[:length], + k_cache, + v_cache, + cu_seqlens_q=cu[:2], + block_table=table[:1], + seqlen_k=seqlen_k[:1], + max_seqlen_q=length, + scale=SCALE, + schedule="monolithic", + fast_path=False, + ) + decode, _ = A.paged_attention_forward( + q_decode, + k_cache, + v_cache, + cu_seqlens_q=cu_decode, + block_table=table_decode, + seqlen_k=seqlen_decode, + max_seqlen_q=1, + scale=SCALE, + schedule="split", + fast_path=True, + ) + assert torch.equal(decode, prefill[positions]) + + +def test_output_buffer_and_padded_rows_are_left_alone(): + torch.manual_seed(3) + k_cache, v_cache, table, _dense = _paged([100, 50]) + q = (torch.randn(2, HQ, D, device=DEV) * 0.5).bfloat16() + cu = torch.tensor([0, 1, 2], dtype=torch.int32, device=DEV) + seqlen_k = torch.tensor([100, 50], dtype=torch.int32, device=DEV) + expected, _ = _run(q, k_cache, v_cache, table, cu, seqlen_k, 1) + buffer = torch.zeros((4, HQ, D), dtype=torch.bfloat16, device=DEV) + out, _ = A.paged_attention_forward( + q, + k_cache, + v_cache, + cu_seqlens_q=cu, + block_table=table, + seqlen_k=seqlen_k, + max_seqlen_q=1, + scale=SCALE, + out=buffer.narrow(0, 0, 2), + ) + assert out.data_ptr() == buffer.data_ptr() + assert torch.equal(buffer[:2], expected) + assert torch.equal(buffer[2:], torch.zeros_like(buffer[2:])) + + +def test_aiter_shaped_entry_point_and_warmup(): + torch.manual_seed(5) + k_cache, v_cache, table, _dense = _paged([64, 64]) + q = (torch.randn(128, HQ, D, device=DEV) * 0.5).bfloat16() + cu = torch.tensor([0, 64, 128], dtype=torch.int32, device=DEV) + seqlen_k = torch.tensor([64, 64], dtype=torch.int32, device=DEV) + indptr = torch.tensor([0, 4, 8], dtype=torch.int32, device=DEV) + out, lse, mask, rng = A.triton_paged_prefill( + q, + k_cache, + v_cache, + cu, + indptr, + table.reshape(-1), + 64, + 64, + 0.0, + SCALE, + 0.0, + False, + True, + -1, + -1, + 0, + True, + False, + block_table=table, + seqlen_k=seqlen_k, + ) + expected, expected_lse = _run(q, k_cache, v_cache, table, cu, seqlen_k, 64) + assert torch.equal(out, expected) + assert torch.equal(lse, expected_lse) + assert mask.numel() == 0 and rng.shape == (2,) + with pytest.raises(ValueError): + A.triton_paged_prefill( + q, + k_cache, + v_cache, + cu, + indptr, + table.reshape(-1), + 64, + 64, + 0.1, + SCALE, + 0.0, + False, + True, + -1, + -1, + 0, + True, + False, + block_table=table, + seqlen_k=seqlen_k, + ) + A.warmup(torch.device(DEV), num_q_heads=HQ, num_kv_heads=HKV, dtype=torch.bfloat16) + + +def test_strict_core_binds_the_triton_contract(monkeypatch): + monkeypatch.setenv("RL_KERNEL_ROCM_ATTENTION_BACKEND", "triton") + monkeypatch.setenv("RL_KERNEL_ROCM_FIXED_PAGED_TILE", "128") + from rl_engine.kernels.ops.rocm.attention import flash_attn as F + + core = F.StrictRocmAiterCKAttentionCore() + assert core.attention_backend == "triton" + assert core.paged_kernel_id == A.CHUNKED_FLASH_ATTENTION_CONTRACT_ID + assert core.paged_entrypoint_id == A.CHUNKED_FLASH_ATTENTION_CONTRACT_ID + assert core.supports_paged_schedule and core.supports_mixed_paged_batches + torch.manual_seed(9) + k_cache, v_cache, table, _dense = _paged([48]) + q = (torch.randn(48, HQ, D, device=DEV) * 0.5).bfloat16() + cu = torch.tensor([0, 48], dtype=torch.int32, device=DEV) + seqlen_k = torch.tensor([48], dtype=torch.int32, device=DEV) + indptr = torch.tensor([0, table.size(1)], dtype=torch.int32, device=DEV) + with torch.inference_mode(): + result = core.forward_paged_varlen_with_lse( + q, + k_cache, + v_cache, + page_table=table, + seqused_k=seqlen_k, + cu_seqlens_q=cu, + kv_indptr=indptr, + max_seqlen_q=48, + max_seqlen_k=48, + causal=True, + scale=SCALE, + ) + expected, expected_lse = _run(q, k_cache, v_cache, table, cu, seqlen_k, 48) + assert torch.equal(result.out, expected) + assert torch.equal(result.lse, expected_lse) + assert result.provenance["forward_entrypoint"] == A.CHUNKED_FLASH_ATTENTION_CONTRACT_ID + assert result.provenance["paged_kernel"] == A.CHUNKED_FLASH_ATTENTION_CONTRACT_ID + # dense training-style forward routes through the same paged contract + dense_q = q.permute(1, 0, 2).unsqueeze(0).contiguous() + dense_k = k_cache[table[0, :3]].reshape(-1, HKV, D).permute(1, 0, 2).unsqueeze(0).contiguous() + dense_v = v_cache[table[0, :3]].reshape(-1, HKV, D).permute(1, 0, 2).unsqueeze(0).contiguous() + positions = torch.arange(48, device=DEV).unsqueeze(0) + dense = core.forward_with_lse( + dense_q, + dense_k, + dense_v, + causal=True, + scale=SCALE, + query_position_ids=positions, + key_position_ids=positions, + ) + assert torch.equal(dense.out.squeeze(0).permute(1, 0, 2), expected) + assert dense.provenance["forward_entrypoint"] == A.CHUNKED_FLASH_ATTENTION_CONTRACT_ID