From 2f6b2e9faf376c68c9bd33e3c5d9e80356cf7681 Mon Sep 17 00:00:00 2001 From: vensen Date: Wed, 9 Sep 2026 06:47:53 -0700 Subject: [PATCH 1/3] perf(rocm): MFMA batch-invariant GEMM and chunked Triton attention for strict R/R Replace the scalar-FMA/BF16-K-tree deterministic GEMM on gfx942 with a pinned-order MFMA kernel (rlkernel.det_gemm.triton_mfma_rocm.v1): fixed 16x16x16 bf16 MFMA with kpack=2, BLOCK_K=64 tiles inside 1024-wide K chunks, FP32 chunk partials combined ascending, single BF16 rounding. Every output element depends only on its own row and column, so tile shape, warp count, weight layout, row count and the split-K decode schedule are bit-identical (verified by tests). Selected through RL_KERNEL_DET_GEMM_BACKEND (auto -> MFMA on gfx942, triton_tree keeps the previous contract). Add a chunked-KV flash attention contract for ROCm (rlkernel.rocm.triton_chunked_flash_attention.v1) that serves Megatron full-sequence forward, vLLM prefill, prefix-cached extend and paged decode from one arithmetic path (fixed 64-key blocks inside 512-token chunks with an ascending exact merge), so all four agree bitwise by construction. Opt in with RL_KERNEL_ROCM_ATTENTION_BACKEND=triton; the CK fixed-M128 path stays default. Vime harness: parameterized PR377-workload runner and summary table, proxy variables stripped from Ray arms, GPU_ARCHS pinned for AITER JIT builds, validator accepts the two Triton contracts, decode-step profiler scripts. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BunaEZQBLE3hDxHyYcb76V --- benchmarks/profile_rocm_rollout_decode.py | 152 ++++ benchmarks/summarize_rollout_trace.py | 86 ++ .../launch_arm.sh | 16 +- examples/vime_rocm_attention_ablation/run.py | 16 + .../run_pr377_workload.py | 289 +++++++ .../summarize_pr377_runs.py | 181 ++++ .../validate_artifacts.py | 18 +- rl_engine/integrations/framework_operators.py | 76 +- .../kernels/ops/rocm/attention/flash_attn.py | 64 +- .../ops/rocm/attention/strict_runtime.py | 7 +- rl_engine/kernels/ops/rocm/matmul/det_gemm.py | 80 +- .../triton/attention/chunked_flash_attn.py | 810 ++++++++++++++++++ .../kernels/ops/triton/matmul/mfma_gemm.py | 486 +++++++++++ tests/test_rocm_mfma_gemm.py | 207 +++++ tests/test_rocm_triton_chunked_attention.py | 422 +++++++++ 15 files changed, 2873 insertions(+), 37 deletions(-) create mode 100644 benchmarks/profile_rocm_rollout_decode.py create mode 100644 benchmarks/summarize_rollout_trace.py create mode 100644 examples/vime_rocm_attention_ablation/run_pr377_workload.py create mode 100644 examples/vime_rocm_attention_ablation/summarize_pr377_runs.py create mode 100644 rl_engine/kernels/ops/triton/attention/chunked_flash_attn.py create mode 100644 rl_engine/kernels/ops/triton/matmul/mfma_gemm.py create mode 100644 tests/test_rocm_mfma_gemm.py create mode 100644 tests/test_rocm_triton_chunked_attention.py 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/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 From 7a9f3b5657ac380ae2419539975498513d18c490 Mon Sep 17 00:00:00 2001 From: vensen Date: Wed, 9 Sep 2026 07:13:48 -0700 Subject: [PATCH 2/3] docs(rocm): add PR 396 performance analysis report Where the strict R/R time went, the MFMA GEMM and chunked attention microbenchmarks, the one-round PR377 workload results for P/P versus both R/R variants, the node-degradation caveat and the remaining gaps. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01BunaEZQBLE3hDxHyYcb76V --- .../results/pr396_rocm_mi300x/report.md | 161 ++++++++++++++++++ 1 file changed, 161 insertions(+) create mode 100644 benchmarks/results/pr396_rocm_mi300x/report.md diff --git a/benchmarks/results/pr396_rocm_mi300x/report.md b/benchmarks/results/pr396_rocm_mi300x/report.md new file mode 100644 index 00000000..c68c5606 --- /dev/null +++ b/benchmarks/results/pr396_rocm_mi300x/report.md @@ -0,0 +1,161 @@ +# 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. Same machine, 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 | + +### 4b. Healthy-node points (before the degradation described below) + +| 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 is needed before the +30-round comparison; the absolute numbers in 4a 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 ~2x native per step even though the GEMMs are within + 1.3x of hipBLASLt and the attention core is faster than CK. The + remaining cost is outside the kernels: eager custom-op boundaries around + the deterministic all-reduce and the packed FFN, plus attention launch + overhead. `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. From e1aaccc6003f709a6ef82ad626cc6501f36df924 Mon Sep 17 00:00:00 2001 From: vensen Date: Wed, 9 Sep 2026 07:43:45 -0700 Subject: [PATCH 3/3] docs(rocm): add healthy-node one-round P/P vs R/R results to the PR 396 report Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Ms5CP2SrVboniULA1vE8s1 --- .../results/pr396_rocm_mi300x/report.md | 43 ++++++++++++++----- 1 file changed, 33 insertions(+), 10 deletions(-) diff --git a/benchmarks/results/pr396_rocm_mi300x/report.md b/benchmarks/results/pr396_rocm_mi300x/report.md index c68c5606..2820cf3c 100644 --- a/benchmarks/results/pr396_rocm_mi300x/report.md +++ b/benchmarks/results/pr396_rocm_mi300x/report.md @@ -98,7 +98,31 @@ 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. Same machine, same conditions (vLLM memory utilization 0.30) +### 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 | |---|---:|---:|:---:| @@ -117,7 +141,7 @@ FULL_AND_PIECEWISE (capture 32), `RL_KERNEL_ROCM_FIXED_PAGED_TILE=128`. | 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 | -### 4b. Healthy-node points (before the degradation described below) +### 4c. Healthy-node points measured before the degradation | Run | rollout | log probs | actor train | step | consistency | |---|---:|---:|---:|---:|---| @@ -138,9 +162,9 @@ 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 is needed before the -30-round comparison; the absolute numbers in 4a are inflated for every arm, -the relative ones are same-condition. +`/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 @@ -150,11 +174,10 @@ strips the proxy variables and pins `GPU_ARCHS`. ## 6. What is still missing for a net win -- Rollout decode is ~2x native per step even though the GEMMs are within - 1.3x of hipBLASLt and the attention core is faster than CK. The - remaining cost is outside the kernels: eager custom-op boundaries around - the deterministic all-reduce and the packed FFN, plus attention launch - overhead. `profile_rocm_rollout_decode.py` + `summarize_rollout_trace.py` +- 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