From 9a991e3b21f72d0070bdb458ad2393d2d713a75b Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Wed, 14 Oct 2026 12:07:18 +0800 Subject: [PATCH 01/26] Add H3 FP32 video input projection with Triton and CUTLASS Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 293 +++++++++++++++ build_tools/extensions.py | 50 +++ .../video_input_projection_fp32_cutlass.cpp | 25 ++ .../video_input_projection_fp32_cutlass.cu | 61 +++ docs/.nav.yml | 1 + docs/operators/video-input-projection-fp32.md | 149 ++++++++ rl_engine/_C_cutlass.pyi | 9 + .../cuda/gemm/video_input_projection_fp32.py | 90 +++++ .../backends/shared/triton/gemm/det_gemm.py | 1 + .../gemm/video_input_projection_fp32.py | 40 ++ rl_engine/ops/gemm/canonical_linear.py | 110 +++++- .../ops/gemm/video_input_projection_fp32.py | 114 ++++++ .../gemm/video_input_projection_fp32.py | 17 + rl_engine/runtime/registry.py | 72 ++++ rl_engine/validation/operators/op_checks.py | 34 +- .../validation/operators/operator_inputs.py | 22 ++ .../validation/operators/operator_specs.py | 36 ++ .../build/test_build_platform_collectives.py | 30 ++ .../gemm/test_video_input_projection_fp32.py | 353 ++++++++++++++++++ tests/runtime/test_kernel_registry.py | 14 + tests/validation/operators/test_op_checks.py | 59 +++ .../operators/test_operator_inputs.py | 16 + 22 files changed, 1571 insertions(+), 25 deletions(-) create mode 100644 benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py create mode 100644 csrc/bindings/video_input_projection_fp32_cutlass.cpp create mode 100644 csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu create mode 100644 docs/operators/video-input-projection-fp32.md create mode 100644 rl_engine/_C_cutlass.pyi create mode 100644 rl_engine/backends/cuda/gemm/video_input_projection_fp32.py create mode 100644 rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py create mode 100644 rl_engine/ops/gemm/video_input_projection_fp32.py create mode 100644 rl_engine/reference/gemm/video_input_projection_fp32.py create mode 100644 tests/ops/gemm/test_video_input_projection_fp32.py diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py new file mode 100644 index 000000000..64b1df48c --- /dev/null +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Measure complete H3 FP32 projection calls after reference accuracy checks.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import platform +import statistics +import shutil +from functools import partial +import subprocess +from pathlib import Path + +import torch + +from rl_engine.contracts.numerical import load_contract +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session + + +def _candidate(name): + if name == "triton": + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + TritonVideoInputProjectionFP32Op, + ) + + return TritonVideoInputProjectionFP32Op() + if name == "cuda": + from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import ( + CUDAVideoInputProjectionFP32Op, + ) + + return CUDAVideoInputProjectionFP32Op() + return None + + +def _tool_version(command): + try: + return subprocess.check_output(command, text=True, timeout=10).strip() + except (OSError, subprocess.SubprocessError): + return None + + +def _forward(op, x, weight, bias, keys): + if op is None: + return torch.nn.functional.linear(x, weight, bias) + return op(x, weight, bias, logical_keys=keys, parameter_id="benchmark.proj_in") + + +def _vjp(op, x, weight, bias, grad, keys): + x, weight, bias = [t.detach().clone().requires_grad_() for t in (x, weight, bias)] + with canonical_backward_session() as session: + output = _forward(op, x, weight, bias, keys) + output.backward(grad) + session.validate_complete() + return output.detach(), x.grad, weight.grad, bias.grad + + +def _accuracy(actual, gold): + judgments = load_contract()["judgments"] + result = {} + for i, (name, a, b) in enumerate(zip(("y", "dx", "dw", "dbias"), actual, gold, strict=True)): + policy = judgments["forward_accuracy" if i == 0 else "gradient_accuracy"]["by_op_class"][ + "reduction" + ]["float32"] + a, b = a.detach().cpu(), b.detach().cpu() + error = (a - b).abs() + limit = policy["atol"] + policy["rtol"] * b.abs() + result[name] = { + "max_abs_error": error.max().item(), + "out_of_tolerance": int((error > limit).sum()), + **policy, + } + return result + + +def _sample(fn, prepare, samples, warmup): + values = [] + for i in range(warmup + samples): + call = prepare() if prepare is not None else fn + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + call() + end.record() + end.synchronize() + if i >= warmup: + values.append(start.elapsed_time(end) * 1000) + return {"median_us": statistics.median(values), "samples_us": values} + + +def _backward_call(op, x, weight, bias, grad, keys): + x.grad = weight.grad = bias.grad = None + cm = canonical_backward_session() + session = cm.__enter__() + try: + output = _forward(op, x, weight, bias, keys) + except BaseException: + cm.__exit__(None, None, None) + raise + + def run(): + try: + output.backward(grad) + session.validate_complete() + finally: + cm.__exit__(None, None, None) + + return run + + +def _forward_backward(op, x, weight, bias, grad, keys): + _backward_call(op, x, weight, bias, grad, keys)() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--backends", + nargs="+", + choices=("pytorch", "triton", "cuda"), + default=["pytorch", "triton", "cuda"], + ) + parser.add_argument("--rows", nargs="+", type=int, default=[1, 128, 1024, 4096, 13104]) + parser.add_argument("--samples", type=int, default=20) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--weights", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if ( + not torch.cuda.is_available() + or torch.version.hip is not None + or torch.cuda.get_device_capability() != (9, 0) + ): + raise RuntimeError("benchmark requires a CUDA SM90 GPU") + if args.samples < 1 or args.warmup < 0 or min(args.rows) < 1: + raise ValueError("rows/samples must be positive and warmup must be non-negative") + torch.backends.cuda.matmul.allow_tf32 = False + torch.manual_seed(420) + weight = torch.randn(5376, 96, device="cuda") * 0.0625 + bias = torch.randn(5376, device="cuda") * 0.125 + fixture_hash = None + fixture_metadata = None + if args.weights: + from safetensors import safe_open + from safetensors.torch import load_file + + with safe_open(str(args.weights), framework="pt") as fixture: + fixture_metadata = fixture.metadata() + tensors = load_file(str(args.weights)) + weight, bias = [ + tensors[f"proj_in.{name}"].cuda().float().contiguous() for name in ("weight", "bias") + ] + fixture_hash = hashlib.sha256(args.weights.read_bytes()).hexdigest() + root = Path(__file__).resolve().parents[3] + sources = [ + "rl_engine/ops/gemm/canonical_linear.py", + "rl_engine/ops/autograd/canonical_backward.py", + "rl_engine/ops/gemm/video_input_projection_fp32.py", + "rl_engine/backends/shared/triton/gemm/det_gemm.py", + "rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py", + "rl_engine/backends/cuda/gemm/video_input_projection_fp32.py", + "csrc/bindings/video_input_projection_fp32_cutlass.cpp", + "csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu", + "build_tools/extensions.py", + ] + report = { + "scope": "single H3 FP32 input projection; eager training and no-grad forward graph", + "environment": { + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "python": platform.python_version(), + "tf32": False, + "driver": _tool_version( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"] + ), + "nvcc": _tool_version( + [shutil.which("nvcc") or "/usr/local/cuda/bin/nvcc", "--version"] + ), + }, + "git_commit": subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=root, text=True + ).strip(), + "git_dirty": bool(subprocess.check_output(["git", "status", "--porcelain"], cwd=root)), + "source_sha256": {p: hashlib.sha256((root / p).read_bytes()).hexdigest() for p in sources}, + "weight_fixture_sha256": fixture_hash, + "weight_fixture_metadata": fixture_metadata, + "timing": "CUDA event interval for complete operator calls, including host launch gaps", + "results": [], + } + for rows in args.rows: + x = torch.randn(rows, 96, device="cuda") * 0.125 + grad = torch.randn(rows, 5376, device="cuda") * 0.0625 + keys = torch.stack( + ( + torch.zeros(rows, dtype=torch.int64, device="cuda"), + torch.arange(rows, device="cuda"), + ), + dim=1, + ) + cx, cw, cb = [t.detach().cpu().requires_grad_() for t in (x, weight, bias)] + cy = torch.nn.functional.linear(cx, cw, cb) + cy.backward(grad.cpu()) + gold = (cy, cx.grad, cw.grad, cb.grad) + for name in args.backends: + op = _candidate(name) + actual = _vjp(op, x, weight, bias, grad, keys) + accuracy = _accuracy(actual, gold) + if op is not None and any(v["out_of_tolerance"] for v in accuracy.values()): + raise RuntimeError(f"{name} failed FP32 accuracy at M={rows}: {accuracy}") + order = torch.arange(rows - 1, -1, -1, device="cuda") + permuted = _vjp( + op, + x[order].contiguous(), + weight, + bias, + grad[order].contiguous(), + keys[order].contiguous(), + ) + mismatches = { + label: int( + ( + a.contiguous().view(torch.int32) + != (b[order] if i < 2 else b).contiguous().view(torch.int32) + ).sum() + ) + for i, (label, a, b) in enumerate( + zip(("y", "dx", "dw", "dbias"), actual, permuted, strict=True) + ) + } + if op is not None and any(mismatches.values()): + raise RuntimeError(f"{name} failed row permutation BI: {mismatches}") + with torch.no_grad(): + forward = _sample( + partial(_forward, op, x, weight, bias, keys), None, args.samples, args.warmup + ) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + expected = _forward(op, x, weight, bias, keys) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = _forward(op, x, weight, bias, keys) + graph.replay() + if not torch.equal(output.view(torch.int32), expected.view(torch.int32)): + raise RuntimeError("forward graph changed output bytes") + graph_time = _sample(graph.replay, None, args.samples, args.warmup) + tx, tw, tb = [t.detach().clone().requires_grad_() for t in (x, weight, bias)] + backward = _sample( + None, partial(_backward_call, op, tx, tw, tb, grad, keys), args.samples, args.warmup + ) + combined = _sample( + partial(_forward_backward, op, tx, tw, tb, grad, keys), + None, + args.samples, + args.warmup, + ) + item = { + "backend": name, + "rows": rows, + "accuracy": accuracy, + "permutation_bit_mismatches": mismatches, + "actual_backend": ( + op.backend_info() + if op is not None + else { + "family": "pytorch", + "operator": "torch.nn.functional.linear", + "strict_bi": False, + } + ), + "forward_eager": forward, + "forward_graph": graph_time, + "backward_eager": backward, + "forward_backward_eager": combined, + } + report["results"].append(item) + print( + f"{name} M={rows}: F={forward['median_us']:.1f} us, " + f"B={backward['median_us']:.1f} us, " + f"F+B={combined['median_us']:.1f} us" + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/build_tools/extensions.py b/build_tools/extensions.py index 670fdb072..3786fd43a 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -2,6 +2,7 @@ # Copyright (c) 2026 RL-Kernel Contributors import importlib.util +import hashlib import os import sysconfig import warnings @@ -106,6 +107,7 @@ def _native_extension_required() -> bool: return ( envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) or envs.env_flag(envs.RL_KERNEL_H3_STRICT) + or envs.env_flag("RL_KERNEL_BUILD_CUTLASS") or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) or envs.env_flag("FORCE_CUDA") @@ -158,6 +160,52 @@ def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: return filtered_flags +_CUTLASS_REVISION = "56ca28adb32ffadd097566545b8f7d83b652e998" +_CUTLASS_HEADERS_SHA256 = "7b5321860c269f8e44eab02e59c32c2d4d0e2fa7a3ef72e0a161de9f58b361fa" + + +def _cutlass_headers_fingerprint(include_root: Path) -> str: + digest = hashlib.sha256() + for path in sorted(include_root.rglob("*")): + if path.is_file() and not path.name.startswith("._"): + digest.update(path.relative_to(include_root).as_posix().encode() + b"\0") + digest.update(path.read_bytes() + b"\0") + return digest.hexdigest() + + +def _cutlass_extension(torch, gpu_extension, link_args): + if torch.version.hip is not None or torch.version.cuda is None: + raise RuntimeError("CUTLASS projection requires a CUDA PyTorch build") + root = os.environ.get("CUTLASS_ROOT", "").strip() + if not root: + raise RuntimeError(f"CUTLASS_ROOT must point to CUTLASS {_CUTLASS_REVISION}") + include = Path(root).resolve() / "include" + if _cutlass_headers_fingerprint(include) != _CUTLASS_HEADERS_SHA256: + raise RuntimeError("CUTLASS include tree does not match the pinned source fingerprint") + return gpu_extension( + name="rl_engine._C_cutlass", + sources=[ + "csrc/bindings/video_input_projection_fp32_cutlass.cpp", + "csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu", + ], + include_dirs=[str(include)], + extra_compile_args={ + "cxx": ["-O3", "-std=c++17", "-fno-fast-math", "-ffp-contract=off"], + "nvcc": [ + "-O3", + "-std=c++17", + "--fmad=false", + "--ftz=false", + "--prec-div=true", + "--prec-sqrt=true", + "--expt-relaxed-constexpr", + "-gencode=arch=compute_90,code=sm_90", + ], + }, + extra_link_args=list(link_args), + ) + + def get_extensions(): h3_flags = envs.h3_strict_cuda_flags() torch, _, GpuExtension = _load_torch_extension_tools() @@ -183,6 +231,8 @@ def get_extensions(): if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") is_rocm = getattr(torch.version, "hip", None) is not None + if envs.env_flag("RL_KERNEL_BUILD_CUTLASS"): + extensions.append(_cutlass_extension(torch, GpuExtension, torch_rpath)) if h3_flags and (is_rocm or _musa_build_available(torch)): raise RuntimeError("RL_KERNEL_H3_STRICT currently qualifies NVIDIA CUDA only") diff --git a/csrc/bindings/video_input_projection_fp32_cutlass.cpp b/csrc/bindings/video_input_projection_fp32_cutlass.cpp new file mode 100644 index 000000000..e4e14a42d --- /dev/null +++ b/csrc/bindings/video_input_projection_fp32_cutlass.cpp @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +#include + +torch::Tensor cutlass_gemm_fp32_tn(torch::Tensor a, torch::Tensor b_transposed); +int cutlass_cuda_version(); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("gemm_fp32_tn", &cutlass_gemm_fp32_tn, + "Fixed FP32 SIMT GEMM: [M,K] times transposed [N,K]"); + m.def("build_info", []() { + pybind11::dict info; + info["cutlass_revision"] = "56ca28adb32ffadd097566545b8f7d83b652e998"; + info["headers_sha256"] = "7b5321860c269f8e44eab02e59c32c2d4d0e2fa7a3ef72e0a161de9f58b361fa"; + info["cuda_version"] = cutlass_cuda_version(); + info["architecture"] = "sm_90"; + info["tile"] = pybind11::make_tuple(64, 64, 8); + info["operator_class"] = "simt"; + info["split_k"] = 1; + info["fmad"] = false; + info["ftz"] = false; + return info; + }); +} diff --git a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu new file mode 100644 index 000000000..d552d3e5d --- /dev/null +++ b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +// Reuse CUTLASS device::Gemm at revision 56ca28adb32ffadd097566545b8f7d83b652e998. +// The dependency's BSD-3-Clause notice is packaged in the Python wrapper. +#include +#include +#include +#include +#include +#include +#include + +#include + +using Gemm = cutlass::gemm::device::Gemm< + float, cutlass::layout::RowMajor, + float, cutlass::layout::ColumnMajor, + float, cutlass::layout::RowMajor, + float, cutlass::arch::OpClassSimt, cutlass::arch::Sm50, + cutlass::gemm::GemmShape<64, 64, 8>, + cutlass::gemm::GemmShape<32, 32, 8>, + cutlass::gemm::GemmShape<1, 1, 1>, + cutlass::epilogue::thread::LinearCombination, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, + 2, 1, 1, false, cutlass::arch::OpMultiplyAdd>; + +at::Tensor cutlass_gemm_fp32_tn(at::Tensor a, at::Tensor b_transposed) { + TORCH_CHECK(a.is_cuda() && b_transposed.is_cuda(), "CUTLASS GEMM requires CUDA tensors"); + TORCH_CHECK(a.device() == b_transposed.device(), "CUTLASS GEMM devices must match"); + TORCH_CHECK(a.scalar_type() == at::kFloat && b_transposed.scalar_type() == at::kFloat, + "CUTLASS GEMM requires FP32"); + TORCH_CHECK(a.dim() == 2 && b_transposed.dim() == 2, + "CUTLASS GEMM requires rank-two matrices"); + TORCH_CHECK(a.is_contiguous() && b_transposed.is_contiguous(), + "CUTLASS GEMM matrices must be contiguous"); + TORCH_CHECK(a.size(1) == b_transposed.size(1), "CUTLASS GEMM K dimensions must match"); + const c10::cuda::CUDAGuard guard(a.device()); + const auto* properties = at::cuda::getDeviceProperties(a.get_device()); + TORCH_CHECK(properties->major == 9 && properties->minor == 0, + "CUTLASS video projection requires CUDA SM90"); + const auto m64 = a.size(0), n64 = b_transposed.size(0), k64 = a.size(1); + TORCH_CHECK(m64 > 0 && n64 > 0 && k64 > 0 && + m64 <= INT_MAX && n64 <= INT_MAX && k64 <= INT_MAX, + "CUTLASS GEMM dimensions must be positive int32 values"); + const int m = static_cast(m64), n = static_cast(n64), k = static_cast(k64); + auto output = at::empty({m, n}, a.options()); + Gemm::Arguments args({m, n, k}, {a.data_ptr(), k}, + {b_transposed.data_ptr(), k}, + {output.data_ptr(), n}, {output.data_ptr(), n}, + {1.0f, 0.0f}, 1); + Gemm gemm; + auto status = gemm.can_implement(args); + TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); + status = gemm(args, nullptr, c10::cuda::getCurrentCUDAStream(a.get_device()).stream()); + TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +int cutlass_cuda_version() { return CUDART_VERSION; } diff --git a/docs/.nav.yml b/docs/.nav.yml index 1037edce5..819b3f31a 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -27,6 +27,7 @@ nav: - operators/ratio-kl.md - operators/pack-and-pad.md - operators/video-patch-pack-unpack.md + - operators/video-input-projection-fp32.md - operators/sampling.md - operators/det-gemm.md - operators/embedding.md diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md new file mode 100644 index 000000000..f3014c0b0 --- /dev/null +++ b/docs/operators/video-input-projection-fp32.md @@ -0,0 +1,149 @@ +# Video Input Projection FP32 + +## Summary + +MiniMax-H3 maps each 96-element video patch to the 5376-wide Transformer residual +space. This operator preserves the FP32 projection boundary of the pinned H3 +reference. The model adapter owns the subsequent cast into the BF16 packed stream. + +The primary implementation reuses RL-Kernel's Triton FP32 GEMM and canonical +parameter-gradient session. A pinned CUTLASS FP32 SIMT GEMM is an explicit CUDA +candidate. Both use the same semantic adapter and autograd implementation. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session + +op = kernel_registry.get_op("video_input_projection_fp32", device=x.device) +with canonical_backward_session() as session: + y = op(x, weight, bias, logical_keys=keys, parameter_id="transformer.proj_in") + y.backward(grad_y) + session.validate_complete() +``` + +For inference, call the same operator under `torch.no_grad()` without row keys. +For parameter gradients, keys identify original logical rows, not their current +physical placement. The same parameter keeps the same `parameter_id` across uses. + +## Backends + +| Backend | Wrapper | Native symbol | Status | +| --- | --- | --- | --- | +| Triton on CUDA | `TritonVideoInputProjectionFP32Op` | `_det_gemm_fp32_kernel` | Default CUDA path; SM90 qualification | +| CUDA / CUTLASS | `CUDAVideoInputProjectionFP32Op` | `_C_cutlass.gemm_fp32_tn` | Explicit candidate; optional build; SM90 | +| PyTorch reference | `NativeVideoInputProjectionFP32Op` | CPU `F.linear` | Accuracy reference | + +CUTLASS revision: `56ca28adb32ffadd097566545b8f7d83b652e998`. The build verifies +the complete include-tree fingerprint. The BSD-3-Clause dependency notice is +preserved in the installed CUDA Python wrapper. + +```bash +git clone https://github.com/NVIDIA/cutlass.git /tmp/h3-cutlass +git -C /tmp/h3-cutlass checkout 56ca28adb32ffadd097566545b8f7d83b652e998 +CUTLASS_ROOT=/tmp/h3-cutlass RL_KERNEL_BUILD_CUTLASS=1 \ + MAX_JOBS=2 python -m pip install --no-build-isolation --no-deps -e . +``` + +The separate `_C_cutlass` extension fixes `--fmad=false`, `--ftz=false`, precise +division/square root, and SM90 code generation. Common `_C` compiler options are +independent. `backend_info()` reads back the candidate's revision and schedule. + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `x` | `[M,96]` or `[B,S,96]` | FP32 | Non-empty, contiguous | +| `weight` | `[5376,96]` | FP32 | Contiguous; same device as `x` | +| `bias` | `[5376]` or `None` | FP32 | Contiguous; same device as `x` | +| `logical_keys` | `[M,2]` or `[M,3]` | int64 | Required for trainable parameters; same device | +| `parameter_id` | Stable module identity | string | Required for trainable parameters | +| Output | `[M,5376]` or `[B,S,5376]` | FP32 | Leading dimensions preserved | + +Keys use `(sample_id, patch_id[, use_id])`. Active identities are unique across +uses of one parameter; a negative first column marks padding. Register every +forward use before starting backward in the session. Fresh backward graphs use +fresh sessions. Parameter freezing registers only the parameters that need VJPs. + +## Dispatch Behavior + +`get_op` selects Triton on CUDA. CUTLASS is selected explicitly through the +semantic catalog or `OP_SPECS` candidate `cuda`. The semantic catalog records +single-device, unsharded support and exact implementation identities. An +unavailable explicit candidate raises an error. CPU/ROCm/MUSA/NPU execution is +outside this contribution's strict support scope. + +## Accuracy + +Forward computes `Y = X @ W.T`, then adds bias in FP32. The upstream gradient `G` +stays FP32, including normalization of non-contiguous autograd inputs: + +- `dX = G @ W` is row-local. +- `dW = G_sorted.T @ X_sorted` follows canonical logical-row order. +- `dbias = ones[1,R] @ G_sorted` uses the same fixed FP32 GEMM bottom layer. + +The existing canonical session collects all contributions, removes padding, and +sorts keys before the final parameter reduction. It returns each parameter VJP +once. Ordinary chunk-by-chunk accumulation is a different arithmetic contract. + +Long K reductions use a fixed midpoint tree with at most 128 terms per leaf. +Each leaf reuses the selected GEMM and every tree node adds in FP32. This bounds +rounding depth for the 5376-term input VJP and long parameter reductions. The +tree depends on logical K and retains its topology across physical row changes. +Existing canonical linear callers keep their original reduction by default. + +Triton fixes the 64x64x32 leaf schedule and disables TF32 and FP fusion. CUTLASS +fixes 64x64x8 FP32 SIMT tiles and a single output owner per leaf, without parallel +split-K. Their leaf arithmetic has separate backend identities. + +The shared WS1 `reduction` FP32 accuracy policy uses `atol=rtol=1e-4` for forward +and gradients. Invariance uses byte equality. Parameter invariance compares the +same logical contribution set and upstream gradients across layout changes. +Adding a genuine training sample changes the expected parameter gradient. + +## Performance Notes + +```bash +PYTHONPATH=. python benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py \ + --backends pytorch triton cuda --rows 1 128 1024 4096 13104 \ + --output /tmp/h3-input-projection.json +``` + +The benchmark first checks independent CPU FP32 forward/backward accuracy and +row-permutation VJPs. It measures eager forward, backward, F+B, and no-grad +forward CUDA Graph. Timings include layout copies, bias, and the canonical +parameter-gradient machinery. Raw samples, environment, source hashes, and +actual backend metadata are written to the explicitly selected local output. +PyTorch is a production performance baseline with separately reported BI results. +The fixed tree adds kernel launches and layout copies to backward; compare full +operator timings before drawing a speedup conclusion. + +## Tests + +```bash +python -m pytest tests/ops/gemm/test_video_input_projection_fp32.py -v +python tools/validation/operators/check_operator.py \ + --op video_input_projection_fp32 --candidate triton --device cuda --dtype fp32 \ + --batch 1 --seq 129 --check-grad --grad-mode random +python tools/validation/operators/check_operator.py \ + --op video_input_projection_fp32 --candidate cuda --device cuda --dtype fp32 \ + --batch 1 --seq 129 --check-grad --grad-mode random +``` + +Set `H3_PROJECTION_WEIGHTS` to a safetensors fixture containing `proj_in.weight` +and `proj_in.bias` from `MiniMaxAI/MiniMax-H3` at revision +`42ed227ee7df40d41602854ae760620d6eb651fe` to run the latent-pack-projection check. +The fixture metadata must include `repository` and `checkpoint_revision` with +these values. The check uses a synthetic `[1,24,13,48,84]` latent with real input +projection weights; it does not load the complete Transformer. +Keep the weights and generated evidence outside tracked source files. + +## Known Limitations + +Qualification targets a single H20 (SM90). Distributed parameter gradients, +interleaved forward/backward microbatch accumulation, higher-order derivatives, +gradient checkpointing, and full training CUDA Graph require further integration +contracts. Training capture fails explicitly; no-grad forward capture is tested. +This operator's evidence is one part of the H3 WS1 operator gate. Full Transformer +and complete RL trajectory acceptance belong to the corresponding model gates. diff --git a/rl_engine/_C_cutlass.pyi b/rl_engine/_C_cutlass.pyi new file mode 100644 index 000000000..a1afa227a --- /dev/null +++ b/rl_engine/_C_cutlass.pyi @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from typing import Any + +import torch + +def gemm_fp32_tn(a: torch.Tensor, b_transposed: torch.Tensor) -> torch.Tensor: ... +def build_info() -> dict[str, Any]: ... diff --git a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py new file mode 100644 index 000000000..0f82ae6c5 --- /dev/null +++ b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 FP32 projection using a pinned CUTLASS SIMT GEMM candidate. + +CUTLASS dependency notice (BSD-3-Clause): +Copyright (c) 2017 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. +3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +""" + +from functools import lru_cache +from importlib import import_module + +import torch + +from rl_engine.ops.gemm.video_input_projection_fp32 import VideoInputProjectionFP32Op + +CUTLASS_REVISION = "56ca28adb32ffadd097566545b8f7d83b652e998" +CUTLASS_HEADERS_SHA256 = "7b5321860c269f8e44eab02e59c32c2d4d0e2fa7a3ef72e0a161de9f58b361fa" + + +@lru_cache(maxsize=1) +def _load_extension(): + try: + extension = import_module("rl_engine._C_cutlass") + except ImportError as exc: + raise RuntimeError( + "build the optional CUTLASS extension with RL_KERNEL_BUILD_CUTLASS=1" + ) from exc + info = extension.build_info() + if ( + info["cutlass_revision"] != CUTLASS_REVISION + or info["headers_sha256"] != CUTLASS_HEADERS_SHA256 + or info["fmad"] + or info["ftz"] + or info["split_k"] != 1 + ): + raise RuntimeError("CUTLASS extension does not match the strict FP32 build contract") + return extension + + +def cutlass_gemm_fp32(a, b): + # The native TN interface represents mathematical B[K,N] as [N,K] storage. + return _load_extension().gemm_fp32_tn(a.contiguous(), b.t().contiguous()) + + +class CUDAVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): + family = "cutlass" + backend_id = "cuda-cutlass-video-input-projection-fp32-v1" + + def __init__(self): + _load_extension() + + def _check_backend(self, x): + if x.device.type != "cuda" or torch.version.hip is not None: + raise ValueError("CUTLASS video input projection requires CUDA") + if torch.cuda.get_device_capability(x.device) != (9, 0): + raise ValueError("CUTLASS video input projection requires CUDA SM90") + + def backend_info(self): + return { + "backend_id": self.backend_id, + "family": "cuda", + "kernel": "gemm_fp32_tn", + "reduction": "fixed_midpoint_tree", + "reduction_leaf_size": 128, + **_load_extension().build_info(), + } diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 7569f815e..558507602 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -745,6 +745,7 @@ def _triton_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: BLOCK_N=_BLOCK_N, BLOCK_K=_BLOCK_K, PROMOTE_INPUTS=True, + **({"enable_fp_fusion": False} if torch.version.hip is None else {}), ) return c diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py new file mode 100644 index 000000000..f594a4b79 --- /dev/null +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 projection through RL-Kernel's fixed FP32 Triton GEMM.""" + +import torch + +from rl_engine.backends.shared.triton.gemm.det_gemm import _TRITON_AVAILABLE +from rl_engine.ops.gemm.video_input_projection_fp32 import VideoInputProjectionFP32Op + + +class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): + family = "triton" + backend_id = "triton-video-input-projection-fp32-v1" + + def __init__(self): + if not _TRITON_AVAILABLE: + raise ImportError("video input projection requires Triton") + + def _check_backend(self, x): + if x.device.type != "cuda" or torch.version.hip is not None: + raise ValueError("Triton video input projection is qualified on CUDA only") + if torch.cuda.get_device_capability(x.device) != (9, 0): + raise ValueError("video input projection requires CUDA SM90") + + def backend_info(self): + import triton + + return { + "backend_id": self.backend_id, + "family": self.family, + "kernel": "_det_gemm_fp32_kernel", + "triton_version": triton.__version__, + "tile": [64, 64, 32], + "tf32": False, + "fp_fusion": False, + "split_k": 1, + "reduction": "fixed_midpoint_tree", + "reduction_leaf_size": 128, + } diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 30175b07a..aa8268024 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -13,7 +13,23 @@ from rl_engine.ops.autograd.canonical_backward import active_session -def _gemm_fp32(a: torch.Tensor, b: torch.Tensor, family: str) -> torch.Tensor: +def _gemm_fp32( + a: torch.Tensor, + b: torch.Tensor, + family: str, + reduction_leaf_size: int | None = None, +) -> torch.Tensor: + # Bound FP32 rounding depth for long reductions. The midpoint topology + # depends only on logical K, so physical row placement cannot change it. + if reduction_leaf_size is not None and a.shape[1] > reduction_leaf_size: + midpoint = a.shape[1] // 2 + lower = _gemm_fp32(a[:, :midpoint], b[:midpoint], family, reduction_leaf_size) + upper = _gemm_fp32(a[:, midpoint:], b[midpoint:], family, reduction_leaf_size) + return lower + upper + if family == "cutlass": + from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import cutlass_gemm_fp32 + + return cutlass_gemm_fp32(a, b) if family == "cuda": return _C.det_gemm_rowwise_fwd_fp32(a.contiguous(), b.contiguous()) if family == "triton": @@ -31,16 +47,46 @@ def _gemm_fp32(a: torch.Tensor, b: torch.Tensor, family: str) -> torch.Tensor: class _CanonicalLinearFn(torch.autograd.Function): @staticmethod - def forward(ctx, a, weight, logical_keys, parameter_id, family): + def forward( + ctx, + a, + weight, + logical_keys, + parameter_id, + family, + bias, + bias_parameter_id, + reduction_leaf_size, + ): session = active_session() - if session is None: + parameter_grad = ctx.needs_input_grad[1] or ctx.needs_input_grad[5] + if parameter_grad and session is None: raise RuntimeError("canonical linear requires an active backward session") + if parameter_grad and logical_keys is None: + raise ValueError("canonical parameter gradients require logical keys") + if bias is not None and str(parameter_id) == str(bias_parameter_id): + raise ValueError("weight and bias parameter identities must differ") ctx.save_for_backward(a, weight) ctx.session = session ctx.parameter_id = str(parameter_id) ctx.family = str(family) - ctx.slot = session.register(ctx.parameter_id, logical_keys) - return _gemm_fp32(a.float(), weight.float().t().contiguous(), ctx.family) + ctx.reduction_leaf_size = reduction_leaf_size + ctx.bias_parameter_id = str(bias_parameter_id) + ctx.bias_dtype = bias.dtype if bias is not None else None + ctx.slot = ( + session.register(ctx.parameter_id, logical_keys) if ctx.needs_input_grad[1] else None + ) + ctx.bias_slot = ( + session.register(ctx.bias_parameter_id, logical_keys) + if ctx.needs_input_grad[5] + else None + ) + rhs = weight.float().t() + # CUTLASS consumes column-major B directly from contiguous [N,K] weight. + if ctx.family != "cutlass": + rhs = rhs.contiguous() + output = _gemm_fp32(a.float(), rhs, ctx.family, ctx.reduction_leaf_size) + return output if bias is None else output + bias.float() @staticmethod def backward(ctx, grad_out): @@ -48,17 +94,36 @@ def backward(ctx, grad_out): grad32 = grad_out.contiguous().float() da = None if ctx.needs_input_grad[0]: - da = _gemm_fp32(grad32, weight.float(), ctx.family).to(a.dtype) + da = _gemm_fp32(grad32, weight.float(), ctx.family, ctx.reduction_leaf_size).to(a.dtype) dweight = None if ctx.needs_input_grad[1]: def reducer(rows, grads): - return _gemm_fp32(grads.float().t().contiguous(), rows.float(), ctx.family).to( - weight.dtype - ) + return _gemm_fp32( + grads.float().t().contiguous(), + rows.float(), + ctx.family, + ctx.reduction_leaf_size, + ).to(weight.dtype) dweight = ctx.session.submit_linear(ctx.parameter_id, ctx.slot, a, grad_out, reducer) + dbias = None + if ctx.needs_input_grad[5]: + + def bias_reducer(grads): + if grads.shape[0] == 0: + raise RuntimeError(f"no active logical rows for {ctx.bias_parameter_id}") + ones = torch.ones((1, grads.shape[0]), device=grads.device, dtype=torch.float32) + return ( + _gemm_fp32(ones, grads.float(), ctx.family, ctx.reduction_leaf_size) + .squeeze(0) + .to(ctx.bias_dtype) + ) + + dbias = ctx.session.submit_rows( + ctx.bias_parameter_id, ctx.bias_slot, grad32, bias_reducer + ) if ctx.family == "triton": record_backward( "det_gemm", @@ -73,18 +138,39 @@ def reducer(rows, grads): impl="ascend_det_gemm_canonical_rowfold", family="ascend", ) - return da, dweight, None, None, None + elif ctx.family == "cutlass": + record_backward( + "video_input_projection_fp32", + kernel_id="rl_engine._C_cutlass.gemm_fp32_tn", + impl="cutlass_fp32_simt_canonical_rowfold", + family="cuda", + ) + return da, dweight, None, None, None, dbias, None, None def canonical_linear_fp32( a: torch.Tensor, weight: torch.Tensor, - logical_keys: torch.Tensor, + logical_keys: torch.Tensor | None, *, parameter_id: str, family: str, + bias: torch.Tensor | None = None, + bias_parameter_id: str | None = None, + reduction_leaf_size: int | None = None, ) -> torch.Tensor: - return _CanonicalLinearFn.apply(a, weight, logical_keys, parameter_id, family) + if reduction_leaf_size is not None and reduction_leaf_size < 1: + raise ValueError("reduction_leaf_size must be positive") + return _CanonicalLinearFn.apply( + a, + weight, + logical_keys, + parameter_id, + family, + bias, + bias_parameter_id if bias_parameter_id is not None else f"{parameter_id}.bias", + reduction_leaf_size, + ) __all__ = ["canonical_linear_fp32"] diff --git a/rl_engine/ops/gemm/video_input_projection_fp32.py b/rl_engine/ops/gemm/video_input_projection_fp32.py new file mode 100644 index 000000000..7c9196504 --- /dev/null +++ b/rl_engine/ops/gemm/video_input_projection_fp32.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""FP32 MiniMax-H3 video input projection with logical-row parameter VJPs.""" + +from __future__ import annotations + +import torch + +from rl_engine.ops.autograd.canonical_backward import active_session +from rl_engine.ops.gemm.canonical_linear import _gemm_fp32, canonical_linear_fp32 + +VIDEO_PATCH_WIDTH = 96 +VIDEO_HIDDEN_WIDTH = 5376 +FP32_REDUCTION_LEAF_SIZE = 128 + + +def validate_projection_inputs(x, weight, bias): + """Validate the public FP32 tensor boundary before backend dispatch.""" + tensors = {"x": x, "weight": weight} + if bias is not None: + tensors["bias"] = bias + for name, tensor in tensors.items(): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype != torch.float32: + raise TypeError(f"{name} must have dtype torch.float32") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if tensor.device != x.device: + raise ValueError("projection tensors must be on the same device") + if x.ndim not in (2, 3) or x.shape[-1] != VIDEO_PATCH_WIDTH or x.numel() == 0: + raise ValueError("x must have non-empty shape [M,96] or [B,S,96]") + if tuple(weight.shape) != (VIDEO_HIDDEN_WIDTH, VIDEO_PATCH_WIDTH): + raise ValueError("weight must have shape [5376,96]") + if bias is not None and tuple(bias.shape) != (VIDEO_HIDDEN_WIDTH,): + raise ValueError("bias must have shape [5376]") + + +def _validate_keys(keys, rows, device, parameter_id, session, weight, bias): + if not isinstance(keys, torch.Tensor): + raise TypeError("parameter gradients require logical_keys") + if keys.dtype != torch.int64 or keys.device != device or not keys.is_contiguous(): + raise ValueError("logical_keys must be contiguous int64 on the projection device") + if keys.ndim != 2 or keys.shape[0] != rows or keys.shape[1] not in (2, 3): + raise ValueError("logical_keys must have shape [M,2] or [M,3]") + if not isinstance(parameter_id, str) or not parameter_id.strip(): + raise ValueError("parameter gradients require a non-empty parameter_id") + # A duplicate active identity leaves tie order dependent on chunk arrival. + for name, needed in ( + ("weight", weight.requires_grad), + ("bias", bias is not None and bias.requires_grad), + ): + if not needed: + continue + previous = session.uses.get(f"{parameter_id}.{name}", []) + if any(use.keys.shape[1] != keys.shape[1] for use in previous): + raise ValueError("logical key columns must agree across parameter uses") + merged = torch.cat([use.keys for use in previous] + [keys], dim=0) + active = merged[merged[:, 0] >= 0] + if torch.unique(active, dim=0).shape[0] != active.shape[0]: + raise ValueError("active logical keys must be unique across parameter uses") + + +class VideoInputProjectionFP32Op: + """Shared semantic adapter; subclasses select one qualified GEMM backend.""" + + family: str + backend_id: str + + def _check_backend(self, x): + raise NotImplementedError + + def __call__(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None): + validate_projection_inputs(x, weight, bias) + self._check_backend(x) + rows = x.reshape(-1, VIDEO_PATCH_WIDTH) + needs_grad = torch.is_grad_enabled() and any( + tensor is not None and tensor.requires_grad for tensor in (x, weight, bias) + ) + if needs_grad: + if x.device.type == "cuda" and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("video projection backward requires eager execution") + if weight.requires_grad or (bias is not None and bias.requires_grad): + session = active_session() + if session is None: + raise RuntimeError("parameter gradients require canonical_backward_session") + _validate_keys( + logical_keys, rows.shape[0], x.device, parameter_id, session, weight, bias + ) + logical_keys = logical_keys.detach().clone() + output = canonical_linear_fp32( + rows, + weight, + logical_keys, + parameter_id=f"{parameter_id}.weight", + family=self.family, + bias=bias, + bias_parameter_id=f"{parameter_id}.bias", + reduction_leaf_size=FP32_REDUCTION_LEAF_SIZE, + ) + else: + rhs = weight.t() + if self.family != "cutlass": + rhs = rhs.contiguous() + output = _gemm_fp32(rows, rhs, self.family, FP32_REDUCTION_LEAF_SIZE) + if bias is not None: + output = output + bias + return output.reshape(*x.shape[:-1], VIDEO_HIDDEN_WIDTH) + + forward = __call__ + + +__all__ = ["VideoInputProjectionFP32Op"] diff --git a/rl_engine/reference/gemm/video_input_projection_fp32.py b/rl_engine/reference/gemm/video_input_projection_fp32.py new file mode 100644 index 000000000..9679cb17b --- /dev/null +++ b/rl_engine/reference/gemm/video_input_projection_fp32.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Independent CPU FP32 Linear reference for the H3 video input head.""" + +import torch.nn.functional as F + +from rl_engine.ops.gemm.video_input_projection_fp32 import validate_projection_inputs + + +class NativeVideoInputProjectionFP32Op: + def forward_fp32(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None): + validate_projection_inputs(x, weight, bias) + # CPU avoids a GPU TF32 or batch-dependent cuBLAS reference path. + return F.linear(x.cpu(), weight.cpu(), None if bias is None else bias.cpu()).to(x.device) + + __call__ = forward_fp32 diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 528c77b98..82f1938f1 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -114,6 +114,13 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # Batch-invariant deterministic GEMM (WS1 #146) CUDA_DET_GEMM = "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp" TRITON_DET_GEMM = "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp" + TRITON_VIDEO_INPUT_PROJECTION_FP32 = ( + "rl_engine.backends.shared.triton.gemm.video_input_projection_fp32." + "TritonVideoInputProjectionFP32Op" + ) + CUDA_VIDEO_INPUT_PROJECTION_FP32 = ( + "rl_engine.backends.cuda.gemm.video_input_projection_fp32.CUDAVideoInputProjectionFP32Op" + ) ASCEND_DET_GEMM = "rl_engine.kernels.ops.ascend.matmul.det_gemm.DetGemmAscendOp" # NON-deterministic reference (torch.matmul); reference/benchmark ONLY, # intentionally excluded from det_gemm dispatch (cuBLAS breaks invariance). @@ -235,6 +242,68 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint="h3-query-chunk-v1; binary profile checked at invocation", ), + OperatorBackendDescriptor( + semantic_op="video_input_projection_fp32", + backend_id="cuda-cutlass-video-input-projection-fp32-v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cuda"}), + supported_dtypes=frozenset({"float32"}), + supported_topologies={ + "rollout": { + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + "training": {"world_size": (1,), "sharding": ("unsharded",)}, + }, + determinism_or_alignment_properties={ + "algorithm": "cutlass_simt_fp32_canonical_parameter_vjp", + "input_width": 96, + "output_width": 5376, + "cuda_sm": 90, + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + "training_mode": "eager", + "parameter_gradient_order": "logical_key", + "reduction": "fixed_midpoint_tree_128", + }, + lifecycle=OperatorLifecycle.REQUEST, + implementation_class_or_factory=OpBackend.CUDA_VIDEO_INPUT_PROJECTION_FP32.value, + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="cutlass-56ca28a-fp32-simt-64x64x8-tree128-h3-v1", + ), + OperatorBackendDescriptor( + semantic_op="video_input_projection_fp32", + backend_id="triton-video-input-projection-fp32-v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cuda"}), + supported_dtypes=frozenset({"float32"}), + supported_topologies={ + "rollout": { + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + "training": {"world_size": (1,), "sharding": ("unsharded",)}, + }, + determinism_or_alignment_properties={ + "algorithm": "fixed_fp32_gemm_canonical_parameter_vjp", + "input_width": 96, + "output_width": 5376, + "cuda_sm": 90, + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + "training_mode": "eager", + "parameter_gradient_order": "logical_key", + "reduction": "fixed_midpoint_tree_128", + }, + lifecycle=OperatorLifecycle.REQUEST, + implementation_class_or_factory=OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32.value, + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="rlkernel-triton-fp32-64x64x32-no-fusion-tree128-h3-v1", + ), OperatorBackendDescriptor( semantic_op="video_patch_pack_unpack", backend_id="pytorch-video-patch-pack-unpack-v1", @@ -649,6 +718,7 @@ def __init__(self): "pack": [OpBackend.PYTORCH_PACK], "video_patch_pack_unpack": [OpBackend.PYTORCH_VIDEO_PATCH_PACK_UNPACK], "det_gemm": [OpBackend.CUDA_DET_GEMM, OpBackend.TRITON_DET_GEMM], + "video_input_projection_fp32": [OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32], "batch_invariant_logp": [ OpBackend.TRITON_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, @@ -1107,6 +1177,8 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: platform = self._platform_for_device(device) if op_type == "h3_full_attention" and platform != "cuda": raise RuntimeError("h3_full_attention requires NVIDIA CUDA; no fallback") + if op_type == "video_input_projection_fp32" and platform != "cuda": + raise RuntimeError("video input projection is qualified on CUDA SM90 only") candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: diff --git a/rl_engine/validation/operators/op_checks.py b/rl_engine/validation/operators/op_checks.py index 5f3aded34..bdec0d9aa 100644 --- a/rl_engine/validation/operators/op_checks.py +++ b/rl_engine/validation/operators/op_checks.py @@ -4,6 +4,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence +from contextlib import AbstractContextManager, nullcontext from dataclasses import asdict, dataclass from typing import Any @@ -40,6 +41,7 @@ class CandidateSpec: backend: str = "unknown" arch_key: str | None = None provenance: BackendProvenance | None = None + backward_context_factory: Callable[[], AbstractContextManager] | None = None @dataclass(frozen=True) @@ -218,23 +220,29 @@ def _run_case_backward( candidate_inputs = _clone_inputs_for_backward(case.inputs, case.grad_input_names) gold_inputs = _clone_inputs_for_backward(case.inputs, case.grad_input_names) - candidate_outputs = _flatten_tensors(_call_candidate(candidate.fn, candidate_inputs)) - gold_outputs = _flatten_tensors(case.gold_fn(**gold_inputs)) + context = ( + candidate.backward_context_factory() + if candidate.backward_context_factory + else nullcontext() + ) # Candidate and gold must use the same upstream gradients; otherwise we # would compare different vector-Jacobian products. # grad_mode="ones" is the old output.sum().backward() smoke path. # grad_mode="random" is closer to training, where dL/doutput is non-uniform. - grad_outputs = _make_grad_outputs(candidate_outputs, grad_mode=grad_mode, seed=grad_seed) - shared_upstreams = [ - grad.to(device=output.device, dtype=output.dtype) - for grad, output in zip(grad_outputs, candidate_outputs, strict=True) - ] - candidate_grads = _backward_grads( - candidate_outputs, - candidate_inputs, - case.grad_input_names, - grad_outputs=shared_upstreams, - ) + with context: + candidate_outputs = _flatten_tensors(_call_candidate(candidate.fn, candidate_inputs)) + grad_outputs = _make_grad_outputs(candidate_outputs, grad_mode=grad_mode, seed=grad_seed) + shared_upstreams = [ + grad.to(device=output.device, dtype=output.dtype) + for grad, output in zip(grad_outputs, candidate_outputs, strict=True) + ] + candidate_grads = _backward_grads( + candidate_outputs, + candidate_inputs, + case.grad_input_names, + grad_outputs=shared_upstreams, + ) + gold_outputs = _flatten_tensors(case.gold_fn(**gold_inputs)) gold_grads = _backward_grads( gold_outputs, gold_inputs, diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index aaebc279f..842142e4f 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -29,6 +29,7 @@ def make_operator_inputs( "qk_norm": _make_qk_norm_inputs, "pack": _make_pack_inputs, "video_patch_pack_unpack": _make_video_patch_pack_unpack_inputs, + "video_input_projection_fp32": _make_video_input_projection_fp32_inputs, "matmul": _make_matmul_inputs, "det_gemm": _make_det_gemm_inputs, "attention": _make_attention_inputs, @@ -62,6 +63,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "qk_norm": f"{batch}x{seq}x{_arg_int(args, 'n_heads', DEFAULT_N_HEADS)}x" f"{_arg_int(args, 'head_dim', DEFAULT_HEAD_DIM)}", "pack": f"{batch}x{seq}x{_normalized_dim(args)}", + "video_input_projection_fp32": f"{batch}x{seq}x96x5376", "matmul": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "det_gemm": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", @@ -85,6 +87,26 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: raise ValueError(f"unsupported operator shape: {op_name}") from exc +def _make_video_input_projection_fp32_inputs(args, dtype, device): + if dtype != torch.float32: + raise TypeError("video input projection validation requires FP32") + batch, seq = _batch_seq(args) + keys = torch.stack( + ( + torch.arange(batch, device=device).repeat_interleave(seq), + torch.arange(seq, device=device).repeat(batch), + ), + dim=1, + ) + return { + "x": _floating_tensor((batch, seq, 96), args, dtype, device, offset=0), + "weight": _floating_tensor((5376, 96), args, dtype, device, offset=1), + "bias": _floating_tensor((5376,), args, dtype, device, offset=2), + "logical_keys": keys, + "parameter_id": "transformer.proj_in", + } + + def _make_rms_norm_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 92995b938..0e7ff39b4 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -5,6 +5,7 @@ import argparse import importlib +from contextlib import contextmanager from dataclasses import dataclass from typing import Any @@ -42,6 +43,27 @@ def _load_object(path: str) -> Any: }, grad_input_names=("q", "k", "v"), ), + "video_input_projection_fp32": OperatorSpec( + name="video_input_projection_fp32", + op_class="reduction", + gold_path="rl_engine.reference.gemm.video_input_projection_fp32.NativeVideoInputProjectionFP32Op", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.reference.gemm.video_input_projection_fp32." + "NativeVideoInputProjectionFP32Op" + ), + "triton": ( + "rl_engine.backends.shared.triton.gemm.video_input_projection_fp32." + "TritonVideoInputProjectionFP32Op" + ), + "cuda": ( + "rl_engine.backends.cuda.gemm.video_input_projection_fp32." + "CUDAVideoInputProjectionFP32Op" + ), + }, + grad_input_names=("x", "weight", "bias"), + ), "rms_norm": OperatorSpec( name="rms_norm", op_class="reduction", @@ -417,6 +439,11 @@ def make_candidate(args: argparse.Namespace) -> CandidateSpec: backend=candidate_name, arch_key=args.arch_key, fn=candidate_op, + backward_context_factory=( + _canonical_validation_context + if args.op == "video_input_projection_fp32" and candidate_name != "pytorch" + else None + ), ) supported = sorted([*spec.candidate_paths, "native"]) @@ -424,3 +451,12 @@ def make_candidate(args: argparse.Namespace) -> CandidateSpec: f"unsupported candidate {args.candidate!r} for op {args.op!r}; " f"supported candidates: {', '.join(supported)}" ) + + +@contextmanager +def _canonical_validation_context(): + from rl_engine.ops.autograd.canonical_backward import canonical_backward_session + + with canonical_backward_session() as session: + yield session + session.validate_complete() diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index b813506b9..271f5b3e4 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -13,8 +13,38 @@ from build_tools.extensions import _resolve_package_data_source, get_cmdclass +def test_optional_cutlass_build_has_independent_strict_flags(monkeypatch, tmp_path): + from build_tools import extensions + + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(torch.version, "cuda", "12.9") + monkeypatch.setenv("CUTLASS_ROOT", str(tmp_path)) + monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", "1") + monkeypatch.setattr( + extensions, "_cutlass_headers_fingerprint", lambda _: extensions._CUTLASS_HEADERS_SHA256 + ) + candidate = extensions._cutlass_extension(torch, lambda **kw: kw, []) + assert candidate["name"] == "rl_engine._C_cutlass" + flags = candidate["extra_compile_args"]["nvcc"] + assert "--fmad=false" in flags and "--ftz=false" in flags + assert "--use_fast_math" not in flags + assert "csrc/bindings/ops.cpp" not in candidate["sources"] + + +def test_cutlass_build_rejects_unpinned_headers(monkeypatch, tmp_path): + from build_tools import extensions + import pytest + + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(torch.version, "cuda", "12.9") + monkeypatch.setenv("CUTLASS_ROOT", str(tmp_path)) + with pytest.raises(RuntimeError, match="fingerprint"): + extensions._cutlass_extension(torch, lambda **kw: kw, []) + + def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: captured: dict[str, Any] = {} + monkeypatch.delenv("RL_KERNEL_BUILD_CUTLASS", raising=False) def fake_setup(**kwargs: Any) -> None: captured.update(kwargs) diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py new file mode 100644 index 000000000..5e831b240 --- /dev/null +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -0,0 +1,353 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 FP32 projection accuracy, actual autograd, and logical-row invariance.""" + +import os +from importlib.util import find_spec + +import pytest +import torch + +from rl_engine.contracts.numerical import load_contract +from rl_engine.ops.autograd.canonical_backward import canonical_backward_session +from rl_engine.ops.gemm.canonical_linear import canonical_linear_fp32 +from rl_engine.ops.gemm.video_input_projection_fp32 import VideoInputProjectionFP32Op + +H20 = ( + torch.cuda.is_available() + and torch.version.hip is None + and torch.cuda.get_device_capability() == (9, 0) +) +BACKENDS = ["triton", "cuda"] + + +def _op(backend): + if backend == "triton": + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + TritonVideoInputProjectionFP32Op, + ) + + return TritonVideoInputProjectionFP32Op() + from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import ( + CUDAVideoInputProjectionFP32Op, + ) + + if find_spec("rl_engine._C_cutlass") is None: + pytest.skip("requires the optional CUTLASS extension") + return CUDAVideoInputProjectionFP32Op() + + +def _keys(rows, device): + return torch.stack( + (torch.zeros(rows, dtype=torch.int64, device=device), torch.arange(rows, device=device)), + dim=1, + ) + + +def _data(rows, device="cuda"): + generator = torch.Generator(device=device).manual_seed(420) + return ( + torch.randn(rows, 96, generator=generator, device=device) * 0.125, + torch.randn(5376, 96, generator=generator, device=device) * 0.0625, + torch.randn(5376, generator=generator, device=device) * 0.125, + torch.randn(rows, 5376, generator=generator, device=device) * 0.0625, + ) + + +def _run(op, data, keys, chunks=None, *, train_weight=True, train_bias=True): + x, weight, bias, grad = [None if t is None else t.detach().clone() for t in data] + x.requires_grad_(True) + weight.requires_grad_(train_weight) + if bias is not None: + bias.requires_grad_(train_bias) + cuts = chunks if chunks is not None else [torch.arange(x.shape[0], device=x.device)] + with canonical_backward_session() as session: + outputs = [ + op( + x.index_select(0, ix), + weight, + bias, + logical_keys=keys.index_select(0, ix), + parameter_id="proj_in", + ) + for ix in cuts + ] + torch.autograd.backward(outputs, [grad.index_select(0, ix) for ix in cuts]) + session.validate_complete() + # Scatter the per-call outputs back to physical row order. + y = torch.empty(x.shape[0], 5376, device=x.device) + for ix, out in zip(cuts, outputs, strict=True): + y.index_copy_(0, ix, out.detach()) + return y, x.grad, weight.grad, None if bias is None else bias.grad + + +def _assert_accuracy(actual, expected, kind): + threshold = load_contract()["judgments"][kind]["by_op_class"]["reduction"]["float32"] + torch.testing.assert_close( + actual.cpu(), expected.cpu(), atol=threshold["atol"], rtol=threshold["rtol"] + ) + + +def _assert_bytes(a, b): + assert torch.equal( + a.detach().contiguous().view(torch.int32), b.detach().contiguous().view(torch.int32) + ) + + +@pytest.mark.parametrize("with_bias", [False, True]) +def test_canonical_linear_bias_and_legacy_call(with_bias): + torch.manual_seed(420) + x = torch.randn(7, 3, requires_grad=True) + weight = torch.randn(5, 3, requires_grad=True) + bias = torch.randn(5, requires_grad=True) if with_bias else None + grad = torch.randn(7, 5) + with canonical_backward_session() as session: + y = canonical_linear_fp32( + x, weight, _keys(7, "cpu"), parameter_id="weight", family="pytorch", bias=bias + ) + y.backward(grad) + session.validate_complete() + torch.testing.assert_close(y, torch.nn.functional.linear(x, weight, bias)) + torch.testing.assert_close(x.grad, grad @ weight) + torch.testing.assert_close(weight.grad, grad.t() @ x) + if with_bias: + torch.testing.assert_close(bias.grad, grad.sum(0)) + + +def test_input_gradient_without_parameter_session(): + x = torch.ones(3, 2, requires_grad=True) + weight = torch.ones(4, 2) + canonical_linear_fp32(x, weight, None, parameter_id="w", family="pytorch").sum().backward() + torch.testing.assert_close(x.grad, torch.full_like(x, 4)) + + +class _CPUAdapter(VideoInputProjectionFP32Op): + family = "pytorch" + + def _check_backend(self, x): + pass + + +@pytest.mark.parametrize("frozen", ["weight", "bias", "both"]) +def test_frozen_parameter_registration(frozen): + result = _run( + _CPUAdapter(), + _data(2, "cpu"), + _keys(2, "cpu"), + train_weight=frozen not in ("weight", "both"), + train_bias=frozen not in ("bias", "both"), + ) + assert (result[2] is None) == (frozen in ("weight", "both")) + assert (result[3] is None) == (frozen in ("bias", "both")) + + +@pytest.mark.parametrize( + "invalid", ["dtype", "shape", "empty", "bias", "noncontiguous", "keys", "duplicates", "session"] +) +def test_invalid_inputs(invalid): + x, weight, bias, _ = _data(2, "cpu") + keys = _keys(2, "cpu") + if invalid == "dtype": + x = x.bfloat16() + elif invalid == "shape": + x = x[:, :95].contiguous() + elif invalid == "empty": + x = x[:0] + elif invalid == "bias": + bias = bias[:10] + elif invalid == "noncontiguous": + x = torch.empty(2, 192)[:, ::2] + elif invalid == "keys": + keys = keys.int() + elif invalid == "duplicates": + keys.zero_() + weight.requires_grad_(True) + with pytest.raises((TypeError, ValueError, RuntimeError)): + if invalid == "session": + _CPUAdapter()(x, weight, bias, logical_keys=keys, parameter_id="p") + else: + with canonical_backward_session(): + _CPUAdapter()(x, weight, bias, logical_keys=keys, parameter_id="p") + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("rows", [1, 9, 17, 65, 129, 257]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_forward_and_actual_gradients(backend, rows, with_bias): + data = list(_data(rows)) + if not with_bias: + data[2] = None + # _run clones optional bias separately to keep the production call identical. + if not with_bias: + x, weight, _, grad = data + x, weight = x.requires_grad_(), weight.requires_grad_() + with canonical_backward_session() as session: + y = _op(backend)(x, weight, logical_keys=_keys(rows, "cuda"), parameter_id="p") + y.backward(grad) + session.validate_complete() + actual = (y, x.grad, weight.grad) + else: + actual = _run(_op(backend), data, _keys(rows, "cuda")) + cx, cw = data[0].detach().cpu().requires_grad_(), data[1].detach().cpu().requires_grad_() + cb = None if not with_bias else data[2].detach().cpu().requires_grad_() + gold = torch.nn.functional.linear(cx, cw, cb) + gold.backward(data[3].cpu()) + expected = (gold, cx.grad, cw.grad) if cb is None else (gold, cx.grad, cw.grad, cb.grad) + for i, (a, b) in enumerate(zip(actual, expected, strict=True)): + _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("rows", [9, 65, 129]) +def test_parameter_bi_permutation_chunk_padding(backend, rows): + op, data, keys = _op(backend), _data(rows), _keys(rows, "cuda") + baseline = _run(op, data, keys) + order = torch.randperm(rows, device="cuda") + permuted = _run( + op, + [t.index_select(0, order) if i in (0, 3) else t for i, t in enumerate(data)], + keys.index_select(0, order), + ) + inverse = order.argsort() + for i in range(4): + _assert_bytes(baseline[i], permuted[i].index_select(0, inverse) if i < 2 else permuted[i]) + chunks = list(torch.tensor_split(order, 3)) + chunked = _run(op, data, keys, chunks) + for a, b in zip(baseline, chunked, strict=True): + _assert_bytes(a, b) + padded = [ + torch.cat((data[0], torch.randn(7, 96, device="cuda"))), + data[1], + data[2], + torch.cat((data[3], torch.randn(7, 5376, device="cuda"))), + ] + padded_keys = torch.cat((keys, torch.full((7, 2), -1, device="cuda", dtype=torch.int64))) + result = _run(op, padded, padded_keys) + for i, (a, b) in enumerate(zip(baseline, result, strict=True)): + _assert_bytes(a, b[:rows] if i < 2 else b) + repeated = _run(op, data, keys) + for a, b in zip(baseline, repeated, strict=True): + _assert_bytes(a, b) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +def test_row_placement_and_unrelated_rows(backend): + op, data = _op(backend), _data(65) + single = _run(op, [data[0][:1], data[1], data[2], data[3][:1]], _keys(1, "cuda")) + data[0][37] = data[0][0] + data[3][37] = data[3][0] + batch = _run(op, data, _keys(65, "cuda")) + _assert_bytes(single[0][0], batch[0][37]) + _assert_bytes(single[1][0], batch[1][37]) + changed = [t.clone() for t in data] + changed[0][0:30] += 3 + changed[3][0:30] -= 2 + result = _run(op, changed, _keys(65, "cuda")) + _assert_bytes(batch[0][37], result[0][37]) + _assert_bytes(batch[1][37], result[1][37]) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +def test_noncontiguous_upstream_and_forward_graph(backend): + op = _op(backend) + x, weight, bias, grad = _data(17) + with canonical_backward_session() as session: + y = op( + x.requires_grad_(), + weight.requires_grad_(), + bias.requires_grad_(), + logical_keys=_keys(17, "cuda"), + parameter_id="p", + ) + y.sum().backward() # Expanded scalar upstream has zero strides. + session.validate_complete() + _assert_accuracy(x.grad, torch.ones_like(grad).cpu() @ weight.cpu(), "gradient_accuracy") + with torch.no_grad(): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + expected = op(x, weight, bias) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = op(x, weight, bias) + graph.replay() + _assert_bytes(output, expected) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +def test_real_weights_patch_path(backend): + path = os.environ.get("H3_PROJECTION_WEIGHTS") + if not path: + pytest.skip("set H3_PROJECTION_WEIGHTS to pinned H3 safetensors fixture") + from safetensors import safe_open + from safetensors.torch import load_file + from rl_engine.reference.packing.video_patch_pack_unpack import NativeVideoPatchPackUnpackOp + + with safe_open(path, framework="pt") as fixture: + metadata = fixture.metadata() + assert metadata["repository"] == "MiniMaxAI/MiniMax-H3" + assert metadata["checkpoint_revision"] == "42ed227ee7df40d41602854ae760620d6eb651fe" + tensors = load_file(path) + weight = tensors["proj_in.weight"].cuda().float().contiguous() + bias = tensors["proj_in.bias"].cuda().float().contiguous() + torch.manual_seed(420) + latent = torch.randn(1, 24, 13, 48, 84, device="cuda", dtype=torch.bfloat16) + x = NativeVideoPatchPackUnpackOp().pack(latent).reshape(-1, 96).float() + grad = torch.randn(x.shape[0], 5376, device="cuda") * 0.0625 + actual = _run(_op(backend), (x, weight, bias, grad), _keys(x.shape[0], "cuda")) + cx, cw, cb = [t.detach().cpu().requires_grad_() for t in (x, weight, bias)] + gold = torch.nn.functional.linear(cx, cw, cb) + gold.backward(grad.cpu()) + for i, (a, b) in enumerate(zip(actual, (gold, cx.grad, cw.grad, cb.grad), strict=True)): + _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") + order = torch.arange(x.shape[0] - 1, -1, -1, device="cuda") + chunked = _run( + _op(backend), + (x, weight, bias, grad), + _keys(x.shape[0], "cuda"), + list(torch.tensor_split(order, 3)), + ) + for a, b in zip(actual, chunked, strict=True): + _assert_bytes(a, b) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +def test_default_cuda_dispatch(): + from rl_engine.runtime.registry import KernelRegistry + + op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v1" + x, weight, bias, _ = _data(3) + with torch.no_grad(): + _assert_bytes(op(x, weight, bias), _op("triton")(x, weight, bias)) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +def test_reused_fp32_gemm(): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + x, weight, _, _ = _data(17) + output = _triton_gemm_fp32(x, weight.t().contiguous()) + _assert_accuracy(output, x.cpu() @ weight.cpu().t(), "forward_accuracy") + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +def test_standard_normal_long_reduction(backend): + # Unscaled inputs expose cancellation error in the 5376-term input VJP. + data = _data(129) + data = tuple(t / scale for t, scale in zip(data, (0.125, 0.0625, 0.125, 0.0625), strict=True)) + actual = _run(_op(backend), data, _keys(129, "cuda")) + cx, cw, cb = [t.detach().cpu().requires_grad_() for t in data[:3]] + gold = torch.nn.functional.linear(cx, cw, cb) + gold.backward(data[3].cpu()) + for i, (a, b) in enumerate(zip(actual, (gold, cx.grad, cw.grad, cb.grad), strict=True)): + _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index b2726e67e..7bdcd03f5 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -7,6 +7,20 @@ from rl_engine.runtime.registry import KernelRegistry, OpBackend +def test_video_input_projection_default_and_strict_scope(): + registry = KernelRegistry() + assert registry._priority_map["cuda"]["video_input_projection_fp32"] == [ + OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 + ] + desc = registry.semantic.backend_descriptor( + "video_input_projection_fp32", "triton-video-input-projection-fp32-v1" + ) + assert desc.supported_devices == frozenset({"cuda"}) + assert desc.supported_dtypes == frozenset({"float32"}) + with pytest.raises(RuntimeError, match="SM90"): + registry.get_op("video_input_projection_fp32", device="cpu") + + def test_rocm_attention_defaults_to_flash_attention(monkeypatch): monkeypatch.delenv("RL_KERNEL_ROCM_ATTN_BACKEND", raising=False) diff --git a/tests/validation/operators/test_op_checks.py b/tests/validation/operators/test_op_checks.py index 9c9b6999b..bd16c169f 100644 --- a/tests/validation/operators/test_op_checks.py +++ b/tests/validation/operators/test_op_checks.py @@ -4,6 +4,7 @@ from __future__ import annotations import argparse +from contextlib import contextmanager import pytest import torch @@ -20,6 +21,64 @@ ) +def test_video_input_projection_reference_forward_and_backward(): + args = _spec_args("video_input_projection_fp32") + case = make_operator_case(args, torch.float32, torch.device("cpu")) + report = run_operator_suite( + "video_input_projection_fp32", + candidates=[make_candidate(args)], + cases=[case], + check_grad=True, + grad_mode="random", + ) + assert report.passed + assert case.grad_input_names == ("x", "weight", "bias") + + +def test_backward_context_surrounds_actual_autograd(): + state = {"active": False, "backward": False} + + @contextmanager + def context(): + state["active"] = True + yield + assert state["backward"] + state["active"] = False + + class Probe(torch.autograd.Function): + @staticmethod + def forward(ctx, logits, token_ids): + assert state["active"] + ctx.save_for_backward(logits, token_ids) + return NativeLogpOp().forward_fp32(logits, token_ids) + + @staticmethod + def backward(ctx, grad): + assert state["active"] + state["backward"] = True + logits, ids = ctx.saved_tensors + with torch.enable_grad(): + x = logits.detach().requires_grad_() + out = NativeLogpOp().forward_fp32(x, ids) + result = torch.autograd.grad(out, x, grad)[0] + return result, None + + report = run_operator_suite( + "logp", + candidates=[ + CandidateSpec( + "probe", + lambda logits, token_ids: Probe.apply(logits, token_ids), + backward_context_factory=context, + ), + ], + cases=[_logp_backward_case("context")], + check_grad=True, + ) + assert report.passed + assert not state["active"] + + def _logp_case(name: str, dtype: torch.dtype, *, seed: int = 0) -> OperatorCase: generator = torch.Generator().manual_seed(seed) logits = torch.randn(2, 8, 257, dtype=dtype, generator=generator) diff --git a/tests/validation/operators/test_operator_inputs.py b/tests/validation/operators/test_operator_inputs.py index 9e5015bf4..938c2d9d5 100644 --- a/tests/validation/operators/test_operator_inputs.py +++ b/tests/validation/operators/test_operator_inputs.py @@ -38,6 +38,22 @@ def _args(**overrides): return argparse.Namespace(**values) +def test_video_input_projection_contract_and_logical_keys(): + args = _args(batch=2, seq=3) + inputs = make_operator_inputs( + "video_input_projection_fp32", args, torch.float32, torch.device("cpu") + ) + assert inputs["x"].shape == (2, 3, 96) + assert inputs["weight"].shape == (5376, 96) + assert inputs["bias"].shape == (5376,) + assert inputs["logical_keys"].tolist() == [[0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]] + assert operator_shape_name("video_input_projection_fp32", args) == "2x3x96x5376" + with pytest.raises(TypeError, match="FP32"): + make_operator_inputs( + "video_input_projection_fp32", args, torch.bfloat16, torch.device("cpu") + ) + + @pytest.mark.parametrize( "op_name", [ From ed7f0dbd4c04dbb4071a74da4a7b4dfd40064066 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Wed, 14 Oct 2026 15:21:22 +0800 Subject: [PATCH 02/26] Reuse shared GPU reduction tree for H3 FP32 input projection Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 4 +- build_tools/extensions.py | 2 +- .../video_input_projection_fp32_cutlass.cpp | 7 ++ .../video_input_projection_fp32_cutlass.cu | 101 ++++++++++++++++++ docs/operators/video-input-projection-fp32.md | 33 ++++-- rl_engine/_C_cutlass.pyi | 8 ++ .../cuda/gemm/video_input_projection_fp32.py | 40 ++++++- .../backends/shared/triton/gemm/det_gemm.py | 60 ++++++++--- .../gemm/video_input_projection_fp32.py | 30 ++++-- rl_engine/ops/gemm/canonical_linear.py | 63 +++++------ .../ops/gemm/video_input_projection_fp32.py | 4 +- rl_engine/runtime/registry.py | 17 +-- .../validation/operators/operator_specs.py | 4 +- .../build/test_build_platform_collectives.py | 3 +- tests/ops/gemm/test_det_gemm.py | 10 +- .../gemm/test_video_input_projection_fp32.py | 57 +++++++++- tests/runtime/test_kernel_registry.py | 2 +- 17 files changed, 354 insertions(+), 91 deletions(-) diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index 64b1df48c..aa144868d 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -9,10 +9,10 @@ import hashlib import json import platform -import statistics import shutil -from functools import partial +import statistics import subprocess +from functools import partial from pathlib import Path import torch diff --git a/build_tools/extensions.py b/build_tools/extensions.py index 3786fd43a..aebba78ee 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -import importlib.util import hashlib +import importlib.util import os import sysconfig import warnings diff --git a/csrc/bindings/video_input_projection_fp32_cutlass.cpp b/csrc/bindings/video_input_projection_fp32_cutlass.cpp index e4e14a42d..60e8d5ed9 100644 --- a/csrc/bindings/video_input_projection_fp32_cutlass.cpp +++ b/csrc/bindings/video_input_projection_fp32_cutlass.cpp @@ -5,10 +5,15 @@ torch::Tensor cutlass_gemm_fp32_tn(torch::Tensor a, torch::Tensor b_transposed); int cutlass_cuda_version(); +void cutlass_gemm_fp32_tree_leaves_tn( + torch::Tensor a, torch::Tensor bt, torch::Tensor workspace, + torch::Tensor starts, torch::Tensor lengths, torch::Tensor nodes); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("gemm_fp32_tn", &cutlass_gemm_fp32_tn, "Fixed FP32 SIMT GEMM: [M,K] times transposed [N,K]"); + m.def("gemm_fp32_tree_leaves_tn", &cutlass_gemm_fp32_tree_leaves_tn, + "Batch all FP32 leaves into a validated canonical tree workspace"); m.def("build_info", []() { pybind11::dict info; info["cutlass_revision"] = "56ca28adb32ffadd097566545b8f7d83b652e998"; @@ -20,6 +25,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { info["split_k"] = 1; info["fmad"] = false; info["ftz"] = false; + info["tree_leaf_size"] = 32; + info["tree_leaf_launcher"] = "cutlass_device_gemm_array"; return info; }); } diff --git a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu index d552d3e5d..d26153e1a 100644 --- a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu +++ b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu @@ -12,6 +12,7 @@ #include #include +#include using Gemm = cutlass::gemm::device::Gemm< float, cutlass::layout::RowMajor, @@ -25,6 +26,106 @@ using Gemm = cutlass::gemm::device::Gemm< cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 2, 1, 1, false, cutlass::arch::OpMultiplyAdd>; +// Match the shared tree plan's 32-term leaves. One batched launch writes all +// leaves into their existing workspace nodes; Triton then merges each level. +constexpr int kTreeLeaf = 32; +using LeafGemm = cutlass::gemm::device::GemmArray< + float, cutlass::layout::RowMajor, + float, cutlass::layout::ColumnMajor, + float, cutlass::layout::RowMajor, + float, cutlass::arch::OpClassSimt, cutlass::arch::Sm50, + cutlass::gemm::GemmShape<64, 64, 8>, + cutlass::gemm::GemmShape<32, 32, 8>, + cutlass::gemm::GemmShape<1, 1, 1>, + cutlass::epilogue::thread::LinearCombination, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 2, 1, 1, cutlass::arch::OpMultiplyAdd>; + +template +__global__ void pack_tree_operand( + const float* input, float* packed, int64_t* pointers, + const int32_t* starts, const int32_t* lengths, const int32_t* nodes, + int64_t rows, int64_t stride_row, int64_t stride_k, int64_t total, + float* workspace, int64_t* output_pointers, int64_t output_elements) { + const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= total) return; + const int64_t leaf_elements = rows * kTreeLeaf; + const int leaf = static_cast(index / leaf_elements); + const int64_t local = index % leaf_elements; + const int64_t row = local / kTreeLeaf; + const int k = static_cast(local % kTreeLeaf); + packed[index] = k < lengths[leaf] + ? input[row * stride_row + (static_cast(starts[leaf]) + k) * stride_k] + : 0.0f; + if (local == 0) { + pointers[leaf] = reinterpret_cast(packed + leaf * leaf_elements); + if constexpr (WriteOutputPointers) { + output_pointers[leaf] = reinterpret_cast( + workspace + static_cast(nodes[leaf]) * output_elements); + } + } +} + +void cutlass_gemm_fp32_tree_leaves_tn( + at::Tensor a, at::Tensor bt, at::Tensor workspace, + at::Tensor starts, at::Tensor lengths, at::Tensor nodes) { + TORCH_CHECK(a.is_cuda() && a.scalar_type() == at::kFloat && a.dim() == 2, + "CUTLASS leaves require rank-two CUDA FP32 A"); + const c10::cuda::CUDAGuard guard(a.device()); + const auto* properties = at::cuda::getDeviceProperties(a.get_device()); + TORCH_CHECK(properties->major == 9 && properties->minor == 0, + "CUTLASS video projection requires CUDA SM90"); + TORCH_CHECK(bt.device() == a.device() && bt.scalar_type() == at::kFloat && bt.dim() == 2, + "CUTLASS leaves require rank-two FP32 Bt on the same device"); + TORCH_CHECK(a.size(1) == bt.size(1) && a.size(0) > 0 && bt.size(0) > 0, + "CUTLASS leaf dimensions must match and be non-empty"); + TORCH_CHECK(a.stride(0) > 0 && a.stride(1) > 0 && bt.stride(0) > 0 && bt.stride(1) > 0, + "CUTLASS leaf operands require positive strides"); + TORCH_CHECK(workspace.device() == a.device() && workspace.scalar_type() == at::kFloat && + workspace.dim() == 3 && workspace.is_contiguous() && + workspace.size(1) == a.size(0) && workspace.size(2) == bt.size(0), + "CUTLASS leaves require contiguous FP32 tree workspace"); + for (const auto& metadata : {starts, lengths, nodes}) { + TORCH_CHECK(metadata.device() == a.device() && metadata.scalar_type() == at::kInt && + metadata.dim() == 1 && metadata.is_contiguous() && + metadata.numel() == starts.numel(), + "CUTLASS tree metadata must be matching contiguous CUDA int32 vectors"); + } + const int64_t leaves = starts.numel(), m = a.size(0), n = bt.size(0); + TORCH_CHECK(leaves > 0 && leaves <= 65535 && m <= INT_MAX && n <= INT_MAX, + "CUTLASS batched leaf grid exceeds its supported dimensions"); + auto packed_a = at::empty({leaves, m, kTreeLeaf}, a.options()); + auto packed_b = at::empty({leaves, n, kTreeLeaf}, a.options()); + auto pointers = at::empty({3, leaves}, a.options().dtype(at::kLong)); + auto* ptr_a = pointers.data_ptr(); + auto* ptr_b = ptr_a + leaves; + auto* ptr_output = ptr_b + leaves; + const auto stream = c10::cuda::getCurrentCUDAStream(a.get_device()).stream(); + constexpr int threads = 256; + pack_tree_operand<<<(packed_a.numel() + threads - 1) / threads, threads, 0, stream>>>( + a.data_ptr(), packed_a.data_ptr(), ptr_a, + starts.data_ptr(), lengths.data_ptr(), nodes.data_ptr(), + m, a.stride(0), a.stride(1), packed_a.numel(), nullptr, nullptr, m * n); + pack_tree_operand<<<(packed_b.numel() + threads - 1) / threads, threads, 0, stream>>>( + bt.data_ptr(), packed_b.data_ptr(), ptr_b, + starts.data_ptr(), lengths.data_ptr(), nodes.data_ptr(), + n, bt.stride(0), bt.stride(1), packed_b.numel(), workspace.data_ptr(), + ptr_output, m * n); + LeafGemm::Arguments args( + {static_cast(m), static_cast(n), kTreeLeaf}, + reinterpret_cast(ptr_a), cutlass::layout::RowMajor(kTreeLeaf), + reinterpret_cast(ptr_b), cutlass::layout::ColumnMajor(kTreeLeaf), + reinterpret_cast(ptr_output), cutlass::layout::RowMajor(n), + reinterpret_cast(ptr_output), cutlass::layout::RowMajor(n), + {1.0f, 0.0f}, static_cast(leaves)); + LeafGemm gemm; + auto status = gemm.can_implement(args); + TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); + status = gemm(args, nullptr, stream); + TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + at::Tensor cutlass_gemm_fp32_tn(at::Tensor a, at::Tensor b_transposed) { TORCH_CHECK(a.is_cuda() && b_transposed.is_cuda(), "CUTLASS GEMM requires CUDA tensors"); TORCH_CHECK(a.device() == b_transposed.device(), "CUTLASS GEMM devices must match"); diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index f3014c0b0..d67511b20 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -8,7 +8,9 @@ reference. The model adapter owns the subsequent cast into the BF16 packed strea The primary implementation reuses RL-Kernel's Triton FP32 GEMM and canonical parameter-gradient session. A pinned CUTLASS FP32 SIMT GEMM is an explicit CUDA -candidate. Both use the same semantic adapter and autograd implementation. +candidate. Both use the same semantic adapter, autograd implementation, and +existing Triton tree plan and GPU level reductions. CUTLASS substitutes the +batched leaf GEMM; its reported execution identity includes the Triton reduction. ## Entry Point @@ -31,8 +33,8 @@ physical placement. The same parameter keeps the same `parameter_id` across uses | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | -| Triton on CUDA | `TritonVideoInputProjectionFP32Op` | `_det_gemm_fp32_kernel` | Default CUDA path; SM90 qualification | -| CUDA / CUTLASS | `CUDAVideoInputProjectionFP32Op` | `_C_cutlass.gemm_fp32_tn` | Explicit candidate; optional build; SM90 | +| Triton on CUDA | `TritonVideoInputProjectionFP32Op` | `_det_gemm_fp32_kernel`, existing tree leaf/reduce kernels | Default CUDA path; SM90 qualification | +| CUDA / CUTLASS + Triton reductions | `CUDAVideoInputProjectionFP32Op` | `_C_cutlass.gemm_fp32_tn`, `_C_cutlass.gemm_fp32_tree_leaves_tn`, shared tree reductions | Explicit candidate; optional build; SM90 | | PyTorch reference | `NativeVideoInputProjectionFP32Op` | CPU `F.linear` | Accuracy reference | CUTLASS revision: `56ca28adb32ffadd097566545b8f7d83b652e998`. The build verifies @@ -87,15 +89,23 @@ The existing canonical session collects all contributions, removes padding, and sorts keys before the final parameter reduction. It returns each parameter VJP once. Ordinary chunk-by-chunk accumulation is a different arithmetic contract. -Long K reductions use a fixed midpoint tree with at most 128 terms per leaf. -Each leaf reuses the selected GEMM and every tree node adds in FP32. This bounds -rounding depth for the 5376-term input VJP and long parameter reductions. The -tree depends on logical K and retains its topology across physical row changes. -Existing canonical linear callers keep their original reduction by default. +Short reductions up to the 96-wide video input use the existing direct FP32 +GEMM. Long reductions reuse the existing midpoint plan with at most 32 terms per +leaf. Workspace, merge nodes, and output stay FP32. The Triton backend launches +all leaves in one grid. CUTLASS packs the same leaf ranges into 32-wide buffers, +pads inactive slots with zero, and uses one `device::GemmArray` launch to write +the plan's leaf nodes. Both then reuse the existing GPU level reductions. + +The tree depends on logical K and retains its topology across physical row +changes. Output-row chunking keeps tree workspace below the existing element +limit (512 MiB for FP32). If one output row exceeds that limit, independent +output columns are chunked with the same K tree. CUTLASS also allocates packed +operands and pointer arrays. Parameter sorting and public validation are included in complete-call +measurements. Existing BF16 tree callers retain their default dtype and schedule. Triton fixes the 64x64x32 leaf schedule and disables TF32 and FP fusion. CUTLASS fixes 64x64x8 FP32 SIMT tiles and a single output owner per leaf, without parallel -split-K. Their leaf arithmetic has separate backend identities. +split-K. Both leaf algorithms have separately versioned backend identities. The shared WS1 `reduction` FP32 accuracy policy uses `atol=rtol=1e-4` for forward and gradients. Invariance uses byte equality. Parameter invariance compares the @@ -116,8 +126,9 @@ forward CUDA Graph. Timings include layout copies, bias, and the canonical parameter-gradient machinery. Raw samples, environment, source hashes, and actual backend metadata are written to the explicitly selected local output. PyTorch is a production performance baseline with separately reported BI results. -The fixed tree adds kernel launches and layout copies to backward; compare full -operator timings before drawing a speedup conclusion. +The shared GPU tree avoids Python recursion per leaf and merge node. Its +workspace traffic, output-row chunking, and CUTLASS operand packing still add +backward work; compare complete operator timings before drawing a speedup claim. ## Tests diff --git a/rl_engine/_C_cutlass.pyi b/rl_engine/_C_cutlass.pyi index a1afa227a..862c08a2e 100644 --- a/rl_engine/_C_cutlass.pyi +++ b/rl_engine/_C_cutlass.pyi @@ -6,4 +6,12 @@ from typing import Any import torch def gemm_fp32_tn(a: torch.Tensor, b_transposed: torch.Tensor) -> torch.Tensor: ... +def gemm_fp32_tree_leaves_tn( + a: torch.Tensor, + b_transposed: torch.Tensor, + workspace: torch.Tensor, + starts: torch.Tensor, + lengths: torch.Tensor, + nodes: torch.Tensor, +) -> None: ... def build_info() -> dict[str, Any]: ... diff --git a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py index 0f82ae6c5..ef921f782 100644 --- a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py @@ -35,7 +35,15 @@ import torch -from rl_engine.ops.gemm.video_input_projection_fp32 import VideoInputProjectionFP32Op +from rl_engine.backends.shared.triton.gemm.det_gemm import ( + _BLOCK_K, + _TRITON_AVAILABLE, + _triton_tree_gemm, +) +from rl_engine.ops.gemm.video_input_projection_fp32 import ( + VIDEO_PATCH_WIDTH, + VideoInputProjectionFP32Op, +) CUTLASS_REVISION = "56ca28adb32ffadd097566545b8f7d83b652e998" CUTLASS_HEADERS_SHA256 = "7b5321860c269f8e44eab02e59c32c2d4d0e2fa7a3ef72e0a161de9f58b361fa" @@ -56,21 +64,38 @@ def _load_extension(): or info["fmad"] or info["ftz"] or info["split_k"] != 1 + or info.get("tree_leaf_size") != _BLOCK_K ): raise RuntimeError("CUTLASS extension does not match the strict FP32 build contract") return extension +def _cutlass_tree_leaves(a, b, workspace, plan): + _load_extension().gemm_fp32_tree_leaves_tn( + a, b.t(), workspace, plan.leaf_starts, plan.leaf_lengths, plan.leaf_nodes + ) + + def cutlass_gemm_fp32(a, b): # The native TN interface represents mathematical B[K,N] as [N,K] storage. - return _load_extension().gemm_fp32_tn(a.contiguous(), b.t().contiguous()) + if a.shape[1] <= VIDEO_PATCH_WIDTH: + return _load_extension().gemm_fp32_tn(a.contiguous(), b.t().contiguous()) + return _triton_tree_gemm( + a, + b, + output_dtype=torch.float32, + preserve_a_strides=True, + fp32_leaf_launcher=_cutlass_tree_leaves, + ) class CUDAVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "cutlass" - backend_id = "cuda-cutlass-video-input-projection-fp32-v1" + backend_id = "cuda-cutlass-video-input-projection-fp32-v2" def __init__(self): + if not _TRITON_AVAILABLE: + raise RuntimeError("CUTLASS projection requires Triton for shared FP32 tree reductions") _load_extension() def _check_backend(self, x): @@ -80,11 +105,16 @@ def _check_backend(self, x): raise ValueError("CUTLASS video input projection requires CUDA SM90") def backend_info(self): + import triton + return { "backend_id": self.backend_id, "family": "cuda", - "kernel": "gemm_fp32_tn", + "kernel": "gemm_fp32_tn / gemm_fp32_tree_leaves_tn", "reduction": "fixed_midpoint_tree", - "reduction_leaf_size": 128, + "reduction_leaf_size": _BLOCK_K, + "reduction_backend": "triton", + "triton_version": triton.__version__, + "leaf_launch": "cutlass_device_gemm_array", **_load_extension().build_info(), } diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 558507602..5f7691189 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -11,7 +11,8 @@ does not match the scalar reference at every BF16 rounding boundary. Autotuning and split-K are intentionally disabled. FP32-output calls preserve -the earlier fixed-order, no-split-K contract and do not use BF16 tree nodes. +the earlier fixed-order, no-split-K contract. The optional CUDA FP32 tree reuses +the canonical plan, leaf kernel, and level reductions with FP32 nodes. """ from __future__ import annotations @@ -773,6 +774,8 @@ def _triton_tree_gemm( out: torch.Tensor | None = None, preserve_a_strides: bool = False, inference_schedule: bool = False, + output_dtype: torch.dtype = torch.bfloat16, + fp32_leaf_launcher=None, ) -> torch.Tensor: if not _TRITON_AVAILABLE: raise RuntimeError("Triton is unavailable") @@ -780,8 +783,16 @@ def _triton_tree_gemm( raise ValueError("Triton deterministic GEMM expects two 2-D tensors") if a.size(1) != b.size(0): raise ValueError(f"Triton deterministic GEMM K mismatch: {a.size(1)} and {b.size(0)}") - if a.dtype != torch.bfloat16 or b.dtype != torch.bfloat16: - raise TypeError("Triton tree GEMM requires BF16 inputs") + if output_dtype not in (torch.bfloat16, torch.float32): + raise TypeError("Triton tree GEMM output must be BF16 or FP32") + fp32_tree = output_dtype == torch.float32 + input_dtype = torch.float32 if fp32_tree else torch.bfloat16 + if a.dtype != input_dtype or b.dtype != input_dtype: + raise TypeError(f"Triton tree GEMM requires {input_dtype} inputs") + if fp32_tree and torch.version.hip is not None: + raise RuntimeError("FP32 tree GEMM is qualified on CUDA only") + if fp32_leaf_launcher is not None and not fp32_tree: + raise ValueError("a custom leaf launcher requires the FP32 tree") if not a.is_cuda or not b.is_cuda or a.device != b.device: raise RuntimeError("Triton tree GEMM inputs must share one CUDA/ROCm device") @@ -796,6 +807,25 @@ def _triton_tree_gemm( plan = _device_tree_plan(k_size, a.device) workspace_elements = plan.host.node_count * m_size * n_size if workspace_elements > _MAX_TREE_WORKSPACE_ELEMENTS and out is None: + if fp32_tree and plan.host.node_count * n_size > _MAX_TREE_WORKSPACE_ELEMENTS: + columns_per_chunk = _MAX_TREE_WORKSPACE_ELEMENTS // plan.host.node_count + if columns_per_chunk < 1: + raise RuntimeError("FP32 tree plan exceeds the workspace budget for one element") + # A one-row bias VJP cannot be split further over M. Splitting + # independent output columns retains the same logical K tree. + chunks = [ + _triton_tree_gemm( + a, + b[:, start : start + columns_per_chunk], + transpose_output=transpose_output, + preserve_a_strides=preserve_a_strides, + inference_schedule=inference_schedule, + output_dtype=output_dtype, + fp32_leaf_launcher=fp32_leaf_launcher, + ) + for start in range(0, n_size, columns_per_chunk) + ] + return torch.cat(chunks, dim=0 if transpose_output else 1) rows_per_chunk = max( 1, _MAX_TREE_WORKSPACE_ELEMENTS // (plan.host.node_count * n_size), @@ -809,20 +839,22 @@ def _triton_tree_gemm( transpose_output=transpose_output, preserve_a_strides=preserve_a_strides, inference_schedule=inference_schedule, + output_dtype=output_dtype, + fp32_leaf_launcher=fp32_leaf_launcher, ) chunks.append(chunk) return torch.cat(chunks, dim=1 if transpose_output else 0) result_shape = (n_size, m_size) if transpose_output else (m_size, n_size) if out is None: - result = torch.empty(result_shape, dtype=torch.bfloat16, device=a.device) + result = torch.empty(result_shape, dtype=output_dtype, device=a.device) else: if tuple(out.shape) != result_shape: raise ValueError( f"Triton tree GEMM output must have shape {result_shape}, " f"got {tuple(out.shape)}" ) - if out.dtype != torch.bfloat16: - raise TypeError(f"Triton tree GEMM output must be BF16, got {out.dtype}") + if out.dtype != output_dtype: + raise TypeError(f"Triton tree GEMM output must be {output_dtype}, got {out.dtype}") if out.device != a.device: raise RuntimeError(f"Triton tree GEMM output must be on {a.device}, got {out.device}") if not out.is_contiguous(): @@ -832,7 +864,7 @@ def _triton_tree_gemm( result = out workspace = torch.empty( (plan.host.node_count, m_size, n_size), - dtype=torch.bfloat16, + dtype=output_dtype, device=a.device, ) device_index = a.device.index if a.device.index is not None else torch.cuda.current_device() @@ -861,7 +893,7 @@ def _triton_tree_gemm( if leaf_config.n_fastest else (len(plan.host.leaf_nodes), tiles_m, tiles_n) ) - direct_root_output = not transpose_output and is_gfx942 + direct_root_output = not transpose_output and (is_gfx942 or fp32_tree) fuse_leaf_reduction = ( _ROCM_FUSE_LEAF_REDUCTION and direct_root_output @@ -869,7 +901,9 @@ def _triton_tree_gemm( and k_size >= 1536 and plan.rocm_leaf_reduction is not None ) - if fuse_leaf_reduction: + if fp32_leaf_launcher is not None: + fp32_leaf_launcher(a, b, workspace, plan) + elif fuse_leaf_reduction: assert plan.rocm_leaf_reduction is not None operation_count = len(plan.host.reduction_levels[0]) fused_leaf_grid = ( @@ -895,6 +929,7 @@ def _triton_tree_gemm( N_FASTEST=leaf_config.n_fastest, WRITE_OUTPUT=len(plan.host.reduction_levels) == 1, num_warps=leaf_config.num_warps, + **({"enable_fp_fusion": False} if fp32_tree else {}), ) else: _det_gemm_kernel[leaf_grid]( @@ -916,6 +951,7 @@ def _triton_tree_gemm( BLOCK_K=_BLOCK_K, N_FASTEST=leaf_config.n_fastest, num_warps=leaf_config.num_warps, + **({"enable_fp_fusion": False} if fp32_tree else {}), ) reduction_block = 256 if fuse_leaf_reduction: @@ -945,7 +981,7 @@ def _triton_tree_gemm( N=n_size, BLOCK=reduction_block, ) - elif direct_root_output and use_inference_schedule: + elif direct_root_output and use_inference_schedule and is_gfx942: blocks = triton.cdiv(m_size * n_size, reduction_block) for pair_index, nodes in enumerate(plan.rocm_fused_reduction_pairs): second_level = pair_index * 2 + 1 @@ -986,8 +1022,8 @@ def _triton_tree_gemm( raise RuntimeError("the final deterministic GEMM tree level must contain one root") if write_final_output: grid = (triton.cdiv(m_size * n_size, reduction_block),) - # direct_root_output is true only for gfx942; CUDA retains its - # established reduction and root-copy path. + # FP32 CUDA also writes the final level directly to output. + # Existing BF16 CUDA keeps its reduction and root-copy path. _det_gemm_tree_reduce_to_output_rocm_kernel[grid]( workspace, result, diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index f594a4b79..95d711e27 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -5,13 +5,27 @@ import torch -from rl_engine.backends.shared.triton.gemm.det_gemm import _TRITON_AVAILABLE -from rl_engine.ops.gemm.video_input_projection_fp32 import VideoInputProjectionFP32Op +from rl_engine.backends.shared.triton.gemm.det_gemm import ( + _BLOCK_K, + _TRITON_AVAILABLE, + _triton_gemm_fp32, + _triton_tree_gemm, +) +from rl_engine.ops.gemm.video_input_projection_fp32 import ( + VIDEO_PATCH_WIDTH, + VideoInputProjectionFP32Op, +) + + +def projection_gemm_fp32(a, b): + if a.shape[1] <= VIDEO_PATCH_WIDTH: + return _triton_gemm_fp32(a, b) + return _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): - family = "triton" - backend_id = "triton-video-input-projection-fp32-v1" + family = "triton_tree" + backend_id = "triton-video-input-projection-fp32-v2" def __init__(self): if not _TRITON_AVAILABLE: @@ -28,13 +42,15 @@ def backend_info(self): return { "backend_id": self.backend_id, - "family": self.family, - "kernel": "_det_gemm_fp32_kernel", + "family": "triton", + "kernel": "_det_gemm_fp32_kernel / _det_gemm_tree_leaf_kernel", "triton_version": triton.__version__, "tile": [64, 64, 32], "tf32": False, "fp_fusion": False, "split_k": 1, "reduction": "fixed_midpoint_tree", - "reduction_leaf_size": 128, + "reduction_leaf_size": _BLOCK_K, + "reduction_backend": "triton", + "leaf_launch": "all_leaves_in_one_grid", } diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index aa8268024..3e572ea0b 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -13,19 +13,13 @@ from rl_engine.ops.autograd.canonical_backward import active_session -def _gemm_fp32( - a: torch.Tensor, - b: torch.Tensor, - family: str, - reduction_leaf_size: int | None = None, -) -> torch.Tensor: - # Bound FP32 rounding depth for long reductions. The midpoint topology - # depends only on logical K, so physical row placement cannot change it. - if reduction_leaf_size is not None and a.shape[1] > reduction_leaf_size: - midpoint = a.shape[1] // 2 - lower = _gemm_fp32(a[:, :midpoint], b[:midpoint], family, reduction_leaf_size) - upper = _gemm_fp32(a[:, midpoint:], b[midpoint:], family, reduction_leaf_size) - return lower + upper +def _gemm_fp32(a: torch.Tensor, b: torch.Tensor, family: str) -> torch.Tensor: + if family == "triton_tree": + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + projection_gemm_fp32, + ) + + return projection_gemm_fp32(a, b) if family == "cutlass": from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import cutlass_gemm_fp32 @@ -56,7 +50,6 @@ def forward( family, bias, bias_parameter_id, - reduction_leaf_size, ): session = active_session() parameter_grad = ctx.needs_input_grad[1] or ctx.needs_input_grad[5] @@ -70,7 +63,6 @@ def forward( ctx.session = session ctx.parameter_id = str(parameter_id) ctx.family = str(family) - ctx.reduction_leaf_size = reduction_leaf_size ctx.bias_parameter_id = str(bias_parameter_id) ctx.bias_dtype = bias.dtype if bias is not None else None ctx.slot = ( @@ -85,7 +77,7 @@ def forward( # CUTLASS consumes column-major B directly from contiguous [N,K] weight. if ctx.family != "cutlass": rhs = rhs.contiguous() - output = _gemm_fp32(a.float(), rhs, ctx.family, ctx.reduction_leaf_size) + output = _gemm_fp32(a.float(), rhs, ctx.family) return output if bias is None else output + bias.float() @staticmethod @@ -94,18 +86,16 @@ def backward(ctx, grad_out): grad32 = grad_out.contiguous().float() da = None if ctx.needs_input_grad[0]: - da = _gemm_fp32(grad32, weight.float(), ctx.family, ctx.reduction_leaf_size).to(a.dtype) + da = _gemm_fp32(grad32, weight.float(), ctx.family).to(a.dtype) dweight = None if ctx.needs_input_grad[1]: def reducer(rows, grads): - return _gemm_fp32( - grads.float().t().contiguous(), - rows.float(), - ctx.family, - ctx.reduction_leaf_size, - ).to(weight.dtype) + lhs = grads.float().t() + if ctx.family not in ("triton_tree", "cutlass"): + lhs = lhs.contiguous() + return _gemm_fp32(lhs, rows.float(), ctx.family).to(weight.dtype) dweight = ctx.session.submit_linear(ctx.parameter_id, ctx.slot, a, grad_out, reducer) dbias = None @@ -115,19 +105,19 @@ def bias_reducer(grads): if grads.shape[0] == 0: raise RuntimeError(f"no active logical rows for {ctx.bias_parameter_id}") ones = torch.ones((1, grads.shape[0]), device=grads.device, dtype=torch.float32) - return ( - _gemm_fp32(ones, grads.float(), ctx.family, ctx.reduction_leaf_size) - .squeeze(0) - .to(ctx.bias_dtype) - ) + return _gemm_fp32(ones, grads.float(), ctx.family).squeeze(0).to(ctx.bias_dtype) dbias = ctx.session.submit_rows( ctx.bias_parameter_id, ctx.bias_slot, grad32, bias_reducer ) - if ctx.family == "triton": + if ctx.family in ("triton", "triton_tree"): record_backward( "det_gemm", - kernel_id="rl_engine.backends.shared.triton.gemm.det_gemm._triton_gemm", + kernel_id=( + "rl_engine.backends.shared.triton.gemm.det_gemm._triton_tree_gemm" + if ctx.family == "triton_tree" + else "rl_engine.backends.shared.triton.gemm.det_gemm._triton_gemm" + ), impl="triton_det_gemm_canonical_rowfold", family="triton", ) @@ -141,11 +131,14 @@ def bias_reducer(grads): elif ctx.family == "cutlass": record_backward( "video_input_projection_fp32", - kernel_id="rl_engine._C_cutlass.gemm_fp32_tn", - impl="cutlass_fp32_simt_canonical_rowfold", + kernel_id=( + "rl_engine._C_cutlass.gemm_fp32_tree_leaves_tn / " + "rl_engine.backends.shared.triton.gemm.det_gemm._triton_tree_gemm" + ), + impl="cutlass_fp32_simt_triton_tree_canonical_rowfold", family="cuda", ) - return da, dweight, None, None, None, dbias, None, None + return da, dweight, None, None, None, dbias, None def canonical_linear_fp32( @@ -157,10 +150,7 @@ def canonical_linear_fp32( family: str, bias: torch.Tensor | None = None, bias_parameter_id: str | None = None, - reduction_leaf_size: int | None = None, ) -> torch.Tensor: - if reduction_leaf_size is not None and reduction_leaf_size < 1: - raise ValueError("reduction_leaf_size must be positive") return _CanonicalLinearFn.apply( a, weight, @@ -169,7 +159,6 @@ def canonical_linear_fp32( family, bias, bias_parameter_id if bias_parameter_id is not None else f"{parameter_id}.bias", - reduction_leaf_size, ) diff --git a/rl_engine/ops/gemm/video_input_projection_fp32.py b/rl_engine/ops/gemm/video_input_projection_fp32.py index 7c9196504..c0d4d0389 100644 --- a/rl_engine/ops/gemm/video_input_projection_fp32.py +++ b/rl_engine/ops/gemm/video_input_projection_fp32.py @@ -12,7 +12,6 @@ VIDEO_PATCH_WIDTH = 96 VIDEO_HIDDEN_WIDTH = 5376 -FP32_REDUCTION_LEAF_SIZE = 128 def validate_projection_inputs(x, weight, bias): @@ -97,13 +96,12 @@ def __call__(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None family=self.family, bias=bias, bias_parameter_id=f"{parameter_id}.bias", - reduction_leaf_size=FP32_REDUCTION_LEAF_SIZE, ) else: rhs = weight.t() if self.family != "cutlass": rhs = rhs.contiguous() - output = _gemm_fp32(rows, rhs, self.family, FP32_REDUCTION_LEAF_SIZE) + output = _gemm_fp32(rows, rhs, self.family) if bias is not None: output = output + bias return output.reshape(*x.shape[:-1], VIDEO_HIDDEN_WIDTH) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 82f1938f1..02d7cade2 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -244,7 +244,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="cuda-cutlass-video-input-projection-fp32-v1", + backend_id="cuda-cutlass-video-input-projection-fp32-v2", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -257,7 +257,8 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "training": {"world_size": (1,), "sharding": ("unsharded",)}, }, determinism_or_alignment_properties={ - "algorithm": "cutlass_simt_fp32_canonical_parameter_vjp", + "algorithm": "cutlass_simt_fp32_triton_tree_canonical_parameter_vjp", + "reduction_backend": "triton", "input_width": 96, "output_width": 5376, "cuda_sm": 90, @@ -266,16 +267,18 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "strict_observable": True, "training_mode": "eager", "parameter_gradient_order": "logical_key", - "reduction": "fixed_midpoint_tree_128", + "reduction": "fixed_midpoint_tree_32", }, lifecycle=OperatorLifecycle.REQUEST, implementation_class_or_factory=OpBackend.CUDA_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, - version_or_build_fingerprint="cutlass-56ca28a-fp32-simt-64x64x8-tree128-h3-v1", + version_or_build_fingerprint=( + "cutlass-56ca28a-fp32-simt-64x64x8-triton-gpu-tree32-h3-v2" + ), ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v1", + backend_id="triton-video-input-projection-fp32-v2", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -297,12 +300,12 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "strict_observable": True, "training_mode": "eager", "parameter_gradient_order": "logical_key", - "reduction": "fixed_midpoint_tree_128", + "reduction": "fixed_midpoint_tree_32", }, lifecycle=OperatorLifecycle.REQUEST, implementation_class_or_factory=OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, - version_or_build_fingerprint="rlkernel-triton-fp32-64x64x32-no-fusion-tree128-h3-v1", + version_or_build_fingerprint="rlkernel-triton-fp32-64x64x32-no-fusion-gpu-tree32-h3-v2", ), OperatorBackendDescriptor( semantic_op="video_patch_pack_unpack", diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 0e7ff39b4..99f8bae32 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -46,7 +46,9 @@ def _load_object(path: str) -> Any: "video_input_projection_fp32": OperatorSpec( name="video_input_projection_fp32", op_class="reduction", - gold_path="rl_engine.reference.gemm.video_input_projection_fp32.NativeVideoInputProjectionFP32Op", + gold_path=( + "rl_engine.reference.gemm.video_input_projection_fp32.NativeVideoInputProjectionFP32Op" + ), gold_method="forward_fp32", candidate_paths={ "pytorch": ( diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index 271f5b3e4..70c45c7cb 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -32,9 +32,10 @@ def test_optional_cutlass_build_has_independent_strict_flags(monkeypatch, tmp_pa def test_cutlass_build_rejects_unpinned_headers(monkeypatch, tmp_path): - from build_tools import extensions import pytest + from build_tools import extensions + monkeypatch.setattr(torch.version, "hip", None) monkeypatch.setattr(torch.version, "cuda", "12.9") monkeypatch.setenv("CUTLASS_ROOT", str(tmp_path)) diff --git a/tests/ops/gemm/test_det_gemm.py b/tests/ops/gemm/test_det_gemm.py index 5b6ad2aed..532491338 100644 --- a/tests/ops/gemm/test_det_gemm.py +++ b/tests/ops/gemm/test_det_gemm.py @@ -28,7 +28,11 @@ _copy_tree_root_kernel, _det_gemm_tree_leaf_kernel, _det_gemm_tree_reduce_kernel, - _det_gemm_tree_reduce_to_output_kernel, + ) + from rl_engine.backends.shared.triton.gemm.det_gemm import ( + _det_gemm_tree_reduce_to_output_rocm_kernel as _det_gemm_tree_reduce_to_output_kernel, + ) + from rl_engine.backends.shared.triton.gemm.det_gemm import ( _device_tree_plan, _gfx942_qwen_tree_leaf_config, _triton_tree_gemm, @@ -147,7 +151,9 @@ def _special_bf16(shape: tuple[int, ...], *, offset: int = 0) -> torch.Tensor: for size in shape: elements *= size indices = (torch.arange(elements, dtype=torch.int64) + offset) % bits.numel() - return bits[indices].view(torch.bfloat16).reshape(shape).to(DEV) + # UInt16 advanced indexing is unavailable on some supported PyTorch builds. + # A signed view keeps every BF16 bit pattern unchanged while indexing. + return bits.view(torch.int16)[indices].view(torch.bfloat16).reshape(shape).to(DEV) _K_TREE_LEAF = 32 diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 5e831b240..8748221ae 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -289,6 +289,7 @@ def test_real_weights_patch_path(backend): pytest.skip("set H3_PROJECTION_WEIGHTS to pinned H3 safetensors fixture") from safetensors import safe_open from safetensors.torch import load_file + from rl_engine.reference.packing.video_patch_pack_unpack import NativeVideoPatchPackUnpackOp with safe_open(path, framework="pt") as fixture: @@ -324,7 +325,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v1" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v2" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op("triton")(x, weight, bias)) @@ -339,6 +340,60 @@ def test_reused_fp32_gemm(): _assert_accuracy(output, x.cpu() @ weight.cpu().t(), "forward_accuracy") +def _fp32_tree_reference(a, b): + # Evaluate the existing plan with independent CPU FP32 scalar arithmetic. + from rl_engine.backends.shared.triton.gemm.det_gemm import _build_tree_plan + + a, b = a.cpu(), b.cpu() + plan = _build_tree_plan(a.shape[1]) + nodes = {} + for start, length, node in zip( + plan.leaf_starts, plan.leaf_lengths, plan.leaf_nodes, strict=True + ): + value = torch.zeros(a.shape[0], b.shape[1]) + for k in range(start, start + length): + value = value + a[:, k, None] * b[k, None, :] + nodes[node] = value + for level in plan.reduction_levels: + for lower, upper, output in level: + nodes[output] = nodes[lower] + nodes[upper] + return nodes[plan.root] + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("k", [31, 33, 129, 5376]) +def test_fp32_tree_preserves_fp32_nodes(k): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_tree_gemm + + torch.manual_seed(420) + a, b = torch.randn(3, k, device="cuda"), torch.randn(k, 5, device="cuda") + output = _triton_tree_gemm(a, b, output_dtype=torch.float32) + assert output.dtype == torch.float32 + _assert_bytes(output.cpu(), _fp32_tree_reference(a, b)) + with torch.inference_mode(): + inference = _triton_tree_gemm(a, b, output_dtype=torch.float32) + _assert_bytes(inference, output) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +def test_fp32_tree_workspace_chunking(backend, monkeypatch): + from rl_engine.backends.shared.triton.gemm import det_gemm + from rl_engine.ops.gemm.canonical_linear import _gemm_fp32 + + _op(backend) # Apply the same optional-extension eligibility as other tests. + family = "triton_tree" if backend == "triton" else "cutlass" + a, b = torch.randn(7, 129, device="cuda"), torch.randn(129, 5, device="cuda") + expected = _gemm_fp32(a, b, family) + plan = det_gemm._build_tree_plan(129) + monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.node_count * 5 * 2) + actual = _gemm_fp32(a, b, family) + _assert_bytes(actual, expected) + monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.node_count * 2) + # Exercise output-column chunking when even one whole row exceeds budget. + _assert_bytes(_gemm_fp32(a, b, family), expected) + + @pytest.mark.skipif(not H20, reason="requires CUDA SM90") @pytest.mark.parametrize("backend", BACKENDS) def test_standard_normal_long_reduction(backend): diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 7bdcd03f5..3b9432659 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,7 +13,7 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v1" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v2" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) From dbfb1a214db67f730e297f9d0550351cb47185d8 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Wed, 14 Oct 2026 23:34:51 +0800 Subject: [PATCH 03/26] Optimize canonical H3 FP32 projection reductions and packing Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 17 ++- .../video_input_projection_fp32_cutlass.cpp | 1 + .../video_input_projection_fp32_cutlass.cu | 74 ++++++++++-- docs/operators/video-input-projection-fp32.md | 25 ++-- .../cuda/gemm/video_input_projection_fp32.py | 4 +- .../backends/shared/triton/gemm/det_gemm.py | 113 ++++++++++++++---- .../gemm/video_input_projection_fp32.py | 10 +- rl_engine/ops/autograd/canonical_backward.py | 83 +++++++++---- rl_engine/ops/gemm/canonical_linear.py | 9 +- .../ops/gemm/video_input_projection_fp32.py | 8 +- rl_engine/runtime/registry.py | 10 +- .../gemm/test_video_input_projection_fp32.py | 108 ++++++++++++++++- tests/runtime/test_kernel_registry.py | 2 +- 13 files changed, 378 insertions(+), 86 deletions(-) diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index aa144868d..e7ea55047 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -39,7 +39,9 @@ def _candidate(name): def _tool_version(command): try: - return subprocess.check_output(command, text=True, timeout=10).strip() + return subprocess.check_output( + command, text=True, stderr=subprocess.DEVNULL, timeout=10 + ).strip() except (OSError, subprocess.SubprocessError): return None @@ -166,6 +168,13 @@ def main(): "csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu", "build_tools/extensions.py", ] + # Devbox worker sync excludes .git; source hashes also identify exported workspaces. + git_commit = _tool_version(["git", "-C", str(root), "rev-parse", "HEAD"]) + git_status = ( + _tool_version(["git", "-C", str(root), "status", "--porcelain"]) + if git_commit is not None + else None + ) report = { "scope": "single H3 FP32 input projection; eager training and no-grad forward graph", "environment": { @@ -181,10 +190,8 @@ def main(): [shutil.which("nvcc") or "/usr/local/cuda/bin/nvcc", "--version"] ), }, - "git_commit": subprocess.check_output( - ["git", "rev-parse", "HEAD"], cwd=root, text=True - ).strip(), - "git_dirty": bool(subprocess.check_output(["git", "status", "--porcelain"], cwd=root)), + "git_commit": git_commit, + "git_dirty": bool(git_status) if git_status is not None else None, "source_sha256": {p: hashlib.sha256((root / p).read_bytes()).hexdigest() for p in sources}, "weight_fixture_sha256": fixture_hash, "weight_fixture_metadata": fixture_metadata, diff --git a/csrc/bindings/video_input_projection_fp32_cutlass.cpp b/csrc/bindings/video_input_projection_fp32_cutlass.cpp index 60e8d5ed9..0afc9ca60 100644 --- a/csrc/bindings/video_input_projection_fp32_cutlass.cpp +++ b/csrc/bindings/video_input_projection_fp32_cutlass.cpp @@ -27,6 +27,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { info["ftz"] = false; info["tree_leaf_size"] = 32; info["tree_leaf_launcher"] = "cutlass_device_gemm_array"; + info["tree_operand_packing"] = "coalesced_transpose32"; return info; }); } diff --git a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu index d26153e1a..dfb986614 100644 --- a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu +++ b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu @@ -66,6 +66,66 @@ __global__ void pack_tree_operand( } } +// The transposed G/X views have contiguous rows and strided K. Read adjacent +// rows together, then transpose in shared memory for the packed row-major leaf. +template +__global__ void pack_tree_operand_transposed( + const float* input, float* packed, int64_t* pointers, + const int32_t* starts, const int32_t* lengths, const int32_t* nodes, + int64_t rows, int64_t stride_k, float* workspace, + int64_t* output_pointers, int64_t output_elements) { + __shared__ float tile[kTreeLeaf][kTreeLeaf + 1]; + const int leaf = blockIdx.y; + const int64_t row_start = static_cast(blockIdx.x) * kTreeLeaf; + const int x = threadIdx.x, y = threadIdx.y; + const int length = lengths[leaf]; + const int64_t start = starts[leaf]; + for (int offset = 0; offset < kTreeLeaf; offset += 8) { + const int k = y + offset; + tile[k][x] = row_start + x < rows && k < length + ? input[row_start + x + (start + k) * stride_k] : 0.0f; + } + __syncthreads(); + for (int offset = 0; offset < kTreeLeaf; offset += 8) { + const int64_t row = row_start + y + offset; + if (row < rows) { + packed[(static_cast(leaf) * rows + row) * kTreeLeaf + x] + = tile[x][y + offset]; + } + } + if (blockIdx.x == 0 && x == 0 && y == 0) { + pointers[leaf] = reinterpret_cast( + packed + static_cast(leaf) * rows * kTreeLeaf); + if constexpr (WriteOutputPointers) { + output_pointers[leaf] = reinterpret_cast( + workspace + static_cast(nodes[leaf]) * output_elements); + } + } +} + +template +void launch_tree_packing( + const at::Tensor& input, const at::Tensor& packed, int64_t* pointers, + const at::Tensor& starts, const at::Tensor& lengths, const at::Tensor& nodes, + float* workspace, int64_t* output_pointers, int64_t output_elements, + cudaStream_t stream) { + if (input.stride(0) == 1 && input.stride(1) > 1) { + const dim3 grid((input.size(0) + kTreeLeaf - 1) / kTreeLeaf, starts.numel()); + pack_tree_operand_transposed<<>>( + input.data_ptr(), packed.data_ptr(), pointers, + starts.data_ptr(), lengths.data_ptr(), nodes.data_ptr(), + input.size(0), input.stride(1), workspace, output_pointers, output_elements); + } else { + constexpr int threads = 256; + pack_tree_operand + <<<(packed.numel() + threads - 1) / threads, threads, 0, stream>>>( + input.data_ptr(), packed.data_ptr(), pointers, + starts.data_ptr(), lengths.data_ptr(), nodes.data_ptr(), + input.size(0), input.stride(0), input.stride(1), packed.numel(), + workspace, output_pointers, output_elements); + } +} + void cutlass_gemm_fp32_tree_leaves_tn( at::Tensor a, at::Tensor bt, at::Tensor workspace, at::Tensor starts, at::Tensor lengths, at::Tensor nodes) { @@ -101,16 +161,10 @@ void cutlass_gemm_fp32_tree_leaves_tn( auto* ptr_b = ptr_a + leaves; auto* ptr_output = ptr_b + leaves; const auto stream = c10::cuda::getCurrentCUDAStream(a.get_device()).stream(); - constexpr int threads = 256; - pack_tree_operand<<<(packed_a.numel() + threads - 1) / threads, threads, 0, stream>>>( - a.data_ptr(), packed_a.data_ptr(), ptr_a, - starts.data_ptr(), lengths.data_ptr(), nodes.data_ptr(), - m, a.stride(0), a.stride(1), packed_a.numel(), nullptr, nullptr, m * n); - pack_tree_operand<<<(packed_b.numel() + threads - 1) / threads, threads, 0, stream>>>( - bt.data_ptr(), packed_b.data_ptr(), ptr_b, - starts.data_ptr(), lengths.data_ptr(), nodes.data_ptr(), - n, bt.stride(0), bt.stride(1), packed_b.numel(), workspace.data_ptr(), - ptr_output, m * n); + launch_tree_packing(a, packed_a, ptr_a, starts, lengths, nodes, + nullptr, nullptr, m * n, stream); + launch_tree_packing(bt, packed_b, ptr_b, starts, lengths, nodes, + workspace.data_ptr(), ptr_output, m * n, stream); LeafGemm::Arguments args( {static_cast(m), static_cast(n), kTreeLeaf}, reinterpret_cast(ptr_a), cutlass::layout::RowMajor(kTreeLeaf), diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index d67511b20..d7521ef22 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -87,23 +87,34 @@ stays FP32, including normalization of non-contiguous autograd inputs: The existing canonical session collects all contributions, removes padding, and sorts keys before the final parameter reduction. It returns each parameter VJP -once. Ordinary chunk-by-chunk accumulation is a different arithmetic contract. +once. Weight and bias share sorting and the sorted upstream gradient when their +complete contribution sequences match. Calls with different contribution sets +normalize each parameter independently. Short reductions up to the 96-wide video input use the existing direct FP32 GEMM. Long reductions reuse the existing midpoint plan with at most 32 terms per -leaf. Workspace, merge nodes, and output stay FP32. The Triton backend launches -all leaves in one grid. CUTLASS packs the same leaf ranges into 32-wide buffers, +leaf. Workspace, merge nodes, and output stay FP32. One-row complete trees fuse each +Triton leaf pair with its first merge. Other complete trees compute leaves in +one grid. Subsequent +merges use the existing two-level kernel and alternating compact workspace +buffers. Ragged trees retain the original per-level schedule. CUTLASS packs +the same leaf ranges into 32-wide buffers, pads inactive slots with zero, and uses one `device::GemmArray` launch to write -the plan's leaf nodes. Both then reuse the existing GPU level reductions. +the plan's leaf nodes. Its transpose-view packer reads adjacent rows together +and transposes a 32x32 tile in shared memory. It uses the same paired level +reductions and compact buffers for complete trees. The tree depends on logical K and retains its topology across physical row changes. Output-row chunking keeps tree workspace below the existing element limit (512 MiB for FP32). If one output row exceeds that limit, independent output columns are chunked with the same K tree. CUTLASS also allocates packed -operands and pointer arrays. Parameter sorting and public validation are included in complete-call -measurements. Existing BF16 tree callers retain their default dtype and schedule. +operands and pointer arrays. Parameter sorting and public validation are included +in complete-call measurements. Existing BF16 tree callers retain their default +dtype and schedule. -Triton fixes the 64x64x32 leaf schedule and disables TF32 and FP fusion. CUTLASS +Triton uses fixed 64x32x32 leaves and 1x128x32 leaves for one-row tree calls, +including bias VJPs. The short forward GEMM retains 64x64x32 tiles. TF32 and FP multiply-add fusion are disabled. Schedule +fusion preserves every FP32 leaf and tree addition boundary. CUTLASS fixes 64x64x8 FP32 SIMT tiles and a single output owner per leaf, without parallel split-K. Both leaf algorithms have separately versioned backend identities. diff --git a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py index ef921f782..fd2f6b53c 100644 --- a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py @@ -65,6 +65,7 @@ def _load_extension(): or info["ftz"] or info["split_k"] != 1 or info.get("tree_leaf_size") != _BLOCK_K + or info.get("tree_operand_packing") != "coalesced_transpose32" ): raise RuntimeError("CUTLASS extension does not match the strict FP32 build contract") return extension @@ -91,7 +92,7 @@ def cutlass_gemm_fp32(a, b): class CUDAVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "cutlass" - backend_id = "cuda-cutlass-video-input-projection-fp32-v2" + backend_id = "cuda-cutlass-video-input-projection-fp32-v3" def __init__(self): if not _TRITON_AVAILABLE: @@ -116,5 +117,6 @@ def backend_info(self): "reduction_backend": "triton", "triton_version": triton.__version__, "leaf_launch": "cutlass_device_gemm_array", + "tree_schedule": "paired_levels_compact_workspace", **_load_extension().build_info(), } diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 5f7691189..c2dc6bdd4 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -41,7 +41,7 @@ # one shot would otherwise create a multi-terabyte virtual tensor on ROCm. # Chunking over M preserves the exact K-tree and BF16 rounding contract. _MAX_TREE_WORKSPACE_ELEMENTS = 128 * 1024 * 1024 -_TREE_PLANS: dict[tuple[int, int], "_DeviceTreePlan"] = {} +_TREE_PLANS: dict[tuple[int, int, bool | None], "_DeviceTreePlan"] = {} _TREE_PLAN_LOCK = threading.Lock() _ROCM_TUNE_DECODE_LEAF = True _ROCM_FUSE_LEAF_REDUCTION = True @@ -194,6 +194,8 @@ class _TreePlan: @dataclass(frozen=True) class _DeviceTreePlan: host: _TreePlan + workspace_node_count: int + fp32_fused: bool leaf_starts: torch.Tensor leaf_lengths: torch.Tensor leaf_nodes: torch.Tensor @@ -259,9 +261,11 @@ def visit(begin: int, end: int) -> tuple[int, int]: ) -def _device_tree_plan(k_size: int, device: torch.device) -> _DeviceTreePlan: +def _device_tree_plan( + k_size: int, device: torch.device, *, fp32_fuse_leaves: bool | None = None +) -> _DeviceTreePlan: device_index = device.index if device.index is not None else torch.cuda.current_device() - key = (device_index, k_size) + key = (device_index, k_size, fp32_fuse_leaves) with _TREE_PLAN_LOCK: cached = _TREE_PLANS.get(key) if cached is not None: @@ -276,10 +280,43 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: for operations in host.reduction_levels: lower_nodes, upper_nodes, output_nodes = zip(*operations, strict=True) levels.append((indices(lower_nodes), indices(upper_nodes), indices(output_nodes))) + # The paired kernels require complete subtrees. Ragged midpoint trees + # retain the original per-level schedule and workspace indices. + leaf_count = len(host.leaf_nodes) + complete = bool(host.reduction_levels) and all( + len(level) * (1 << (height + 1)) == leaf_count + for height, level in enumerate(host.reduction_levels) + ) + fp32_fused = fp32_fuse_leaves is not None and complete + node_map = {node: node for node in range(host.node_count)} + workspace_node_count = host.node_count + if fp32_fused: + # Alternate disjoint buffers for the retained levels. Each launch + # reads its previous buffer completely before the next launch. + first_height = 0 if fp32_fuse_leaves else -1 + first_nodes = ( + tuple(op[2] for op in host.reduction_levels[0]) + if fp32_fuse_leaves + else host.leaf_nodes + ) + next_count = max(1, len(first_nodes) // 4) + workspace_node_count = len(first_nodes) + next_count + for i, node in enumerate(first_nodes): + node_map[node] = i + for step, height in enumerate( + range(first_height + 2, len(host.reduction_levels), 2), start=1 + ): + base = len(first_nodes) if step % 2 else 0 + for i, (_, _, node) in enumerate(host.reduction_levels[height]): + node_map[node] = base + i + + def mapped_indices(values): + return indices(tuple(node_map[node] for node in values)) + fused_pairs = [] leaf_reduction = None fused_pairs_after_leaf = [] - if _device_arch(device_index) == "gfx942": + if _device_arch(device_index) == "gfx942" or fp32_fused: def build_fused_pairs(start_level: int): result = [] @@ -296,21 +333,23 @@ def build_fused_pairs(start_level: int): node0, node1, node2, node3 = zip(*grandchildren, strict=True) result.append( ( - indices(node0), - indices(node1), - indices(node2), - indices(node3), - indices(tuple(outputs)), + mapped_indices(node0), + mapped_indices(node1), + mapped_indices(node2), + mapped_indices(node3), + mapped_indices(tuple(outputs)), ) ) return result - fused_pairs = build_fused_pairs(0) - fused_pairs_after_leaf = build_fused_pairs(1) + if not fp32_fused or not fp32_fuse_leaves: + fused_pairs = build_fused_pairs(0) + if not fp32_fused or fp32_fuse_leaves: + fused_pairs_after_leaf = build_fused_pairs(1) if host.reduction_levels: first_level = host.reduction_levels[0] paired_nodes = {node for op in first_level for node in op[:2]} - if paired_nodes == set(host.leaf_nodes): + if paired_nodes == set(host.leaf_nodes) and (not fp32_fused or fp32_fuse_leaves): leaf_metadata = { node: (start, length) for start, length, node in zip( @@ -335,14 +374,23 @@ def build_fused_pairs(start_level: int): indices(tuple(lower_lengths)), indices(tuple(upper_starts)), indices(tuple(upper_lengths)), - indices(tuple(outputs)), + mapped_indices(tuple(outputs)), ) result = _DeviceTreePlan( host=host, + workspace_node_count=workspace_node_count, + fp32_fused=fp32_fused, leaf_starts=indices(host.leaf_starts), leaf_lengths=indices(host.leaf_lengths), - leaf_nodes=indices(host.leaf_nodes), - reduction_levels=tuple(levels), + leaf_nodes=mapped_indices(host.leaf_nodes), + reduction_levels=( + tuple( + tuple(mapped_indices(tuple(op[column] for op in level)) for column in range(3)) + for level in host.reduction_levels + ) + if fp32_fused + else tuple(levels) + ), rocm_fused_reduction_pairs=tuple(fused_pairs), rocm_leaf_reduction=leaf_reduction, rocm_fused_reduction_pairs_after_leaf=tuple(fused_pairs_after_leaf), @@ -804,11 +852,19 @@ def _triton_tree_gemm( b = b.contiguous() m_size, k_size = a.shape n_size = b.size(1) - plan = _device_tree_plan(k_size, a.device) - workspace_elements = plan.host.node_count * m_size * n_size + plan = _device_tree_plan( + k_size, + a.device, + fp32_fuse_leaves=( + (fp32_leaf_launcher is None and m_size == 1) + if fp32_tree and not transpose_output + else None + ), + ) + workspace_elements = plan.workspace_node_count * m_size * n_size if workspace_elements > _MAX_TREE_WORKSPACE_ELEMENTS and out is None: - if fp32_tree and plan.host.node_count * n_size > _MAX_TREE_WORKSPACE_ELEMENTS: - columns_per_chunk = _MAX_TREE_WORKSPACE_ELEMENTS // plan.host.node_count + if fp32_tree and plan.workspace_node_count * n_size > _MAX_TREE_WORKSPACE_ELEMENTS: + columns_per_chunk = _MAX_TREE_WORKSPACE_ELEMENTS // plan.workspace_node_count if columns_per_chunk < 1: raise RuntimeError("FP32 tree plan exceeds the workspace budget for one element") # A one-row bias VJP cannot be split further over M. Splitting @@ -828,7 +884,7 @@ def _triton_tree_gemm( return torch.cat(chunks, dim=0 if transpose_output else 1) rows_per_chunk = max( 1, - _MAX_TREE_WORKSPACE_ELEMENTS // (plan.host.node_count * n_size), + _MAX_TREE_WORKSPACE_ELEMENTS // (plan.workspace_node_count * n_size), ) chunks = [] for start in range(0, m_size, rows_per_chunk): @@ -863,7 +919,7 @@ def _triton_tree_gemm( raise ValueError("Triton tree GEMM output buffer must not require gradients") result = out workspace = torch.empty( - (plan.host.node_count, m_size, n_size), + (plan.workspace_node_count, m_size, n_size), dtype=output_dtype, device=a.device, ) @@ -878,6 +934,12 @@ def _triton_tree_gemm( transpose_output=transpose_output, preserve_a_strides=preserve_a_strides, ) + if fp32_tree and fp32_leaf_launcher is None: + # Offline H20 sweep: 64x32 avoids both masked N=96 columns and the + # register spills of paired 64x64 leaves. One-row calls retain fusion. + leaf_config = ( + _TreeLeafConfig(1, 128, 4, True) if m_size == 1 else _TreeLeafConfig(64, 32, 4, False) + ) if ( _ROCM_TUNE_DECODE_LEAF and is_gfx942 @@ -895,10 +957,11 @@ def _triton_tree_gemm( ) direct_root_output = not transpose_output and (is_gfx942 or fp32_tree) fuse_leaf_reduction = ( - _ROCM_FUSE_LEAF_REDUCTION + (plan.fp32_fused or (_ROCM_FUSE_LEAF_REDUCTION and is_gfx942)) and direct_root_output - and use_inference_schedule - and k_size >= 1536 + and (use_inference_schedule or plan.fp32_fused) + and (k_size >= 1536 or plan.fp32_fused) + and (fp32_leaf_launcher is None) and plan.rocm_leaf_reduction is not None ) if fp32_leaf_launcher is not None: @@ -981,7 +1044,7 @@ def _triton_tree_gemm( N=n_size, BLOCK=reduction_block, ) - elif direct_root_output and use_inference_schedule and is_gfx942: + elif direct_root_output and ((use_inference_schedule and is_gfx942) or plan.fp32_fused): blocks = triton.cdiv(m_size * n_size, reduction_block) for pair_index, nodes in enumerate(plan.rocm_fused_reduction_pairs): second_level = pair_index * 2 + 1 diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index 95d711e27..66c5d9d47 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -25,7 +25,7 @@ def projection_gemm_fp32(a, b): class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton_tree" - backend_id = "triton-video-input-projection-fp32-v2" + backend_id = "triton-video-input-projection-fp32-v3" def __init__(self): if not _TRITON_AVAILABLE: @@ -43,7 +43,8 @@ def backend_info(self): return { "backend_id": self.backend_id, "family": "triton", - "kernel": "_det_gemm_fp32_kernel / _det_gemm_tree_leaf_kernel", + "kernel": "_det_gemm_fp32_kernel / _det_gemm_tree_leaf_kernel / " + "_det_gemm_tree_leaf_reduce_rocm_kernel", "triton_version": triton.__version__, "tile": [64, 64, 32], "tf32": False, @@ -52,5 +53,8 @@ def backend_info(self): "reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, "reduction_backend": "triton", - "leaf_launch": "all_leaves_in_one_grid", + "leaf_launch": "paired_leaves_for_one_row_complete_trees", + "tree_tile": [64, 32, 32], + "bias_tree_tile": [1, 128, 32], + "tree_schedule": "paired_levels_compact_workspace", } diff --git a/rl_engine/ops/autograd/canonical_backward.py b/rl_engine/ops/autograd/canonical_backward.py index 75ddd0269..e295b2eb9 100644 --- a/rl_engine/ops/autograd/canonical_backward.py +++ b/rl_engine/ops/autograd/canonical_backward.py @@ -5,6 +5,7 @@ from __future__ import annotations +from collections import Counter from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass, field @@ -18,6 +19,7 @@ @dataclass class _Use: keys: torch.Tensor + contribution: object | None = None rows: torch.Tensor | None = None grads: torch.Tensor | None = None @@ -29,16 +31,59 @@ class CanonicalBackwardSession: uses: dict[str, list[_Use]] = field(default_factory=dict) received: dict[str, int] = field(default_factory=dict) backward_started: bool = False - - def register(self, parameter_id: str, keys: torch.Tensor) -> int: + _normalized: dict[tuple[object, ...], tuple[torch.Tensor, torch.Tensor]] = field( + default_factory=dict, repr=False + ) + _shareable: set[tuple[object, ...]] | None = field(default=None, repr=False) + + def register( + self, parameter_id: str, keys: torch.Tensor, *, contribution: object | None = None + ) -> int: if self.backward_started: raise RuntimeError("cannot register canonical rows after backward started") if keys.dim() != 2 or keys.shape[1] not in (2, 3): raise ValueError("logical keys must have shape [rows, 2] or [rows, 3]") slot = len(self.uses.setdefault(parameter_id, [])) - self.uses[parameter_id].append(_Use(keys=keys.detach())) + self.uses[parameter_id].append(_Use(keys=keys.detach(), contribution=contribution)) return slot + @staticmethod + def _concatenate(values: list[torch.Tensor]) -> torch.Tensor: + return values[0] if len(values) == 1 else torch.cat(values, dim=0) + + def _normalize( + self, parameter_id: str, values: list[torch.Tensor] + ) -> tuple[torch.Tensor, torch.Tensor]: + entries = self.uses[parameter_id] + # A token names one forward's shared weight/bias contribution. Matching + # complete token sequences prove that both parameters receive the same + # keys and upstream gradients, without a device equality check. + signature = tuple(item.contribution for item in entries) + if self._shareable is None: + # Forward registration is complete before backward. Build this once + # so a model's parameters do not each scan the complete session. + signatures = Counter( + tuple(item.contribution for item in other) + for other in self.uses.values() + if all(item.contribution is not None for item in other) + ) + self._shareable = {key for key, count in signatures.items() if count == 2} + shared = signature in self._shareable + if shared and signature in self._normalized: + return self._normalized.pop(signature) + keys = self._concatenate([item.keys for item in entries]) + valid = torch.nonzero(keys[:, 0] >= 0, as_tuple=True)[0] + active_keys = keys.index_select(0, valid) + order = torch.arange(valid.numel(), device=keys.device) + for column in range(keys.shape[1] - 1, -1, -1): + column_values = active_keys.index_select(0, order)[:, column] + order = order.index_select(0, torch.argsort(column_values, stable=True)) + order = valid.index_select(0, order) + normalized = order, self._concatenate(values).index_select(0, order) + if shared: + self._normalized[signature] = normalized + return normalized + def submit_linear( self, parameter_id: str, @@ -60,20 +105,13 @@ def submit_linear( self.received[parameter_id] = count if count != len(entries): return None - all_keys = torch.cat([item.keys for item in entries], dim=0) - all_rows = torch.cat([item.rows for item in entries if item.rows is not None], dim=0) - all_grads = torch.cat([item.grads for item in entries if item.grads is not None], dim=0) - valid = all_keys[:, 0] >= 0 - all_keys = all_keys[valid] - all_rows = all_rows[valid] - all_grads = all_grads[valid] - if all_keys.shape[0] == 0: + order, grads = self._normalize( + parameter_id, [item.grads for item in entries if item.grads is not None] + ) + if order.numel() == 0: raise RuntimeError(f"no active logical rows for {parameter_id}") - order = torch.arange(all_keys.shape[0], device=all_keys.device) - for column in range(all_keys.shape[1] - 1, -1, -1): - values = all_keys.index_select(0, order)[:, column] - order = order.index_select(0, torch.argsort(values, stable=True)) - return reducer(all_rows.index_select(0, order), all_grads.index_select(0, order)) + rows = self._concatenate([item.rows for item in entries if item.rows is not None]) + return reducer(rows.index_select(0, order), grads) def submit_rows( self, @@ -94,15 +132,10 @@ def submit_rows( self.received[parameter_id] = count if count != len(entries): return None - all_keys = torch.cat([item.keys for item in entries], dim=0) - all_rows = torch.cat([item.rows for item in entries if item.rows is not None], dim=0) - valid = all_keys[:, 0] >= 0 - all_keys, all_rows = all_keys[valid], all_rows[valid] - order = torch.arange(all_keys.shape[0], device=all_keys.device) - for column in range(all_keys.shape[1] - 1, -1, -1): - values = all_keys.index_select(0, order)[:, column] - order = order.index_select(0, torch.argsort(values, stable=True)) - return reducer(all_rows.index_select(0, order)) + _, rows = self._normalize( + parameter_id, [item.rows for item in entries if item.rows is not None] + ) + return reducer(rows) def validate_complete(self) -> None: missing = { diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 3e572ea0b..59c66f876 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -65,11 +65,14 @@ def forward( ctx.family = str(family) ctx.bias_parameter_id = str(bias_parameter_id) ctx.bias_dtype = bias.dtype if bias is not None else None + contribution = object() if ctx.needs_input_grad[1] and ctx.needs_input_grad[5] else None ctx.slot = ( - session.register(ctx.parameter_id, logical_keys) if ctx.needs_input_grad[1] else None + session.register(ctx.parameter_id, logical_keys, contribution=contribution) + if ctx.needs_input_grad[1] + else None ) ctx.bias_slot = ( - session.register(ctx.bias_parameter_id, logical_keys) + session.register(ctx.bias_parameter_id, logical_keys, contribution=contribution) if ctx.needs_input_grad[5] else None ) @@ -97,7 +100,7 @@ def reducer(rows, grads): lhs = lhs.contiguous() return _gemm_fp32(lhs, rows.float(), ctx.family).to(weight.dtype) - dweight = ctx.session.submit_linear(ctx.parameter_id, ctx.slot, a, grad_out, reducer) + dweight = ctx.session.submit_linear(ctx.parameter_id, ctx.slot, a, grad32, reducer) dbias = None if ctx.needs_input_grad[5]: diff --git a/rl_engine/ops/gemm/video_input_projection_fp32.py b/rl_engine/ops/gemm/video_input_projection_fp32.py index c0d4d0389..e31f76c8b 100644 --- a/rl_engine/ops/gemm/video_input_projection_fp32.py +++ b/rl_engine/ops/gemm/video_input_projection_fp32.py @@ -46,6 +46,7 @@ def _validate_keys(keys, rows, device, parameter_id, session, weight, bias): if not isinstance(parameter_id, str) or not parameter_id.strip(): raise ValueError("parameter gradients require a non-empty parameter_id") # A duplicate active identity leaves tie order dependent on chunk arrival. + checked = set() for name, needed in ( ("weight", weight.requires_grad), ("bias", bias is not None and bias.requires_grad), @@ -53,9 +54,14 @@ def _validate_keys(keys, rows, device, parameter_id, session, weight, bias): if not needed: continue previous = session.uses.get(f"{parameter_id}.{name}", []) + signature = tuple(use.contribution for use in previous) + if all(token is not None for token in signature): + if signature in checked: + continue + checked.add(signature) if any(use.keys.shape[1] != keys.shape[1] for use in previous): raise ValueError("logical key columns must agree across parameter uses") - merged = torch.cat([use.keys for use in previous] + [keys], dim=0) + merged = torch.cat([use.keys for use in previous] + [keys], dim=0) if previous else keys active = merged[merged[:, 0] >= 0] if torch.unique(active, dim=0).shape[0] != active.shape[0]: raise ValueError("active logical keys must be unique across parameter uses") diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 02d7cade2..43c231b70 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -244,7 +244,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="cuda-cutlass-video-input-projection-fp32-v2", + backend_id="cuda-cutlass-video-input-projection-fp32-v3", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -273,12 +273,12 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: implementation_class_or_factory=OpBackend.CUDA_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( - "cutlass-56ca28a-fp32-simt-64x64x8-triton-gpu-tree32-h3-v2" + "cutlass-56ca28a-fp32-simt-64x64x8-triton-fused-tree32-coalesced-pack-h3-v3" ), ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v2", + backend_id="triton-video-input-projection-fp32-v3", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -305,7 +305,9 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: lifecycle=OperatorLifecycle.REQUEST, implementation_class_or_factory=OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, - version_or_build_fingerprint="rlkernel-triton-fp32-64x64x32-no-fusion-gpu-tree32-h3-v2", + version_or_build_fingerprint=( + "rlkernel-triton-fp32-no-fma-paired-tree32-leaf64x32-bias1x128-h3-v3" + ), ), OperatorBackendDescriptor( semantic_op="video_patch_pack_unpack", diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 8748221ae..11bc57088 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -325,7 +325,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v2" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v3" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op("triton")(x, weight, bias)) @@ -406,3 +406,109 @@ def test_standard_normal_long_reduction(backend): gold.backward(data[3].cpu()) for i, (a, b) in enumerate(zip(actual, (gold, cx.grad, cw.grad, cb.grad), strict=True)): _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") + + +@pytest.mark.parametrize("bias_uses", ["all", "partial", "distinct"]) +def test_canonical_bias_sharing_requires_matching_contributions(bias_uses, monkeypatch): + from rl_engine.ops.autograd.canonical_backward import CanonicalBackwardSession + + torch.manual_seed(421) + x = torch.randn(7, 3, requires_grad=True) + weight = torch.randn(5, 3, requires_grad=True) + biases = [torch.randn(5, requires_grad=True) for _ in range(2)] + grad = torch.randn(7, 5) + indices = [torch.tensor([4, 1, 6]), torch.tensor([0, 5, 2, 3])] + keys = _keys(7, "cpu") + keys[2, 0] = -1 + normalized = [] + original = CanonicalBackwardSession._normalize + + def observe(self, parameter_id, values): + result = original(self, parameter_id, values) + normalized.append((parameter_id, result[1])) + return result + + monkeypatch.setattr(CanonicalBackwardSession, "_normalize", observe) + with canonical_backward_session() as session: + outputs = [] + for use, ix in enumerate(indices): + bias = ( + None + if bias_uses == "partial" and use == 1 + else biases[use if bias_uses == "distinct" else 0] + ) + outputs.append( + canonical_linear_fp32( + x[ix], + weight, + keys[ix], + parameter_id="weight", + family="pytorch", + bias=bias, + bias_parameter_id=f"bias.{use if bias_uses == 'distinct' else 0}", + ) + ) + torch.autograd.backward(outputs, [grad[ix] for ix in indices]) + session.validate_complete() + assert not session._normalized + active = keys[:, 0] >= 0 + torch.testing.assert_close(weight.grad, grad[active].t() @ x.detach()[active]) + torch.testing.assert_close(x.grad, grad @ weight.detach()) + for use, bias in enumerate(biases): + if use == 1 and bias_uses != "distinct": + assert bias.grad is None + continue + ix = torch.arange(7) if bias_uses == "all" else indices[use] + ix = ix[active[ix]] + torch.testing.assert_close(bias.grad, grad[ix].sum(0)) + if bias_uses == "all": + assert len(normalized) == 2 + assert normalized[0][1] is normalized[1][1] + else: + assert all(a[1] is not b[1] for i, a in enumerate(normalized) for b in normalized[i + 1 :]) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("k", [64, 65, 128, 1024, 13104]) +def test_fp32_fused_tree_and_transposed_packing_match_scalar_reference(backend, k): + from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import _cutlass_tree_leaves + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_tree_gemm + + _op(backend) + torch.manual_seed(422) + # Both operands exercise the strided transpose packing, including tail rows. + a = torch.randn(k, 3, device="cuda").t() + b = torch.randn(k, 5, device="cuda") + + def gemm(a, b): + return _triton_tree_gemm( + a, + b, + output_dtype=torch.float32, + preserve_a_strides=True, + fp32_leaf_launcher=_cutlass_tree_leaves if backend == "cuda" else None, + ) + + result = gemm(a, b) + _assert_bytes(result.cpu(), _fp32_tree_reference(a, b)) + _assert_bytes(gemm(a[:1], b), result[:1]) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +def test_compact_fp32_tree_workspace_preserves_chunked_outputs(backend, monkeypatch): + from rl_engine.backends.shared.triton.gemm import det_gemm + from rl_engine.ops.gemm.canonical_linear import _gemm_fp32 + + _op(backend) + family = "triton_tree" if backend == "triton" else "cutlass" + a, b = torch.randn(7, 1024, device="cuda"), torch.randn(1024, 5, device="cuda") + expected = _gemm_fp32(a, b, family) + plan = det_gemm._device_tree_plan(1024, a.device, fp32_fuse_leaves=False) + assert plan.fp32_fused + assert plan.workspace_node_count < plan.host.node_count + monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.workspace_node_count * 10) + _assert_bytes(_gemm_fp32(a, b, family), expected) + monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.workspace_node_count * 2) + _assert_bytes(_gemm_fp32(a, b, family), expected) diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 3b9432659..5ae57e9d5 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,7 +13,7 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v2" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v3" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) From ac95108f436eb3ca53128bc6043439ec99e143a0 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Thu, 15 Oct 2026 11:34:01 +0800 Subject: [PATCH 04/26] Optimize H3 FP32 projection leaf loops and bias paths Signed-off-by: walker-ai <2398833647@qq.com> --- docs/operators/video-input-projection-fp32.md | 15 +++- .../cuda/gemm/video_input_projection_fp32.py | 6 +- .../backends/shared/triton/gemm/det_gemm.py | 34 +++++++-- .../gemm/video_input_projection_fp32.py | 11 ++- rl_engine/ops/gemm/canonical_linear.py | 25 +++++-- .../ops/gemm/video_input_projection_fp32.py | 4 +- rl_engine/runtime/registry.py | 14 ++-- .../gemm/test_video_input_projection_fp32.py | 74 ++++++++++++++++++- tests/runtime/test_kernel_registry.py | 2 +- 9 files changed, 152 insertions(+), 33 deletions(-) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index d7521ef22..3ffcd8cb4 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -10,7 +10,8 @@ The primary implementation reuses RL-Kernel's Triton FP32 GEMM and canonical parameter-gradient session. A pinned CUTLASS FP32 SIMT GEMM is an explicit CUDA candidate. Both use the same semantic adapter, autograd implementation, and existing Triton tree plan and GPU level reductions. CUTLASS substitutes the -batched leaf GEMM; its reported execution identity includes the Triton reduction. +batched leaf GEMM for dX/dW. Long bias VJPs reuse the one-row Triton leaves; +the reported execution identity includes both paths. ## Entry Point @@ -78,12 +79,16 @@ outside this contribution's strict support scope. ## Accuracy -Forward computes `Y = X @ W.T`, then adds bias in FP32. The upstream gradient `G` +Forward computes `Y = X @ W.T + bias` in FP32. The Triton forward adds bias +in the existing GEMM epilogue, preserving the rounded FP32 GEMM result before +the bias addition. The upstream gradient `G` stays FP32, including normalization of non-contiguous autograd inputs: - `dX = G @ W` is row-local. - `dW = G_sorted.T @ X_sorted` follows canonical logical-row order. -- `dbias = ones[1,R] @ G_sorted` uses the same fixed FP32 GEMM bottom layer. +- `dbias = ones[1,R] @ G_sorted` uses the fixed FP32 GEMM bottom layer. + Both backends reuse the one-row Triton tree for reductions longer than 96; + the CUTLASS candidate retains its direct GEMM for shorter reductions. The existing canonical session collects all contributions, removes padding, and sorts keys before the final parameter reduction. It returns each parameter VJP @@ -93,7 +98,9 @@ normalize each parameter independently. Short reductions up to the 96-wide video input use the existing direct FP32 GEMM. Long reductions reuse the existing midpoint plan with at most 32 terms per -leaf. Workspace, merge nodes, and output stay FP32. One-row complete trees fuse each +leaf. The Triton scalar loop stops at the largest logical leaf length +(21 for K=5376; 26 for K=13104), preserving the midpoint tree and ascending +products. Workspace, merge nodes, and output stay FP32. One-row complete trees fuse each Triton leaf pair with its first merge. Other complete trees compute leaves in one grid. Subsequent merges use the existing two-level kernel and alternating compact workspace diff --git a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py index fd2f6b53c..3d2799242 100644 --- a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py @@ -92,7 +92,7 @@ def cutlass_gemm_fp32(a, b): class CUDAVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "cutlass" - backend_id = "cuda-cutlass-video-input-projection-fp32-v3" + backend_id = "cuda-cutlass-video-input-projection-fp32-v4" def __init__(self): if not _TRITON_AVAILABLE: @@ -111,10 +111,12 @@ def backend_info(self): return { "backend_id": self.backend_id, "family": "cuda", - "kernel": "gemm_fp32_tn / gemm_fp32_tree_leaves_tn", + "kernel": "gemm_fp32_tn / gemm_fp32_tree_leaves_tn / " + "_det_gemm_tree_leaf_kernel / _det_gemm_tree_leaf_reduce_rocm_kernel", "reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, "reduction_backend": "triton", + "bias_reduction_backend": "triton_for_k_gt_96", "triton_version": triton.__version__, "leaf_launch": "cutlass_device_gemm_array", "tree_schedule": "paired_levels_compact_workspace", diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index c2dc6bdd4..b8f3c24aa 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -420,6 +420,7 @@ def _det_gemm_tree_leaf_kernel( BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, N_FASTEST: tl.constexpr, + LEAF_LOOP_BOUND: tl.constexpr = _BLOCK_K, ): if N_FASTEST: pid_n = tl.program_id(0) @@ -439,9 +440,10 @@ def _det_gemm_tree_leaf_kernel( leaf_node = tl.load(leaf_nodes_ptr + leaf).to(tl.int64) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) # Keep the leaf's ascending scalar FMA order identical to the native - # gfx942 correctness kernel. A tl.dot/MFMA leaf is topology-stable but + # gfx942 correctness kernel. FP32 launches trim only the masked tail. + # A tl.dot/MFMA leaf is topology-stable but # differs from the scalar reference at rare BF16 rounding boundaries. - for offset in tl.static_range(0, BLOCK_K): + for offset in tl.static_range(0, LEAF_LOOP_BOUND): k_offset = leaf_start + offset active = offset < leaf_length a = tl.load( @@ -491,6 +493,7 @@ def _det_gemm_tree_leaf_reduce_rocm_kernel( BLOCK_K: tl.constexpr, N_FASTEST: tl.constexpr, WRITE_OUTPUT: tl.constexpr, + LEAF_LOOP_BOUND: tl.constexpr = _BLOCK_K, ): if N_FASTEST: pid_n = tl.program_id(0) @@ -508,7 +511,7 @@ def _det_gemm_tree_leaf_reduce_rocm_kernel( upper_length = tl.load(upper_lengths_ptr + pair).to(tl.int64) lower_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for offset in tl.static_range(0, BLOCK_K): + for offset in tl.static_range(0, LEAF_LOOP_BOUND): k_offset = lower_start + offset active = offset < lower_length a = tl.load( @@ -526,7 +529,7 @@ def _det_gemm_tree_leaf_reduce_rocm_kernel( lower = lower_acc.to(workspace_ptr.dtype.element_ty) upper_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for offset in tl.static_range(0, BLOCK_K): + for offset in tl.static_range(0, LEAF_LOOP_BOUND): k_offset = upper_start + offset active = offset < upper_length a = tl.load( @@ -730,6 +733,7 @@ def _det_gemm_fp32_kernel( a_ptr, b_ptr, c_ptr, + bias_ptr, M, N, K, @@ -743,6 +747,7 @@ def _det_gemm_fp32_kernel( BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, PROMOTE_INPUTS: tl.constexpr, + HAS_BIAS: tl.constexpr, ): # One program = one output tile, walks the whole K in fixed order. # No split-K -> K-accumulation order independent of M -> batch-invariant. @@ -765,13 +770,18 @@ def _det_gemm_fp32_kernel( acc += tl.dot(a, b, allow_tf32=False) a_ptrs += BLOCK_K * stride_ak b_ptrs += BLOCK_K * stride_bk + if HAS_BIAS: + bias = tl.load(bias_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32) + acc = acc + bias[None, :] c = acc.to(c_ptr.dtype.element_ty) c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) tl.store(c_ptrs, c, mask=mask) -def _triton_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: +def _triton_gemm_fp32( + a: torch.Tensor, b: torch.Tensor, *, bias: torch.Tensor | None = None +) -> torch.Tensor: a, b = a.contiguous(), b.contiguous() M, K = a.shape _, N = b.shape @@ -781,6 +791,7 @@ def _triton_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: a, b, c, + bias, M, N, K, @@ -794,6 +805,7 @@ def _triton_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: BLOCK_N=_BLOCK_N, BLOCK_K=_BLOCK_K, PROMOTE_INPUTS=True, + HAS_BIAS=bias is not None, **({"enable_fp_fusion": False} if torch.version.hip is None else {}), ) return c @@ -992,7 +1004,11 @@ def _triton_tree_gemm( N_FASTEST=leaf_config.n_fastest, WRITE_OUTPUT=len(plan.host.reduction_levels) == 1, num_warps=leaf_config.num_warps, - **({"enable_fp_fusion": False} if fp32_tree else {}), + **( + {"enable_fp_fusion": False, "LEAF_LOOP_BOUND": max(plan.host.leaf_lengths)} + if fp32_tree + else {} + ), ) else: _det_gemm_kernel[leaf_grid]( @@ -1014,7 +1030,11 @@ def _triton_tree_gemm( BLOCK_K=_BLOCK_K, N_FASTEST=leaf_config.n_fastest, num_warps=leaf_config.num_warps, - **({"enable_fp_fusion": False} if fp32_tree else {}), + **( + {"enable_fp_fusion": False, "LEAF_LOOP_BOUND": max(plan.host.leaf_lengths)} + if fp32_tree + else {} + ), ) reduction_block = 256 if fuse_leaf_reduction: diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index 66c5d9d47..dee6b3935 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -17,15 +17,16 @@ ) -def projection_gemm_fp32(a, b): +def projection_gemm_fp32(a, b, *, bias=None): if a.shape[1] <= VIDEO_PATCH_WIDTH: - return _triton_gemm_fp32(a, b) - return _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) + return _triton_gemm_fp32(a, b, bias=bias) + output = _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) + return output if bias is None else output + bias class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton_tree" - backend_id = "triton-video-input-projection-fp32-v3" + backend_id = "triton-video-input-projection-fp32-v4" def __init__(self): if not _TRITON_AVAILABLE: @@ -49,6 +50,8 @@ def backend_info(self): "tile": [64, 64, 32], "tf32": False, "fp_fusion": False, + "forward_bias": "fp32_epilogue", + "leaf_loop_bound": "max_logical_leaf_length", "split_k": 1, "reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 59c66f876..8fe329dbb 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -13,13 +13,17 @@ from rl_engine.ops.autograd.canonical_backward import active_session -def _gemm_fp32(a: torch.Tensor, b: torch.Tensor, family: str) -> torch.Tensor: +def _gemm_fp32( + a: torch.Tensor, b: torch.Tensor, family: str, *, bias: torch.Tensor | None = None +) -> torch.Tensor: if family == "triton_tree": from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( projection_gemm_fp32, ) - return projection_gemm_fp32(a, b) + return projection_gemm_fp32(a, b, bias=bias) + if bias is not None: + return _gemm_fp32(a, b, family) + bias if family == "cutlass": from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import cutlass_gemm_fp32 @@ -80,8 +84,9 @@ def forward( # CUTLASS consumes column-major B directly from contiguous [N,K] weight. if ctx.family != "cutlass": rhs = rhs.contiguous() - output = _gemm_fp32(a.float(), rhs, ctx.family) - return output if bias is None else output + bias.float() + return _gemm_fp32( + a.float(), rhs, ctx.family, bias=bias.float() if bias is not None else None + ) @staticmethod def backward(ctx, grad_out): @@ -108,7 +113,15 @@ def bias_reducer(grads): if grads.shape[0] == 0: raise RuntimeError(f"no active logical rows for {ctx.bias_parameter_id}") ones = torch.ones((1, grads.shape[0]), device=grads.device, dtype=torch.float32) - return _gemm_fp32(ones, grads.float(), ctx.family).squeeze(0).to(ctx.bias_dtype) + # Reuse the existing one-row Triton leaves for long bias VJPs. + # CUTLASS retains its direct GEMM for short reductions. + family = ctx.family + if family == "cutlass": + from rl_engine.ops.gemm.video_input_projection_fp32 import VIDEO_PATCH_WIDTH + + if grads.shape[0] > VIDEO_PATCH_WIDTH: + family = "triton_tree" + return _gemm_fp32(ones, grads.float(), family).squeeze(0).to(ctx.bias_dtype) dbias = ctx.session.submit_rows( ctx.bias_parameter_id, ctx.bias_slot, grad32, bias_reducer @@ -138,7 +151,7 @@ def bias_reducer(grads): "rl_engine._C_cutlass.gemm_fp32_tree_leaves_tn / " "rl_engine.backends.shared.triton.gemm.det_gemm._triton_tree_gemm" ), - impl="cutlass_fp32_simt_triton_tree_canonical_rowfold", + impl="cutlass_fp32_simt_triton_tree_bias_canonical_rowfold", family="cuda", ) return da, dweight, None, None, None, dbias, None diff --git a/rl_engine/ops/gemm/video_input_projection_fp32.py b/rl_engine/ops/gemm/video_input_projection_fp32.py index e31f76c8b..404bd49f1 100644 --- a/rl_engine/ops/gemm/video_input_projection_fp32.py +++ b/rl_engine/ops/gemm/video_input_projection_fp32.py @@ -107,9 +107,7 @@ def __call__(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None rhs = weight.t() if self.family != "cutlass": rhs = rhs.contiguous() - output = _gemm_fp32(rows, rhs, self.family) - if bias is not None: - output = output + bias + output = _gemm_fp32(rows, rhs, self.family, bias=bias) return output.reshape(*x.shape[:-1], VIDEO_HIDDEN_WIDTH) forward = __call__ diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 43c231b70..269db8187 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -244,7 +244,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="cuda-cutlass-video-input-projection-fp32-v3", + backend_id="cuda-cutlass-video-input-projection-fp32-v4", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -257,8 +257,9 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "training": {"world_size": (1,), "sharding": ("unsharded",)}, }, determinism_or_alignment_properties={ - "algorithm": "cutlass_simt_fp32_triton_tree_canonical_parameter_vjp", + "algorithm": "cutlass_simt_fp32_triton_tree_bias_canonical_parameter_vjp", "reduction_backend": "triton", + "bias_reduction_backend": "triton_for_k_gt_96", "input_width": 96, "output_width": 5376, "cuda_sm": 90, @@ -273,12 +274,13 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: implementation_class_or_factory=OpBackend.CUDA_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( - "cutlass-56ca28a-fp32-simt-64x64x8-triton-fused-tree32-coalesced-pack-h3-v3" + "cutlass-56ca28a-fp32-simt-64x64x8-triton-fused-tree32-" + "coalesced-pack-triton-bias-h3-v4" ), ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v3", + backend_id="triton-video-input-projection-fp32-v4", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -292,6 +294,8 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: }, determinism_or_alignment_properties={ "algorithm": "fixed_fp32_gemm_canonical_parameter_vjp", + "forward_bias": "fp32_epilogue", + "leaf_loop_bound": "max_logical_leaf_length", "input_width": 96, "output_width": 5376, "cuda_sm": 90, @@ -306,7 +310,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: implementation_class_or_factory=OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( - "rlkernel-triton-fp32-no-fma-paired-tree32-leaf64x32-bias1x128-h3-v3" + "rlkernel-triton-fp32-no-fma-paired-tree32-leaf64x32-short-loop-bias-epilogue-h3-v4" ), ), OperatorBackendDescriptor( diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 11bc57088..9218c6279 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -325,7 +325,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v3" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v4" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op("triton")(x, weight, bias)) @@ -512,3 +512,75 @@ def test_compact_fp32_tree_workspace_preserves_chunked_outputs(backend, monkeypa _assert_bytes(_gemm_fp32(a, b, family), expected) monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.workspace_node_count * 2) _assert_bytes(_gemm_fp32(a, b, family), expected) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("rows", [1, 3]) +@pytest.mark.parametrize("k", [129, 5376, 13104]) +def test_shortened_fp32_leaves_preserve_padded_arithmetic(rows, k, monkeypatch): + from rl_engine.backends.shared.triton.gemm import det_gemm + + class PaddedLeaves: + def __init__(self, kernel): + self.kernel = kernel + + def __getitem__(self, grid): + def launch(*args, **kwargs): + kwargs["LEAF_LOOP_BOUND"] = det_gemm._BLOCK_K + return self.kernel[grid](*args, **kwargs) + + return launch + + torch.manual_seed(423) + a = torch.randn(k, rows, device="cuda").t() + b = torch.randn(k, 5, device="cuda") + inputs = [(a, b)] + # Include signed zero, gradual underflow, and non-finite propagation. + for value in (-0.0, torch.finfo(torch.float32).tiny / 8, float("inf"), float("nan")): + inputs.append((torch.full_like(a, value), b)) + actual = [ + det_gemm._triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) + for a, b in inputs + ] + for name in ("_det_gemm_kernel", "_det_gemm_tree_leaf_reduce_rocm_kernel"): + monkeypatch.setattr(det_gemm, name, PaddedLeaves(getattr(det_gemm, name))) + for (a, b), output in zip(inputs, actual, strict=True): + padded = det_gemm._triton_tree_gemm( + a, b, output_dtype=torch.float32, preserve_a_strides=True + ) + _assert_bytes(output, padded) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("rows", [1, 17, 129]) +def test_triton_forward_bias_fusion_preserves_fp32_rounding(rows): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + x, weight, bias, _ = _data(rows) + rhs = weight.t().contiguous() + expected = _triton_gemm_fp32(x, rhs) + bias + actual = _triton_gemm_fp32(x, rhs, bias=bias) + _assert_bytes(actual, expected) + with torch.no_grad(): + _assert_bytes(_op("triton")(x, weight, bias), expected) + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("frozen", ["weight", "bias", "both"]) +def test_backend_frozen_parameters_preserve_trainable_gradients(backend, frozen): + op = _op(backend) + data, keys = _data(129), _keys(129, "cuda") + expected = _run(op, data, keys) + actual = _run( + op, + data, + keys, + train_weight=frozen not in ("weight", "both"), + train_bias=frozen not in ("bias", "both"), + ) + for i, (value, gold) in enumerate(zip(actual, expected, strict=True)): + if i == 2 and frozen in ("weight", "both") or i == 3 and frozen in ("bias", "both"): + assert value is None + else: + _assert_bytes(value, gold) diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 5ae57e9d5..a8a110587 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,7 +13,7 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v3" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v4" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) From 51d78211c1ee71200848c8fd1656cf82e5090f0d Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Thu, 15 Oct 2026 16:02:20 +0800 Subject: [PATCH 05/26] Optimize H3 FP32 projection chunks and CUTLASS layouts Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 5 +- .../video_input_projection_fp32_cutlass.cpp | 10 +- .../video_input_projection_fp32_cutlass.cu | 258 +++++++----------- docs/operators/video-input-projection-fp32.md | 91 +++--- rl_engine/_C_cutlass.pyi | 5 +- .../cuda/gemm/video_input_projection_fp32.py | 28 +- .../backends/shared/triton/gemm/det_gemm.py | 88 +++++- .../gemm/video_input_projection_fp32.py | 19 +- rl_engine/ops/gemm/canonical_linear.py | 5 +- .../gemm/video_input_projection_fp32.py | 41 ++- rl_engine/runtime/registry.py | 10 +- .../gemm/test_video_input_projection_fp32.py | 46 +++- tests/runtime/test_kernel_registry.py | 2 +- 13 files changed, 378 insertions(+), 230 deletions(-) diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index e7ea55047..56845e8b3 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -19,6 +19,7 @@ from rl_engine.contracts.numerical import load_contract from rl_engine.ops.autograd.canonical_backward import canonical_backward_session +from rl_engine.reference.gemm.video_input_projection_fp32 import NativeVideoInputProjectionFP32Op def _candidate(name): @@ -164,6 +165,7 @@ def main(): "rl_engine/backends/shared/triton/gemm/det_gemm.py", "rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py", "rl_engine/backends/cuda/gemm/video_input_projection_fp32.py", + "rl_engine/reference/gemm/video_input_projection_fp32.py", "csrc/bindings/video_input_projection_fp32_cutlass.cpp", "csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu", "build_tools/extensions.py", @@ -177,6 +179,7 @@ def main(): ) report = { "scope": "single H3 FP32 input projection; eager training and no-grad forward graph", + "reference_id": NativeVideoInputProjectionFP32Op.reference_id, "environment": { "gpu": torch.cuda.get_device_name(), "torch": torch.__version__, @@ -209,7 +212,7 @@ def main(): dim=1, ) cx, cw, cb = [t.detach().cpu().requires_grad_() for t in (x, weight, bias)] - cy = torch.nn.functional.linear(cx, cw, cb) + cy = NativeVideoInputProjectionFP32Op()(cx, cw, cb) cy.backward(grad.cpu()) gold = (cy, cx.grad, cw.grad, cb.grad) for name in args.backends: diff --git a/csrc/bindings/video_input_projection_fp32_cutlass.cpp b/csrc/bindings/video_input_projection_fp32_cutlass.cpp index 0afc9ca60..9923611d2 100644 --- a/csrc/bindings/video_input_projection_fp32_cutlass.cpp +++ b/csrc/bindings/video_input_projection_fp32_cutlass.cpp @@ -5,15 +5,15 @@ torch::Tensor cutlass_gemm_fp32_tn(torch::Tensor a, torch::Tensor b_transposed); int cutlass_cuda_version(); -void cutlass_gemm_fp32_tree_leaves_tn( +void cutlass_gemm_fp32_tree_leaves_strided( torch::Tensor a, torch::Tensor bt, torch::Tensor workspace, - torch::Tensor starts, torch::Tensor lengths, torch::Tensor nodes); + torch::Tensor starts, torch::Tensor nodes, torch::Tensor leaf_ids, int64_t length); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("gemm_fp32_tn", &cutlass_gemm_fp32_tn, "Fixed FP32 SIMT GEMM: [M,K] times transposed [N,K]"); - m.def("gemm_fp32_tree_leaves_tn", &cutlass_gemm_fp32_tree_leaves_tn, - "Batch all FP32 leaves into a validated canonical tree workspace"); + m.def("gemm_fp32_tree_leaves_strided", &cutlass_gemm_fp32_tree_leaves_strided, + "Evaluate one actual-K leaf group directly from matrix views"); m.def("build_info", []() { pybind11::dict info; info["cutlass_revision"] = "56ca28adb32ffadd097566545b8f7d83b652e998"; @@ -27,7 +27,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { info["ftz"] = false; info["tree_leaf_size"] = 32; info["tree_leaf_launcher"] = "cutlass_device_gemm_array"; - info["tree_operand_packing"] = "coalesced_transpose32"; + info["tree_operand_packing"] = "direct_strides_grouped_actual_k"; return info; }); } diff --git a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu index dfb986614..3e942634d 100644 --- a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu +++ b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu @@ -26,159 +26,7 @@ using Gemm = cutlass::gemm::device::Gemm< cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 2, 1, 1, false, cutlass::arch::OpMultiplyAdd>; -// Match the shared tree plan's 32-term leaves. One batched launch writes all -// leaves into their existing workspace nodes; Triton then merges each level. constexpr int kTreeLeaf = 32; -using LeafGemm = cutlass::gemm::device::GemmArray< - float, cutlass::layout::RowMajor, - float, cutlass::layout::ColumnMajor, - float, cutlass::layout::RowMajor, - float, cutlass::arch::OpClassSimt, cutlass::arch::Sm50, - cutlass::gemm::GemmShape<64, 64, 8>, - cutlass::gemm::GemmShape<32, 32, 8>, - cutlass::gemm::GemmShape<1, 1, 1>, - cutlass::epilogue::thread::LinearCombination, - cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, - 2, 1, 1, cutlass::arch::OpMultiplyAdd>; - -template -__global__ void pack_tree_operand( - const float* input, float* packed, int64_t* pointers, - const int32_t* starts, const int32_t* lengths, const int32_t* nodes, - int64_t rows, int64_t stride_row, int64_t stride_k, int64_t total, - float* workspace, int64_t* output_pointers, int64_t output_elements) { - const int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (index >= total) return; - const int64_t leaf_elements = rows * kTreeLeaf; - const int leaf = static_cast(index / leaf_elements); - const int64_t local = index % leaf_elements; - const int64_t row = local / kTreeLeaf; - const int k = static_cast(local % kTreeLeaf); - packed[index] = k < lengths[leaf] - ? input[row * stride_row + (static_cast(starts[leaf]) + k) * stride_k] - : 0.0f; - if (local == 0) { - pointers[leaf] = reinterpret_cast(packed + leaf * leaf_elements); - if constexpr (WriteOutputPointers) { - output_pointers[leaf] = reinterpret_cast( - workspace + static_cast(nodes[leaf]) * output_elements); - } - } -} - -// The transposed G/X views have contiguous rows and strided K. Read adjacent -// rows together, then transpose in shared memory for the packed row-major leaf. -template -__global__ void pack_tree_operand_transposed( - const float* input, float* packed, int64_t* pointers, - const int32_t* starts, const int32_t* lengths, const int32_t* nodes, - int64_t rows, int64_t stride_k, float* workspace, - int64_t* output_pointers, int64_t output_elements) { - __shared__ float tile[kTreeLeaf][kTreeLeaf + 1]; - const int leaf = blockIdx.y; - const int64_t row_start = static_cast(blockIdx.x) * kTreeLeaf; - const int x = threadIdx.x, y = threadIdx.y; - const int length = lengths[leaf]; - const int64_t start = starts[leaf]; - for (int offset = 0; offset < kTreeLeaf; offset += 8) { - const int k = y + offset; - tile[k][x] = row_start + x < rows && k < length - ? input[row_start + x + (start + k) * stride_k] : 0.0f; - } - __syncthreads(); - for (int offset = 0; offset < kTreeLeaf; offset += 8) { - const int64_t row = row_start + y + offset; - if (row < rows) { - packed[(static_cast(leaf) * rows + row) * kTreeLeaf + x] - = tile[x][y + offset]; - } - } - if (blockIdx.x == 0 && x == 0 && y == 0) { - pointers[leaf] = reinterpret_cast( - packed + static_cast(leaf) * rows * kTreeLeaf); - if constexpr (WriteOutputPointers) { - output_pointers[leaf] = reinterpret_cast( - workspace + static_cast(nodes[leaf]) * output_elements); - } - } -} - -template -void launch_tree_packing( - const at::Tensor& input, const at::Tensor& packed, int64_t* pointers, - const at::Tensor& starts, const at::Tensor& lengths, const at::Tensor& nodes, - float* workspace, int64_t* output_pointers, int64_t output_elements, - cudaStream_t stream) { - if (input.stride(0) == 1 && input.stride(1) > 1) { - const dim3 grid((input.size(0) + kTreeLeaf - 1) / kTreeLeaf, starts.numel()); - pack_tree_operand_transposed<<>>( - input.data_ptr(), packed.data_ptr(), pointers, - starts.data_ptr(), lengths.data_ptr(), nodes.data_ptr(), - input.size(0), input.stride(1), workspace, output_pointers, output_elements); - } else { - constexpr int threads = 256; - pack_tree_operand - <<<(packed.numel() + threads - 1) / threads, threads, 0, stream>>>( - input.data_ptr(), packed.data_ptr(), pointers, - starts.data_ptr(), lengths.data_ptr(), nodes.data_ptr(), - input.size(0), input.stride(0), input.stride(1), packed.numel(), - workspace, output_pointers, output_elements); - } -} - -void cutlass_gemm_fp32_tree_leaves_tn( - at::Tensor a, at::Tensor bt, at::Tensor workspace, - at::Tensor starts, at::Tensor lengths, at::Tensor nodes) { - TORCH_CHECK(a.is_cuda() && a.scalar_type() == at::kFloat && a.dim() == 2, - "CUTLASS leaves require rank-two CUDA FP32 A"); - const c10::cuda::CUDAGuard guard(a.device()); - const auto* properties = at::cuda::getDeviceProperties(a.get_device()); - TORCH_CHECK(properties->major == 9 && properties->minor == 0, - "CUTLASS video projection requires CUDA SM90"); - TORCH_CHECK(bt.device() == a.device() && bt.scalar_type() == at::kFloat && bt.dim() == 2, - "CUTLASS leaves require rank-two FP32 Bt on the same device"); - TORCH_CHECK(a.size(1) == bt.size(1) && a.size(0) > 0 && bt.size(0) > 0, - "CUTLASS leaf dimensions must match and be non-empty"); - TORCH_CHECK(a.stride(0) > 0 && a.stride(1) > 0 && bt.stride(0) > 0 && bt.stride(1) > 0, - "CUTLASS leaf operands require positive strides"); - TORCH_CHECK(workspace.device() == a.device() && workspace.scalar_type() == at::kFloat && - workspace.dim() == 3 && workspace.is_contiguous() && - workspace.size(1) == a.size(0) && workspace.size(2) == bt.size(0), - "CUTLASS leaves require contiguous FP32 tree workspace"); - for (const auto& metadata : {starts, lengths, nodes}) { - TORCH_CHECK(metadata.device() == a.device() && metadata.scalar_type() == at::kInt && - metadata.dim() == 1 && metadata.is_contiguous() && - metadata.numel() == starts.numel(), - "CUTLASS tree metadata must be matching contiguous CUDA int32 vectors"); - } - const int64_t leaves = starts.numel(), m = a.size(0), n = bt.size(0); - TORCH_CHECK(leaves > 0 && leaves <= 65535 && m <= INT_MAX && n <= INT_MAX, - "CUTLASS batched leaf grid exceeds its supported dimensions"); - auto packed_a = at::empty({leaves, m, kTreeLeaf}, a.options()); - auto packed_b = at::empty({leaves, n, kTreeLeaf}, a.options()); - auto pointers = at::empty({3, leaves}, a.options().dtype(at::kLong)); - auto* ptr_a = pointers.data_ptr(); - auto* ptr_b = ptr_a + leaves; - auto* ptr_output = ptr_b + leaves; - const auto stream = c10::cuda::getCurrentCUDAStream(a.get_device()).stream(); - launch_tree_packing(a, packed_a, ptr_a, starts, lengths, nodes, - nullptr, nullptr, m * n, stream); - launch_tree_packing(bt, packed_b, ptr_b, starts, lengths, nodes, - workspace.data_ptr(), ptr_output, m * n, stream); - LeafGemm::Arguments args( - {static_cast(m), static_cast(n), kTreeLeaf}, - reinterpret_cast(ptr_a), cutlass::layout::RowMajor(kTreeLeaf), - reinterpret_cast(ptr_b), cutlass::layout::ColumnMajor(kTreeLeaf), - reinterpret_cast(ptr_output), cutlass::layout::RowMajor(n), - reinterpret_cast(ptr_output), cutlass::layout::RowMajor(n), - {1.0f, 0.0f}, static_cast(leaves)); - LeafGemm gemm; - auto status = gemm.can_implement(args); - TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); - status = gemm(args, nullptr, stream); - TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} at::Tensor cutlass_gemm_fp32_tn(at::Tensor a, at::Tensor b_transposed) { TORCH_CHECK(a.is_cuda() && b_transposed.is_cuda(), "CUTLASS GEMM requires CUDA tensors"); @@ -214,3 +62,109 @@ at::Tensor cutlass_gemm_fp32_tn(at::Tensor a, at::Tensor b_transposed) { } int cutlass_cuda_version() { return CUDART_VERSION; } + +template +using StridedLeafGemm = cutlass::gemm::device::GemmArray< + float, LayoutA, float, LayoutB, float, cutlass::layout::RowMajor, + float, cutlass::arch::OpClassSimt, cutlass::arch::Sm50, + cutlass::gemm::GemmShape<64, 64, 8>, + cutlass::gemm::GemmShape<32, 32, 8>, + cutlass::gemm::GemmShape<1, 1, 1>, + cutlass::epilogue::thread::LinearCombination, + cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, + 2, 1, 1, cutlass::arch::OpMultiplyAdd>; + +__global__ void tree_leaf_pointers( + const float* a, const float* bt, float* workspace, int64_t* pointers, + const int32_t* starts, const int32_t* nodes, const int32_t* leaf_ids, + int64_t count, int64_t stride_ak, int64_t stride_btk, int64_t elements) { + const int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= count) return; + const int leaf = leaf_ids[index]; + const int64_t start = starts[leaf]; + pointers[index] = reinterpret_cast(a + start * stride_ak); + pointers[count + index] = reinterpret_cast(bt + start * stride_btk); + pointers[2 * count + index] = reinterpret_cast( + workspace + static_cast(nodes[leaf]) * elements); +} + +template +void launch_strided_leaves( + const at::Tensor& a, const at::Tensor& bt, const at::Tensor& pointers, + int64_t count, int length, int64_t lda, int64_t ldb, cudaStream_t stream) { + using Leaf = StridedLeafGemm; + const auto* addresses = pointers.data_ptr(); + typename Leaf::Arguments args( + {static_cast(a.size(0)), static_cast(bt.size(0)), length}, + reinterpret_cast(addresses), LayoutA(lda), + reinterpret_cast(addresses + count), LayoutB(ldb), + reinterpret_cast(addresses + 2 * count), + cutlass::layout::RowMajor(bt.size(0)), + reinterpret_cast(addresses + 2 * count), + cutlass::layout::RowMajor(bt.size(0)), {1.0f, 0.0f}, static_cast(count)); + Leaf gemm; + auto status = gemm.can_implement(args); + TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); + status = gemm(args, nullptr, stream); + TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); +} + +void cutlass_gemm_fp32_tree_leaves_strided( + at::Tensor a, at::Tensor bt, at::Tensor workspace, at::Tensor starts, + at::Tensor nodes, at::Tensor leaf_ids, int64_t length) { + TORCH_CHECK(a.is_cuda() && a.scalar_type() == at::kFloat && a.dim() == 2, + "CUTLASS leaves require rank-two CUDA FP32 A"); + const c10::cuda::CUDAGuard guard(a.device()); + const auto* properties = at::cuda::getDeviceProperties(a.get_device()); + TORCH_CHECK(properties->major == 9 && properties->minor == 0, + "CUTLASS video projection requires CUDA SM90"); + TORCH_CHECK(bt.device() == a.device() && bt.scalar_type() == at::kFloat && + bt.dim() == 2 && a.size(1) == bt.size(1), + "CUTLASS leaf dimensions, dtypes and devices must match"); + TORCH_CHECK(a.size(0) > 0 && bt.size(0) > 0 && length > 0 && length <= kTreeLeaf && + a.size(0) <= INT_MAX && bt.size(0) <= INT_MAX, + "CUTLASS leaf sizes exceed the supported dimensions"); + TORCH_CHECK(workspace.device() == a.device() && workspace.scalar_type() == at::kFloat && + workspace.dim() == 3 && workspace.is_contiguous() && + workspace.size(1) == a.size(0) && workspace.size(2) == bt.size(0), + "CUTLASS leaves require contiguous FP32 tree workspace"); + for (const auto& metadata : {starts, nodes, leaf_ids}) { + TORCH_CHECK(metadata.device() == a.device() && metadata.scalar_type() == at::kInt && + metadata.dim() == 1 && metadata.is_contiguous(), + "CUTLASS leaf metadata must be contiguous CUDA int32 vectors"); + } + TORCH_CHECK(starts.numel() == nodes.numel() && leaf_ids.numel() > 0 && + leaf_ids.numel() <= 65535, + "CUTLASS leaf metadata sizes are inconsistent"); + const bool a_row = a.stride(1) == 1; + const bool b_column = bt.stride(1) == 1; + TORCH_CHECK((a_row || a.stride(0) == 1) && + (b_column || bt.stride(0) == 1) && + a.stride(0) > 0 && a.stride(1) > 0 && + bt.stride(0) > 0 && bt.stride(1) > 0, + "CUTLASS leaves require a row-major or column-major view"); + const int64_t lda = a_row ? a.stride(0) : a.stride(1); + const int64_t ldb = b_column ? bt.stride(0) : bt.stride(1); + TORCH_CHECK(lda <= INT_MAX && ldb <= INT_MAX, + "CUTLASS leaf leading dimensions exceed int32"); + const int64_t count = leaf_ids.numel(); + auto pointers = at::empty({3, count}, a.options().dtype(at::kLong)); + const auto stream = c10::cuda::getCurrentCUDAStream(a.get_device()).stream(); + tree_leaf_pointers<<<(count + 255) / 256, 256, 0, stream>>>( + a.data_ptr(), bt.data_ptr(), workspace.data_ptr(), + pointers.data_ptr(), starts.data_ptr(), nodes.data_ptr(), + leaf_ids.data_ptr(), count, a.stride(1), bt.stride(1), + a.size(0) * bt.size(0)); + using Row = cutlass::layout::RowMajor; + using Column = cutlass::layout::ColumnMajor; + if (a_row && b_column) { + launch_strided_leaves(a, bt, pointers, count, length, lda, ldb, stream); + } else if (a_row) { + launch_strided_leaves(a, bt, pointers, count, length, lda, ldb, stream); + } else if (b_column) { + launch_strided_leaves(a, bt, pointers, count, length, lda, ldb, stream); + } else { + launch_strided_leaves(a, bt, pointers, count, length, lda, ldb, stream); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 3ffcd8cb4..f82b2072b 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -7,11 +7,11 @@ space. This operator preserves the FP32 projection boundary of the pinned H3 reference. The model adapter owns the subsequent cast into the BF16 packed stream. The primary implementation reuses RL-Kernel's Triton FP32 GEMM and canonical -parameter-gradient session. A pinned CUTLASS FP32 SIMT GEMM is an explicit CUDA -candidate. Both use the same semantic adapter, autograd implementation, and -existing Triton tree plan and GPU level reductions. CUTLASS substitutes the -batched leaf GEMM for dX/dW. Long bias VJPs reuse the one-row Triton leaves; -the reported execution identity includes both paths. +parameter-gradient session. Input and weight VJPs use fixed 128-term K chunks +and one balanced FP32 merge. A pinned CUTLASS FP32 SIMT GEMM is an explicit +CUDA candidate. It retains the midpoint tree and reads each leaf directly +from matrix views. Both paths share the semantic adapter and canonical session. +Long bias VJPs retain the one-row Triton tree. ## Entry Point @@ -34,9 +34,9 @@ physical placement. The same parameter keeps the same `parameter_id` across uses | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | -| Triton on CUDA | `TritonVideoInputProjectionFP32Op` | `_det_gemm_fp32_kernel`, existing tree leaf/reduce kernels | Default CUDA path; SM90 qualification | -| CUDA / CUTLASS + Triton reductions | `CUDAVideoInputProjectionFP32Op` | `_C_cutlass.gemm_fp32_tn`, `_C_cutlass.gemm_fp32_tree_leaves_tn`, shared tree reductions | Explicit candidate; optional build; SM90 | -| PyTorch reference | `NativeVideoInputProjectionFP32Op` | CPU `F.linear` | Accuracy reference | +| Triton on CUDA | `TritonVideoInputProjectionFP32Op` | `_det_gemm_fp32_kernel`, `_det_gemm_fp32_chunk_merge_kernel`, bias tree kernels | Default CUDA path; SM90 qualification | +| CUDA / CUTLASS + Triton reductions | `CUDAVideoInputProjectionFP32Op` | `_C_cutlass.gemm_fp32_tn`, `_C_cutlass.gemm_fp32_tree_leaves_strided`, shared tree reductions | Explicit candidate; optional build; SM90 | +| PyTorch reference | `NativeVideoInputProjectionFP32Op` | CPU `F.linear`; FP32 product/sum VJPs | Accuracy reference | CUTLASS revision: `56ca28adb32ffadd097566545b8f7d83b652e998`. The build verifies the complete include-tree fingerprint. The BSD-3-Clause dependency notice is @@ -97,33 +97,47 @@ complete contribution sequences match. Calls with different contribution sets normalize each parameter independently. Short reductions up to the 96-wide video input use the existing direct FP32 -GEMM. Long reductions reuse the existing midpoint plan with at most 32 terms per -leaf. The Triton scalar loop stops at the largest logical leaf length -(21 for K=5376; 26 for K=13104), preserving the midpoint tree and ascending -products. Workspace, merge nodes, and output stay FP32. One-row complete trees fuse each -Triton leaf pair with its first merge. Other complete trees compute leaves in -one grid. Subsequent -merges use the existing two-level kernel and alternating compact workspace -buffers. Ragged trees retain the original per-level schedule. CUTLASS packs -the same leaf ranges into 32-wide buffers, -pads inactive slots with zero, and uses one `device::GemmArray` launch to write -the plan's leaf nodes. Its transpose-view packer reads adjacent rows together -and transposes a 32x32 tile in shared memory. It uses the same paired level -reductions and compact buffers for complete trees. - -The tree depends on logical K and retains its topology across physical row -changes. Output-row chunking keeps tree workspace below the existing element -limit (512 MiB for FP32). If one output row exceeds that limit, independent -output columns are chunked with the same K tree. CUTLASS also allocates packed -operands and pointer arrays. Parameter sorting and public validation are included -in complete-call measurements. Existing BF16 tree callers retain their default -dtype and schedule. - -Triton uses fixed 64x32x32 leaves and 1x128x32 leaves for one-row tree calls, -including bias VJPs. The short forward GEMM retains 64x64x32 tiles. TF32 and FP multiply-add fusion are disabled. Schedule -fusion preserves every FP32 leaf and tree addition boundary. CUTLASS -fixes 64x64x8 FP32 SIMT tiles and a single output owner per leaf, without parallel -split-K. Both leaf algorithms have separately versioned backend identities. +GEMM. The Triton input and weight VJPs reuse that kernel with fixed 128-term +chunks. Each chunk computes four ascending 32-term FP32 dots from zero. Chunk +partials remain FP32. A second grid uses a fixed balanced reduction over the +next power-of-two chunk count, with zero padding. The merge uses 16 output +elements per program and four warps. These settings define the v5 arithmetic +contract. The K partition does not depend on M, row placement, or occupancy. +The VJP kernels consume positive-stride views directly, including `G.T`. + +The fixed-chunk design follows the scheduling separation in merged PR #396. +Its AMD BF16 MFMA instructions are not used by this CUDA FP32 implementation. +The FP32 GEMM body is reused from the helper present in merged PR #325. +The short forward path and long bias tree retain their earlier arithmetic. + +CUTLASS retains the existing midpoint plan with at most 32 terms per leaf. +Leaves are grouped by actual K length. A pointer kernel selects each leaf's +original operand range, and a fixed `device::GemmArray` reads row-major or +column-major views directly. This applies the layout approach from merged +PR #343. It avoids packed operand buffers while retaining the FP32 midpoint +merge. The actual-K grouping and layouts are recorded in `backend_info()`. +CUTLASS still uses Triton for tree merging and long bias VJPs. + +Both schedules cap their FP32 partial workspace at 512 MiB by chunking output +rows. The midpoint implementation also supports output-column chunking when +one complete row exceeds its budget. CUTLASS allocates pointer arrays for each +leaf group. Parameter sorting and public validation are included in complete +operator measurements. Existing BF16 callers retain their dtype and schedule. +TF32 and compiler multiply-add fusion remain disabled. +The reused Triton `tl.dot` lowers its IEEE FP32 inner products to fixed SIMT +`fma.rn.f32` instructions. The compiler fusion flag controls surrounding +expressions; it does not remove the intrinsic dot's FMA instructions. This +arithmetic is included in the Triton contract and runtime metadata. The pinned +CUTLASS SIMT build uses separate FP32 multiply and add instructions. + +The independent reference keeps the forward definition as CPU FP32 `F.linear`. +Its analytic VJPs form FP32 products and use CPU PyTorch `sum` over the complete +K dimension. Only output rows and columns are tiled to bound temporary memory. +It does not reuse candidate K partitions or GPU kernels. This avoids long CPU +GEMM accumulation errors observed in the 13104-row unscaled regression. +The reference identity is `cpu-fp32-linear-products-sum-v2`. FP64 is used only +in a small reference-validation test and diagnostic evidence, not by the +reference implementation or candidate kernels. The shared WS1 `reduction` FP32 accuracy policy uses `atol=rtol=1e-4` for forward and gradients. Invariance uses byte equality. Parameter invariance compares the @@ -144,9 +158,10 @@ forward CUDA Graph. Timings include layout copies, bias, and the canonical parameter-gradient machinery. Raw samples, environment, source hashes, and actual backend metadata are written to the explicitly selected local output. PyTorch is a production performance baseline with separately reported BI results. -The shared GPU tree avoids Python recursion per leaf and merge node. Its -workspace traffic, output-row chunking, and CUTLASS operand packing still add -backward work; compare complete operator timings before drawing a speedup claim. +The Triton VJPs use one partial-GEMM launch and one merge launch when the +workspace fits. CUTLASS retains tree merge launches but removes operand packing. +Compare complete operator timings, including sorting, before drawing a speedup +claim. The benchmark records the reference identity and its source hash. ## Tests diff --git a/rl_engine/_C_cutlass.pyi b/rl_engine/_C_cutlass.pyi index 862c08a2e..de2152e41 100644 --- a/rl_engine/_C_cutlass.pyi +++ b/rl_engine/_C_cutlass.pyi @@ -6,12 +6,13 @@ from typing import Any import torch def gemm_fp32_tn(a: torch.Tensor, b_transposed: torch.Tensor) -> torch.Tensor: ... -def gemm_fp32_tree_leaves_tn( +def gemm_fp32_tree_leaves_strided( a: torch.Tensor, b_transposed: torch.Tensor, workspace: torch.Tensor, starts: torch.Tensor, - lengths: torch.Tensor, nodes: torch.Tensor, + leaf_ids: torch.Tensor, + length: int, ) -> None: ... def build_info() -> dict[str, Any]: ... diff --git a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py index 3d2799242..32e20952b 100644 --- a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py @@ -65,15 +65,31 @@ def _load_extension(): or info["ftz"] or info["split_k"] != 1 or info.get("tree_leaf_size") != _BLOCK_K - or info.get("tree_operand_packing") != "coalesced_transpose32" + or info.get("tree_operand_packing") != "direct_strides_grouped_actual_k" ): raise RuntimeError("CUTLASS extension does not match the strict FP32 build contract") return extension def _cutlass_tree_leaves(a, b, workspace, plan): - _load_extension().gemm_fp32_tree_leaves_tn( - a, b.t(), workspace, plan.leaf_starts, plan.leaf_lengths, plan.leaf_nodes + for length, leaf_ids in _leaf_groups(plan.host.leaf_lengths, a.device): + _load_extension().gemm_fp32_tree_leaves_strided( + a, b.t(), workspace, plan.leaf_starts, plan.leaf_nodes, leaf_ids, length + ) + + +@lru_cache(maxsize=None) +def _leaf_groups(lengths, device): + return tuple( + ( + length, + torch.tensor( + [index for index, actual in enumerate(lengths) if actual == length], + device=device, + dtype=torch.int32, + ), + ) + for length in sorted(set(lengths)) ) @@ -92,7 +108,7 @@ def cutlass_gemm_fp32(a, b): class CUDAVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "cutlass" - backend_id = "cuda-cutlass-video-input-projection-fp32-v4" + backend_id = "cuda-cutlass-video-input-projection-fp32-v5" def __init__(self): if not _TRITON_AVAILABLE: @@ -111,8 +127,8 @@ def backend_info(self): return { "backend_id": self.backend_id, "family": "cuda", - "kernel": "gemm_fp32_tn / gemm_fp32_tree_leaves_tn / " - "_det_gemm_tree_leaf_kernel / _det_gemm_tree_leaf_reduce_rocm_kernel", + "kernel": "gemm_fp32_tn / gemm_fp32_tree_leaves_strided / " + "_det_gemm_kernel / _det_gemm_tree_leaf_reduce_rocm_kernel", "reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, "reduction_backend": "triton", diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index b8f3c24aa..c5db36fc5 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -36,6 +36,7 @@ # Pinned. NOT autotuned (autotune picks per-shape configs -> breaks invariance). _BLOCK_M, _BLOCK_N, _BLOCK_K = 64, 64, 32 +_FP32_CHUNK_K = 128 # Keep the intermediate tree workspace bounded. A VLLM warmup can present a # very large flattened token dimension; allocating ``node_count * M * N`` in # one shot would otherwise create a multi-terabyte virtual tensor on ROCm. @@ -748,17 +749,21 @@ def _det_gemm_fp32_kernel( BLOCK_K: tl.constexpr, PROMOTE_INPUTS: tl.constexpr, HAS_BIAS: tl.constexpr, + CHUNK_K: tl.constexpr = 0, ): - # One program = one output tile, walks the whole K in fixed order. - # No split-K -> K-accumulation order independent of M -> batch-invariant. + # One program walks the whole K or one declared fixed-width chunk. + # The K order does not depend on output rows or their placement. pid_m, pid_n = tl.program_id(0), tl.program_id(1) offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - offs_k = tl.arange(0, BLOCK_K) + chunk = tl.program_id(2) if CHUNK_K else 0 + start_k = chunk * CHUNK_K + offs_k = start_k + tl.arange(0, BLOCK_K) a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak) b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_K)): + reduction_size = CHUNK_K if CHUNK_K else K + for k in range(0, tl.cdiv(reduction_size, BLOCK_K)): k_rem = K - k * BLOCK_K a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < k_rem) a = tl.load(a_ptrs, mask=a_mask, other=0.0) @@ -774,10 +779,83 @@ def _det_gemm_fp32_kernel( bias = tl.load(bias_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32) acc = acc + bias[None, :] c = acc.to(c_ptr.dtype.element_ty) - c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn + c_ptrs = c_ptr + chunk * M * N + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) tl.store(c_ptrs, c, mask=mask) + @triton.jit + def _det_gemm_fp32_chunk_merge_kernel( + partial_ptr, + output_ptr, + ELEMENTS, + CHUNKS: tl.constexpr, + PADDED_CHUNKS: tl.constexpr, + BLOCK_C: tl.constexpr, + ): + offsets = tl.program_id(0) * BLOCK_C + tl.arange(0, BLOCK_C) + chunks = tl.arange(0, PADDED_CHUNKS) + values = tl.load( + partial_ptr + chunks[:, None] * ELEMENTS + offsets[None, :], + mask=(chunks[:, None] < CHUNKS) & (offsets[None, :] < ELEMENTS), + other=0.0, + ) + # Pinned padded width and launch configuration define this FP32 tree. + total = tl.sum(values, axis=0) + tl.store(output_ptr + offsets, total, mask=offsets < ELEMENTS) + + +def _triton_chunked_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Reuse the FP32 GEMM in fixed K chunks and merge partials in one grid.""" + m_size, k_size = a.shape + n_size = b.shape[1] + chunks = triton.cdiv(k_size, _FP32_CHUNK_K) + rows_per_chunk = _MAX_TREE_WORKSPACE_ELEMENTS // (chunks * n_size) + if rows_per_chunk < 1: + raise RuntimeError("FP32 chunk GEMM exceeds the workspace budget for one row") + if m_size > rows_per_chunk: + return torch.cat( + [ + _triton_chunked_gemm_fp32(a[start : start + rows_per_chunk], b) + for start in range(0, m_size, rows_per_chunk) + ], + dim=0, + ) + partial = torch.empty((chunks, m_size, n_size), device=a.device, dtype=torch.float32) + output = torch.empty((m_size, n_size), device=a.device, dtype=torch.float32) + grid = (triton.cdiv(m_size, _BLOCK_M), triton.cdiv(n_size, _BLOCK_N), chunks) + _det_gemm_fp32_kernel[grid]( + a, + b, + partial, + None, + m_size, + n_size, + k_size, + *a.stride(), + *b.stride(), + n_size, + 1, + BLOCK_M=_BLOCK_M, + BLOCK_N=_BLOCK_N, + BLOCK_K=_BLOCK_K, + PROMOTE_INPUTS=True, + HAS_BIAS=False, + CHUNK_K=_FP32_CHUNK_K, + enable_fp_fusion=False, + num_warps=4, + ) + _det_gemm_fp32_chunk_merge_kernel[(triton.cdiv(m_size * n_size, 16),)]( + partial, + output, + m_size * n_size, + CHUNKS=chunks, + PADDED_CHUNKS=triton.next_power_of_2(chunks), + BLOCK_C=16, + enable_fp_fusion=False, + num_warps=4, + ) + return output + def _triton_gemm_fp32( a: torch.Tensor, b: torch.Tensor, *, bias: torch.Tensor | None = None diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index dee6b3935..e39b88942 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -8,6 +8,7 @@ from rl_engine.backends.shared.triton.gemm.det_gemm import ( _BLOCK_K, _TRITON_AVAILABLE, + _triton_chunked_gemm_fp32, _triton_gemm_fp32, _triton_tree_gemm, ) @@ -20,13 +21,16 @@ def projection_gemm_fp32(a, b, *, bias=None): if a.shape[1] <= VIDEO_PATCH_WIDTH: return _triton_gemm_fp32(a, b, bias=bias) - output = _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) + if b.shape[1] == VIDEO_PATCH_WIDTH: + output = _triton_chunked_gemm_fp32(a, b) + else: + output = _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) return output if bias is None else output + bias class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton_tree" - backend_id = "triton-video-input-projection-fp32-v4" + backend_id = "triton-video-input-projection-fp32-v5" def __init__(self): if not _TRITON_AVAILABLE: @@ -44,16 +48,21 @@ def backend_info(self): return { "backend_id": self.backend_id, "family": "triton", - "kernel": "_det_gemm_fp32_kernel / _det_gemm_tree_leaf_kernel / " + "kernel": "_det_gemm_fp32_kernel / _det_gemm_fp32_chunk_merge_kernel / " + "_det_gemm_kernel / " "_det_gemm_tree_leaf_reduce_rocm_kernel", "triton_version": triton.__version__, "tile": [64, 64, 32], "tf32": False, "fp_fusion": False, + "dot_arithmetic": "ieee_fp32_simt_fma", "forward_bias": "fp32_epilogue", "leaf_loop_bound": "max_logical_leaf_length", - "split_k": 1, - "reduction": "fixed_midpoint_tree", + "split_k": "fixed_128_term_vjp_chunks", + "reduction": "fixed_128_chunks_balanced_fp32_merge", + "vjp_chunk_k": 128, + "vjp_merge_tile": 16, + "bias_reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, "reduction_backend": "triton", "leaf_launch": "paired_leaves_for_one_row_complete_trees", diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 8fe329dbb..b0bd9a08e 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -130,7 +130,8 @@ def bias_reducer(grads): record_backward( "det_gemm", kernel_id=( - "rl_engine.backends.shared.triton.gemm.det_gemm._triton_tree_gemm" + "rl_engine.backends.shared.triton.gemm.det_gemm." + "_triton_chunked_gemm_fp32 / _triton_tree_gemm" if ctx.family == "triton_tree" else "rl_engine.backends.shared.triton.gemm.det_gemm._triton_gemm" ), @@ -148,7 +149,7 @@ def bias_reducer(grads): record_backward( "video_input_projection_fp32", kernel_id=( - "rl_engine._C_cutlass.gemm_fp32_tree_leaves_tn / " + "rl_engine._C_cutlass.gemm_fp32_tree_leaves_strided / " "rl_engine.backends.shared.triton.gemm.det_gemm._triton_tree_gemm" ), impl="cutlass_fp32_simt_triton_tree_bias_canonical_rowfold", diff --git a/rl_engine/reference/gemm/video_input_projection_fp32.py b/rl_engine/reference/gemm/video_input_projection_fp32.py index 9679cb17b..98f8f01c6 100644 --- a/rl_engine/reference/gemm/video_input_projection_fp32.py +++ b/rl_engine/reference/gemm/video_input_projection_fp32.py @@ -3,15 +3,54 @@ """Independent CPU FP32 Linear reference for the H3 video input head.""" +import torch import torch.nn.functional as F from rl_engine.ops.gemm.video_input_projection_fp32 import validate_projection_inputs +def _reference_matmul_fp32(a, b): + """Evaluate each CPU dot with independent FP32 products and torch.sum.""" + a, bt = a.cpu().contiguous(), b.t().cpu().contiguous() + output = torch.empty((a.shape[0], bt.shape[0]), dtype=torch.float32) + # Tile only the output. Every torch.sum consumes the full logical K. + # This bounds temporary memory without copying a candidate's K partition. + for row in range(0, a.shape[0], 16): + for column in range(0, bt.shape[0], 16): + products = a[row : row + 16, None, :] * bt[None, column : column + 16, :] + output[row : row + 16, column : column + 16] = products.sum(dim=-1) + return output + + +class _ReferenceLinearFn(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, bias): + ctx.save_for_backward(x, weight) + ctx.has_bias = bias is not None + return F.linear(x, weight, bias) + + @staticmethod + def backward(ctx, grad_output): + x, weight = ctx.saved_tensors + grad = grad_output.reshape(-1, weight.shape[0]).contiguous().float() + rows = x.reshape(-1, weight.shape[1]) + dx = _reference_matmul_fp32(grad, weight).reshape_as(x) if ctx.needs_input_grad[0] else None + dw = _reference_matmul_fp32(grad.t(), rows) if ctx.needs_input_grad[1] else None + dbias = grad.sum(dim=0) if ctx.has_bias and ctx.needs_input_grad[2] else None + return dx, dw, dbias + + class NativeVideoInputProjectionFP32Op: + reference_id = "cpu-fp32-linear-products-sum-v2" + def forward_fp32(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None): validate_projection_inputs(x, weight, bias) # CPU avoids a GPU TF32 or batch-dependent cuBLAS reference path. - return F.linear(x.cpu(), weight.cpu(), None if bias is None else bias.cpu()).to(x.device) + # Long CPU GEMM VJPs can exceed the shared FP32 accuracy threshold. + # Use the analytic Linear VJP with PyTorch's independent FP32 sum. + with torch.autocast(device_type="cpu", enabled=False): + return _ReferenceLinearFn.apply( + x.cpu(), weight.cpu(), None if bias is None else bias.cpu() + ).to(x.device) __call__ = forward_fp32 diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 269db8187..405a30847 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -244,7 +244,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="cuda-cutlass-video-input-projection-fp32-v4", + backend_id="cuda-cutlass-video-input-projection-fp32-v5", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -275,12 +275,12 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( "cutlass-56ca28a-fp32-simt-64x64x8-triton-fused-tree32-" - "coalesced-pack-triton-bias-h3-v4" + "strided-actual-k-triton-bias-h3-v5" ), ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v4", + backend_id="triton-video-input-projection-fp32-v5", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -304,13 +304,13 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "strict_observable": True, "training_mode": "eager", "parameter_gradient_order": "logical_key", - "reduction": "fixed_midpoint_tree_32", + "reduction": "fixed_128_chunks_balanced_fp32_merge", }, lifecycle=OperatorLifecycle.REQUEST, implementation_class_or_factory=OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( - "rlkernel-triton-fp32-no-fma-paired-tree32-leaf64x32-short-loop-bias-epilogue-h3-v4" + "rlkernel-triton-ieee-fp32-dot-chunk128-merge16-strided-bias-epilogue-h3-v5" ), ), OperatorBackendDescriptor( diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 9218c6279..c6ab8e508 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -13,6 +13,7 @@ from rl_engine.ops.autograd.canonical_backward import canonical_backward_session from rl_engine.ops.gemm.canonical_linear import canonical_linear_fp32 from rl_engine.ops.gemm.video_input_projection_fp32 import VideoInputProjectionFP32Op +from rl_engine.reference.gemm.video_input_projection_fp32 import NativeVideoInputProjectionFP32Op H20 = ( torch.cuda.is_available() @@ -192,7 +193,7 @@ def test_forward_and_actual_gradients(backend, rows, with_bias): actual = _run(_op(backend), data, _keys(rows, "cuda")) cx, cw = data[0].detach().cpu().requires_grad_(), data[1].detach().cpu().requires_grad_() cb = None if not with_bias else data[2].detach().cpu().requires_grad_() - gold = torch.nn.functional.linear(cx, cw, cb) + gold = NativeVideoInputProjectionFP32Op()(cx, cw, cb) gold.backward(data[3].cpu()) expected = (gold, cx.grad, cw.grad) if cb is None else (gold, cx.grad, cw.grad, cb.grad) for i, (a, b) in enumerate(zip(actual, expected, strict=True)): @@ -305,7 +306,7 @@ def test_real_weights_patch_path(backend): grad = torch.randn(x.shape[0], 5376, device="cuda") * 0.0625 actual = _run(_op(backend), (x, weight, bias, grad), _keys(x.shape[0], "cuda")) cx, cw, cb = [t.detach().cpu().requires_grad_() for t in (x, weight, bias)] - gold = torch.nn.functional.linear(cx, cw, cb) + gold = NativeVideoInputProjectionFP32Op()(cx, cw, cb) gold.backward(grad.cpu()) for i, (a, b) in enumerate(zip(actual, (gold, cx.grad, cw.grad, cb.grad), strict=True)): _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") @@ -325,7 +326,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v4" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v5" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op("triton")(x, weight, bias)) @@ -396,18 +397,49 @@ def test_fp32_tree_workspace_chunking(backend, monkeypatch): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") @pytest.mark.parametrize("backend", BACKENDS) -def test_standard_normal_long_reduction(backend): +@pytest.mark.parametrize("rows", [129, 13104]) +def test_standard_normal_long_reduction(backend, rows): # Unscaled inputs expose cancellation error in the 5376-term input VJP. - data = _data(129) + data = _data(rows) data = tuple(t / scale for t, scale in zip(data, (0.125, 0.0625, 0.125, 0.0625), strict=True)) - actual = _run(_op(backend), data, _keys(129, "cuda")) + actual = _run(_op(backend), data, _keys(rows, "cuda")) cx, cw, cb = [t.detach().cpu().requires_grad_() for t in data[:3]] - gold = torch.nn.functional.linear(cx, cw, cb) + gold = NativeVideoInputProjectionFP32Op()(cx, cw, cb) gold.backward(data[3].cpu()) for i, (a, b) in enumerate(zip(actual, (gold, cx.grad, cw.grad, cb.grad), strict=True)): _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") +def test_cpu_fp32_reference_long_dot_accuracy(): + from rl_engine.reference.gemm.video_input_projection_fp32 import _reference_matmul_fp32 + + generator = torch.Generator().manual_seed(420) + a = torch.randn(3, 13104, generator=generator) + b = torch.randn(13104, 5, generator=generator) + expected = (a.double() @ b.double()).float() + actual = _reference_matmul_fp32(a, b) + assert actual.dtype == torch.float32 + _assert_accuracy(actual, expected, "gradient_accuracy") + + +@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("k", [129, 5376, 13104]) +def test_fp32_chunked_gemm_strides_and_workspace(k, monkeypatch): + from rl_engine.backends.shared.triton.gemm import det_gemm + from rl_engine.reference.gemm.video_input_projection_fp32 import _reference_matmul_fp32 + + torch.manual_seed(420) + a = torch.randn(k, 7, device="cuda").t() + b = torch.randn(k, 96, device="cuda") + expected = det_gemm._triton_chunked_gemm_fp32(a.contiguous(), b) + actual = det_gemm._triton_chunked_gemm_fp32(a, b) + _assert_bytes(actual, expected) + _assert_accuracy(actual, _reference_matmul_fp32(a, b), "gradient_accuracy") + chunks = (k + 127) // 128 + monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", chunks * 96 * 2) + _assert_bytes(det_gemm._triton_chunked_gemm_fp32(a, b), expected) + + @pytest.mark.parametrize("bias_uses", ["all", "partial", "distinct"]) def test_canonical_bias_sharing_requires_matching_contributions(bias_uses, monkeypatch): from rl_engine.ops.autograd.canonical_backward import CanonicalBackwardSession diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index a8a110587..41789b98b 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,7 +13,7 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v4" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v5" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) From ca13589cc204eddba1de598251fdaa3486e28b1b Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Thu, 15 Oct 2026 17:47:40 +0800 Subject: [PATCH 06/26] Optimize H3 FP32 projection chunks and Triton tiles Signed-off-by: walker-ai <2398833647@qq.com> --- .../video_input_projection_fp32_cutlass.cpp | 13 ++-- .../video_input_projection_fp32_cutlass.cu | 64 +++++++++---------- docs/operators/video-input-projection-fp32.md | 42 ++++++------ rl_engine/_C_cutlass.pyi | 4 +- .../cuda/gemm/video_input_projection_fp32.py | 36 +++++++++-- .../backends/shared/triton/gemm/det_gemm.py | 63 +++++++++++------- .../gemm/video_input_projection_fp32.py | 8 ++- rl_engine/ops/gemm/canonical_linear.py | 6 +- rl_engine/runtime/registry.py | 19 ++++-- .../gemm/test_video_input_projection_fp32.py | 18 ++++-- tests/runtime/test_kernel_registry.py | 8 ++- 11 files changed, 172 insertions(+), 109 deletions(-) diff --git a/csrc/bindings/video_input_projection_fp32_cutlass.cpp b/csrc/bindings/video_input_projection_fp32_cutlass.cpp index 9923611d2..2b2d3c806 100644 --- a/csrc/bindings/video_input_projection_fp32_cutlass.cpp +++ b/csrc/bindings/video_input_projection_fp32_cutlass.cpp @@ -5,15 +5,15 @@ torch::Tensor cutlass_gemm_fp32_tn(torch::Tensor a, torch::Tensor b_transposed); int cutlass_cuda_version(); -void cutlass_gemm_fp32_tree_leaves_strided( +void cutlass_gemm_fp32_partials_strided( torch::Tensor a, torch::Tensor bt, torch::Tensor workspace, - torch::Tensor starts, torch::Tensor nodes, torch::Tensor leaf_ids, int64_t length); + torch::Tensor starts, torch::Tensor nodes, torch::Tensor partial_ids, int64_t length); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("gemm_fp32_tn", &cutlass_gemm_fp32_tn, "Fixed FP32 SIMT GEMM: [M,K] times transposed [N,K]"); - m.def("gemm_fp32_tree_leaves_strided", &cutlass_gemm_fp32_tree_leaves_strided, - "Evaluate one actual-K leaf group directly from matrix views"); + m.def("gemm_fp32_partials_strided", &cutlass_gemm_fp32_partials_strided, + "Evaluate one actual-K partial group directly from matrix views"); m.def("build_info", []() { pybind11::dict info; info["cutlass_revision"] = "56ca28adb32ffadd097566545b8f7d83b652e998"; @@ -26,8 +26,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { info["fmad"] = false; info["ftz"] = false; info["tree_leaf_size"] = 32; - info["tree_leaf_launcher"] = "cutlass_device_gemm_array"; - info["tree_operand_packing"] = "direct_strides_grouped_actual_k"; + info["vjp_chunk_k"] = 128; + info["partial_launch"] = "cutlass_device_gemm_array"; + info["partial_operand_layout"] = "direct_strides_grouped_actual_k"; return info; }); } diff --git a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu index 3e942634d..e9fe81d3c 100644 --- a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu +++ b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu @@ -26,7 +26,7 @@ using Gemm = cutlass::gemm::device::Gemm< cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 2, 1, 1, false, cutlass::arch::OpMultiplyAdd>; -constexpr int kTreeLeaf = 32; +constexpr int kMaxPartialK = 128; at::Tensor cutlass_gemm_fp32_tn(at::Tensor a, at::Tensor b_transposed) { TORCH_CHECK(a.is_cuda() && b_transposed.is_cuda(), "CUTLASS GEMM requires CUDA tensors"); @@ -64,7 +64,7 @@ at::Tensor cutlass_gemm_fp32_tn(at::Tensor a, at::Tensor b_transposed) { int cutlass_cuda_version() { return CUDART_VERSION; } template -using StridedLeafGemm = cutlass::gemm::device::GemmArray< +using StridedPartialGemm = cutlass::gemm::device::GemmArray< float, LayoutA, float, LayoutB, float, cutlass::layout::RowMajor, float, cutlass::arch::OpClassSimt, cutlass::arch::Sm50, cutlass::gemm::GemmShape<64, 64, 8>, @@ -74,27 +74,27 @@ using StridedLeafGemm = cutlass::gemm::device::GemmArray< cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, 2, 1, 1, cutlass::arch::OpMultiplyAdd>; -__global__ void tree_leaf_pointers( +__global__ void partial_pointers( const float* a, const float* bt, float* workspace, int64_t* pointers, - const int32_t* starts, const int32_t* nodes, const int32_t* leaf_ids, + const int32_t* starts, const int32_t* nodes, const int32_t* partial_ids, int64_t count, int64_t stride_ak, int64_t stride_btk, int64_t elements) { const int index = blockIdx.x * blockDim.x + threadIdx.x; if (index >= count) return; - const int leaf = leaf_ids[index]; - const int64_t start = starts[leaf]; + const int partial = partial_ids[index]; + const int64_t start = starts[partial]; pointers[index] = reinterpret_cast(a + start * stride_ak); pointers[count + index] = reinterpret_cast(bt + start * stride_btk); pointers[2 * count + index] = reinterpret_cast( - workspace + static_cast(nodes[leaf]) * elements); + workspace + static_cast(nodes[partial]) * elements); } template -void launch_strided_leaves( +void launch_strided_partials( const at::Tensor& a, const at::Tensor& bt, const at::Tensor& pointers, int64_t count, int length, int64_t lda, int64_t ldb, cudaStream_t stream) { - using Leaf = StridedLeafGemm; + using Partial = StridedPartialGemm; const auto* addresses = pointers.data_ptr(); - typename Leaf::Arguments args( + typename Partial::Arguments args( {static_cast(a.size(0)), static_cast(bt.size(0)), length}, reinterpret_cast(addresses), LayoutA(lda), reinterpret_cast(addresses + count), LayoutB(ldb), @@ -102,69 +102,69 @@ void launch_strided_leaves( cutlass::layout::RowMajor(bt.size(0)), reinterpret_cast(addresses + 2 * count), cutlass::layout::RowMajor(bt.size(0)), {1.0f, 0.0f}, static_cast(count)); - Leaf gemm; + Partial gemm; auto status = gemm.can_implement(args); TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); status = gemm(args, nullptr, stream); TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); } -void cutlass_gemm_fp32_tree_leaves_strided( +void cutlass_gemm_fp32_partials_strided( at::Tensor a, at::Tensor bt, at::Tensor workspace, at::Tensor starts, - at::Tensor nodes, at::Tensor leaf_ids, int64_t length) { + at::Tensor nodes, at::Tensor partial_ids, int64_t length) { TORCH_CHECK(a.is_cuda() && a.scalar_type() == at::kFloat && a.dim() == 2, - "CUTLASS leaves require rank-two CUDA FP32 A"); + "CUTLASS partials require rank-two CUDA FP32 A"); const c10::cuda::CUDAGuard guard(a.device()); const auto* properties = at::cuda::getDeviceProperties(a.get_device()); TORCH_CHECK(properties->major == 9 && properties->minor == 0, "CUTLASS video projection requires CUDA SM90"); TORCH_CHECK(bt.device() == a.device() && bt.scalar_type() == at::kFloat && bt.dim() == 2 && a.size(1) == bt.size(1), - "CUTLASS leaf dimensions, dtypes and devices must match"); - TORCH_CHECK(a.size(0) > 0 && bt.size(0) > 0 && length > 0 && length <= kTreeLeaf && + "CUTLASS partial dimensions, dtypes and devices must match"); + TORCH_CHECK(a.size(0) > 0 && bt.size(0) > 0 && length > 0 && length <= kMaxPartialK && a.size(0) <= INT_MAX && bt.size(0) <= INT_MAX, - "CUTLASS leaf sizes exceed the supported dimensions"); + "CUTLASS partial sizes exceed the supported dimensions"); TORCH_CHECK(workspace.device() == a.device() && workspace.scalar_type() == at::kFloat && workspace.dim() == 3 && workspace.is_contiguous() && workspace.size(1) == a.size(0) && workspace.size(2) == bt.size(0), - "CUTLASS leaves require contiguous FP32 tree workspace"); - for (const auto& metadata : {starts, nodes, leaf_ids}) { + "CUTLASS partials require contiguous FP32 tree workspace"); + for (const auto& metadata : {starts, nodes, partial_ids}) { TORCH_CHECK(metadata.device() == a.device() && metadata.scalar_type() == at::kInt && metadata.dim() == 1 && metadata.is_contiguous(), - "CUTLASS leaf metadata must be contiguous CUDA int32 vectors"); + "CUTLASS partial metadata must be contiguous CUDA int32 vectors"); } - TORCH_CHECK(starts.numel() == nodes.numel() && leaf_ids.numel() > 0 && - leaf_ids.numel() <= 65535, - "CUTLASS leaf metadata sizes are inconsistent"); + TORCH_CHECK(starts.numel() == nodes.numel() && partial_ids.numel() > 0 && + partial_ids.numel() <= 65535, + "CUTLASS partial metadata sizes are inconsistent"); const bool a_row = a.stride(1) == 1; const bool b_column = bt.stride(1) == 1; TORCH_CHECK((a_row || a.stride(0) == 1) && (b_column || bt.stride(0) == 1) && a.stride(0) > 0 && a.stride(1) > 0 && bt.stride(0) > 0 && bt.stride(1) > 0, - "CUTLASS leaves require a row-major or column-major view"); + "CUTLASS partials require a row-major or column-major view"); const int64_t lda = a_row ? a.stride(0) : a.stride(1); const int64_t ldb = b_column ? bt.stride(0) : bt.stride(1); TORCH_CHECK(lda <= INT_MAX && ldb <= INT_MAX, - "CUTLASS leaf leading dimensions exceed int32"); - const int64_t count = leaf_ids.numel(); + "CUTLASS partial leading dimensions exceed int32"); + const int64_t count = partial_ids.numel(); auto pointers = at::empty({3, count}, a.options().dtype(at::kLong)); const auto stream = c10::cuda::getCurrentCUDAStream(a.get_device()).stream(); - tree_leaf_pointers<<<(count + 255) / 256, 256, 0, stream>>>( + partial_pointers<<<(count + 255) / 256, 256, 0, stream>>>( a.data_ptr(), bt.data_ptr(), workspace.data_ptr(), pointers.data_ptr(), starts.data_ptr(), nodes.data_ptr(), - leaf_ids.data_ptr(), count, a.stride(1), bt.stride(1), + partial_ids.data_ptr(), count, a.stride(1), bt.stride(1), a.size(0) * bt.size(0)); using Row = cutlass::layout::RowMajor; using Column = cutlass::layout::ColumnMajor; if (a_row && b_column) { - launch_strided_leaves(a, bt, pointers, count, length, lda, ldb, stream); + launch_strided_partials(a, bt, pointers, count, length, lda, ldb, stream); } else if (a_row) { - launch_strided_leaves(a, bt, pointers, count, length, lda, ldb, stream); + launch_strided_partials(a, bt, pointers, count, length, lda, ldb, stream); } else if (b_column) { - launch_strided_leaves(a, bt, pointers, count, length, lda, ldb, stream); + launch_strided_partials(a, bt, pointers, count, length, lda, ldb, stream); } else { - launch_strided_leaves(a, bt, pointers, count, length, lda, ldb, stream); + launch_strided_partials(a, bt, pointers, count, length, lda, ldb, stream); } C10_CUDA_KERNEL_LAUNCH_CHECK(); } diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index f82b2072b..62a33442d 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -9,8 +9,9 @@ reference. The model adapter owns the subsequent cast into the BF16 packed strea The primary implementation reuses RL-Kernel's Triton FP32 GEMM and canonical parameter-gradient session. Input and weight VJPs use fixed 128-term K chunks and one balanced FP32 merge. A pinned CUTLASS FP32 SIMT GEMM is an explicit -CUDA candidate. It retains the midpoint tree and reads each leaf directly -from matrix views. Both paths share the semantic adapter and canonical session. +CUDA candidate. It uses the same fixed-chunk scheduler and balanced merge, +and reads each partial directly from matrix views. Both paths share the +semantic adapter and canonical session. Long bias VJPs retain the one-row Triton tree. ## Entry Point @@ -35,7 +36,7 @@ physical placement. The same parameter keeps the same `parameter_id` across uses | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | Triton on CUDA | `TritonVideoInputProjectionFP32Op` | `_det_gemm_fp32_kernel`, `_det_gemm_fp32_chunk_merge_kernel`, bias tree kernels | Default CUDA path; SM90 qualification | -| CUDA / CUTLASS + Triton reductions | `CUDAVideoInputProjectionFP32Op` | `_C_cutlass.gemm_fp32_tn`, `_C_cutlass.gemm_fp32_tree_leaves_strided`, shared tree reductions | Explicit candidate; optional build; SM90 | +| CUDA / CUTLASS + Triton reductions | `CUDAVideoInputProjectionFP32Op` | `_C_cutlass.gemm_fp32_tn`, `_C_cutlass.gemm_fp32_partials_strided`, shared chunk merge and bias tree | Explicit candidate; optional build; SM90 | | PyTorch reference | `NativeVideoInputProjectionFP32Op` | CPU `F.linear`; FP32 product/sum VJPs | Accuracy reference | CUTLASS revision: `56ca28adb32ffadd097566545b8f7d83b652e998`. The build verifies @@ -101,8 +102,10 @@ GEMM. The Triton input and weight VJPs reuse that kernel with fixed 128-term chunks. Each chunk computes four ascending 32-term FP32 dots from zero. Chunk partials remain FP32. A second grid uses a fixed balanced reduction over the next power-of-two chunk count, with zero padding. The merge uses 16 output -elements per program and four warps. These settings define the v5 arithmetic -contract. The K partition does not depend on M, row placement, or occupancy. +elements per program and four warps. VJP GEMMs use fixed 128-by-32 output tiles, +32-term dot tiles, and four warps. The forward retains its 64-by-64 output tile. +These settings define the v6 implementation identity. The K partition and VJP +tile do not depend on M, row placement, or occupancy. The VJP kernels consume positive-stride views directly, including `G.T`. The fixed-chunk design follows the scheduling separation in merged PR #396. @@ -110,18 +113,20 @@ Its AMD BF16 MFMA instructions are not used by this CUDA FP32 implementation. The FP32 GEMM body is reused from the helper present in merged PR #325. The short forward path and long bias tree retain their earlier arithmetic. -CUTLASS retains the existing midpoint plan with at most 32 terms per leaf. -Leaves are grouped by actual K length. A pointer kernel selects each leaf's -original operand range, and a fixed `device::GemmArray` reads row-major or -column-major views directly. This applies the layout approach from merged -PR #343. It avoids packed operand buffers while retaining the FP32 midpoint -merge. The actual-K grouping and layouts are recorded in `backend_info()`. -CUTLASS still uses Triton for tree merging and long bias VJPs. - -Both schedules cap their FP32 partial workspace at 512 MiB by chunking output -rows. The midpoint implementation also supports output-column chunking when -one complete row exceeds its budget. CUTLASS allocates pointer arrays for each -leaf group. Parameter sorting and public validation are included in complete +CUTLASS also uses fixed 128-term chunks for input and weight VJPs. Partials +are grouped by actual K length: 128 for complete chunks, and the remaining +length for the tail. A pointer kernel selects each partial's original operand +range, and the pinned `device::GemmArray` reads row-major or column-major views +directly. This applies the layout approach from merged PR #343. The shared +chunk scheduler calls this SDK launcher and then the same balanced Triton merge. +The SDK retains its separate FP32 multiply/add arithmetic inside each partial. +This schedule defines the CUTLASS v6 identity. Its actual-K grouping and layouts +are recorded in `backend_info()`. Long bias VJPs retain the original Triton tree. + +Both VJP schedules cap their FP32 partial workspace at 512 MiB by chunking +output rows. The retained midpoint helper also supports output-column chunking +when one complete row exceeds its budget. CUTLASS allocates pointer arrays for +each partial group. Parameter sorting and public validation are included in complete operator measurements. Existing BF16 callers retain their dtype and schedule. TF32 and compiler multiply-add fusion remain disabled. The reused Triton `tl.dot` lowers its IEEE FP32 inner products to fixed SIMT @@ -159,7 +164,8 @@ parameter-gradient machinery. Raw samples, environment, source hashes, and actual backend metadata are written to the explicitly selected local output. PyTorch is a production performance baseline with separately reported BI results. The Triton VJPs use one partial-GEMM launch and one merge launch when the -workspace fits. CUTLASS retains tree merge launches but removes operand packing. +workspace fits. CUTLASS launches the SDK partial groups, their pointer metadata, +and one balanced merge. A tail can require a separate SDK launch. Compare complete operator timings, including sorting, before drawing a speedup claim. The benchmark records the reference identity and its source hash. diff --git a/rl_engine/_C_cutlass.pyi b/rl_engine/_C_cutlass.pyi index de2152e41..890103a31 100644 --- a/rl_engine/_C_cutlass.pyi +++ b/rl_engine/_C_cutlass.pyi @@ -6,13 +6,13 @@ from typing import Any import torch def gemm_fp32_tn(a: torch.Tensor, b_transposed: torch.Tensor) -> torch.Tensor: ... -def gemm_fp32_tree_leaves_strided( +def gemm_fp32_partials_strided( a: torch.Tensor, b_transposed: torch.Tensor, workspace: torch.Tensor, starts: torch.Tensor, nodes: torch.Tensor, - leaf_ids: torch.Tensor, + partial_ids: torch.Tensor, length: int, ) -> None: ... def build_info() -> dict[str, Any]: ... diff --git a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py index 32e20952b..6cf7f7c21 100644 --- a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py @@ -37,7 +37,9 @@ from rl_engine.backends.shared.triton.gemm.det_gemm import ( _BLOCK_K, + _FP32_CHUNK_K, _TRITON_AVAILABLE, + _triton_chunked_gemm_fp32, _triton_tree_gemm, ) from rl_engine.ops.gemm.video_input_projection_fp32 import ( @@ -65,7 +67,8 @@ def _load_extension(): or info["ftz"] or info["split_k"] != 1 or info.get("tree_leaf_size") != _BLOCK_K - or info.get("tree_operand_packing") != "direct_strides_grouped_actual_k" + or info.get("vjp_chunk_k") != _FP32_CHUNK_K + or info.get("partial_operand_layout") != "direct_strides_grouped_actual_k" ): raise RuntimeError("CUTLASS extension does not match the strict FP32 build contract") return extension @@ -73,7 +76,7 @@ def _load_extension(): def _cutlass_tree_leaves(a, b, workspace, plan): for length, leaf_ids in _leaf_groups(plan.host.leaf_lengths, a.device): - _load_extension().gemm_fp32_tree_leaves_strided( + _load_extension().gemm_fp32_partials_strided( a, b.t(), workspace, plan.leaf_starts, plan.leaf_nodes, leaf_ids, length ) @@ -93,10 +96,26 @@ def _leaf_groups(lengths, device): ) +@lru_cache(maxsize=None) +def _chunk_metadata(k_size, device): + full, tail = divmod(k_size, _FP32_CHUNK_K) + lengths = (_FP32_CHUNK_K,) * full + ((tail,) if tail else ()) + nodes = torch.arange(len(lengths), device=device, dtype=torch.int32) + return nodes * _FP32_CHUNK_K, nodes, _leaf_groups(lengths, device) + + +def _cutlass_chunks(a, b, partial): + starts, nodes, groups = _chunk_metadata(a.shape[1], a.device) + for length, ids in groups: + _load_extension().gemm_fp32_partials_strided(a, b.t(), partial, starts, nodes, ids, length) + + def cutlass_gemm_fp32(a, b): # The native TN interface represents mathematical B[K,N] as [N,K] storage. if a.shape[1] <= VIDEO_PATCH_WIDTH: return _load_extension().gemm_fp32_tn(a.contiguous(), b.t().contiguous()) + if b.shape[1] == VIDEO_PATCH_WIDTH: + return _triton_chunked_gemm_fp32(a, b, fp32_chunk_launcher=_cutlass_chunks) return _triton_tree_gemm( a, b, @@ -108,7 +127,7 @@ def cutlass_gemm_fp32(a, b): class CUDAVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "cutlass" - backend_id = "cuda-cutlass-video-input-projection-fp32-v5" + backend_id = "cuda-cutlass-video-input-projection-fp32-v6" def __init__(self): if not _TRITON_AVAILABLE: @@ -127,14 +146,17 @@ def backend_info(self): return { "backend_id": self.backend_id, "family": "cuda", - "kernel": "gemm_fp32_tn / gemm_fp32_tree_leaves_strided / " + "kernel": "gemm_fp32_tn / gemm_fp32_partials_strided / " + "_det_gemm_fp32_chunk_merge_kernel / " "_det_gemm_kernel / _det_gemm_tree_leaf_reduce_rocm_kernel", - "reduction": "fixed_midpoint_tree", + "reduction": "fixed_128_chunks_balanced_fp32_merge", + "vjp_chunk_k": _FP32_CHUNK_K, + "vjp_merge_tile": 16, + "bias_reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, "reduction_backend": "triton", "bias_reduction_backend": "triton_for_k_gt_96", "triton_version": triton.__version__, - "leaf_launch": "cutlass_device_gemm_array", - "tree_schedule": "paired_levels_compact_workspace", + "bias_tree_schedule": "paired_levels_compact_workspace", **_load_extension().build_info(), } diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index c5db36fc5..e3571d21f 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -804,7 +804,13 @@ def _det_gemm_fp32_chunk_merge_kernel( tl.store(output_ptr + offsets, total, mask=offsets < ELEMENTS) -def _triton_chunked_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: +def _triton_chunked_gemm_fp32( + a: torch.Tensor, + b: torch.Tensor, + *, + tile: tuple[int, int, int] = (_BLOCK_M, _BLOCK_N, 4), + fp32_chunk_launcher=None, +) -> torch.Tensor: """Reuse the FP32 GEMM in fixed K chunks and merge partials in one grid.""" m_size, k_size = a.shape n_size = b.shape[1] @@ -815,35 +821,44 @@ def _triton_chunked_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: if m_size > rows_per_chunk: return torch.cat( [ - _triton_chunked_gemm_fp32(a[start : start + rows_per_chunk], b) + _triton_chunked_gemm_fp32( + a[start : start + rows_per_chunk], + b, + tile=tile, + fp32_chunk_launcher=fp32_chunk_launcher, + ) for start in range(0, m_size, rows_per_chunk) ], dim=0, ) partial = torch.empty((chunks, m_size, n_size), device=a.device, dtype=torch.float32) output = torch.empty((m_size, n_size), device=a.device, dtype=torch.float32) - grid = (triton.cdiv(m_size, _BLOCK_M), triton.cdiv(n_size, _BLOCK_N), chunks) - _det_gemm_fp32_kernel[grid]( - a, - b, - partial, - None, - m_size, - n_size, - k_size, - *a.stride(), - *b.stride(), - n_size, - 1, - BLOCK_M=_BLOCK_M, - BLOCK_N=_BLOCK_N, - BLOCK_K=_BLOCK_K, - PROMOTE_INPUTS=True, - HAS_BIAS=False, - CHUNK_K=_FP32_CHUNK_K, - enable_fp_fusion=False, - num_warps=4, - ) + if fp32_chunk_launcher is None: + block_m, block_n, num_warps = tile + grid = (triton.cdiv(m_size, block_m), triton.cdiv(n_size, block_n), chunks) + _det_gemm_fp32_kernel[grid]( + a, + b, + partial, + None, + m_size, + n_size, + k_size, + *a.stride(), + *b.stride(), + n_size, + 1, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=_BLOCK_K, + PROMOTE_INPUTS=True, + HAS_BIAS=False, + CHUNK_K=_FP32_CHUNK_K, + enable_fp_fusion=False, + num_warps=num_warps, + ) + else: + fp32_chunk_launcher(a, b, partial) _det_gemm_fp32_chunk_merge_kernel[(triton.cdiv(m_size * n_size, 16),)]( partial, output, diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index e39b88942..ae2bc4486 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -17,12 +17,14 @@ VideoInputProjectionFP32Op, ) +_VJP_TILE = (128, 32, 4) + def projection_gemm_fp32(a, b, *, bias=None): if a.shape[1] <= VIDEO_PATCH_WIDTH: return _triton_gemm_fp32(a, b, bias=bias) if b.shape[1] == VIDEO_PATCH_WIDTH: - output = _triton_chunked_gemm_fp32(a, b) + output = _triton_chunked_gemm_fp32(a, b, tile=_VJP_TILE) else: output = _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) return output if bias is None else output + bias @@ -30,7 +32,7 @@ def projection_gemm_fp32(a, b, *, bias=None): class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton_tree" - backend_id = "triton-video-input-projection-fp32-v5" + backend_id = "triton-video-input-projection-fp32-v6" def __init__(self): if not _TRITON_AVAILABLE: @@ -53,6 +55,8 @@ def backend_info(self): "_det_gemm_tree_leaf_reduce_rocm_kernel", "triton_version": triton.__version__, "tile": [64, 64, 32], + "vjp_tile": [_VJP_TILE[0], _VJP_TILE[1], _BLOCK_K], + "vjp_num_warps": _VJP_TILE[2], "tf32": False, "fp_fusion": False, "dot_arithmetic": "ieee_fp32_simt_fma", diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index b0bd9a08e..9e1f4da81 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -149,10 +149,10 @@ def bias_reducer(grads): record_backward( "video_input_projection_fp32", kernel_id=( - "rl_engine._C_cutlass.gemm_fp32_tree_leaves_strided / " - "rl_engine.backends.shared.triton.gemm.det_gemm._triton_tree_gemm" + "rl_engine._C_cutlass.gemm_fp32_partials_strided / " + "rl_engine.backends.shared.triton.gemm.det_gemm._triton_chunked_gemm_fp32" ), - impl="cutlass_fp32_simt_triton_tree_bias_canonical_rowfold", + impl="cutlass_fp32_simt_chunk128_triton_merge_bias_canonical_rowfold", family="cuda", ) return da, dweight, None, None, None, dbias, None diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 405a30847..00f467411 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -244,7 +244,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="cuda-cutlass-video-input-projection-fp32-v5", + backend_id="cuda-cutlass-video-input-projection-fp32-v6", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -257,7 +257,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "training": {"world_size": (1,), "sharding": ("unsharded",)}, }, determinism_or_alignment_properties={ - "algorithm": "cutlass_simt_fp32_triton_tree_bias_canonical_parameter_vjp", + "algorithm": "cutlass_simt_fp32_chunk128_triton_merge_bias_canonical_parameter_vjp", "reduction_backend": "triton", "bias_reduction_backend": "triton_for_k_gt_96", "input_width": 96, @@ -268,19 +268,21 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "strict_observable": True, "training_mode": "eager", "parameter_gradient_order": "logical_key", - "reduction": "fixed_midpoint_tree_32", + "reduction": "fixed_128_chunks_balanced_fp32_merge", + "vjp_chunk_k": 128, + "bias_reduction": "fixed_midpoint_tree_32", }, lifecycle=OperatorLifecycle.REQUEST, implementation_class_or_factory=OpBackend.CUDA_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( - "cutlass-56ca28a-fp32-simt-64x64x8-triton-fused-tree32-" - "strided-actual-k-triton-bias-h3-v5" + "cutlass-56ca28a-fp32-simt-64x64x8-chunk128-merge16-" + "strided-actual-k-triton-bias-h3-v6" ), ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v5", + backend_id="triton-video-input-projection-fp32-v6", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -305,12 +307,15 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "training_mode": "eager", "parameter_gradient_order": "logical_key", "reduction": "fixed_128_chunks_balanced_fp32_merge", + "vjp_tile": (128, 32, 32), + "vjp_num_warps": 4, }, lifecycle=OperatorLifecycle.REQUEST, implementation_class_or_factory=OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( - "rlkernel-triton-ieee-fp32-dot-chunk128-merge16-strided-bias-epilogue-h3-v5" + "rlkernel-triton-ieee-fp32-dot-chunk128-tile128x32-merge16-" + "strided-bias-epilogue-h3-v6" ), ), OperatorBackendDescriptor( diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index c6ab8e508..356c70dfd 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -202,7 +202,7 @@ def test_forward_and_actual_gradients(backend, rows, with_bias): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") @pytest.mark.parametrize("backend", BACKENDS) -@pytest.mark.parametrize("rows", [9, 65, 129]) +@pytest.mark.parametrize("rows", [9, 65, 127, 128, 129, 257]) def test_parameter_bi_permutation_chunk_padding(backend, rows): op, data, keys = _op(backend), _data(rows), _keys(rows, "cuda") baseline = _run(op, data, keys) @@ -326,7 +326,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v5" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v6" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op("triton")(x, weight, bias)) @@ -423,21 +423,25 @@ def test_cpu_fp32_reference_long_dot_accuracy(): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("k", [129, 5376, 13104]) -def test_fp32_chunked_gemm_strides_and_workspace(k, monkeypatch): +@pytest.mark.parametrize("backend", BACKENDS) +@pytest.mark.parametrize("k", [127, 128, 129, 5376, 13104]) +def test_fp32_chunked_gemm_strides_and_workspace(backend, k, monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm + from rl_engine.ops.gemm.canonical_linear import _gemm_fp32 from rl_engine.reference.gemm.video_input_projection_fp32 import _reference_matmul_fp32 torch.manual_seed(420) + _op(backend) + family = "triton_tree" if backend == "triton" else "cutlass" a = torch.randn(k, 7, device="cuda").t() b = torch.randn(k, 96, device="cuda") - expected = det_gemm._triton_chunked_gemm_fp32(a.contiguous(), b) - actual = det_gemm._triton_chunked_gemm_fp32(a, b) + expected = _gemm_fp32(a.contiguous(), b, family) + actual = _gemm_fp32(a, b, family) _assert_bytes(actual, expected) _assert_accuracy(actual, _reference_matmul_fp32(a, b), "gradient_accuracy") chunks = (k + 127) // 128 monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", chunks * 96 * 2) - _assert_bytes(det_gemm._triton_chunked_gemm_fp32(a, b), expected) + _assert_bytes(_gemm_fp32(a, b, family), expected) @pytest.mark.parametrize("bias_uses", ["all", "partial", "distinct"]) diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 41789b98b..5243c57a7 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,10 +13,16 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v5" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v6" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) + assert desc.determinism_or_alignment_properties["vjp_tile"] == (128, 32, 32) + cutlass = registry.semantic.backend_descriptor( + "video_input_projection_fp32", "cuda-cutlass-video-input-projection-fp32-v6" + ) + assert cutlass.determinism_or_alignment_properties["vjp_chunk_k"] == 128 + assert cutlass.determinism_or_alignment_properties["bias_reduction"] == "fixed_midpoint_tree_32" with pytest.raises(RuntimeError, match="SM90"): registry.get_op("video_input_projection_fp32", device="cpu") From 541c737d633b134ea1c046c40341aaf02bee56a3 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Fri, 16 Oct 2026 12:16:50 +0800 Subject: [PATCH 07/26] Fuse H3 CUTLASS bias and widen FP32 chunk merges Signed-off-by: walker-ai <2398833647@qq.com> --- .../video_input_projection_fp32_cutlass.cpp | 6 +++++- .../video_input_projection_fp32_cutlass.cu | 16 ++++++++++++--- docs/operators/video-input-projection-fp32.md | 13 +++++++++--- rl_engine/_C_cutlass.pyi | 4 +++- .../cuda/gemm/video_input_projection_fp32.py | 11 ++++++---- .../backends/shared/triton/gemm/det_gemm.py | 4 ++-- .../gemm/video_input_projection_fp32.py | 4 ++-- rl_engine/ops/gemm/canonical_linear.py | 6 +++--- rl_engine/runtime/registry.py | 13 ++++++------ .../gemm/test_video_input_projection_fp32.py | 20 +++++++++++++------ tests/runtime/test_kernel_registry.py | 4 ++-- 11 files changed, 68 insertions(+), 33 deletions(-) diff --git a/csrc/bindings/video_input_projection_fp32_cutlass.cpp b/csrc/bindings/video_input_projection_fp32_cutlass.cpp index 2b2d3c806..eb803874c 100644 --- a/csrc/bindings/video_input_projection_fp32_cutlass.cpp +++ b/csrc/bindings/video_input_projection_fp32_cutlass.cpp @@ -2,8 +2,10 @@ // Copyright (c) 2026 RL-Kernel Contributors #include +#include -torch::Tensor cutlass_gemm_fp32_tn(torch::Tensor a, torch::Tensor b_transposed); +torch::Tensor cutlass_gemm_fp32_tn( + torch::Tensor a, torch::Tensor b_transposed, std::optional bias); int cutlass_cuda_version(); void cutlass_gemm_fp32_partials_strided( torch::Tensor a, torch::Tensor bt, torch::Tensor workspace, @@ -11,6 +13,7 @@ void cutlass_gemm_fp32_partials_strided( PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("gemm_fp32_tn", &cutlass_gemm_fp32_tn, + pybind11::arg("a"), pybind11::arg("b_transposed"), pybind11::arg("bias") = pybind11::none(), "Fixed FP32 SIMT GEMM: [M,K] times transposed [N,K]"); m.def("gemm_fp32_partials_strided", &cutlass_gemm_fp32_partials_strided, "Evaluate one actual-K partial group directly from matrix views"); @@ -29,6 +32,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { info["vjp_chunk_k"] = 128; info["partial_launch"] = "cutlass_device_gemm_array"; info["partial_operand_layout"] = "direct_strides_grouped_actual_k"; + info["forward_bias"] = "linear_combination_broadcast_source"; return info; }); } diff --git a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu index e9fe81d3c..7aa9b9412 100644 --- a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu +++ b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -28,7 +29,8 @@ using Gemm = cutlass::gemm::device::Gemm< constexpr int kMaxPartialK = 128; -at::Tensor cutlass_gemm_fp32_tn(at::Tensor a, at::Tensor b_transposed) { +at::Tensor cutlass_gemm_fp32_tn( + at::Tensor a, at::Tensor b_transposed, std::optional bias) { TORCH_CHECK(a.is_cuda() && b_transposed.is_cuda(), "CUTLASS GEMM requires CUDA tensors"); TORCH_CHECK(a.device() == b_transposed.device(), "CUTLASS GEMM devices must match"); TORCH_CHECK(a.scalar_type() == at::kFloat && b_transposed.scalar_type() == at::kFloat, @@ -47,11 +49,19 @@ at::Tensor cutlass_gemm_fp32_tn(at::Tensor a, at::Tensor b_transposed) { m64 <= INT_MAX && n64 <= INT_MAX && k64 <= INT_MAX, "CUTLASS GEMM dimensions must be positive int32 values"); const int m = static_cast(m64), n = static_cast(n64), k = static_cast(k64); + if (bias.has_value()) { + TORCH_CHECK(bias->device() == a.device() && bias->scalar_type() == at::kFloat && + bias->dim() == 1 && bias->size(0) == n && bias->is_contiguous(), + "CUTLASS GEMM bias must be a contiguous FP32 vector on the input device"); + } auto output = at::empty({m, n}, a.options()); + // A zero source row stride broadcasts bias through the SDK epilogue. + // --fmad=false retains the rounded FP32 GEMM followed by FP32 addition. + const float* source = bias.has_value() ? bias->data_ptr() : output.data_ptr(); Gemm::Arguments args({m, n, k}, {a.data_ptr(), k}, {b_transposed.data_ptr(), k}, - {output.data_ptr(), n}, {output.data_ptr(), n}, - {1.0f, 0.0f}, 1); + {source, bias.has_value() ? 0 : n}, {output.data_ptr(), n}, + {1.0f, bias.has_value() ? 1.0f : 0.0f}, 1); Gemm gemm; auto status = gemm.can_implement(args); TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 62a33442d..4804f9bfb 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -101,10 +101,10 @@ Short reductions up to the 96-wide video input use the existing direct FP32 GEMM. The Triton input and weight VJPs reuse that kernel with fixed 128-term chunks. Each chunk computes four ascending 32-term FP32 dots from zero. Chunk partials remain FP32. A second grid uses a fixed balanced reduction over the -next power-of-two chunk count, with zero padding. The merge uses 16 output +next power-of-two chunk count, with zero padding. The merge uses 64 output elements per program and four warps. VJP GEMMs use fixed 128-by-32 output tiles, 32-term dot tiles, and four warps. The forward retains its 64-by-64 output tile. -These settings define the v6 implementation identity. The K partition and VJP +These settings define the v7 implementation identity. The K partition and VJP tile do not depend on M, row placement, or occupancy. The VJP kernels consume positive-stride views directly, including `G.T`. @@ -112,6 +112,9 @@ The fixed-chunk design follows the scheduling separation in merged PR #396. Its AMD BF16 MFMA instructions are not used by this CUDA FP32 implementation. The FP32 GEMM body is reused from the helper present in merged PR #325. The short forward path and long bias tree retain their earlier arithmetic. +The wider merge changes the reduction instruction schedule. Both backends pin +64 output elements and four warps for the v7 identity. Accuracy and logical-row +invariance must be checked independently of bitwise agreement with older versions. CUTLASS also uses fixed 128-term chunks for input and weight VJPs. Partials are grouped by actual K length: 128 for complete chunks, and the remaining @@ -120,8 +123,12 @@ range, and the pinned `device::GemmArray` reads row-major or column-major views directly. This applies the layout approach from merged PR #343. The shared chunk scheduler calls this SDK launcher and then the same balanced Triton merge. The SDK retains its separate FP32 multiply/add arithmetic inside each partial. -This schedule defines the CUTLASS v6 identity. Its actual-K grouping and layouts +This schedule defines the CUTLASS v7 identity. Its actual-K grouping and layouts are recorded in `backend_info()`. Long bias VJPs retain the original Triton tree. +The forward reuses CUTLASS `LinearCombination` with source row stride zero to +broadcast bias. It adds bias after the FP32 accumulator in the same epilogue. +The build keeps `--fmad=false`. Tests compare this output bitwise with a separate +GEMM followed by FP32 bias addition. Both VJP schedules cap their FP32 partial workspace at 512 MiB by chunking output rows. The retained midpoint helper also supports output-column chunking diff --git a/rl_engine/_C_cutlass.pyi b/rl_engine/_C_cutlass.pyi index 890103a31..ff59a01d2 100644 --- a/rl_engine/_C_cutlass.pyi +++ b/rl_engine/_C_cutlass.pyi @@ -5,7 +5,9 @@ from typing import Any import torch -def gemm_fp32_tn(a: torch.Tensor, b_transposed: torch.Tensor) -> torch.Tensor: ... +def gemm_fp32_tn( + a: torch.Tensor, b_transposed: torch.Tensor, bias: torch.Tensor | None = None +) -> torch.Tensor: ... def gemm_fp32_partials_strided( a: torch.Tensor, b_transposed: torch.Tensor, diff --git a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py index 6cf7f7c21..ba4b18637 100644 --- a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py @@ -69,6 +69,7 @@ def _load_extension(): or info.get("tree_leaf_size") != _BLOCK_K or info.get("vjp_chunk_k") != _FP32_CHUNK_K or info.get("partial_operand_layout") != "direct_strides_grouped_actual_k" + or info.get("forward_bias") != "linear_combination_broadcast_source" ): raise RuntimeError("CUTLASS extension does not match the strict FP32 build contract") return extension @@ -110,10 +111,12 @@ def _cutlass_chunks(a, b, partial): _load_extension().gemm_fp32_partials_strided(a, b.t(), partial, starts, nodes, ids, length) -def cutlass_gemm_fp32(a, b): +def cutlass_gemm_fp32(a, b, *, bias=None): # The native TN interface represents mathematical B[K,N] as [N,K] storage. if a.shape[1] <= VIDEO_PATCH_WIDTH: - return _load_extension().gemm_fp32_tn(a.contiguous(), b.t().contiguous()) + return _load_extension().gemm_fp32_tn(a.contiguous(), b.t().contiguous(), bias) + if bias is not None: + return cutlass_gemm_fp32(a, b) + bias if b.shape[1] == VIDEO_PATCH_WIDTH: return _triton_chunked_gemm_fp32(a, b, fp32_chunk_launcher=_cutlass_chunks) return _triton_tree_gemm( @@ -127,7 +130,7 @@ def cutlass_gemm_fp32(a, b): class CUDAVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "cutlass" - backend_id = "cuda-cutlass-video-input-projection-fp32-v6" + backend_id = "cuda-cutlass-video-input-projection-fp32-v7" def __init__(self): if not _TRITON_AVAILABLE: @@ -151,7 +154,7 @@ def backend_info(self): "_det_gemm_kernel / _det_gemm_tree_leaf_reduce_rocm_kernel", "reduction": "fixed_128_chunks_balanced_fp32_merge", "vjp_chunk_k": _FP32_CHUNK_K, - "vjp_merge_tile": 16, + "vjp_merge_tile": 64, "bias_reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, "reduction_backend": "triton", diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index e3571d21f..4e3b97470 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -859,13 +859,13 @@ def _triton_chunked_gemm_fp32( ) else: fp32_chunk_launcher(a, b, partial) - _det_gemm_fp32_chunk_merge_kernel[(triton.cdiv(m_size * n_size, 16),)]( + _det_gemm_fp32_chunk_merge_kernel[(triton.cdiv(m_size * n_size, 64),)]( partial, output, m_size * n_size, CHUNKS=chunks, PADDED_CHUNKS=triton.next_power_of_2(chunks), - BLOCK_C=16, + BLOCK_C=64, enable_fp_fusion=False, num_warps=4, ) diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index ae2bc4486..7b5ec6dcf 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -32,7 +32,7 @@ def projection_gemm_fp32(a, b, *, bias=None): class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton_tree" - backend_id = "triton-video-input-projection-fp32-v6" + backend_id = "triton-video-input-projection-fp32-v7" def __init__(self): if not _TRITON_AVAILABLE: @@ -65,7 +65,7 @@ def backend_info(self): "split_k": "fixed_128_term_vjp_chunks", "reduction": "fixed_128_chunks_balanced_fp32_merge", "vjp_chunk_k": 128, - "vjp_merge_tile": 16, + "vjp_merge_tile": 64, "bias_reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, "reduction_backend": "triton", diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 9e1f4da81..6cf33b7f0 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -22,12 +22,12 @@ def _gemm_fp32( ) return projection_gemm_fp32(a, b, bias=bias) - if bias is not None: - return _gemm_fp32(a, b, family) + bias if family == "cutlass": from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import cutlass_gemm_fp32 - return cutlass_gemm_fp32(a, b) + return cutlass_gemm_fp32(a, b, bias=bias) + if bias is not None: + return _gemm_fp32(a, b, family) + bias if family == "cuda": return _C.det_gemm_rowwise_fwd_fp32(a.contiguous(), b.contiguous()) if family == "triton": diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 00f467411..3658b3e05 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -244,7 +244,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="cuda-cutlass-video-input-projection-fp32-v6", + backend_id="cuda-cutlass-video-input-projection-fp32-v7", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -258,6 +258,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: }, determinism_or_alignment_properties={ "algorithm": "cutlass_simt_fp32_chunk128_triton_merge_bias_canonical_parameter_vjp", + "forward_bias": "linear_combination_broadcast_source", "reduction_backend": "triton", "bias_reduction_backend": "triton_for_k_gt_96", "input_width": 96, @@ -276,13 +277,13 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: implementation_class_or_factory=OpBackend.CUDA_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( - "cutlass-56ca28a-fp32-simt-64x64x8-chunk128-merge16-" - "strided-actual-k-triton-bias-h3-v6" + "cutlass-56ca28a-fp32-simt-64x64x8-chunk128-merge64-" + "strided-actual-k-bias-epilogue-triton-bias-h3-v7" ), ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v6", + backend_id="triton-video-input-projection-fp32-v7", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -314,8 +315,8 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: implementation_class_or_factory=OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( - "rlkernel-triton-ieee-fp32-dot-chunk128-tile128x32-merge16-" - "strided-bias-epilogue-h3-v6" + "rlkernel-triton-ieee-fp32-dot-chunk128-tile128x32-merge64-" + "strided-bias-epilogue-h3-v7" ), ), OperatorBackendDescriptor( diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 356c70dfd..029a38050 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -326,7 +326,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v6" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v7" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op("triton")(x, weight, bias)) @@ -588,17 +588,25 @@ def launch(*args, **kwargs): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("rows", [1, 17, 129]) -def test_triton_forward_bias_fusion_preserves_fp32_rounding(rows): +def test_forward_bias_fusion_preserves_fp32_rounding(backend, rows): from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 x, weight, bias, _ = _data(rows) - rhs = weight.t().contiguous() - expected = _triton_gemm_fp32(x, rhs) + bias - actual = _triton_gemm_fp32(x, rhs, bias=bias) + if backend == "cuda": + from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import cutlass_gemm_fp32 + + gemm = cutlass_gemm_fp32 + rhs = weight.t() + else: + gemm = _triton_gemm_fp32 + rhs = weight.t().contiguous() + expected = gemm(x, rhs) + bias + actual = gemm(x, rhs, bias=bias) _assert_bytes(actual, expected) with torch.no_grad(): - _assert_bytes(_op("triton")(x, weight, bias), expected) + _assert_bytes(_op(backend)(x, weight, bias), expected) @pytest.mark.skipif(not H20, reason="requires CUDA SM90") diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 5243c57a7..6d7ea48ea 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,13 +13,13 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v6" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v7" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) assert desc.determinism_or_alignment_properties["vjp_tile"] == (128, 32, 32) cutlass = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "cuda-cutlass-video-input-projection-fp32-v6" + "video_input_projection_fp32", "cuda-cutlass-video-input-projection-fp32-v7" ) assert cutlass.determinism_or_alignment_properties["vjp_chunk_k"] == 128 assert cutlass.determinism_or_alignment_properties["bias_reduction"] == "fixed_midpoint_tree_32" From 96b39cfd68205cfc1bb473b3bc7dcff0bd2be840 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Fri, 16 Oct 2026 15:03:55 +0800 Subject: [PATCH 08/26] Optimize H3 Triton FP32 GEMM and canonical row reuse Signed-off-by: walker-ai <2398833647@qq.com> --- docs/operators/video-input-projection-fp32.md | 21 +++- .../backends/shared/triton/gemm/det_gemm.py | 71 +++++++++++- .../gemm/video_input_projection_fp32.py | 13 ++- rl_engine/ops/autograd/canonical_backward.py | 61 ++++++++-- rl_engine/ops/gemm/canonical_linear.py | 14 ++- rl_engine/runtime/registry.py | 7 +- .../gemm/test_video_input_projection_fp32.py | 109 +++++++++++++++++- tests/runtime/test_kernel_registry.py | 2 +- 8 files changed, 272 insertions(+), 26 deletions(-) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 4804f9bfb..3233b2054 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -97,6 +97,15 @@ once. Weight and bias share sorting and the sorted upstream gradient when their complete contribution sequences match. Calls with different contribution sets normalize each parameter independently. +The H3 Triton path explicitly enables verified canonical-order reuse for at least +4096 submitted rows. Smaller inputs retain stable normalization to avoid the +proof overhead. At each eligible normalization, a small GPU check proves that every row is active and adjacent +keys strictly increase in lexicographic order. This includes block boundaries. +When the proof succeeds, the session reuses contiguous gradients without sorting +or gathering the large matrix. Otherwise it runs the original stable sort. +The proof reads the current keys and does not cache a caller's ordering claim. +Other canonical-session callers retain stable normalization by default. + Short reductions up to the 96-wide video input use the existing direct FP32 GEMM. The Triton input and weight VJPs reuse that kernel with fixed 128-term chunks. Each chunk computes four ascending 32-term FP32 dots from zero. Chunk @@ -104,17 +113,21 @@ partials remain FP32. A second grid uses a fixed balanced reduction over the next power-of-two chunk count, with zero padding. The merge uses 64 output elements per program and four warps. VJP GEMMs use fixed 128-by-32 output tiles, 32-term dot tiles, and four warps. The forward retains its 64-by-64 output tile. -These settings define the v7 implementation identity. The K partition and VJP +These settings define the v8 Triton implementation identity. The K partition and VJP tile do not depend on M, row placement, or occupancy. The VJP kernels consume positive-stride views directly, including `G.T`. +The Triton v8 forward omits K or N load predicates only when the corresponding +32-term or 64-column tiles are complete. It retains M predicates for every row. +Column-major A views with unit row stride use a fixed `maxnreg=96` VJP allocation. +Other VJP layouts retain the default allocation. These choices preserve the dot +tile and K accumulation order. No configuration depends on physical row count. The fixed-chunk design follows the scheduling separation in merged PR #396. Its AMD BF16 MFMA instructions are not used by this CUDA FP32 implementation. The FP32 GEMM body is reused from the helper present in merged PR #325. The short forward path and long bias tree retain their earlier arithmetic. -The wider merge changes the reduction instruction schedule. Both backends pin -64 output elements and four warps for the v7 identity. Accuracy and logical-row -invariance must be checked independently of bitwise agreement with older versions. +The merge layout defines its reduction instruction schedule. Both backends pin +64 output elements and four warps as part of their implementation identities. CUTLASS also uses fixed 128-term chunks for input and weight VJPs. Partials are grouped by actual K length: 128 for complete chunks, and the remaining diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 4e3b97470..1a6df3885 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -729,6 +729,36 @@ def _copy_tree_root_transposed_kernel( mask=tl.trans(mask), ) + @triton.jit + def _canonical_key_order_kernel( + keys_ptr, + flags_ptr, + ROWS, + stride_row, + stride_col, + COLUMNS: tl.constexpr, + BLOCK: tl.constexpr, + ): + rows = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + first = tl.load(keys_ptr + rows * stride_row, mask=rows < ROWS, other=-1) + greater = tl.full((BLOCK,), False, tl.int1) + equal = tl.full((BLOCK,), True, tl.int1) + for column in tl.static_range(COLUMNS): + current = tl.load( + keys_ptr + rows * stride_row + column * stride_col, + mask=rows < ROWS, + other=0, + ) + previous = tl.load( + keys_ptr + (rows - 1) * stride_row + column * stride_col, + mask=(rows > 0) & (rows < ROWS), + other=0, + ) + greater = greater | (equal & (current > previous)) + equal = equal & (current == previous) + valid = (rows >= ROWS) | ((first >= 0) & ((rows == 0) | greater)) + tl.store(flags_ptr + tl.program_id(0), tl.min(valid.to(tl.int32), axis=0)) + @triton.jit def _det_gemm_fp32_kernel( a_ptr, @@ -750,6 +780,8 @@ def _det_gemm_fp32_kernel( PROMOTE_INPUTS: tl.constexpr, HAS_BIAS: tl.constexpr, CHUNK_K: tl.constexpr = 0, + EVEN_K: tl.constexpr = False, + EVEN_N: tl.constexpr = False, ): # One program walks the whole K or one declared fixed-width chunk. # The K order does not depend on output rows or their placement. @@ -765,9 +797,15 @@ def _det_gemm_fp32_kernel( reduction_size = CHUNK_K if CHUNK_K else K for k in range(0, tl.cdiv(reduction_size, BLOCK_K)): k_rem = K - k * BLOCK_K - a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < k_rem) + a_mask = offs_m[:, None] < M + if not EVEN_K: + a_mask = a_mask & (offs_k[None, :] < k_rem) a = tl.load(a_ptrs, mask=a_mask, other=0.0) - b_mask = (offs_k[:, None] < k_rem) & (offs_n[None, :] < N) + b_mask = tl.full((BLOCK_K, BLOCK_N), True, tl.int1) + if not EVEN_K: + b_mask = b_mask & (offs_k[:, None] < k_rem) + if not EVEN_N: + b_mask = b_mask & (offs_n[None, :] < N) b = tl.load(b_ptrs, mask=b_mask, other=0.0) if PROMOTE_INPUTS: a = a.to(tl.float32) @@ -804,11 +842,30 @@ def _det_gemm_fp32_chunk_merge_kernel( tl.store(output_ptr + offsets, total, mask=offsets < ELEMENTS) +def _triton_keys_are_canonical(keys: torch.Tensor) -> bool: + """Check the current key contents before reusing physical row order.""" + if keys.shape[0] == 0: + return True + blocks = triton.cdiv(keys.shape[0], 256) + flags = torch.empty((blocks,), device=keys.device, dtype=torch.int32) + _canonical_key_order_kernel[(blocks,)]( + keys, + flags, + keys.shape[0], + *keys.stride(), + COLUMNS=keys.shape[1], + BLOCK=256, + num_warps=4, + ) + return bool(torch.all(flags)) + + def _triton_chunked_gemm_fp32( a: torch.Tensor, b: torch.Tensor, *, tile: tuple[int, int, int] = (_BLOCK_M, _BLOCK_N, 4), + maxnreg: int | None = None, fp32_chunk_launcher=None, ) -> torch.Tensor: """Reuse the FP32 GEMM in fixed K chunks and merge partials in one grid.""" @@ -825,6 +882,7 @@ def _triton_chunked_gemm_fp32( a[start : start + rows_per_chunk], b, tile=tile, + maxnreg=maxnreg, fp32_chunk_launcher=fp32_chunk_launcher, ) for start in range(0, m_size, rows_per_chunk) @@ -856,6 +914,7 @@ def _triton_chunked_gemm_fp32( CHUNK_K=_FP32_CHUNK_K, enable_fp_fusion=False, num_warps=num_warps, + **({"maxnreg": maxnreg} if maxnreg is not None else {}), ) else: fp32_chunk_launcher(a, b, partial) @@ -873,7 +932,11 @@ def _triton_chunked_gemm_fp32( def _triton_gemm_fp32( - a: torch.Tensor, b: torch.Tensor, *, bias: torch.Tensor | None = None + a: torch.Tensor, + b: torch.Tensor, + *, + bias: torch.Tensor | None = None, + aligned_tiles: bool = False, ) -> torch.Tensor: a, b = a.contiguous(), b.contiguous() M, K = a.shape @@ -899,6 +962,8 @@ def _triton_gemm_fp32( BLOCK_K=_BLOCK_K, PROMOTE_INPUTS=True, HAS_BIAS=bias is not None, + EVEN_K=aligned_tiles and K % _BLOCK_K == 0, + EVEN_N=aligned_tiles and N % _BLOCK_N == 0, **({"enable_fp_fusion": False} if torch.version.hip is None else {}), ) return c diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index 7b5ec6dcf..a6a0425a0 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -18,13 +18,17 @@ ) _VJP_TILE = (128, 32, 4) +_VJP_MAXNREG = 96 def projection_gemm_fp32(a, b, *, bias=None): if a.shape[1] <= VIDEO_PATCH_WIDTH: - return _triton_gemm_fp32(a, b, bias=bias) + return _triton_gemm_fp32(a, b, bias=bias, aligned_tiles=True) if b.shape[1] == VIDEO_PATCH_WIDTH: - output = _triton_chunked_gemm_fp32(a, b, tile=_VJP_TILE) + # Pin register allocation for column-major A without changing dot order. + output = _triton_chunked_gemm_fp32( + a, b, tile=_VJP_TILE, maxnreg=_VJP_MAXNREG if a.stride(0) == 1 else None + ) else: output = _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) return output if bias is None else output + bias @@ -32,7 +36,7 @@ def projection_gemm_fp32(a, b, *, bias=None): class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton_tree" - backend_id = "triton-video-input-projection-fp32-v7" + backend_id = "triton-video-input-projection-fp32-v8" def __init__(self): if not _TRITON_AVAILABLE: @@ -57,6 +61,9 @@ def backend_info(self): "tile": [64, 64, 32], "vjp_tile": [_VJP_TILE[0], _VJP_TILE[1], _BLOCK_K], "vjp_num_warps": _VJP_TILE[2], + "vjp_maxnreg_unit_row_stride": _VJP_MAXNREG, + "forward_aligned_tiles": True, + "canonical_rows": "live_strict_order_proof_or_stable_sort", "tf32": False, "fp_fusion": False, "dot_arithmetic": "ieee_fp32_simt_fma", diff --git a/rl_engine/ops/autograd/canonical_backward.py b/rl_engine/ops/autograd/canonical_backward.py index e295b2eb9..d8dceaf6d 100644 --- a/rl_engine/ops/autograd/canonical_backward.py +++ b/rl_engine/ops/autograd/canonical_backward.py @@ -14,6 +14,26 @@ import torch WeightReducer = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] +_CANONICAL_ORDER_REUSE_MIN_ROWS = 4096 + + +def _keys_are_canonical(keys: torch.Tensor) -> bool: + """Prove that physical rows already have the unpadded logical order.""" + if keys.device.type == "cuda" and torch.version.hip is None: + from rl_engine.backends.shared.triton.gemm.det_gemm import ( + _TRITON_AVAILABLE, + _triton_keys_are_canonical, + ) + + if _TRITON_AVAILABLE: + return _triton_keys_are_canonical(keys) + before, after = keys[:-1], keys[1:] + greater = after[:, 0] > before[:, 0] + equal = after[:, 0] == before[:, 0] + for column in range(1, keys.shape[1]): + greater = greater | (equal & (after[:, column] > before[:, column])) + equal = equal & (after[:, column] == before[:, column]) + return bool(torch.all(keys[:, 0] >= 0) & torch.all(greater)) @dataclass @@ -22,6 +42,7 @@ class _Use: contribution: object | None = None rows: torch.Tensor | None = None grads: torch.Tensor | None = None + reuse_canonical_order: bool = False @dataclass @@ -37,14 +58,25 @@ class CanonicalBackwardSession: _shareable: set[tuple[object, ...]] | None = field(default=None, repr=False) def register( - self, parameter_id: str, keys: torch.Tensor, *, contribution: object | None = None + self, + parameter_id: str, + keys: torch.Tensor, + *, + contribution: object | None = None, + reuse_canonical_order: bool = False, ) -> int: if self.backward_started: raise RuntimeError("cannot register canonical rows after backward started") if keys.dim() != 2 or keys.shape[1] not in (2, 3): raise ValueError("logical keys must have shape [rows, 2] or [rows, 3]") slot = len(self.uses.setdefault(parameter_id, [])) - self.uses[parameter_id].append(_Use(keys=keys.detach(), contribution=contribution)) + self.uses[parameter_id].append( + _Use( + keys=keys.detach(), + contribution=contribution, + reuse_canonical_order=reuse_canonical_order, + ) + ) return slot @staticmethod @@ -72,14 +104,23 @@ def _normalize( if shared and signature in self._normalized: return self._normalized.pop(signature) keys = self._concatenate([item.keys for item in entries]) - valid = torch.nonzero(keys[:, 0] >= 0, as_tuple=True)[0] - active_keys = keys.index_select(0, valid) - order = torch.arange(valid.numel(), device=keys.device) - for column in range(keys.shape[1] - 1, -1, -1): - column_values = active_keys.index_select(0, order)[:, column] - order = order.index_select(0, torch.argsort(column_values, stable=True)) - order = valid.index_select(0, order) - normalized = order, self._concatenate(values).index_select(0, order) + reuse_order = ( + keys.shape[0] >= _CANONICAL_ORDER_REUSE_MIN_ROWS + and all(item.reuse_canonical_order for item in entries) + and _keys_are_canonical(keys) + ) + if reuse_order: + order = torch.arange(keys.shape[0], device=keys.device) + normalized = order, self._concatenate(values).contiguous() + else: + valid = torch.nonzero(keys[:, 0] >= 0, as_tuple=True)[0] + active_keys = keys.index_select(0, valid) + order = torch.arange(valid.numel(), device=keys.device) + for column in range(keys.shape[1] - 1, -1, -1): + column_values = active_keys.index_select(0, order)[:, column] + order = order.index_select(0, torch.argsort(column_values, stable=True)) + order = valid.index_select(0, order) + normalized = order, self._concatenate(values).index_select(0, order) if shared: self._normalized[signature] = normalized return normalized diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 6cf33b7f0..6df3d1b8e 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -71,12 +71,22 @@ def forward( ctx.bias_dtype = bias.dtype if bias is not None else None contribution = object() if ctx.needs_input_grad[1] and ctx.needs_input_grad[5] else None ctx.slot = ( - session.register(ctx.parameter_id, logical_keys, contribution=contribution) + session.register( + ctx.parameter_id, + logical_keys, + contribution=contribution, + reuse_canonical_order=ctx.family == "triton_tree", + ) if ctx.needs_input_grad[1] else None ) ctx.bias_slot = ( - session.register(ctx.bias_parameter_id, logical_keys, contribution=contribution) + session.register( + ctx.bias_parameter_id, + logical_keys, + contribution=contribution, + reuse_canonical_order=ctx.family == "triton_tree", + ) if ctx.needs_input_grad[5] else None ) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 3658b3e05..e50dcb9d0 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -283,7 +283,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v7", + backend_id="triton-video-input-projection-fp32-v8", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -310,13 +310,16 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "reduction": "fixed_128_chunks_balanced_fp32_merge", "vjp_tile": (128, 32, 32), "vjp_num_warps": 4, + "vjp_maxnreg_unit_row_stride": 96, + "forward_aligned_tiles": True, + "canonical_rows": "live_strict_order_proof_or_stable_sort", }, lifecycle=OperatorLifecycle.REQUEST, implementation_class_or_factory=OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( "rlkernel-triton-ieee-fp32-dot-chunk128-tile128x32-merge64-" - "strided-bias-epilogue-h3-v7" + "strided-aligned-forward-reg96-canonical-proof-h3-v8" ), ), OperatorBackendDescriptor( diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 029a38050..494b0dad8 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -23,6 +23,113 @@ BACKENDS = ["triton", "cuda"] +@pytest.mark.parametrize( + "device", + ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not H20, reason="requires CUDA SM90"))], +) +@pytest.mark.parametrize( + "key_rows", + [ + [], + [[0, 0]], + [[0, 0], [0, 1], [1, 0]], + [[0, 1], [0, 0], [1, 0]], + [[0, 0], [-1, 0], [1, 0]], + [[0, 0], [0, 0]], + [[0, 2**63 - 1], [1, -(2**63)]], + [[0, 0, 1], [0, 1, -1], [1, 0, 0]], + [[0, 0, 1], [0, 0, 0], [1, 0, 0]], + [[0, 0, 1], [0, 0, 1]], + ], +) +def test_canonical_order_proof_preserves_stable_normalization(key_rows, device, monkeypatch): + from rl_engine.ops.autograd.canonical_backward import ( + CanonicalBackwardSession, + _keys_are_canonical, + ) + + monkeypatch.setattr( + "rl_engine.ops.autograd.canonical_backward._CANONICAL_ORDER_REUSE_MIN_ROWS", 0 + ) + columns = len(key_rows[0]) if key_rows else 2 + keys = torch.tensor(key_rows, dtype=torch.int64, device=device).reshape(-1, columns) + values = torch.arange(4 * len(key_rows), device=device).reshape(4, len(key_rows)).t() + expected_order = sorted( + (i for i, key in enumerate(key_rows) if key[0] >= 0), key=lambda i: key_rows[i] + ) + ordered = expected_order == list(range(len(key_rows))) + # Equal keys still require the stable fallback, even in physical order. + strict = ordered and len({tuple(key) for key in key_rows}) == len(key_rows) + assert _keys_are_canonical(keys) == strict + if not strict: + assert not _keys_are_canonical(keys) + else: + + def unexpected_sort(*args, **kwargs): + raise AssertionError("canonical keys must reuse their physical order") + + monkeypatch.setattr(torch, "argsort", unexpected_sort) + session = CanonicalBackwardSession() + session.register("weight", keys, reuse_canonical_order=True) + order, actual = session._normalize("weight", [values]) + expected = values.index_select( + 0, torch.tensor(expected_order, dtype=torch.int64, device=device) + ) + assert order.tolist() == expected_order + assert torch.equal(actual, expected) + assert actual.is_contiguous() + + +def test_canonical_order_proof_is_recomputed_after_key_changes(monkeypatch): + from rl_engine.ops.autograd.canonical_backward import CanonicalBackwardSession + + monkeypatch.setattr( + "rl_engine.ops.autograd.canonical_backward._CANONICAL_ORDER_REUSE_MIN_ROWS", 0 + ) + keys = torch.tensor([[0, 0], [0, 1]], dtype=torch.int64) + values = torch.tensor([[1.0], [2.0]]) + session = CanonicalBackwardSession() + session.register("weight", keys, reuse_canonical_order=True) + _, first = session._normalize("weight", [values]) + keys.copy_(keys.flip(0)) + _, second = session._normalize("weight", [values]) + assert torch.equal(first, values) + assert torch.equal(second, values.flip(0)) + + +def test_canonical_order_reuse_requires_explicit_registration(monkeypatch): + from rl_engine.ops.autograd import canonical_backward as canonical + + def unexpected_proof(keys): + raise AssertionError("existing callers must retain stable normalization") + + monkeypatch.setattr(canonical, "_keys_are_canonical", unexpected_proof) + monkeypatch.setattr(canonical, "_CANONICAL_ORDER_REUSE_MIN_ROWS", 0) + session = canonical.CanonicalBackwardSession() + keys = torch.tensor([[0, 0], [0, 1]], dtype=torch.int64) + values = torch.tensor([[1.0], [2.0]]) + session.register("weight", keys) + order, actual = session._normalize("weight", [values]) + assert order.tolist() == [0, 1] + assert torch.equal(actual, values) + + +def test_small_canonical_uses_retain_stable_normalization(monkeypatch): + from rl_engine.ops.autograd import canonical_backward as canonical + + def unexpected_proof(keys): + raise AssertionError("small inputs must avoid the order proof overhead") + + monkeypatch.setattr(canonical, "_keys_are_canonical", unexpected_proof) + session = canonical.CanonicalBackwardSession() + keys = torch.tensor([[0, 0], [0, 1]], dtype=torch.int64) + values = torch.tensor([[1.0], [2.0]]) + session.register("weight", keys, reuse_canonical_order=True) + order, actual = session._normalize("weight", [values]) + assert order.tolist() == [0, 1] + assert torch.equal(actual, values) + + def _op(backend): if backend == "triton": from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( @@ -326,7 +433,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v7" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v8" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op("triton")(x, weight, bias)) diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 6d7ea48ea..d315eb9da 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,7 +13,7 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v7" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v8" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) From e9ebceec531e094e8c1ffc0c8d9eb7703cdb834b Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Sat, 17 Oct 2026 12:39:25 +0800 Subject: [PATCH 09/26] Remove H3 FP32 projection CUTLASS candidate Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 14 +- build_tools/extensions.py | 50 ----- .../video_input_projection_fp32_cutlass.cpp | 38 ---- .../video_input_projection_fp32_cutlass.cu | 180 ------------------ docs/operators/video-input-projection-fp32.md | 68 ++----- rl_engine/_C_cutlass.pyi | 20 -- .../cuda/gemm/video_input_projection_fp32.py | 165 ---------------- .../backends/shared/triton/gemm/det_gemm.py | 71 +++---- rl_engine/ops/gemm/canonical_linear.py | 29 +-- .../ops/gemm/video_input_projection_fp32.py | 4 +- rl_engine/runtime/registry.py | 42 ---- .../validation/operators/operator_specs.py | 4 - .../build/test_build_platform_collectives.py | 31 --- .../gemm/test_video_input_projection_fp32.py | 100 ++++------ tests/runtime/test_kernel_registry.py | 5 - 15 files changed, 80 insertions(+), 741 deletions(-) delete mode 100644 csrc/bindings/video_input_projection_fp32_cutlass.cpp delete mode 100644 csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu delete mode 100644 rl_engine/_C_cutlass.pyi delete mode 100644 rl_engine/backends/cuda/gemm/video_input_projection_fp32.py diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index 56845e8b3..28fabdb4c 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -29,12 +29,6 @@ def _candidate(name): ) return TritonVideoInputProjectionFP32Op() - if name == "cuda": - from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import ( - CUDAVideoInputProjectionFP32Op, - ) - - return CUDAVideoInputProjectionFP32Op() return None @@ -123,8 +117,8 @@ def main(): parser.add_argument( "--backends", nargs="+", - choices=("pytorch", "triton", "cuda"), - default=["pytorch", "triton", "cuda"], + choices=("pytorch", "triton"), + default=["pytorch", "triton"], ) parser.add_argument("--rows", nargs="+", type=int, default=[1, 128, 1024, 4096, 13104]) parser.add_argument("--samples", type=int, default=20) @@ -164,11 +158,7 @@ def main(): "rl_engine/ops/gemm/video_input_projection_fp32.py", "rl_engine/backends/shared/triton/gemm/det_gemm.py", "rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py", - "rl_engine/backends/cuda/gemm/video_input_projection_fp32.py", "rl_engine/reference/gemm/video_input_projection_fp32.py", - "csrc/bindings/video_input_projection_fp32_cutlass.cpp", - "csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu", - "build_tools/extensions.py", ] # Devbox worker sync excludes .git; source hashes also identify exported workspaces. git_commit = _tool_version(["git", "-C", str(root), "rev-parse", "HEAD"]) diff --git a/build_tools/extensions.py b/build_tools/extensions.py index aebba78ee..670fdb072 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -import hashlib import importlib.util import os import sysconfig @@ -107,7 +106,6 @@ def _native_extension_required() -> bool: return ( envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) or envs.env_flag(envs.RL_KERNEL_H3_STRICT) - or envs.env_flag("RL_KERNEL_BUILD_CUTLASS") or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) or envs.env_flag("FORCE_CUDA") @@ -160,52 +158,6 @@ def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: return filtered_flags -_CUTLASS_REVISION = "56ca28adb32ffadd097566545b8f7d83b652e998" -_CUTLASS_HEADERS_SHA256 = "7b5321860c269f8e44eab02e59c32c2d4d0e2fa7a3ef72e0a161de9f58b361fa" - - -def _cutlass_headers_fingerprint(include_root: Path) -> str: - digest = hashlib.sha256() - for path in sorted(include_root.rglob("*")): - if path.is_file() and not path.name.startswith("._"): - digest.update(path.relative_to(include_root).as_posix().encode() + b"\0") - digest.update(path.read_bytes() + b"\0") - return digest.hexdigest() - - -def _cutlass_extension(torch, gpu_extension, link_args): - if torch.version.hip is not None or torch.version.cuda is None: - raise RuntimeError("CUTLASS projection requires a CUDA PyTorch build") - root = os.environ.get("CUTLASS_ROOT", "").strip() - if not root: - raise RuntimeError(f"CUTLASS_ROOT must point to CUTLASS {_CUTLASS_REVISION}") - include = Path(root).resolve() / "include" - if _cutlass_headers_fingerprint(include) != _CUTLASS_HEADERS_SHA256: - raise RuntimeError("CUTLASS include tree does not match the pinned source fingerprint") - return gpu_extension( - name="rl_engine._C_cutlass", - sources=[ - "csrc/bindings/video_input_projection_fp32_cutlass.cpp", - "csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu", - ], - include_dirs=[str(include)], - extra_compile_args={ - "cxx": ["-O3", "-std=c++17", "-fno-fast-math", "-ffp-contract=off"], - "nvcc": [ - "-O3", - "-std=c++17", - "--fmad=false", - "--ftz=false", - "--prec-div=true", - "--prec-sqrt=true", - "--expt-relaxed-constexpr", - "-gencode=arch=compute_90,code=sm_90", - ], - }, - extra_link_args=list(link_args), - ) - - def get_extensions(): h3_flags = envs.h3_strict_cuda_flags() torch, _, GpuExtension = _load_torch_extension_tools() @@ -231,8 +183,6 @@ def get_extensions(): if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") is_rocm = getattr(torch.version, "hip", None) is not None - if envs.env_flag("RL_KERNEL_BUILD_CUTLASS"): - extensions.append(_cutlass_extension(torch, GpuExtension, torch_rpath)) if h3_flags and (is_rocm or _musa_build_available(torch)): raise RuntimeError("RL_KERNEL_H3_STRICT currently qualifies NVIDIA CUDA only") diff --git a/csrc/bindings/video_input_projection_fp32_cutlass.cpp b/csrc/bindings/video_input_projection_fp32_cutlass.cpp deleted file mode 100644 index eb803874c..000000000 --- a/csrc/bindings/video_input_projection_fp32_cutlass.cpp +++ /dev/null @@ -1,38 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright (c) 2026 RL-Kernel Contributors - -#include -#include - -torch::Tensor cutlass_gemm_fp32_tn( - torch::Tensor a, torch::Tensor b_transposed, std::optional bias); -int cutlass_cuda_version(); -void cutlass_gemm_fp32_partials_strided( - torch::Tensor a, torch::Tensor bt, torch::Tensor workspace, - torch::Tensor starts, torch::Tensor nodes, torch::Tensor partial_ids, int64_t length); - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("gemm_fp32_tn", &cutlass_gemm_fp32_tn, - pybind11::arg("a"), pybind11::arg("b_transposed"), pybind11::arg("bias") = pybind11::none(), - "Fixed FP32 SIMT GEMM: [M,K] times transposed [N,K]"); - m.def("gemm_fp32_partials_strided", &cutlass_gemm_fp32_partials_strided, - "Evaluate one actual-K partial group directly from matrix views"); - m.def("build_info", []() { - pybind11::dict info; - info["cutlass_revision"] = "56ca28adb32ffadd097566545b8f7d83b652e998"; - info["headers_sha256"] = "7b5321860c269f8e44eab02e59c32c2d4d0e2fa7a3ef72e0a161de9f58b361fa"; - info["cuda_version"] = cutlass_cuda_version(); - info["architecture"] = "sm_90"; - info["tile"] = pybind11::make_tuple(64, 64, 8); - info["operator_class"] = "simt"; - info["split_k"] = 1; - info["fmad"] = false; - info["ftz"] = false; - info["tree_leaf_size"] = 32; - info["vjp_chunk_k"] = 128; - info["partial_launch"] = "cutlass_device_gemm_array"; - info["partial_operand_layout"] = "direct_strides_grouped_actual_k"; - info["forward_bias"] = "linear_combination_broadcast_source"; - return info; - }); -} diff --git a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu b/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu deleted file mode 100644 index 7aa9b9412..000000000 --- a/csrc/cuda/gemm/video_input_projection_fp32_cutlass.cu +++ /dev/null @@ -1,180 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright (c) 2026 RL-Kernel Contributors - -// Reuse CUTLASS device::Gemm at revision 56ca28adb32ffadd097566545b8f7d83b652e998. -// The dependency's BSD-3-Clause notice is packaged in the Python wrapper. -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -using Gemm = cutlass::gemm::device::Gemm< - float, cutlass::layout::RowMajor, - float, cutlass::layout::ColumnMajor, - float, cutlass::layout::RowMajor, - float, cutlass::arch::OpClassSimt, cutlass::arch::Sm50, - cutlass::gemm::GemmShape<64, 64, 8>, - cutlass::gemm::GemmShape<32, 32, 8>, - cutlass::gemm::GemmShape<1, 1, 1>, - cutlass::epilogue::thread::LinearCombination, - cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, - 2, 1, 1, false, cutlass::arch::OpMultiplyAdd>; - -constexpr int kMaxPartialK = 128; - -at::Tensor cutlass_gemm_fp32_tn( - at::Tensor a, at::Tensor b_transposed, std::optional bias) { - TORCH_CHECK(a.is_cuda() && b_transposed.is_cuda(), "CUTLASS GEMM requires CUDA tensors"); - TORCH_CHECK(a.device() == b_transposed.device(), "CUTLASS GEMM devices must match"); - TORCH_CHECK(a.scalar_type() == at::kFloat && b_transposed.scalar_type() == at::kFloat, - "CUTLASS GEMM requires FP32"); - TORCH_CHECK(a.dim() == 2 && b_transposed.dim() == 2, - "CUTLASS GEMM requires rank-two matrices"); - TORCH_CHECK(a.is_contiguous() && b_transposed.is_contiguous(), - "CUTLASS GEMM matrices must be contiguous"); - TORCH_CHECK(a.size(1) == b_transposed.size(1), "CUTLASS GEMM K dimensions must match"); - const c10::cuda::CUDAGuard guard(a.device()); - const auto* properties = at::cuda::getDeviceProperties(a.get_device()); - TORCH_CHECK(properties->major == 9 && properties->minor == 0, - "CUTLASS video projection requires CUDA SM90"); - const auto m64 = a.size(0), n64 = b_transposed.size(0), k64 = a.size(1); - TORCH_CHECK(m64 > 0 && n64 > 0 && k64 > 0 && - m64 <= INT_MAX && n64 <= INT_MAX && k64 <= INT_MAX, - "CUTLASS GEMM dimensions must be positive int32 values"); - const int m = static_cast(m64), n = static_cast(n64), k = static_cast(k64); - if (bias.has_value()) { - TORCH_CHECK(bias->device() == a.device() && bias->scalar_type() == at::kFloat && - bias->dim() == 1 && bias->size(0) == n && bias->is_contiguous(), - "CUTLASS GEMM bias must be a contiguous FP32 vector on the input device"); - } - auto output = at::empty({m, n}, a.options()); - // A zero source row stride broadcasts bias through the SDK epilogue. - // --fmad=false retains the rounded FP32 GEMM followed by FP32 addition. - const float* source = bias.has_value() ? bias->data_ptr() : output.data_ptr(); - Gemm::Arguments args({m, n, k}, {a.data_ptr(), k}, - {b_transposed.data_ptr(), k}, - {source, bias.has_value() ? 0 : n}, {output.data_ptr(), n}, - {1.0f, bias.has_value() ? 1.0f : 0.0f}, 1); - Gemm gemm; - auto status = gemm.can_implement(args); - TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); - status = gemm(args, nullptr, c10::cuda::getCurrentCUDAStream(a.get_device()).stream()); - TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - return output; -} - -int cutlass_cuda_version() { return CUDART_VERSION; } - -template -using StridedPartialGemm = cutlass::gemm::device::GemmArray< - float, LayoutA, float, LayoutB, float, cutlass::layout::RowMajor, - float, cutlass::arch::OpClassSimt, cutlass::arch::Sm50, - cutlass::gemm::GemmShape<64, 64, 8>, - cutlass::gemm::GemmShape<32, 32, 8>, - cutlass::gemm::GemmShape<1, 1, 1>, - cutlass::epilogue::thread::LinearCombination, - cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle, - 2, 1, 1, cutlass::arch::OpMultiplyAdd>; - -__global__ void partial_pointers( - const float* a, const float* bt, float* workspace, int64_t* pointers, - const int32_t* starts, const int32_t* nodes, const int32_t* partial_ids, - int64_t count, int64_t stride_ak, int64_t stride_btk, int64_t elements) { - const int index = blockIdx.x * blockDim.x + threadIdx.x; - if (index >= count) return; - const int partial = partial_ids[index]; - const int64_t start = starts[partial]; - pointers[index] = reinterpret_cast(a + start * stride_ak); - pointers[count + index] = reinterpret_cast(bt + start * stride_btk); - pointers[2 * count + index] = reinterpret_cast( - workspace + static_cast(nodes[partial]) * elements); -} - -template -void launch_strided_partials( - const at::Tensor& a, const at::Tensor& bt, const at::Tensor& pointers, - int64_t count, int length, int64_t lda, int64_t ldb, cudaStream_t stream) { - using Partial = StridedPartialGemm; - const auto* addresses = pointers.data_ptr(); - typename Partial::Arguments args( - {static_cast(a.size(0)), static_cast(bt.size(0)), length}, - reinterpret_cast(addresses), LayoutA(lda), - reinterpret_cast(addresses + count), LayoutB(ldb), - reinterpret_cast(addresses + 2 * count), - cutlass::layout::RowMajor(bt.size(0)), - reinterpret_cast(addresses + 2 * count), - cutlass::layout::RowMajor(bt.size(0)), {1.0f, 0.0f}, static_cast(count)); - Partial gemm; - auto status = gemm.can_implement(args); - TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); - status = gemm(args, nullptr, stream); - TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status)); -} - -void cutlass_gemm_fp32_partials_strided( - at::Tensor a, at::Tensor bt, at::Tensor workspace, at::Tensor starts, - at::Tensor nodes, at::Tensor partial_ids, int64_t length) { - TORCH_CHECK(a.is_cuda() && a.scalar_type() == at::kFloat && a.dim() == 2, - "CUTLASS partials require rank-two CUDA FP32 A"); - const c10::cuda::CUDAGuard guard(a.device()); - const auto* properties = at::cuda::getDeviceProperties(a.get_device()); - TORCH_CHECK(properties->major == 9 && properties->minor == 0, - "CUTLASS video projection requires CUDA SM90"); - TORCH_CHECK(bt.device() == a.device() && bt.scalar_type() == at::kFloat && - bt.dim() == 2 && a.size(1) == bt.size(1), - "CUTLASS partial dimensions, dtypes and devices must match"); - TORCH_CHECK(a.size(0) > 0 && bt.size(0) > 0 && length > 0 && length <= kMaxPartialK && - a.size(0) <= INT_MAX && bt.size(0) <= INT_MAX, - "CUTLASS partial sizes exceed the supported dimensions"); - TORCH_CHECK(workspace.device() == a.device() && workspace.scalar_type() == at::kFloat && - workspace.dim() == 3 && workspace.is_contiguous() && - workspace.size(1) == a.size(0) && workspace.size(2) == bt.size(0), - "CUTLASS partials require contiguous FP32 tree workspace"); - for (const auto& metadata : {starts, nodes, partial_ids}) { - TORCH_CHECK(metadata.device() == a.device() && metadata.scalar_type() == at::kInt && - metadata.dim() == 1 && metadata.is_contiguous(), - "CUTLASS partial metadata must be contiguous CUDA int32 vectors"); - } - TORCH_CHECK(starts.numel() == nodes.numel() && partial_ids.numel() > 0 && - partial_ids.numel() <= 65535, - "CUTLASS partial metadata sizes are inconsistent"); - const bool a_row = a.stride(1) == 1; - const bool b_column = bt.stride(1) == 1; - TORCH_CHECK((a_row || a.stride(0) == 1) && - (b_column || bt.stride(0) == 1) && - a.stride(0) > 0 && a.stride(1) > 0 && - bt.stride(0) > 0 && bt.stride(1) > 0, - "CUTLASS partials require a row-major or column-major view"); - const int64_t lda = a_row ? a.stride(0) : a.stride(1); - const int64_t ldb = b_column ? bt.stride(0) : bt.stride(1); - TORCH_CHECK(lda <= INT_MAX && ldb <= INT_MAX, - "CUTLASS partial leading dimensions exceed int32"); - const int64_t count = partial_ids.numel(); - auto pointers = at::empty({3, count}, a.options().dtype(at::kLong)); - const auto stream = c10::cuda::getCurrentCUDAStream(a.get_device()).stream(); - partial_pointers<<<(count + 255) / 256, 256, 0, stream>>>( - a.data_ptr(), bt.data_ptr(), workspace.data_ptr(), - pointers.data_ptr(), starts.data_ptr(), nodes.data_ptr(), - partial_ids.data_ptr(), count, a.stride(1), bt.stride(1), - a.size(0) * bt.size(0)); - using Row = cutlass::layout::RowMajor; - using Column = cutlass::layout::ColumnMajor; - if (a_row && b_column) { - launch_strided_partials(a, bt, pointers, count, length, lda, ldb, stream); - } else if (a_row) { - launch_strided_partials(a, bt, pointers, count, length, lda, ldb, stream); - } else if (b_column) { - launch_strided_partials(a, bt, pointers, count, length, lda, ldb, stream); - } else { - launch_strided_partials(a, bt, pointers, count, length, lda, ldb, stream); - } - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 3233b2054..450e6e826 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -6,12 +6,9 @@ MiniMax-H3 maps each 96-element video patch to the 5376-wide Transformer residua space. This operator preserves the FP32 projection boundary of the pinned H3 reference. The model adapter owns the subsequent cast into the BF16 packed stream. -The primary implementation reuses RL-Kernel's Triton FP32 GEMM and canonical +The implementation reuses RL-Kernel's Triton FP32 GEMM and canonical parameter-gradient session. Input and weight VJPs use fixed 128-term K chunks -and one balanced FP32 merge. A pinned CUTLASS FP32 SIMT GEMM is an explicit -CUDA candidate. It uses the same fixed-chunk scheduler and balanced merge, -and reads each partial directly from matrix views. Both paths share the -semantic adapter and canonical session. +and one balanced FP32 merge. Long bias VJPs retain the one-row Triton tree. ## Entry Point @@ -36,24 +33,8 @@ physical placement. The same parameter keeps the same `parameter_id` across uses | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | Triton on CUDA | `TritonVideoInputProjectionFP32Op` | `_det_gemm_fp32_kernel`, `_det_gemm_fp32_chunk_merge_kernel`, bias tree kernels | Default CUDA path; SM90 qualification | -| CUDA / CUTLASS + Triton reductions | `CUDAVideoInputProjectionFP32Op` | `_C_cutlass.gemm_fp32_tn`, `_C_cutlass.gemm_fp32_partials_strided`, shared chunk merge and bias tree | Explicit candidate; optional build; SM90 | | PyTorch reference | `NativeVideoInputProjectionFP32Op` | CPU `F.linear`; FP32 product/sum VJPs | Accuracy reference | -CUTLASS revision: `56ca28adb32ffadd097566545b8f7d83b652e998`. The build verifies -the complete include-tree fingerprint. The BSD-3-Clause dependency notice is -preserved in the installed CUDA Python wrapper. - -```bash -git clone https://github.com/NVIDIA/cutlass.git /tmp/h3-cutlass -git -C /tmp/h3-cutlass checkout 56ca28adb32ffadd097566545b8f7d83b652e998 -CUTLASS_ROOT=/tmp/h3-cutlass RL_KERNEL_BUILD_CUTLASS=1 \ - MAX_JOBS=2 python -m pip install --no-build-isolation --no-deps -e . -``` - -The separate `_C_cutlass` extension fixes `--fmad=false`, `--ftz=false`, precise -division/square root, and SM90 code generation. Common `_C` compiler options are -independent. `backend_info()` reads back the candidate's revision and schedule. - ## Tensor Contract | Argument | Shape | Dtype | Requirements | @@ -72,8 +53,8 @@ fresh sessions. Parameter freezing registers only the parameters that need VJPs. ## Dispatch Behavior -`get_op` selects Triton on CUDA. CUTLASS is selected explicitly through the -semantic catalog or `OP_SPECS` candidate `cuda`. The semantic catalog records +`get_op` selects Triton on CUDA. `OP_SPECS` exposes candidate `triton` and the +independent PyTorch accuracy reference. The semantic catalog records single-device, unsharded support and exact implementation identities. An unavailable explicit candidate raises an error. CPU/ROCm/MUSA/NPU execution is outside this contribution's strict support scope. @@ -87,9 +68,8 @@ stays FP32, including normalization of non-contiguous autograd inputs: - `dX = G @ W` is row-local. - `dW = G_sorted.T @ X_sorted` follows canonical logical-row order. -- `dbias = ones[1,R] @ G_sorted` uses the fixed FP32 GEMM bottom layer. - Both backends reuse the one-row Triton tree for reductions longer than 96; - the CUTLASS candidate retains its direct GEMM for shorter reductions. +- `dbias = ones[1,R] @ G_sorted` uses the fixed FP32 GEMM for reductions up to 96 + and the one-row Triton tree for longer reductions. The existing canonical session collects all contributions, removes padding, and sorts keys before the final parameter reduction. It returns each parameter VJP @@ -126,34 +106,18 @@ The fixed-chunk design follows the scheduling separation in merged PR #396. Its AMD BF16 MFMA instructions are not used by this CUDA FP32 implementation. The FP32 GEMM body is reused from the helper present in merged PR #325. The short forward path and long bias tree retain their earlier arithmetic. -The merge layout defines its reduction instruction schedule. Both backends pin -64 output elements and four warps as part of their implementation identities. - -CUTLASS also uses fixed 128-term chunks for input and weight VJPs. Partials -are grouped by actual K length: 128 for complete chunks, and the remaining -length for the tail. A pointer kernel selects each partial's original operand -range, and the pinned `device::GemmArray` reads row-major or column-major views -directly. This applies the layout approach from merged PR #343. The shared -chunk scheduler calls this SDK launcher and then the same balanced Triton merge. -The SDK retains its separate FP32 multiply/add arithmetic inside each partial. -This schedule defines the CUTLASS v7 identity. Its actual-K grouping and layouts -are recorded in `backend_info()`. Long bias VJPs retain the original Triton tree. -The forward reuses CUTLASS `LinearCombination` with source row stride zero to -broadcast bias. It adds bias after the FP32 accumulator in the same epilogue. -The build keeps `--fmad=false`. Tests compare this output bitwise with a separate -GEMM followed by FP32 bias addition. - -Both VJP schedules cap their FP32 partial workspace at 512 MiB by chunking +The merge layout defines its reduction instruction schedule. The implementation +pins 64 output elements and four warps as part of its implementation identity. + +The VJP schedule caps its FP32 partial workspace at 512 MiB by chunking output rows. The retained midpoint helper also supports output-column chunking -when one complete row exceeds its budget. CUTLASS allocates pointer arrays for -each partial group. Parameter sorting and public validation are included in complete +when one complete row exceeds its budget. Parameter sorting and public validation are included in complete operator measurements. Existing BF16 callers retain their dtype and schedule. TF32 and compiler multiply-add fusion remain disabled. The reused Triton `tl.dot` lowers its IEEE FP32 inner products to fixed SIMT `fma.rn.f32` instructions. The compiler fusion flag controls surrounding expressions; it does not remove the intrinsic dot's FMA instructions. This -arithmetic is included in the Triton contract and runtime metadata. The pinned -CUTLASS SIMT build uses separate FP32 multiply and add instructions. +arithmetic is included in the Triton contract and runtime metadata. The independent reference keeps the forward definition as CPU FP32 `F.linear`. Its analytic VJPs form FP32 products and use CPU PyTorch `sum` over the complete @@ -173,7 +137,7 @@ Adding a genuine training sample changes the expected parameter gradient. ```bash PYTHONPATH=. python benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py \ - --backends pytorch triton cuda --rows 1 128 1024 4096 13104 \ + --backends pytorch triton --rows 1 128 1024 4096 13104 \ --output /tmp/h3-input-projection.json ``` @@ -184,8 +148,7 @@ parameter-gradient machinery. Raw samples, environment, source hashes, and actual backend metadata are written to the explicitly selected local output. PyTorch is a production performance baseline with separately reported BI results. The Triton VJPs use one partial-GEMM launch and one merge launch when the -workspace fits. CUTLASS launches the SDK partial groups, their pointer metadata, -and one balanced merge. A tail can require a separate SDK launch. +workspace fits. Compare complete operator timings, including sorting, before drawing a speedup claim. The benchmark records the reference identity and its source hash. @@ -196,9 +159,6 @@ python -m pytest tests/ops/gemm/test_video_input_projection_fp32.py -v python tools/validation/operators/check_operator.py \ --op video_input_projection_fp32 --candidate triton --device cuda --dtype fp32 \ --batch 1 --seq 129 --check-grad --grad-mode random -python tools/validation/operators/check_operator.py \ - --op video_input_projection_fp32 --candidate cuda --device cuda --dtype fp32 \ - --batch 1 --seq 129 --check-grad --grad-mode random ``` Set `H3_PROJECTION_WEIGHTS` to a safetensors fixture containing `proj_in.weight` diff --git a/rl_engine/_C_cutlass.pyi b/rl_engine/_C_cutlass.pyi deleted file mode 100644 index ff59a01d2..000000000 --- a/rl_engine/_C_cutlass.pyi +++ /dev/null @@ -1,20 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -from typing import Any - -import torch - -def gemm_fp32_tn( - a: torch.Tensor, b_transposed: torch.Tensor, bias: torch.Tensor | None = None -) -> torch.Tensor: ... -def gemm_fp32_partials_strided( - a: torch.Tensor, - b_transposed: torch.Tensor, - workspace: torch.Tensor, - starts: torch.Tensor, - nodes: torch.Tensor, - partial_ids: torch.Tensor, - length: int, -) -> None: ... -def build_info() -> dict[str, Any]: ... diff --git a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py b/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py deleted file mode 100644 index ba4b18637..000000000 --- a/rl_engine/backends/cuda/gemm/video_input_projection_fp32.py +++ /dev/null @@ -1,165 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""H3 FP32 projection using a pinned CUTLASS SIMT GEMM candidate. - -CUTLASS dependency notice (BSD-3-Clause): -Copyright (c) 2017 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are met: - -1. Redistributions of source code must retain the above copyright notice, this - list of conditions and the following disclaimer. -2. Redistributions in binary form must reproduce the above copyright notice, - this list of conditions and the following disclaimer in the documentation - and/or other materials provided with the distribution. -3. Neither the name of the copyright holder nor the names of its contributors - may be used to endorse or promote products derived from this software without - specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" -AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE -DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR -ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES -(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; -LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON -ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS -SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -""" - -from functools import lru_cache -from importlib import import_module - -import torch - -from rl_engine.backends.shared.triton.gemm.det_gemm import ( - _BLOCK_K, - _FP32_CHUNK_K, - _TRITON_AVAILABLE, - _triton_chunked_gemm_fp32, - _triton_tree_gemm, -) -from rl_engine.ops.gemm.video_input_projection_fp32 import ( - VIDEO_PATCH_WIDTH, - VideoInputProjectionFP32Op, -) - -CUTLASS_REVISION = "56ca28adb32ffadd097566545b8f7d83b652e998" -CUTLASS_HEADERS_SHA256 = "7b5321860c269f8e44eab02e59c32c2d4d0e2fa7a3ef72e0a161de9f58b361fa" - - -@lru_cache(maxsize=1) -def _load_extension(): - try: - extension = import_module("rl_engine._C_cutlass") - except ImportError as exc: - raise RuntimeError( - "build the optional CUTLASS extension with RL_KERNEL_BUILD_CUTLASS=1" - ) from exc - info = extension.build_info() - if ( - info["cutlass_revision"] != CUTLASS_REVISION - or info["headers_sha256"] != CUTLASS_HEADERS_SHA256 - or info["fmad"] - or info["ftz"] - or info["split_k"] != 1 - or info.get("tree_leaf_size") != _BLOCK_K - or info.get("vjp_chunk_k") != _FP32_CHUNK_K - or info.get("partial_operand_layout") != "direct_strides_grouped_actual_k" - or info.get("forward_bias") != "linear_combination_broadcast_source" - ): - raise RuntimeError("CUTLASS extension does not match the strict FP32 build contract") - return extension - - -def _cutlass_tree_leaves(a, b, workspace, plan): - for length, leaf_ids in _leaf_groups(plan.host.leaf_lengths, a.device): - _load_extension().gemm_fp32_partials_strided( - a, b.t(), workspace, plan.leaf_starts, plan.leaf_nodes, leaf_ids, length - ) - - -@lru_cache(maxsize=None) -def _leaf_groups(lengths, device): - return tuple( - ( - length, - torch.tensor( - [index for index, actual in enumerate(lengths) if actual == length], - device=device, - dtype=torch.int32, - ), - ) - for length in sorted(set(lengths)) - ) - - -@lru_cache(maxsize=None) -def _chunk_metadata(k_size, device): - full, tail = divmod(k_size, _FP32_CHUNK_K) - lengths = (_FP32_CHUNK_K,) * full + ((tail,) if tail else ()) - nodes = torch.arange(len(lengths), device=device, dtype=torch.int32) - return nodes * _FP32_CHUNK_K, nodes, _leaf_groups(lengths, device) - - -def _cutlass_chunks(a, b, partial): - starts, nodes, groups = _chunk_metadata(a.shape[1], a.device) - for length, ids in groups: - _load_extension().gemm_fp32_partials_strided(a, b.t(), partial, starts, nodes, ids, length) - - -def cutlass_gemm_fp32(a, b, *, bias=None): - # The native TN interface represents mathematical B[K,N] as [N,K] storage. - if a.shape[1] <= VIDEO_PATCH_WIDTH: - return _load_extension().gemm_fp32_tn(a.contiguous(), b.t().contiguous(), bias) - if bias is not None: - return cutlass_gemm_fp32(a, b) + bias - if b.shape[1] == VIDEO_PATCH_WIDTH: - return _triton_chunked_gemm_fp32(a, b, fp32_chunk_launcher=_cutlass_chunks) - return _triton_tree_gemm( - a, - b, - output_dtype=torch.float32, - preserve_a_strides=True, - fp32_leaf_launcher=_cutlass_tree_leaves, - ) - - -class CUDAVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): - family = "cutlass" - backend_id = "cuda-cutlass-video-input-projection-fp32-v7" - - def __init__(self): - if not _TRITON_AVAILABLE: - raise RuntimeError("CUTLASS projection requires Triton for shared FP32 tree reductions") - _load_extension() - - def _check_backend(self, x): - if x.device.type != "cuda" or torch.version.hip is not None: - raise ValueError("CUTLASS video input projection requires CUDA") - if torch.cuda.get_device_capability(x.device) != (9, 0): - raise ValueError("CUTLASS video input projection requires CUDA SM90") - - def backend_info(self): - import triton - - return { - "backend_id": self.backend_id, - "family": "cuda", - "kernel": "gemm_fp32_tn / gemm_fp32_partials_strided / " - "_det_gemm_fp32_chunk_merge_kernel / " - "_det_gemm_kernel / _det_gemm_tree_leaf_reduce_rocm_kernel", - "reduction": "fixed_128_chunks_balanced_fp32_merge", - "vjp_chunk_k": _FP32_CHUNK_K, - "vjp_merge_tile": 64, - "bias_reduction": "fixed_midpoint_tree", - "reduction_leaf_size": _BLOCK_K, - "reduction_backend": "triton", - "bias_reduction_backend": "triton_for_k_gt_96", - "triton_version": triton.__version__, - "bias_tree_schedule": "paired_levels_compact_workspace", - **_load_extension().build_info(), - } diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 1a6df3885..7127ec623 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -866,7 +866,6 @@ def _triton_chunked_gemm_fp32( *, tile: tuple[int, int, int] = (_BLOCK_M, _BLOCK_N, 4), maxnreg: int | None = None, - fp32_chunk_launcher=None, ) -> torch.Tensor: """Reuse the FP32 GEMM in fixed K chunks and merge partials in one grid.""" m_size, k_size = a.shape @@ -883,7 +882,6 @@ def _triton_chunked_gemm_fp32( b, tile=tile, maxnreg=maxnreg, - fp32_chunk_launcher=fp32_chunk_launcher, ) for start in range(0, m_size, rows_per_chunk) ], @@ -891,33 +889,30 @@ def _triton_chunked_gemm_fp32( ) partial = torch.empty((chunks, m_size, n_size), device=a.device, dtype=torch.float32) output = torch.empty((m_size, n_size), device=a.device, dtype=torch.float32) - if fp32_chunk_launcher is None: - block_m, block_n, num_warps = tile - grid = (triton.cdiv(m_size, block_m), triton.cdiv(n_size, block_n), chunks) - _det_gemm_fp32_kernel[grid]( - a, - b, - partial, - None, - m_size, - n_size, - k_size, - *a.stride(), - *b.stride(), - n_size, - 1, - BLOCK_M=block_m, - BLOCK_N=block_n, - BLOCK_K=_BLOCK_K, - PROMOTE_INPUTS=True, - HAS_BIAS=False, - CHUNK_K=_FP32_CHUNK_K, - enable_fp_fusion=False, - num_warps=num_warps, - **({"maxnreg": maxnreg} if maxnreg is not None else {}), - ) - else: - fp32_chunk_launcher(a, b, partial) + block_m, block_n, num_warps = tile + grid = (triton.cdiv(m_size, block_m), triton.cdiv(n_size, block_n), chunks) + _det_gemm_fp32_kernel[grid]( + a, + b, + partial, + None, + m_size, + n_size, + k_size, + *a.stride(), + *b.stride(), + n_size, + 1, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=_BLOCK_K, + PROMOTE_INPUTS=True, + HAS_BIAS=False, + CHUNK_K=_FP32_CHUNK_K, + enable_fp_fusion=False, + num_warps=num_warps, + **({"maxnreg": maxnreg} if maxnreg is not None else {}), + ) _det_gemm_fp32_chunk_merge_kernel[(triton.cdiv(m_size * n_size, 64),)]( partial, output, @@ -993,7 +988,6 @@ def _triton_tree_gemm( preserve_a_strides: bool = False, inference_schedule: bool = False, output_dtype: torch.dtype = torch.bfloat16, - fp32_leaf_launcher=None, ) -> torch.Tensor: if not _TRITON_AVAILABLE: raise RuntimeError("Triton is unavailable") @@ -1009,8 +1003,6 @@ def _triton_tree_gemm( raise TypeError(f"Triton tree GEMM requires {input_dtype} inputs") if fp32_tree and torch.version.hip is not None: raise RuntimeError("FP32 tree GEMM is qualified on CUDA only") - if fp32_leaf_launcher is not None and not fp32_tree: - raise ValueError("a custom leaf launcher requires the FP32 tree") if not a.is_cuda or not b.is_cuda or a.device != b.device: raise RuntimeError("Triton tree GEMM inputs must share one CUDA/ROCm device") @@ -1025,11 +1017,7 @@ def _triton_tree_gemm( plan = _device_tree_plan( k_size, a.device, - fp32_fuse_leaves=( - (fp32_leaf_launcher is None and m_size == 1) - if fp32_tree and not transpose_output - else None - ), + fp32_fuse_leaves=(m_size == 1) if fp32_tree and not transpose_output else None, ) workspace_elements = plan.workspace_node_count * m_size * n_size if workspace_elements > _MAX_TREE_WORKSPACE_ELEMENTS and out is None: @@ -1047,7 +1035,6 @@ def _triton_tree_gemm( preserve_a_strides=preserve_a_strides, inference_schedule=inference_schedule, output_dtype=output_dtype, - fp32_leaf_launcher=fp32_leaf_launcher, ) for start in range(0, n_size, columns_per_chunk) ] @@ -1066,7 +1053,6 @@ def _triton_tree_gemm( preserve_a_strides=preserve_a_strides, inference_schedule=inference_schedule, output_dtype=output_dtype, - fp32_leaf_launcher=fp32_leaf_launcher, ) chunks.append(chunk) return torch.cat(chunks, dim=1 if transpose_output else 0) @@ -1104,7 +1090,7 @@ def _triton_tree_gemm( transpose_output=transpose_output, preserve_a_strides=preserve_a_strides, ) - if fp32_tree and fp32_leaf_launcher is None: + if fp32_tree: # Offline H20 sweep: 64x32 avoids both masked N=96 columns and the # register spills of paired 64x64 leaves. One-row calls retain fusion. leaf_config = ( @@ -1131,12 +1117,9 @@ def _triton_tree_gemm( and direct_root_output and (use_inference_schedule or plan.fp32_fused) and (k_size >= 1536 or plan.fp32_fused) - and (fp32_leaf_launcher is None) and plan.rocm_leaf_reduction is not None ) - if fp32_leaf_launcher is not None: - fp32_leaf_launcher(a, b, workspace, plan) - elif fuse_leaf_reduction: + if fuse_leaf_reduction: assert plan.rocm_leaf_reduction is not None operation_count = len(plan.host.reduction_levels[0]) fused_leaf_grid = ( diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 6df3d1b8e..c87e0b381 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -22,10 +22,6 @@ def _gemm_fp32( ) return projection_gemm_fp32(a, b, bias=bias) - if family == "cutlass": - from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import cutlass_gemm_fp32 - - return cutlass_gemm_fp32(a, b, bias=bias) if bias is not None: return _gemm_fp32(a, b, family) + bias if family == "cuda": @@ -91,9 +87,7 @@ def forward( else None ) rhs = weight.float().t() - # CUTLASS consumes column-major B directly from contiguous [N,K] weight. - if ctx.family != "cutlass": - rhs = rhs.contiguous() + rhs = rhs.contiguous() return _gemm_fp32( a.float(), rhs, ctx.family, bias=bias.float() if bias is not None else None ) @@ -111,7 +105,7 @@ def backward(ctx, grad_out): def reducer(rows, grads): lhs = grads.float().t() - if ctx.family not in ("triton_tree", "cutlass"): + if ctx.family != "triton_tree": lhs = lhs.contiguous() return _gemm_fp32(lhs, rows.float(), ctx.family).to(weight.dtype) @@ -124,14 +118,7 @@ def bias_reducer(grads): raise RuntimeError(f"no active logical rows for {ctx.bias_parameter_id}") ones = torch.ones((1, grads.shape[0]), device=grads.device, dtype=torch.float32) # Reuse the existing one-row Triton leaves for long bias VJPs. - # CUTLASS retains its direct GEMM for short reductions. - family = ctx.family - if family == "cutlass": - from rl_engine.ops.gemm.video_input_projection_fp32 import VIDEO_PATCH_WIDTH - - if grads.shape[0] > VIDEO_PATCH_WIDTH: - family = "triton_tree" - return _gemm_fp32(ones, grads.float(), family).squeeze(0).to(ctx.bias_dtype) + return _gemm_fp32(ones, grads.float(), ctx.family).squeeze(0).to(ctx.bias_dtype) dbias = ctx.session.submit_rows( ctx.bias_parameter_id, ctx.bias_slot, grad32, bias_reducer @@ -155,16 +142,6 @@ def bias_reducer(grads): impl="ascend_det_gemm_canonical_rowfold", family="ascend", ) - elif ctx.family == "cutlass": - record_backward( - "video_input_projection_fp32", - kernel_id=( - "rl_engine._C_cutlass.gemm_fp32_partials_strided / " - "rl_engine.backends.shared.triton.gemm.det_gemm._triton_chunked_gemm_fp32" - ), - impl="cutlass_fp32_simt_chunk128_triton_merge_bias_canonical_rowfold", - family="cuda", - ) return da, dweight, None, None, None, dbias, None diff --git a/rl_engine/ops/gemm/video_input_projection_fp32.py b/rl_engine/ops/gemm/video_input_projection_fp32.py index 404bd49f1..37aab820d 100644 --- a/rl_engine/ops/gemm/video_input_projection_fp32.py +++ b/rl_engine/ops/gemm/video_input_projection_fp32.py @@ -104,9 +104,7 @@ def __call__(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None bias_parameter_id=f"{parameter_id}.bias", ) else: - rhs = weight.t() - if self.family != "cutlass": - rhs = rhs.contiguous() + rhs = weight.t().contiguous() output = _gemm_fp32(rows, rhs, self.family, bias=bias) return output.reshape(*x.shape[:-1], VIDEO_HIDDEN_WIDTH) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index e50dcb9d0..b88ccf152 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -118,9 +118,6 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): "rl_engine.backends.shared.triton.gemm.video_input_projection_fp32." "TritonVideoInputProjectionFP32Op" ) - CUDA_VIDEO_INPUT_PROJECTION_FP32 = ( - "rl_engine.backends.cuda.gemm.video_input_projection_fp32.CUDAVideoInputProjectionFP32Op" - ) ASCEND_DET_GEMM = "rl_engine.kernels.ops.ascend.matmul.det_gemm.DetGemmAscendOp" # NON-deterministic reference (torch.matmul); reference/benchmark ONLY, # intentionally excluded from det_gemm dispatch (cuBLAS breaks invariance). @@ -242,45 +239,6 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint="h3-query-chunk-v1; binary profile checked at invocation", ), - OperatorBackendDescriptor( - semantic_op="video_input_projection_fp32", - backend_id="cuda-cutlass-video-input-projection-fp32-v7", - supported_targets=frozenset({"rollout", "training"}), - supported_devices=frozenset({"cuda"}), - supported_dtypes=frozenset({"float32"}), - supported_topologies={ - "rollout": { - "world_size": (1,), - "tensor_parallel_size": (1,), - "context_parallel_size": (1,), - }, - "training": {"world_size": (1,), "sharding": ("unsharded",)}, - }, - determinism_or_alignment_properties={ - "algorithm": "cutlass_simt_fp32_chunk128_triton_merge_bias_canonical_parameter_vjp", - "forward_bias": "linear_combination_broadcast_source", - "reduction_backend": "triton", - "bias_reduction_backend": "triton_for_k_gt_96", - "input_width": 96, - "output_width": 5376, - "cuda_sm": 90, - "batch_invariant": True, - "deterministic": True, - "strict_observable": True, - "training_mode": "eager", - "parameter_gradient_order": "logical_key", - "reduction": "fixed_128_chunks_balanced_fp32_merge", - "vjp_chunk_k": 128, - "bias_reduction": "fixed_midpoint_tree_32", - }, - lifecycle=OperatorLifecycle.REQUEST, - implementation_class_or_factory=OpBackend.CUDA_VIDEO_INPUT_PROJECTION_FP32.value, - fallback_policy=OperatorFallbackPolicy.ERROR, - version_or_build_fingerprint=( - "cutlass-56ca28a-fp32-simt-64x64x8-chunk128-merge64-" - "strided-actual-k-bias-epilogue-triton-bias-h3-v7" - ), - ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", backend_id="triton-video-input-projection-fp32-v8", diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 99f8bae32..76f4e0a58 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -59,10 +59,6 @@ def _load_object(path: str) -> Any: "rl_engine.backends.shared.triton.gemm.video_input_projection_fp32." "TritonVideoInputProjectionFP32Op" ), - "cuda": ( - "rl_engine.backends.cuda.gemm.video_input_projection_fp32." - "CUDAVideoInputProjectionFP32Op" - ), }, grad_input_names=("x", "weight", "bias"), ), diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index 70c45c7cb..b813506b9 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -13,39 +13,8 @@ from build_tools.extensions import _resolve_package_data_source, get_cmdclass -def test_optional_cutlass_build_has_independent_strict_flags(monkeypatch, tmp_path): - from build_tools import extensions - - monkeypatch.setattr(torch.version, "hip", None) - monkeypatch.setattr(torch.version, "cuda", "12.9") - monkeypatch.setenv("CUTLASS_ROOT", str(tmp_path)) - monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", "1") - monkeypatch.setattr( - extensions, "_cutlass_headers_fingerprint", lambda _: extensions._CUTLASS_HEADERS_SHA256 - ) - candidate = extensions._cutlass_extension(torch, lambda **kw: kw, []) - assert candidate["name"] == "rl_engine._C_cutlass" - flags = candidate["extra_compile_args"]["nvcc"] - assert "--fmad=false" in flags and "--ftz=false" in flags - assert "--use_fast_math" not in flags - assert "csrc/bindings/ops.cpp" not in candidate["sources"] - - -def test_cutlass_build_rejects_unpinned_headers(monkeypatch, tmp_path): - import pytest - - from build_tools import extensions - - monkeypatch.setattr(torch.version, "hip", None) - monkeypatch.setattr(torch.version, "cuda", "12.9") - monkeypatch.setenv("CUTLASS_ROOT", str(tmp_path)) - with pytest.raises(RuntimeError, match="fingerprint"): - extensions._cutlass_extension(torch, lambda **kw: kw, []) - - def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: captured: dict[str, Any] = {} - monkeypatch.delenv("RL_KERNEL_BUILD_CUTLASS", raising=False) def fake_setup(**kwargs: Any) -> None: captured.update(kwargs) diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 494b0dad8..f141759c5 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -4,7 +4,6 @@ """H3 FP32 projection accuracy, actual autograd, and logical-row invariance.""" import os -from importlib.util import find_spec import pytest import torch @@ -20,7 +19,6 @@ and torch.version.hip is None and torch.cuda.get_device_capability() == (9, 0) ) -BACKENDS = ["triton", "cuda"] @pytest.mark.parametrize( @@ -130,20 +128,12 @@ def unexpected_proof(keys): assert torch.equal(actual, values) -def _op(backend): - if backend == "triton": - from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( - TritonVideoInputProjectionFP32Op, - ) - - return TritonVideoInputProjectionFP32Op() - from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import ( - CUDAVideoInputProjectionFP32Op, +def _op(): + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + TritonVideoInputProjectionFP32Op, ) - if find_spec("rl_engine._C_cutlass") is None: - pytest.skip("requires the optional CUTLASS extension") - return CUDAVideoInputProjectionFP32Op() + return TritonVideoInputProjectionFP32Op() def _keys(rows, device): @@ -280,10 +270,9 @@ def test_invalid_inputs(invalid): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("rows", [1, 9, 17, 65, 129, 257]) @pytest.mark.parametrize("with_bias", [False, True]) -def test_forward_and_actual_gradients(backend, rows, with_bias): +def test_forward_and_actual_gradients(rows, with_bias): data = list(_data(rows)) if not with_bias: data[2] = None @@ -292,12 +281,12 @@ def test_forward_and_actual_gradients(backend, rows, with_bias): x, weight, _, grad = data x, weight = x.requires_grad_(), weight.requires_grad_() with canonical_backward_session() as session: - y = _op(backend)(x, weight, logical_keys=_keys(rows, "cuda"), parameter_id="p") + y = _op()(x, weight, logical_keys=_keys(rows, "cuda"), parameter_id="p") y.backward(grad) session.validate_complete() actual = (y, x.grad, weight.grad) else: - actual = _run(_op(backend), data, _keys(rows, "cuda")) + actual = _run(_op(), data, _keys(rows, "cuda")) cx, cw = data[0].detach().cpu().requires_grad_(), data[1].detach().cpu().requires_grad_() cb = None if not with_bias else data[2].detach().cpu().requires_grad_() gold = NativeVideoInputProjectionFP32Op()(cx, cw, cb) @@ -308,10 +297,9 @@ def test_forward_and_actual_gradients(backend, rows, with_bias): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("rows", [9, 65, 127, 128, 129, 257]) -def test_parameter_bi_permutation_chunk_padding(backend, rows): - op, data, keys = _op(backend), _data(rows), _keys(rows, "cuda") +def test_parameter_bi_permutation_chunk_padding(rows): + op, data, keys = _op(), _data(rows), _keys(rows, "cuda") baseline = _run(op, data, keys) order = torch.randperm(rows, device="cuda") permuted = _run( @@ -342,9 +330,8 @@ def test_parameter_bi_permutation_chunk_padding(backend, rows): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) -def test_row_placement_and_unrelated_rows(backend): - op, data = _op(backend), _data(65) +def test_row_placement_and_unrelated_rows(): + op, data = _op(), _data(65) single = _run(op, [data[0][:1], data[1], data[2], data[3][:1]], _keys(1, "cuda")) data[0][37] = data[0][0] data[3][37] = data[3][0] @@ -360,9 +347,8 @@ def test_row_placement_and_unrelated_rows(backend): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) -def test_noncontiguous_upstream_and_forward_graph(backend): - op = _op(backend) +def test_noncontiguous_upstream_and_forward_graph(): + op = _op() x, weight, bias, grad = _data(17) with canonical_backward_session() as session: y = op( @@ -390,8 +376,7 @@ def test_noncontiguous_upstream_and_forward_graph(backend): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) -def test_real_weights_patch_path(backend): +def test_real_weights_patch_path(): path = os.environ.get("H3_PROJECTION_WEIGHTS") if not path: pytest.skip("set H3_PROJECTION_WEIGHTS to pinned H3 safetensors fixture") @@ -411,7 +396,7 @@ def test_real_weights_patch_path(backend): latent = torch.randn(1, 24, 13, 48, 84, device="cuda", dtype=torch.bfloat16) x = NativeVideoPatchPackUnpackOp().pack(latent).reshape(-1, 96).float() grad = torch.randn(x.shape[0], 5376, device="cuda") * 0.0625 - actual = _run(_op(backend), (x, weight, bias, grad), _keys(x.shape[0], "cuda")) + actual = _run(_op(), (x, weight, bias, grad), _keys(x.shape[0], "cuda")) cx, cw, cb = [t.detach().cpu().requires_grad_() for t in (x, weight, bias)] gold = NativeVideoInputProjectionFP32Op()(cx, cw, cb) gold.backward(grad.cpu()) @@ -419,7 +404,7 @@ def test_real_weights_patch_path(backend): _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") order = torch.arange(x.shape[0] - 1, -1, -1, device="cuda") chunked = _run( - _op(backend), + _op(), (x, weight, bias, grad), _keys(x.shape[0], "cuda"), list(torch.tensor_split(order, 3)), @@ -436,7 +421,7 @@ def test_default_cuda_dispatch(): assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v8" x, weight, bias, _ = _data(3) with torch.no_grad(): - _assert_bytes(op(x, weight, bias), _op("triton")(x, weight, bias)) + _assert_bytes(op(x, weight, bias), _op()(x, weight, bias)) @pytest.mark.skipif(not H20, reason="requires CUDA SM90") @@ -484,13 +469,11 @@ def test_fp32_tree_preserves_fp32_nodes(k): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) -def test_fp32_tree_workspace_chunking(backend, monkeypatch): +def test_fp32_tree_workspace_chunking(monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm from rl_engine.ops.gemm.canonical_linear import _gemm_fp32 - _op(backend) # Apply the same optional-extension eligibility as other tests. - family = "triton_tree" if backend == "triton" else "cutlass" + family = "triton_tree" a, b = torch.randn(7, 129, device="cuda"), torch.randn(129, 5, device="cuda") expected = _gemm_fp32(a, b, family) plan = det_gemm._build_tree_plan(129) @@ -503,13 +486,12 @@ def test_fp32_tree_workspace_chunking(backend, monkeypatch): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("rows", [129, 13104]) -def test_standard_normal_long_reduction(backend, rows): +def test_standard_normal_long_reduction(rows): # Unscaled inputs expose cancellation error in the 5376-term input VJP. data = _data(rows) data = tuple(t / scale for t, scale in zip(data, (0.125, 0.0625, 0.125, 0.0625), strict=True)) - actual = _run(_op(backend), data, _keys(rows, "cuda")) + actual = _run(_op(), data, _keys(rows, "cuda")) cx, cw, cb = [t.detach().cpu().requires_grad_() for t in data[:3]] gold = NativeVideoInputProjectionFP32Op()(cx, cw, cb) gold.backward(data[3].cpu()) @@ -530,16 +512,14 @@ def test_cpu_fp32_reference_long_dot_accuracy(): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("k", [127, 128, 129, 5376, 13104]) -def test_fp32_chunked_gemm_strides_and_workspace(backend, k, monkeypatch): +def test_fp32_chunked_gemm_strides_and_workspace(k, monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm from rl_engine.ops.gemm.canonical_linear import _gemm_fp32 from rl_engine.reference.gemm.video_input_projection_fp32 import _reference_matmul_fp32 torch.manual_seed(420) - _op(backend) - family = "triton_tree" if backend == "triton" else "cutlass" + family = "triton_tree" a = torch.randn(k, 7, device="cuda").t() b = torch.randn(k, 96, device="cuda") expected = _gemm_fp32(a.contiguous(), b, family) @@ -612,15 +592,12 @@ def observe(self, parameter_id, values): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("k", [64, 65, 128, 1024, 13104]) -def test_fp32_fused_tree_and_transposed_packing_match_scalar_reference(backend, k): - from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import _cutlass_tree_leaves +def test_fp32_fused_tree_and_strided_inputs_match_scalar_reference(k): from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_tree_gemm - _op(backend) torch.manual_seed(422) - # Both operands exercise the strided transpose packing, including tail rows. + # Exercise a transposed A view and tail rows. a = torch.randn(k, 3, device="cuda").t() b = torch.randn(k, 5, device="cuda") @@ -630,7 +607,6 @@ def gemm(a, b): b, output_dtype=torch.float32, preserve_a_strides=True, - fp32_leaf_launcher=_cutlass_tree_leaves if backend == "cuda" else None, ) result = gemm(a, b) @@ -639,13 +615,11 @@ def gemm(a, b): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) -def test_compact_fp32_tree_workspace_preserves_chunked_outputs(backend, monkeypatch): +def test_compact_fp32_tree_workspace_preserves_chunked_outputs(monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm from rl_engine.ops.gemm.canonical_linear import _gemm_fp32 - _op(backend) - family = "triton_tree" if backend == "triton" else "cutlass" + family = "triton_tree" a, b = torch.randn(7, 1024, device="cuda"), torch.randn(1024, 5, device="cuda") expected = _gemm_fp32(a, b, family) plan = det_gemm._device_tree_plan(1024, a.device, fp32_fuse_leaves=False) @@ -695,32 +669,24 @@ def launch(*args, **kwargs): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("rows", [1, 17, 129]) -def test_forward_bias_fusion_preserves_fp32_rounding(backend, rows): +def test_forward_bias_fusion_preserves_fp32_rounding(rows): from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 x, weight, bias, _ = _data(rows) - if backend == "cuda": - from rl_engine.backends.cuda.gemm.video_input_projection_fp32 import cutlass_gemm_fp32 - - gemm = cutlass_gemm_fp32 - rhs = weight.t() - else: - gemm = _triton_gemm_fp32 - rhs = weight.t().contiguous() + gemm = _triton_gemm_fp32 + rhs = weight.t().contiguous() expected = gemm(x, rhs) + bias actual = gemm(x, rhs, bias=bias) _assert_bytes(actual, expected) with torch.no_grad(): - _assert_bytes(_op(backend)(x, weight, bias), expected) + _assert_bytes(_op()(x, weight, bias), expected) @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -@pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("frozen", ["weight", "bias", "both"]) -def test_backend_frozen_parameters_preserve_trainable_gradients(backend, frozen): - op = _op(backend) +def test_backend_frozen_parameters_preserve_trainable_gradients(frozen): + op = _op() data, keys = _data(129), _keys(129, "cuda") expected = _run(op, data, keys) actual = _run( diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index d315eb9da..437e6988b 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -18,11 +18,6 @@ def test_video_input_projection_default_and_strict_scope(): assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) assert desc.determinism_or_alignment_properties["vjp_tile"] == (128, 32, 32) - cutlass = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "cuda-cutlass-video-input-projection-fp32-v7" - ) - assert cutlass.determinism_or_alignment_properties["vjp_chunk_k"] == 128 - assert cutlass.determinism_or_alignment_properties["bias_reduction"] == "fixed_midpoint_tree_32" with pytest.raises(RuntimeError, match="SM90"): registry.get_op("video_input_projection_fp32", device="cpu") From 0a152fb2d9b609dfe6b8947329d72e57f832a869 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Sat, 17 Oct 2026 14:49:26 +0800 Subject: [PATCH 10/26] Separate canonical linear backend from projection matmul Signed-off-by: walker-ai <2398833647@qq.com> --- docs/operators/video-input-projection-fp32.md | 9 ++ .../gemm/video_input_projection_fp32.py | 4 +- rl_engine/ops/gemm/canonical_linear.py | 77 ++++++--- .../ops/gemm/video_input_projection_fp32.py | 15 +- .../gemm/test_video_input_projection_fp32.py | 146 ++++++++++++++++-- 5 files changed, 208 insertions(+), 43 deletions(-) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 450e6e826..d673f09c1 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -59,6 +59,15 @@ single-device, unsharded support and exact implementation identities. An unavailable explicit candidate raises an error. CPU/ROCm/MUSA/NPU execution is outside this contribution's strict support scope. +The wrapper reuses `canonical_linear_fp32` with `family="triton"` and +`matmul_op=projection_gemm_fp32`. The family identifies the backend. The callable +selects the existing forward and VJP GEMM schedules and handles their input +strides. Forward, dX, dW, and dbias use the same callable. Existing canonical +linear callers omit `matmul_op` and retain their original backend dispatch. +The H3 wrapper separately enables `reuse_canonical_order`; the session verifies +the keys before it skips sorting. Backward provenance records the actual +matmul entry point. + ## Accuracy Forward computes `Y = X @ W.T + bias` in FP32. The Triton forward adds bias diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index a6a0425a0..16bb3b582 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -35,7 +35,9 @@ def projection_gemm_fp32(a, b, *, bias=None): class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): - family = "triton_tree" + family = "triton" + matmul_op = staticmethod(projection_gemm_fp32) + reuse_canonical_order = True backend_id = "triton-video-input-projection-fp32-v8" def __init__(self): diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index c87e0b381..7c79a458f 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -5,6 +5,8 @@ from __future__ import annotations +from collections.abc import Callable + import torch from rl_engine.backends.extension import _C @@ -14,14 +16,15 @@ def _gemm_fp32( - a: torch.Tensor, b: torch.Tensor, family: str, *, bias: torch.Tensor | None = None + a: torch.Tensor, + b: torch.Tensor, + family: str, + *, + bias: torch.Tensor | None = None, + matmul_op: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor: - if family == "triton_tree": - from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( - projection_gemm_fp32, - ) - - return projection_gemm_fp32(a, b, bias=bias) + if matmul_op is not None: + return matmul_op(a, b) if bias is None else matmul_op(a, b, bias=bias) if bias is not None: return _gemm_fp32(a, b, family) + bias if family == "cuda": @@ -50,6 +53,8 @@ def forward( family, bias, bias_parameter_id, + matmul_op, + reuse_canonical_order, ): session = active_session() parameter_grad = ctx.needs_input_grad[1] or ctx.needs_input_grad[5] @@ -63,6 +68,7 @@ def forward( ctx.session = session ctx.parameter_id = str(parameter_id) ctx.family = str(family) + ctx.matmul_op = matmul_op ctx.bias_parameter_id = str(bias_parameter_id) ctx.bias_dtype = bias.dtype if bias is not None else None contribution = object() if ctx.needs_input_grad[1] and ctx.needs_input_grad[5] else None @@ -71,7 +77,7 @@ def forward( ctx.parameter_id, logical_keys, contribution=contribution, - reuse_canonical_order=ctx.family == "triton_tree", + reuse_canonical_order=reuse_canonical_order, ) if ctx.needs_input_grad[1] else None @@ -81,15 +87,17 @@ def forward( ctx.bias_parameter_id, logical_keys, contribution=contribution, - reuse_canonical_order=ctx.family == "triton_tree", + reuse_canonical_order=reuse_canonical_order, ) if ctx.needs_input_grad[5] else None ) - rhs = weight.float().t() - rhs = rhs.contiguous() return _gemm_fp32( - a.float(), rhs, ctx.family, bias=bias.float() if bias is not None else None + a.float(), + weight.float().t().contiguous(), + ctx.family, + bias=bias.float() if bias is not None else None, + matmul_op=ctx.matmul_op, ) @staticmethod @@ -98,16 +106,18 @@ def backward(ctx, grad_out): grad32 = grad_out.contiguous().float() da = None if ctx.needs_input_grad[0]: - da = _gemm_fp32(grad32, weight.float(), ctx.family).to(a.dtype) + da = _gemm_fp32(grad32, weight.float(), ctx.family, matmul_op=ctx.matmul_op).to(a.dtype) dweight = None if ctx.needs_input_grad[1]: def reducer(rows, grads): lhs = grads.float().t() - if ctx.family != "triton_tree": + if ctx.matmul_op is None: lhs = lhs.contiguous() - return _gemm_fp32(lhs, rows.float(), ctx.family).to(weight.dtype) + return _gemm_fp32(lhs, rows.float(), ctx.family, matmul_op=ctx.matmul_op).to( + weight.dtype + ) dweight = ctx.session.submit_linear(ctx.parameter_id, ctx.slot, a, grad32, reducer) dbias = None @@ -117,21 +127,28 @@ def bias_reducer(grads): if grads.shape[0] == 0: raise RuntimeError(f"no active logical rows for {ctx.bias_parameter_id}") ones = torch.ones((1, grads.shape[0]), device=grads.device, dtype=torch.float32) - # Reuse the existing one-row Triton leaves for long bias VJPs. - return _gemm_fp32(ones, grads.float(), ctx.family).squeeze(0).to(ctx.bias_dtype) + return ( + _gemm_fp32(ones, grads.float(), ctx.family, matmul_op=ctx.matmul_op) + .squeeze(0) + .to(ctx.bias_dtype) + ) dbias = ctx.session.submit_rows( ctx.bias_parameter_id, ctx.bias_slot, grad32, bias_reducer ) - if ctx.family in ("triton", "triton_tree"): + if ctx.matmul_op is not None: + module = getattr(ctx.matmul_op, "__module__", type(ctx.matmul_op).__module__) + qualname = getattr(ctx.matmul_op, "__qualname__", type(ctx.matmul_op).__qualname__) record_backward( "det_gemm", - kernel_id=( - "rl_engine.backends.shared.triton.gemm.det_gemm." - "_triton_chunked_gemm_fp32 / _triton_tree_gemm" - if ctx.family == "triton_tree" - else "rl_engine.backends.shared.triton.gemm.det_gemm._triton_gemm" - ), + kernel_id=f"{module}.{qualname}", + impl=f"{ctx.family}_det_gemm_canonical_rowfold", + family=ctx.family, + ) + elif ctx.family == "triton": + record_backward( + "det_gemm", + kernel_id="rl_engine.backends.shared.triton.gemm.det_gemm._triton_gemm", impl="triton_det_gemm_canonical_rowfold", family="triton", ) @@ -142,7 +159,7 @@ def bias_reducer(grads): impl="ascend_det_gemm_canonical_rowfold", family="ascend", ) - return da, dweight, None, None, None, dbias, None + return da, dweight, None, None, None, dbias, None, None, None def canonical_linear_fp32( @@ -154,7 +171,15 @@ def canonical_linear_fp32( family: str, bias: torch.Tensor | None = None, bias_parameter_id: str | None = None, + matmul_op: Callable[..., torch.Tensor] | None = None, + reuse_canonical_order: bool = False, ) -> torch.Tensor: + """Apply a canonical linear using the default backend or an FP32 matmul callable. + + A custom matmul handles input strides and accepts ``bias=`` when bias is + supplied. The same callable computes forward, dX, dW, and dbias. Canonical + order reuse requests a runtime proof before skipping stable sorting. + """ return _CanonicalLinearFn.apply( a, weight, @@ -163,6 +188,8 @@ def canonical_linear_fp32( family, bias, bias_parameter_id if bias_parameter_id is not None else f"{parameter_id}.bias", + matmul_op, + reuse_canonical_order, ) diff --git a/rl_engine/ops/gemm/video_input_projection_fp32.py b/rl_engine/ops/gemm/video_input_projection_fp32.py index 37aab820d..8b2c6849e 100644 --- a/rl_engine/ops/gemm/video_input_projection_fp32.py +++ b/rl_engine/ops/gemm/video_input_projection_fp32.py @@ -5,6 +5,8 @@ from __future__ import annotations +from collections.abc import Callable + import torch from rl_engine.ops.autograd.canonical_backward import active_session @@ -72,6 +74,8 @@ class VideoInputProjectionFP32Op: family: str backend_id: str + matmul_op: Callable[..., torch.Tensor] | None = None + reuse_canonical_order: bool = False def _check_backend(self, x): raise NotImplementedError @@ -102,10 +106,17 @@ def __call__(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None family=self.family, bias=bias, bias_parameter_id=f"{parameter_id}.bias", + matmul_op=self.matmul_op, + reuse_canonical_order=self.reuse_canonical_order, ) else: - rhs = weight.t().contiguous() - output = _gemm_fp32(rows, rhs, self.family, bias=bias) + output = _gemm_fp32( + rows, + weight.t().contiguous(), + self.family, + bias=bias, + matmul_op=self.matmul_op, + ) return output.reshape(*x.shape[:-1], VIDEO_HIDDEN_WIDTH) forward = __call__ diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index f141759c5..8c077197f 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -213,6 +213,119 @@ def test_canonical_linear_bias_and_legacy_call(with_bias): torch.testing.assert_close(bias.grad, grad.sum(0)) +@pytest.mark.parametrize("with_bias", [False, True]) +@pytest.mark.parametrize("reuse_order", [False, True]) +def test_canonical_linear_custom_matmul_covers_all_vjps(with_bias, reuse_order, monkeypatch): + from rl_engine.ops.autograd.backward_runtime import ( + reset_backward_runtime, + snapshot_backward_runtime, + ) + from rl_engine.ops.autograd.canonical_backward import CanonicalBackwardSession + + torch.manual_seed(424) + x = torch.randn(7, 3, requires_grad=True) + weight = torch.randn(5, 3, requires_grad=True) + bias = torch.randn(5, requires_grad=True) if with_bias else None + grad = torch.randn(7, 5) + calls, registrations = [], [] + original_register = CanonicalBackwardSession.register + + def register(self, parameter_id, keys, **kwargs): + registrations.append(kwargs["reuse_canonical_order"]) + return original_register(self, parameter_id, keys, **kwargs) + + def gemm(a, b, *, bias=None): + calls.append((a.shape, b.shape, a.stride(0), bias is not None)) + value = a @ b + return value if bias is None else value + bias + + def two_argument_gemm(a, b): + return gemm(a, b) + + matmul_op = gemm if with_bias else two_argument_gemm + monkeypatch.setattr(CanonicalBackwardSession, "register", register) + reset_backward_runtime() + try: + with canonical_backward_session() as session: + y = canonical_linear_fp32( + x, + weight, + _keys(7, "cpu"), + parameter_id="weight", + family="pytorch", + bias=bias, + matmul_op=matmul_op, + reuse_canonical_order=reuse_order, + ) + y.backward(grad) + session.validate_complete() + torch.testing.assert_close(y, torch.nn.functional.linear(x, weight, bias)) + torch.testing.assert_close(x.grad, grad @ weight.detach()) + torch.testing.assert_close(weight.grad, grad.t() @ x.detach()) + expected_shapes = [((7, 3), (3, 5)), ((7, 5), (5, 3)), ((5, 7), (7, 3))] + if with_bias: + torch.testing.assert_close(bias.grad, grad.sum(0)) + expected_shapes.append(((1, 7), (7, 5))) + assert [(a, b) for a, b, _, _ in calls] == expected_shapes + assert calls[0][3] == with_bias + assert all(not call[3] for call in calls[1:]) + assert calls[2][2] == 1 + assert registrations == [reuse_order] * (2 if with_bias else 1) + provenance = snapshot_backward_runtime()["det_gemm"] + assert provenance["kernel_id"] == f"{matmul_op.__module__}.{matmul_op.__qualname__}" + assert provenance["family"] == "pytorch" + finally: + reset_backward_runtime() + + +def test_canonical_linear_default_triton_retains_existing_dispatch(monkeypatch): + from rl_engine.ops.autograd.backward_runtime import ( + reset_backward_runtime, + snapshot_backward_runtime, + ) + from rl_engine.ops.gemm import canonical_linear + + calls = [] + + def existing_gemm(a, b, *, output_dtype): + assert output_dtype == torch.float32 + calls.append((a.shape, b.shape)) + return a @ b + + monkeypatch.setattr(canonical_linear, "_triton_gemm", existing_gemm) + x = torch.randn(7, 3, requires_grad=True) + weight = torch.randn(5, 3, requires_grad=True) + grad = torch.randn(7, 5) + reset_backward_runtime() + try: + with canonical_backward_session() as session: + output = canonical_linear_fp32( + x, weight, _keys(7, "cpu"), parameter_id="weight", family="triton" + ) + output.backward(grad) + session.validate_complete() + torch.testing.assert_close(output, x @ weight.t()) + torch.testing.assert_close(x.grad, grad @ weight.detach()) + torch.testing.assert_close(weight.grad, grad.t() @ x.detach()) + assert calls == [((7, 3), (3, 5)), ((7, 5), (5, 3)), ((5, 7), (7, 3))] + assert snapshot_backward_runtime()["det_gemm"]["kernel_id"] == ( + "rl_engine.backends.shared.triton.gemm.det_gemm._triton_gemm" + ) + finally: + reset_backward_runtime() + + +def test_video_projection_declares_backend_and_matmul_separately(): + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + TritonVideoInputProjectionFP32Op, + projection_gemm_fp32, + ) + + assert TritonVideoInputProjectionFP32Op.family == "triton" + assert TritonVideoInputProjectionFP32Op.matmul_op is projection_gemm_fp32 + assert TritonVideoInputProjectionFP32Op.reuse_canonical_order + + def test_input_gradient_without_parameter_session(): x = torch.ones(3, 2, requires_grad=True) weight = torch.ones(4, 2) @@ -471,18 +584,19 @@ def test_fp32_tree_preserves_fp32_nodes(k): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") def test_fp32_tree_workspace_chunking(monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm - from rl_engine.ops.gemm.canonical_linear import _gemm_fp32 + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + projection_gemm_fp32, + ) - family = "triton_tree" a, b = torch.randn(7, 129, device="cuda"), torch.randn(129, 5, device="cuda") - expected = _gemm_fp32(a, b, family) + expected = projection_gemm_fp32(a, b) plan = det_gemm._build_tree_plan(129) monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.node_count * 5 * 2) - actual = _gemm_fp32(a, b, family) + actual = projection_gemm_fp32(a, b) _assert_bytes(actual, expected) monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.node_count * 2) # Exercise output-column chunking when even one whole row exceeds budget. - _assert_bytes(_gemm_fp32(a, b, family), expected) + _assert_bytes(projection_gemm_fp32(a, b), expected) @pytest.mark.skipif(not H20, reason="requires CUDA SM90") @@ -515,20 +629,21 @@ def test_cpu_fp32_reference_long_dot_accuracy(): @pytest.mark.parametrize("k", [127, 128, 129, 5376, 13104]) def test_fp32_chunked_gemm_strides_and_workspace(k, monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm - from rl_engine.ops.gemm.canonical_linear import _gemm_fp32 + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + projection_gemm_fp32, + ) from rl_engine.reference.gemm.video_input_projection_fp32 import _reference_matmul_fp32 torch.manual_seed(420) - family = "triton_tree" a = torch.randn(k, 7, device="cuda").t() b = torch.randn(k, 96, device="cuda") - expected = _gemm_fp32(a.contiguous(), b, family) - actual = _gemm_fp32(a, b, family) + expected = projection_gemm_fp32(a.contiguous(), b) + actual = projection_gemm_fp32(a, b) _assert_bytes(actual, expected) _assert_accuracy(actual, _reference_matmul_fp32(a, b), "gradient_accuracy") chunks = (k + 127) // 128 monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", chunks * 96 * 2) - _assert_bytes(_gemm_fp32(a, b, family), expected) + _assert_bytes(projection_gemm_fp32(a, b), expected) @pytest.mark.parametrize("bias_uses", ["all", "partial", "distinct"]) @@ -617,18 +732,19 @@ def gemm(a, b): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") def test_compact_fp32_tree_workspace_preserves_chunked_outputs(monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm - from rl_engine.ops.gemm.canonical_linear import _gemm_fp32 + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + projection_gemm_fp32, + ) - family = "triton_tree" a, b = torch.randn(7, 1024, device="cuda"), torch.randn(1024, 5, device="cuda") - expected = _gemm_fp32(a, b, family) + expected = projection_gemm_fp32(a, b) plan = det_gemm._device_tree_plan(1024, a.device, fp32_fuse_leaves=False) assert plan.fp32_fused assert plan.workspace_node_count < plan.host.node_count monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.workspace_node_count * 10) - _assert_bytes(_gemm_fp32(a, b, family), expected) + _assert_bytes(projection_gemm_fp32(a, b), expected) monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.workspace_node_count * 2) - _assert_bytes(_gemm_fp32(a, b, family), expected) + _assert_bytes(projection_gemm_fp32(a, b), expected) @pytest.mark.skipif(not H20, reason="requires CUDA SM90") From d38dff510e83ea35588cd6bd2696c92597c816e3 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Sat, 17 Oct 2026 15:54:42 +0800 Subject: [PATCH 11/26] Remove caller-controlled canonical ordering option Signed-off-by: walker-ai <2398833647@qq.com> --- docs/operators/video-input-projection-fp32.md | 22 ++++++++---------- .../gemm/video_input_projection_fp32.py | 3 +-- rl_engine/ops/autograd/canonical_backward.py | 6 ++--- rl_engine/ops/gemm/canonical_linear.py | 23 ++++--------------- .../ops/gemm/video_input_projection_fp32.py | 2 -- rl_engine/runtime/registry.py | 4 ++-- .../gemm/test_video_input_projection_fp32.py | 23 ++++++++----------- tests/runtime/test_kernel_registry.py | 6 ++++- 8 files changed, 35 insertions(+), 54 deletions(-) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index d673f09c1..f542980e8 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -64,8 +64,7 @@ The wrapper reuses `canonical_linear_fp32` with `family="triton"` and selects the existing forward and VJP GEMM schedules and handles their input strides. Forward, dX, dW, and dbias use the same callable. Existing canonical linear callers omit `matmul_op` and retain their original backend dispatch. -The H3 wrapper separately enables `reuse_canonical_order`; the session verifies -the keys before it skips sorting. Backward provenance records the actual +On CUDA, the session internally checks eligible keys before skipping sorting. Backward provenance records the actual matmul entry point. ## Accuracy @@ -86,14 +85,13 @@ once. Weight and bias share sorting and the sorted upstream gradient when their complete contribution sequences match. Calls with different contribution sets normalize each parameter independently. -The H3 Triton path explicitly enables verified canonical-order reuse for at least -4096 submitted rows. Smaller inputs retain stable normalization to avoid the -proof overhead. At each eligible normalization, a small GPU check proves that every row is active and adjacent -keys strictly increase in lexicographic order. This includes block boundaries. -When the proof succeeds, the session reuses contiguous gradients without sorting -or gathering the large matrix. Otherwise it runs the original stable sort. -The proof reads the current keys and does not cache a caller's ordering claim. -Other canonical-session callers retain stable normalization by default. +For at least 4096 submitted CUDA rows, the session checks that every row is +active and adjacent keys strictly increase in logical order. This includes +block boundaries. When the check succeeds, the session reuses contiguous +gradients without sorting or gathering the large matrix. Otherwise it runs the +existing stable sort. The check reads the current keys on each normalization. +Callers supply logical keys and do not select this optimization. Smaller inputs, +CPU, and other device paths retain stable sorting. Short reductions up to the 96-wide video input use the existing direct FP32 GEMM. The Triton input and weight VJPs reuse that kernel with fixed 128-term @@ -102,10 +100,10 @@ partials remain FP32. A second grid uses a fixed balanced reduction over the next power-of-two chunk count, with zero padding. The merge uses 64 output elements per program and four warps. VJP GEMMs use fixed 128-by-32 output tiles, 32-term dot tiles, and four warps. The forward retains its 64-by-64 output tile. -These settings define the v8 Triton implementation identity. The K partition and VJP +These settings define the v9 Triton implementation identity. The K partition and VJP tile do not depend on M, row placement, or occupancy. The VJP kernels consume positive-stride views directly, including `G.T`. -The Triton v8 forward omits K or N load predicates only when the corresponding +The Triton forward omits K or N load predicates only when the corresponding 32-term or 64-column tiles are complete. It retains M predicates for every row. Column-major A views with unit row stride use a fixed `maxnreg=96` VJP allocation. Other VJP layouts retain the default allocation. These choices preserve the dot diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index 16bb3b582..244f08d4b 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -37,8 +37,7 @@ def projection_gemm_fp32(a, b, *, bias=None): class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton" matmul_op = staticmethod(projection_gemm_fp32) - reuse_canonical_order = True - backend_id = "triton-video-input-projection-fp32-v8" + backend_id = "triton-video-input-projection-fp32-v9" def __init__(self): if not _TRITON_AVAILABLE: diff --git a/rl_engine/ops/autograd/canonical_backward.py b/rl_engine/ops/autograd/canonical_backward.py index d8dceaf6d..4ae7819cf 100644 --- a/rl_engine/ops/autograd/canonical_backward.py +++ b/rl_engine/ops/autograd/canonical_backward.py @@ -42,7 +42,6 @@ class _Use: contribution: object | None = None rows: torch.Tensor | None = None grads: torch.Tensor | None = None - reuse_canonical_order: bool = False @dataclass @@ -63,7 +62,6 @@ def register( keys: torch.Tensor, *, contribution: object | None = None, - reuse_canonical_order: bool = False, ) -> int: if self.backward_started: raise RuntimeError("cannot register canonical rows after backward started") @@ -74,7 +72,6 @@ def register( _Use( keys=keys.detach(), contribution=contribution, - reuse_canonical_order=reuse_canonical_order, ) ) return slot @@ -106,7 +103,8 @@ def _normalize( keys = self._concatenate([item.keys for item in entries]) reuse_order = ( keys.shape[0] >= _CANONICAL_ORDER_REUSE_MIN_ROWS - and all(item.reuse_canonical_order for item in entries) + and keys.device.type == "cuda" + and torch.version.hip is None and _keys_are_canonical(keys) ) if reuse_order: diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 7c79a458f..08a402b1a 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -54,7 +54,6 @@ def forward( bias, bias_parameter_id, matmul_op, - reuse_canonical_order, ): session = active_session() parameter_grad = ctx.needs_input_grad[1] or ctx.needs_input_grad[5] @@ -73,22 +72,12 @@ def forward( ctx.bias_dtype = bias.dtype if bias is not None else None contribution = object() if ctx.needs_input_grad[1] and ctx.needs_input_grad[5] else None ctx.slot = ( - session.register( - ctx.parameter_id, - logical_keys, - contribution=contribution, - reuse_canonical_order=reuse_canonical_order, - ) + session.register(ctx.parameter_id, logical_keys, contribution=contribution) if ctx.needs_input_grad[1] else None ) ctx.bias_slot = ( - session.register( - ctx.bias_parameter_id, - logical_keys, - contribution=contribution, - reuse_canonical_order=reuse_canonical_order, - ) + session.register(ctx.bias_parameter_id, logical_keys, contribution=contribution) if ctx.needs_input_grad[5] else None ) @@ -159,7 +148,7 @@ def bias_reducer(grads): impl="ascend_det_gemm_canonical_rowfold", family="ascend", ) - return da, dweight, None, None, None, dbias, None, None, None + return da, dweight, None, None, None, dbias, None, None def canonical_linear_fp32( @@ -172,13 +161,12 @@ def canonical_linear_fp32( bias: torch.Tensor | None = None, bias_parameter_id: str | None = None, matmul_op: Callable[..., torch.Tensor] | None = None, - reuse_canonical_order: bool = False, ) -> torch.Tensor: """Apply a canonical linear using the default backend or an FP32 matmul callable. A custom matmul handles input strides and accepts ``bias=`` when bias is - supplied. The same callable computes forward, dX, dW, and dbias. Canonical - order reuse requests a runtime proof before skipping stable sorting. + supplied. The same callable computes forward, dX, dW, and dbias. Parameter + gradients reuse the canonical session's stable logical-row sorting. """ return _CanonicalLinearFn.apply( a, @@ -189,7 +177,6 @@ def canonical_linear_fp32( bias, bias_parameter_id if bias_parameter_id is not None else f"{parameter_id}.bias", matmul_op, - reuse_canonical_order, ) diff --git a/rl_engine/ops/gemm/video_input_projection_fp32.py b/rl_engine/ops/gemm/video_input_projection_fp32.py index 8b2c6849e..a19a36d7f 100644 --- a/rl_engine/ops/gemm/video_input_projection_fp32.py +++ b/rl_engine/ops/gemm/video_input_projection_fp32.py @@ -75,7 +75,6 @@ class VideoInputProjectionFP32Op: family: str backend_id: str matmul_op: Callable[..., torch.Tensor] | None = None - reuse_canonical_order: bool = False def _check_backend(self, x): raise NotImplementedError @@ -107,7 +106,6 @@ def __call__(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None bias=bias, bias_parameter_id=f"{parameter_id}.bias", matmul_op=self.matmul_op, - reuse_canonical_order=self.reuse_canonical_order, ) else: output = _gemm_fp32( diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index b88ccf152..6d42c2df6 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -241,7 +241,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v8", + backend_id="triton-video-input-projection-fp32-v9", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -277,7 +277,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( "rlkernel-triton-ieee-fp32-dot-chunk128-tile128x32-merge64-" - "strided-aligned-forward-reg96-canonical-proof-h3-v8" + "strided-aligned-forward-reg96-auto-canonical-proof-h3-v9" ), ), OperatorBackendDescriptor( diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 8c077197f..437f1b1d6 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -61,14 +61,14 @@ def test_canonical_order_proof_preserves_stable_normalization(key_rows, device, assert _keys_are_canonical(keys) == strict if not strict: assert not _keys_are_canonical(keys) - else: + elif device == "cuda": def unexpected_sort(*args, **kwargs): raise AssertionError("canonical keys must reuse their physical order") monkeypatch.setattr(torch, "argsort", unexpected_sort) session = CanonicalBackwardSession() - session.register("weight", keys, reuse_canonical_order=True) + session.register("weight", keys) order, actual = session._normalize("weight", [values]) expected = values.index_select( 0, torch.tensor(expected_order, dtype=torch.int64, device=device) @@ -78,7 +78,7 @@ def unexpected_sort(*args, **kwargs): assert actual.is_contiguous() -def test_canonical_order_proof_is_recomputed_after_key_changes(monkeypatch): +def test_canonical_normalization_reads_changed_keys(monkeypatch): from rl_engine.ops.autograd.canonical_backward import CanonicalBackwardSession monkeypatch.setattr( @@ -87,7 +87,7 @@ def test_canonical_order_proof_is_recomputed_after_key_changes(monkeypatch): keys = torch.tensor([[0, 0], [0, 1]], dtype=torch.int64) values = torch.tensor([[1.0], [2.0]]) session = CanonicalBackwardSession() - session.register("weight", keys, reuse_canonical_order=True) + session.register("weight", keys) _, first = session._normalize("weight", [values]) keys.copy_(keys.flip(0)) _, second = session._normalize("weight", [values]) @@ -95,7 +95,7 @@ def test_canonical_order_proof_is_recomputed_after_key_changes(monkeypatch): assert torch.equal(second, values.flip(0)) -def test_canonical_order_reuse_requires_explicit_registration(monkeypatch): +def test_cpu_callers_retain_stable_normalization(monkeypatch): from rl_engine.ops.autograd import canonical_backward as canonical def unexpected_proof(keys): @@ -122,7 +122,7 @@ def unexpected_proof(keys): session = canonical.CanonicalBackwardSession() keys = torch.tensor([[0, 0], [0, 1]], dtype=torch.int64) values = torch.tensor([[1.0], [2.0]]) - session.register("weight", keys, reuse_canonical_order=True) + session.register("weight", keys) order, actual = session._normalize("weight", [values]) assert order.tolist() == [0, 1] assert torch.equal(actual, values) @@ -214,8 +214,7 @@ def test_canonical_linear_bias_and_legacy_call(with_bias): @pytest.mark.parametrize("with_bias", [False, True]) -@pytest.mark.parametrize("reuse_order", [False, True]) -def test_canonical_linear_custom_matmul_covers_all_vjps(with_bias, reuse_order, monkeypatch): +def test_canonical_linear_custom_matmul_covers_all_vjps(with_bias, monkeypatch): from rl_engine.ops.autograd.backward_runtime import ( reset_backward_runtime, snapshot_backward_runtime, @@ -231,7 +230,7 @@ def test_canonical_linear_custom_matmul_covers_all_vjps(with_bias, reuse_order, original_register = CanonicalBackwardSession.register def register(self, parameter_id, keys, **kwargs): - registrations.append(kwargs["reuse_canonical_order"]) + registrations.append(parameter_id) return original_register(self, parameter_id, keys, **kwargs) def gemm(a, b, *, bias=None): @@ -255,7 +254,6 @@ def two_argument_gemm(a, b): family="pytorch", bias=bias, matmul_op=matmul_op, - reuse_canonical_order=reuse_order, ) y.backward(grad) session.validate_complete() @@ -270,7 +268,7 @@ def two_argument_gemm(a, b): assert calls[0][3] == with_bias assert all(not call[3] for call in calls[1:]) assert calls[2][2] == 1 - assert registrations == [reuse_order] * (2 if with_bias else 1) + assert registrations == (["weight", "weight.bias"] if with_bias else ["weight"]) provenance = snapshot_backward_runtime()["det_gemm"] assert provenance["kernel_id"] == f"{matmul_op.__module__}.{matmul_op.__qualname__}" assert provenance["family"] == "pytorch" @@ -323,7 +321,6 @@ def test_video_projection_declares_backend_and_matmul_separately(): assert TritonVideoInputProjectionFP32Op.family == "triton" assert TritonVideoInputProjectionFP32Op.matmul_op is projection_gemm_fp32 - assert TritonVideoInputProjectionFP32Op.reuse_canonical_order def test_input_gradient_without_parameter_session(): @@ -531,7 +528,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v8" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v9" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op()(x, weight, bias)) diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 437e6988b..bfe1a44a7 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,11 +13,15 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v8" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v9" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) assert desc.determinism_or_alignment_properties["vjp_tile"] == (128, 32, 32) + assert ( + desc.determinism_or_alignment_properties["canonical_rows"] + == "live_strict_order_proof_or_stable_sort" + ) with pytest.raises(RuntimeError, match="SM90"): registry.get_op("video_input_projection_fp32", device="cpu") From b6be6024fa2a5790346bfc571d67b6a7036ca746 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Mon, 19 Oct 2026 11:51:57 +0800 Subject: [PATCH 12/26] Fix H3 FP32 chunk accuracy and extend regression coverage Balance the two 64-term halves inside each existing 128-term chunk. Use default register allocation and update the backend identity to v10. Reject non-finite benchmark comparisons and document matmul reuse. Extend random accuracy and logical-row invariance coverage. Single H20: 198 regression tests and 25 existing Triton GEMM tests pass. Historical checks: 48 accuracy and 30 BI cases pass. All 18 Qwen3 call comparisons preserve output and gradient bytes. The 13104-row H3 eager F+B median changes from 2.613 ms to 2.705 ms. Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 6 ++- docs/operators/video-input-projection-fp32.md | 23 ++++++--- .../backends/shared/triton/gemm/det_gemm.py | 13 +++-- .../gemm/video_input_projection_fp32.py | 11 ++-- rl_engine/ops/gemm/canonical_linear.py | 3 ++ rl_engine/runtime/registry.py | 9 ++-- .../gemm/test_video_input_projection_fp32.py | 51 +++++++++++++++++-- tests/runtime/test_kernel_registry.py | 2 +- 8 files changed, 88 insertions(+), 30 deletions(-) diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index 28fabdb4c..1070a3a25 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -66,9 +66,11 @@ def _accuracy(actual, gold): a, b = a.detach().cpu(), b.detach().cpu() error = (a - b).abs() limit = policy["atol"] + policy["rtol"] * b.abs() + finite = torch.isfinite(a) & torch.isfinite(b) result[name] = { - "max_abs_error": error.max().item(), - "out_of_tolerance": int((error > limit).sum()), + "max_abs_error": error.max().item() if torch.isfinite(error).all() else None, + "out_of_tolerance": int(((error > limit) | ~finite).sum()), + "non_finite_elements": int((~finite).sum()), **policy, } return result diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index f542980e8..283b96b87 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -64,8 +64,11 @@ The wrapper reuses `canonical_linear_fp32` with `family="triton"` and selects the existing forward and VJP GEMM schedules and handles their input strides. Forward, dX, dW, and dbias use the same callable. Existing canonical linear callers omit `matmul_op` and retain their original backend dispatch. -On CUDA, the session internally checks eligible keys before skipping sorting. Backward provenance records the actual -matmul entry point. +The existing `canonical_row_lm_head` uses the same separation: `family` records +the backend, and `matmul_op` supplies the GEMM implementation. Qwen3 passes its +selected `det_gemm.forward_accum_fp32` to that function. +On CUDA, the session internally checks eligible keys before skipping sorting. +Backward provenance records the actual matmul entry point. ## Accuracy @@ -95,19 +98,20 @@ CPU, and other device paths retain stable sorting. Short reductions up to the 96-wide video input use the existing direct FP32 GEMM. The Triton input and weight VJPs reuse that kernel with fixed 128-term -chunks. Each chunk computes four ascending 32-term FP32 dots from zero. Chunk +chunks. Each chunk computes four ascending 32-term FP32 dots from zero. It adds the +first pair and the second pair independently, then combines both 64-term halves. +This balanced sum reduces cancellation error in the long weight VJP. Chunk partials remain FP32. A second grid uses a fixed balanced reduction over the next power-of-two chunk count, with zero padding. The merge uses 64 output elements per program and four warps. VJP GEMMs use fixed 128-by-32 output tiles, 32-term dot tiles, and four warps. The forward retains its 64-by-64 output tile. -These settings define the v9 Triton implementation identity. The K partition and VJP +These settings define the v10 Triton implementation identity. The K partition and VJP tile do not depend on M, row placement, or occupancy. The VJP kernels consume positive-stride views directly, including `G.T`. The Triton forward omits K or N load predicates only when the corresponding 32-term or 64-column tiles are complete. It retains M predicates for every row. -Column-major A views with unit row stride use a fixed `maxnreg=96` VJP allocation. -Other VJP layouts retain the default allocation. These choices preserve the dot -tile and K accumulation order. No configuration depends on physical row count. +VJP kernels use the compiler's default register allocation for every input layout. +The K accumulation order is fixed. No configuration depends on physical row count. The fixed-chunk design follows the scheduling separation in merged PR #396. Its AMD BF16 MFMA instructions are not used by this CUDA FP32 implementation. @@ -168,6 +172,11 @@ python tools/validation/operators/check_operator.py \ --batch 1 --seq 129 --check-grad --grad-mode random ``` +The standard-normal regression covers seeds 420 through 423 at 129 and 13104 +rows. It checks CPU FP32 accuracy and byte equality across permutations, +microbatches, padding, and repeated calls. The benchmark rejects NaN or Inf in +either the candidate or reference, including gradients. + Set `H3_PROJECTION_WEIGHTS` to a safetensors fixture containing `proj_in.weight` and `proj_in.bias` from `MiniMaxAI/MiniMax-H3` at revision `42ed227ee7df40d41602854ae760620d6eb651fe` to run the latent-pack-projection check. diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 7127ec623..5416d16c5 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -795,6 +795,7 @@ def _det_gemm_fp32_kernel( b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) reduction_size = CHUNK_K if CHUNK_K else K + second_half = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(reduction_size, BLOCK_K)): k_rem = K - k * BLOCK_K a_mask = offs_m[:, None] < M @@ -810,9 +811,16 @@ def _det_gemm_fp32_kernel( if PROMOTE_INPUTS: a = a.to(tl.float32) b = b.to(tl.float32) - acc += tl.dot(a, b, allow_tf32=False) + dot = tl.dot(a, b, allow_tf32=False) + # Shorten the FP32 accumulation chain inside each fixed chunk. + if CHUNK_K and k >= tl.cdiv(CHUNK_K, 2 * BLOCK_K): + second_half += dot + else: + acc += dot a_ptrs += BLOCK_K * stride_ak b_ptrs += BLOCK_K * stride_bk + if CHUNK_K: + acc += second_half if HAS_BIAS: bias = tl.load(bias_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32) acc = acc + bias[None, :] @@ -865,7 +873,6 @@ def _triton_chunked_gemm_fp32( b: torch.Tensor, *, tile: tuple[int, int, int] = (_BLOCK_M, _BLOCK_N, 4), - maxnreg: int | None = None, ) -> torch.Tensor: """Reuse the FP32 GEMM in fixed K chunks and merge partials in one grid.""" m_size, k_size = a.shape @@ -881,7 +888,6 @@ def _triton_chunked_gemm_fp32( a[start : start + rows_per_chunk], b, tile=tile, - maxnreg=maxnreg, ) for start in range(0, m_size, rows_per_chunk) ], @@ -911,7 +917,6 @@ def _triton_chunked_gemm_fp32( CHUNK_K=_FP32_CHUNK_K, enable_fp_fusion=False, num_warps=num_warps, - **({"maxnreg": maxnreg} if maxnreg is not None else {}), ) _det_gemm_fp32_chunk_merge_kernel[(triton.cdiv(m_size * n_size, 64),)]( partial, diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index 244f08d4b..93df5d01c 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -18,17 +18,13 @@ ) _VJP_TILE = (128, 32, 4) -_VJP_MAXNREG = 96 def projection_gemm_fp32(a, b, *, bias=None): if a.shape[1] <= VIDEO_PATCH_WIDTH: return _triton_gemm_fp32(a, b, bias=bias, aligned_tiles=True) if b.shape[1] == VIDEO_PATCH_WIDTH: - # Pin register allocation for column-major A without changing dot order. - output = _triton_chunked_gemm_fp32( - a, b, tile=_VJP_TILE, maxnreg=_VJP_MAXNREG if a.stride(0) == 1 else None - ) + output = _triton_chunked_gemm_fp32(a, b, tile=_VJP_TILE) else: output = _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) return output if bias is None else output + bias @@ -37,7 +33,7 @@ def projection_gemm_fp32(a, b, *, bias=None): class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton" matmul_op = staticmethod(projection_gemm_fp32) - backend_id = "triton-video-input-projection-fp32-v9" + backend_id = "triton-video-input-projection-fp32-v10" def __init__(self): if not _TRITON_AVAILABLE: @@ -62,7 +58,7 @@ def backend_info(self): "tile": [64, 64, 32], "vjp_tile": [_VJP_TILE[0], _VJP_TILE[1], _BLOCK_K], "vjp_num_warps": _VJP_TILE[2], - "vjp_maxnreg_unit_row_stride": _VJP_MAXNREG, + "vjp_register_allocation": "compiler_default", "forward_aligned_tiles": True, "canonical_rows": "live_strict_order_proof_or_stable_sort", "tf32": False, @@ -73,6 +69,7 @@ def backend_info(self): "split_k": "fixed_128_term_vjp_chunks", "reduction": "fixed_128_chunks_balanced_fp32_merge", "vjp_chunk_k": 128, + "vjp_chunk_accumulation": "two_64_term_halves", "vjp_merge_tile": 64, "bias_reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 08a402b1a..751820dfe 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -164,6 +164,9 @@ def canonical_linear_fp32( ) -> torch.Tensor: """Apply a canonical linear using the default backend or an FP32 matmul callable. + Without ``matmul_op``, ``family`` selects the existing backend dispatch. + With a callable, ``family`` identifies its backend for provenance, as in + ``canonical_row_lm_head``. A custom matmul handles input strides and accepts ``bias=`` when bias is supplied. The same callable computes forward, dX, dW, and dbias. Parameter gradients reuse the canonical session's stable logical-row sorting. diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 6d42c2df6..5cd07ba7a 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -241,7 +241,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v9", + backend_id="triton-video-input-projection-fp32-v10", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -268,7 +268,8 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "reduction": "fixed_128_chunks_balanced_fp32_merge", "vjp_tile": (128, 32, 32), "vjp_num_warps": 4, - "vjp_maxnreg_unit_row_stride": 96, + "vjp_register_allocation": "compiler_default", + "vjp_chunk_accumulation": "two_64_term_halves", "forward_aligned_tiles": True, "canonical_rows": "live_strict_order_proof_or_stable_sort", }, @@ -276,8 +277,8 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: implementation_class_or_factory=OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32.value, fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( - "rlkernel-triton-ieee-fp32-dot-chunk128-tile128x32-merge64-" - "strided-aligned-forward-reg96-auto-canonical-proof-h3-v9" + "rlkernel-triton-ieee-fp32-dot-chunk128-halves64-tile128x32-merge64-" + "strided-aligned-forward-auto-canonical-proof-h3-v10" ), ), OperatorBackendDescriptor( diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 437f1b1d6..5a8095de3 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -3,6 +3,7 @@ """H3 FP32 projection accuracy, actual autograd, and logical-row invariance.""" +import json import os import pytest @@ -143,8 +144,8 @@ def _keys(rows, device): ) -def _data(rows, device="cuda"): - generator = torch.Generator(device=device).manual_seed(420) +def _data(rows, device="cuda", *, seed=420): + generator = torch.Generator(device=device).manual_seed(seed) return ( torch.randn(rows, 96, generator=generator, device=device) * 0.125, torch.randn(5376, 96, generator=generator, device=device) * 0.0625, @@ -193,6 +194,39 @@ def _assert_bytes(a, b): ) +@pytest.mark.parametrize("output", ["y", "dx", "dw", "dbias"]) +@pytest.mark.parametrize( + "candidate,reference,violations,non_finite", + [ + (1.0, 1.0, 0, 0), + (2.0, 1.0, 1, 0), + (float("nan"), 1.0, 1, 1), + (1.0, float("nan"), 1, 1), + (float("inf"), 1.0, 1, 1), + (1.0, float("inf"), 1, 1), + (float("inf"), float("inf"), 1, 1), + (-float("inf"), 1.0, 1, 1), + ], +) +def test_benchmark_accuracy_rejects_non_finite( + output, candidate, reference, violations, non_finite +): + from benchmarks.operators.gemm.benchmark_video_input_projection_fp32 import _accuracy + + names = ("y", "dx", "dw", "dbias") + actual = [torch.ones(2) for _ in names] + gold = [torch.ones(2) for _ in names] + actual[names.index(output)][0] = candidate + gold[names.index(output)][0] = reference + report = _accuracy(actual, gold) + for name in names: + assert report[name]["out_of_tolerance"] == (violations if name == output else 0) + assert report[name]["non_finite_elements"] == (non_finite if name == output else 0) + if non_finite: + assert report[output]["max_abs_error"] is None + json.dumps(report, allow_nan=False) + + @pytest.mark.parametrize("with_bias", [False, True]) def test_canonical_linear_bias_and_legacy_call(with_bias): torch.manual_seed(420) @@ -410,6 +444,11 @@ def test_forward_and_actual_gradients(rows, with_bias): @pytest.mark.parametrize("rows", [9, 65, 127, 128, 129, 257]) def test_parameter_bi_permutation_chunk_padding(rows): op, data, keys = _op(), _data(rows), _keys(rows, "cuda") + _check_parameter_bi(op, data, keys) + + +def _check_parameter_bi(op, data, keys): + rows = data[0].shape[0] baseline = _run(op, data, keys) order = torch.randperm(rows, device="cuda") permuted = _run( @@ -528,7 +567,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v9" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v10" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op()(x, weight, bias)) @@ -598,9 +637,10 @@ def test_fp32_tree_workspace_chunking(monkeypatch): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") @pytest.mark.parametrize("rows", [129, 13104]) -def test_standard_normal_long_reduction(rows): +@pytest.mark.parametrize("seed", [420, 421, 422, 423]) +def test_standard_normal_long_reduction(rows, seed): # Unscaled inputs expose cancellation error in the 5376-term input VJP. - data = _data(rows) + data = _data(rows, seed=seed) data = tuple(t / scale for t, scale in zip(data, (0.125, 0.0625, 0.125, 0.0625), strict=True)) actual = _run(_op(), data, _keys(rows, "cuda")) cx, cw, cb = [t.detach().cpu().requires_grad_() for t in data[:3]] @@ -608,6 +648,7 @@ def test_standard_normal_long_reduction(rows): gold.backward(data[3].cpu()) for i, (a, b) in enumerate(zip(actual, (gold, cx.grad, cw.grad, cb.grad), strict=True)): _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") + _check_parameter_bi(_op(), data, _keys(rows, "cuda")) def test_cpu_fp32_reference_long_dot_accuracy(): diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index bfe1a44a7..f97c7c3ce 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,7 +13,7 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v9" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v10" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) From 3fc97a0cfaa9ba1a3e4778ec83e8cc2abf918563 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Mon, 19 Oct 2026 15:34:02 +0800 Subject: [PATCH 13/26] Simplify H3 FP32 bias tree scheduling Signed-off-by: walker-ai <2398833647@qq.com> --- docs/operators/video-input-projection-fp32.md | 2 + .../backends/shared/triton/gemm/det_gemm.py | 44 +++++++------------ .../gemm/video_input_projection_fp32.py | 1 - .../gemm/test_video_input_projection_fp32.py | 24 ++++++---- 4 files changed, 35 insertions(+), 36 deletions(-) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 283b96b87..f874e8400 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -10,6 +10,8 @@ The implementation reuses RL-Kernel's Triton FP32 GEMM and canonical parameter-gradient session. Input and weight VJPs use fixed 128-term K chunks and one balanced FP32 merge. Long bias VJPs retain the one-row Triton tree. +Only this one-row FP32 tree uses compact paired-leaf scheduling. Other FP32 tree +shapes reuse the existing default tile and per-level schedule. ## Entry Point diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 5416d16c5..a080831ca 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -10,9 +10,10 @@ a separately versioned arithmetic contract because its internal accumulation does not match the scalar reference at every BF16 rounding boundary. -Autotuning and split-K are intentionally disabled. FP32-output calls preserve -the earlier fixed-order, no-split-K contract. The optional CUDA FP32 tree reuses -the canonical plan, leaf kernel, and level reductions with FP32 nodes. +Autotuning and adaptive split-K are intentionally disabled. Default FP32-output +calls preserve the earlier fixed-order, no-split-K contract. Explicit FP32 VJP +calls use fixed 128-term chunks. The optional CUDA FP32 tree reuses the canonical +plan, leaf kernel, and level reductions with FP32 nodes. """ from __future__ import annotations @@ -42,7 +43,7 @@ # one shot would otherwise create a multi-terabyte virtual tensor on ROCm. # Chunking over M preserves the exact K-tree and BF16 rounding contract. _MAX_TREE_WORKSPACE_ELEMENTS = 128 * 1024 * 1024 -_TREE_PLANS: dict[tuple[int, int, bool | None], "_DeviceTreePlan"] = {} +_TREE_PLANS: dict[tuple[int, int, bool], "_DeviceTreePlan"] = {} _TREE_PLAN_LOCK = threading.Lock() _ROCM_TUNE_DECODE_LEAF = True _ROCM_FUSE_LEAF_REDUCTION = True @@ -263,7 +264,7 @@ def visit(begin: int, end: int) -> tuple[int, int]: def _device_tree_plan( - k_size: int, device: torch.device, *, fp32_fuse_leaves: bool | None = None + k_size: int, device: torch.device, *, fp32_fuse_leaves: bool = False ) -> _DeviceTreePlan: device_index = device.index if device.index is not None else torch.cuda.current_device() key = (device_index, k_size, fp32_fuse_leaves) @@ -288,25 +289,18 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: len(level) * (1 << (height + 1)) == leaf_count for height, level in enumerate(host.reduction_levels) ) - fp32_fused = fp32_fuse_leaves is not None and complete + fp32_fused = fp32_fuse_leaves and complete node_map = {node: node for node in range(host.node_count)} workspace_node_count = host.node_count if fp32_fused: # Alternate disjoint buffers for the retained levels. Each launch # reads its previous buffer completely before the next launch. - first_height = 0 if fp32_fuse_leaves else -1 - first_nodes = ( - tuple(op[2] for op in host.reduction_levels[0]) - if fp32_fuse_leaves - else host.leaf_nodes - ) + first_nodes = tuple(op[2] for op in host.reduction_levels[0]) next_count = max(1, len(first_nodes) // 4) workspace_node_count = len(first_nodes) + next_count for i, node in enumerate(first_nodes): node_map[node] = i - for step, height in enumerate( - range(first_height + 2, len(host.reduction_levels), 2), start=1 - ): + for step, height in enumerate(range(2, len(host.reduction_levels), 2), start=1): base = len(first_nodes) if step % 2 else 0 for i, (_, _, node) in enumerate(host.reduction_levels[height]): node_map[node] = base + i @@ -343,14 +337,13 @@ def build_fused_pairs(start_level: int): ) return result - if not fp32_fused or not fp32_fuse_leaves: + if not fp32_fused: fused_pairs = build_fused_pairs(0) - if not fp32_fused or fp32_fuse_leaves: - fused_pairs_after_leaf = build_fused_pairs(1) + fused_pairs_after_leaf = build_fused_pairs(1) if host.reduction_levels: first_level = host.reduction_levels[0] paired_nodes = {node for op in first_level for node in op[:2]} - if paired_nodes == set(host.leaf_nodes) and (not fp32_fused or fp32_fuse_leaves): + if paired_nodes == set(host.leaf_nodes): leaf_metadata = { node: (start, length) for start, length, node in zip( @@ -1022,7 +1015,7 @@ def _triton_tree_gemm( plan = _device_tree_plan( k_size, a.device, - fp32_fuse_leaves=(m_size == 1) if fp32_tree and not transpose_output else None, + fp32_fuse_leaves=fp32_tree and m_size == 1 and not transpose_output, ) workspace_elements = plan.workspace_node_count * m_size * n_size if workspace_elements > _MAX_TREE_WORKSPACE_ELEMENTS and out is None: @@ -1095,12 +1088,9 @@ def _triton_tree_gemm( transpose_output=transpose_output, preserve_a_strides=preserve_a_strides, ) - if fp32_tree: - # Offline H20 sweep: 64x32 avoids both masked N=96 columns and the - # register spills of paired 64x64 leaves. One-row calls retain fusion. - leaf_config = ( - _TreeLeafConfig(1, 128, 4, True) if m_size == 1 else _TreeLeafConfig(64, 32, 4, False) - ) + if fp32_tree and m_size == 1: + # The H3 bias VJP has one output row. Other shapes use the existing tile. + leaf_config = _TreeLeafConfig(1, 128, 4, True) if ( _ROCM_TUNE_DECODE_LEAF and is_gfx942 @@ -1210,7 +1200,7 @@ def _triton_tree_gemm( N=n_size, BLOCK=reduction_block, ) - elif direct_root_output and ((use_inference_schedule and is_gfx942) or plan.fp32_fused): + elif direct_root_output and use_inference_schedule and is_gfx942: blocks = triton.cdiv(m_size * n_size, reduction_block) for pair_index, nodes in enumerate(plan.rocm_fused_reduction_pairs): second_level = pair_index * 2 + 1 diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index 93df5d01c..5da3fc5a8 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -75,7 +75,6 @@ def backend_info(self): "reduction_leaf_size": _BLOCK_K, "reduction_backend": "triton", "leaf_launch": "paired_leaves_for_one_row_complete_trees", - "tree_tile": [64, 32, 32], "bias_tree_tile": [1, 128, 32], "tree_schedule": "paired_levels_compact_workspace", } diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 5a8095de3..82beca83a 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -746,11 +746,11 @@ def observe(self, parameter_id, values): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") @pytest.mark.parametrize("k", [64, 65, 128, 1024, 13104]) -def test_fp32_fused_tree_and_strided_inputs_match_scalar_reference(k): +def test_fp32_tree_default_schedule_and_bias_match_scalar_reference(k): from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_tree_gemm torch.manual_seed(422) - # Exercise a transposed A view and tail rows. + # Multirow FP32 calls use the original schedule; bias retains one-row fusion. a = torch.randn(k, 3, device="cuda").t() b = torch.randn(k, 5, device="cuda") @@ -768,19 +768,27 @@ def gemm(a, b): @pytest.mark.skipif(not H20, reason="requires CUDA SM90") -def test_compact_fp32_tree_workspace_preserves_chunked_outputs(monkeypatch): +@pytest.mark.parametrize("k", [128, 1024, 13104]) +def test_compact_bias_tree_workspace_preserves_chunked_outputs(k, monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( projection_gemm_fp32, ) - a, b = torch.randn(7, 1024, device="cuda"), torch.randn(1024, 5, device="cuda") - expected = projection_gemm_fp32(a, b) - plan = det_gemm._device_tree_plan(1024, a.device, fp32_fuse_leaves=False) + a, b = torch.ones(1, k, device="cuda"), torch.randn(k, 5, device="cuda") + original_plan = det_gemm._device_tree_plan(k, a.device) + assert not original_plan.fp32_fused + assert original_plan.workspace_node_count == original_plan.host.node_count + plan = det_gemm._device_tree_plan(k, a.device, fp32_fuse_leaves=True) assert plan.fp32_fused assert plan.workspace_node_count < plan.host.node_count - monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.workspace_node_count * 10) - _assert_bytes(projection_gemm_fp32(a, b), expected) + node_arrays = [plan.rocm_leaf_reduction[-1]] + node_arrays.extend( + nodes for pair in plan.rocm_fused_reduction_pairs_after_leaf for nodes in pair + ) + for nodes in node_arrays: + assert bool(torch.all((nodes >= 0) & (nodes < plan.workspace_node_count))) + expected = projection_gemm_fp32(a, b) monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.workspace_node_count * 2) _assert_bytes(projection_gemm_fp32(a, b), expected) From fc0b332e58db5d6a634c32b5b1d06f890d8af1c8 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Mon, 19 Oct 2026 15:55:29 +0800 Subject: [PATCH 14/26] Align H3 projection naming and documentation style Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 2 +- docs/operators/video-input-projection-fp32.md | 35 ++++++++++------- .../validation/operators/operator_inputs.py | 4 +- .../gemm/test_video_input_projection_fp32.py | 39 ++++++++++--------- 4 files changed, 47 insertions(+), 33 deletions(-) diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index 1070a3a25..4e5a9b66a 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -162,7 +162,7 @@ def main(): "rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py", "rl_engine/reference/gemm/video_input_projection_fp32.py", ] - # Devbox worker sync excludes .git; source hashes also identify exported workspaces. + # Source hashes also identify exported workspaces without Git metadata. git_commit = _tool_version(["git", "-C", str(root), "rev-parse", "HEAD"]) git_status = ( _tool_version(["git", "-C", str(root), "status", "--porcelain"]) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index f874e8400..1aace14f5 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -8,16 +8,15 @@ reference. The model adapter owns the subsequent cast into the BF16 packed strea The implementation reuses RL-Kernel's Triton FP32 GEMM and canonical parameter-gradient session. Input and weight VJPs use fixed 128-term K chunks -and one balanced FP32 merge. -Long bias VJPs retain the one-row Triton tree. -Only this one-row FP32 tree uses compact paired-leaf scheduling. Other FP32 tree -shapes reuse the existing default tile and per-level schedule. +and one balanced FP32 merge. Long bias VJPs retain the one-row Triton tree. +Only this one-row FP32 tree uses compact paired-leaf scheduling. Other FP32 +tree shapes reuse the existing default tile and per-level schedule. ## Entry Point ```python -from rl_engine.runtime.registry import kernel_registry from rl_engine.ops.autograd.canonical_backward import canonical_backward_session +from rl_engine.runtime.registry import kernel_registry op = kernel_registry.get_op("video_input_projection_fp32", device=x.device) with canonical_backward_session() as session: @@ -74,16 +73,20 @@ Backward provenance records the actual matmul entry point. ## Accuracy +### Forward and Gradients + Forward computes `Y = X @ W.T + bias` in FP32. The Triton forward adds bias in the existing GEMM epilogue, preserving the rounded FP32 GEMM result before -the bias addition. The upstream gradient `G` -stays FP32, including normalization of non-contiguous autograd inputs: +the bias addition. The upstream gradient `G` stays FP32, including normalization +of non-contiguous autograd inputs: - `dX = G @ W` is row-local. - `dW = G_sorted.T @ X_sorted` follows canonical logical-row order. - `dbias = ones[1,R] @ G_sorted` uses the fixed FP32 GEMM for reductions up to 96 and the one-row Triton tree for longer reductions. +### Canonical Parameter Gradients + The existing canonical session collects all contributions, removes padding, and sorts keys before the final parameter reduction. It returns each parameter VJP once. Weight and bias share sorting and the sorted upstream gradient when their @@ -98,17 +101,19 @@ existing stable sort. The check reads the current keys on each normalization. Callers supply logical keys and do not select this optimization. Smaller inputs, CPU, and other device paths retain stable sorting. +### Fixed FP32 Accumulation + Short reductions up to the 96-wide video input use the existing direct FP32 GEMM. The Triton input and weight VJPs reuse that kernel with fixed 128-term -chunks. Each chunk computes four ascending 32-term FP32 dots from zero. It adds the -first pair and the second pair independently, then combines both 64-term halves. +chunks. Each chunk computes four ascending 32-term FP32 dots from zero. It adds +the first pair and the second pair independently, then combines both 64-term halves. This balanced sum reduces cancellation error in the long weight VJP. Chunk partials remain FP32. A second grid uses a fixed balanced reduction over the next power-of-two chunk count, with zero padding. The merge uses 64 output elements per program and four warps. VJP GEMMs use fixed 128-by-32 output tiles, 32-term dot tiles, and four warps. The forward retains its 64-by-64 output tile. -These settings define the v10 Triton implementation identity. The K partition and VJP -tile do not depend on M, row placement, or occupancy. +These settings define the v10 Triton implementation identity. The K partition +and VJP tile do not depend on M, row placement, or occupancy. The VJP kernels consume positive-stride views directly, including `G.T`. The Triton forward omits K or N load predicates only when the corresponding 32-term or 64-column tiles are complete. It retains M predicates for every row. @@ -124,14 +129,16 @@ pins 64 output elements and four warps as part of its implementation identity. The VJP schedule caps its FP32 partial workspace at 512 MiB by chunking output rows. The retained midpoint helper also supports output-column chunking -when one complete row exceeds its budget. Parameter sorting and public validation are included in complete -operator measurements. Existing BF16 callers retain their dtype and schedule. +when one complete row exceeds its budget. Existing BF16 callers retain their +dtype and schedule. TF32 and compiler multiply-add fusion remain disabled. The reused Triton `tl.dot` lowers its IEEE FP32 inner products to fixed SIMT `fma.rn.f32` instructions. The compiler fusion flag controls surrounding expressions; it does not remove the intrinsic dot's FMA instructions. This arithmetic is included in the Triton contract and runtime metadata. +### Independent Reference + The independent reference keeps the forward definition as CPU FP32 `F.linear`. Its analytic VJPs form FP32 products and use CPU PyTorch `sum` over the complete K dimension. Only output rows and columns are tiled to bound temporary memory. @@ -141,6 +148,8 @@ The reference identity is `cpu-fp32-linear-products-sum-v2`. FP64 is used only in a small reference-validation test and diagnostic evidence, not by the reference implementation or candidate kernels. +### Acceptance Criteria + The shared WS1 `reduction` FP32 accuracy policy uses `atol=rtol=1e-4` for forward and gradients. Invariance uses byte equality. Parameter invariance compares the same logical contribution set and upstream gradients across layout changes. diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index 842142e4f..2e82ca410 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -87,7 +87,9 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: raise ValueError(f"unsupported operator shape: {op_name}") from exc -def _make_video_input_projection_fp32_inputs(args, dtype, device): +def _make_video_input_projection_fp32_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: if dtype != torch.float32: raise TypeError("video input projection validation requires FP32") batch, seq = _batch_seq(args) diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 82beca83a..ae97c3743 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -15,7 +15,7 @@ from rl_engine.ops.gemm.video_input_projection_fp32 import VideoInputProjectionFP32Op from rl_engine.reference.gemm.video_input_projection_fp32 import NativeVideoInputProjectionFP32Op -H20 = ( +HAS_SM90 = ( torch.cuda.is_available() and torch.version.hip is None and torch.cuda.get_device_capability() == (9, 0) @@ -24,7 +24,10 @@ @pytest.mark.parametrize( "device", - ["cpu", pytest.param("cuda", marks=pytest.mark.skipif(not H20, reason="requires CUDA SM90"))], + [ + "cpu", + pytest.param("cuda", marks=pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90")), + ], ) @pytest.mark.parametrize( "key_rows", @@ -413,7 +416,7 @@ def test_invalid_inputs(invalid): _CPUAdapter()(x, weight, bias, logical_keys=keys, parameter_id="p") -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("rows", [1, 9, 17, 65, 129, 257]) @pytest.mark.parametrize("with_bias", [False, True]) def test_forward_and_actual_gradients(rows, with_bias): @@ -440,7 +443,7 @@ def test_forward_and_actual_gradients(rows, with_bias): _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("rows", [9, 65, 127, 128, 129, 257]) def test_parameter_bi_permutation_chunk_padding(rows): op, data, keys = _op(), _data(rows), _keys(rows, "cuda") @@ -478,7 +481,7 @@ def _check_parameter_bi(op, data, keys): _assert_bytes(a, b) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") def test_row_placement_and_unrelated_rows(): op, data = _op(), _data(65) single = _run(op, [data[0][:1], data[1], data[2], data[3][:1]], _keys(1, "cuda")) @@ -495,7 +498,7 @@ def test_row_placement_and_unrelated_rows(): _assert_bytes(batch[1][37], result[1][37]) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") def test_noncontiguous_upstream_and_forward_graph(): op = _op() x, weight, bias, grad = _data(17) @@ -524,7 +527,7 @@ def test_noncontiguous_upstream_and_forward_graph(): _assert_bytes(output, expected) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") def test_real_weights_patch_path(): path = os.environ.get("H3_PROJECTION_WEIGHTS") if not path: @@ -562,7 +565,7 @@ def test_real_weights_patch_path(): _assert_bytes(a, b) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry @@ -573,7 +576,7 @@ def test_default_cuda_dispatch(): _assert_bytes(op(x, weight, bias), _op()(x, weight, bias)) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") def test_reused_fp32_gemm(): from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 @@ -602,7 +605,7 @@ def _fp32_tree_reference(a, b): return nodes[plan.root] -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("k", [31, 33, 129, 5376]) def test_fp32_tree_preserves_fp32_nodes(k): from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_tree_gemm @@ -617,7 +620,7 @@ def test_fp32_tree_preserves_fp32_nodes(k): _assert_bytes(inference, output) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") def test_fp32_tree_workspace_chunking(monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( @@ -635,7 +638,7 @@ def test_fp32_tree_workspace_chunking(monkeypatch): _assert_bytes(projection_gemm_fp32(a, b), expected) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("rows", [129, 13104]) @pytest.mark.parametrize("seed", [420, 421, 422, 423]) def test_standard_normal_long_reduction(rows, seed): @@ -663,7 +666,7 @@ def test_cpu_fp32_reference_long_dot_accuracy(): _assert_accuracy(actual, expected, "gradient_accuracy") -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("k", [127, 128, 129, 5376, 13104]) def test_fp32_chunked_gemm_strides_and_workspace(k, monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm @@ -744,7 +747,7 @@ def observe(self, parameter_id, values): assert all(a[1] is not b[1] for i, a in enumerate(normalized) for b in normalized[i + 1 :]) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("k", [64, 65, 128, 1024, 13104]) def test_fp32_tree_default_schedule_and_bias_match_scalar_reference(k): from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_tree_gemm @@ -767,7 +770,7 @@ def gemm(a, b): _assert_bytes(gemm(a[:1], b), result[:1]) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("k", [128, 1024, 13104]) def test_compact_bias_tree_workspace_preserves_chunked_outputs(k, monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm @@ -793,7 +796,7 @@ def test_compact_bias_tree_workspace_preserves_chunked_outputs(k, monkeypatch): _assert_bytes(projection_gemm_fp32(a, b), expected) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("rows", [1, 3]) @pytest.mark.parametrize("k", [129, 5376, 13104]) def test_shortened_fp32_leaves_preserve_padded_arithmetic(rows, k, monkeypatch): @@ -830,7 +833,7 @@ def launch(*args, **kwargs): _assert_bytes(output, padded) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("rows", [1, 17, 129]) def test_forward_bias_fusion_preserves_fp32_rounding(rows): from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 @@ -845,7 +848,7 @@ def test_forward_bias_fusion_preserves_fp32_rounding(rows): _assert_bytes(_op()(x, weight, bias), expected) -@pytest.mark.skipif(not H20, reason="requires CUDA SM90") +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize("frozen", ["weight", "bias", "both"]) def test_backend_frozen_parameters_preserve_trainable_gradients(frozen): op = _op() From 19df90dc4ccf33c146527df4c12d821387d53751 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Mon, 19 Oct 2026 19:32:12 +0800 Subject: [PATCH 15/26] Clarify H3 projection block configuration names Signed-off-by: walker-ai <2398833647@qq.com> --- .../gemm/video_input_projection_fp32.py | 23 +++++++++++-------- rl_engine/ops/autograd/canonical_backward.py | 3 +++ rl_engine/runtime/registry.py | 8 +++---- tests/runtime/test_kernel_registry.py | 3 ++- 4 files changed, 23 insertions(+), 14 deletions(-) diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index 5da3fc5a8..9546d7596 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -17,14 +17,19 @@ VideoInputProjectionFP32Op, ) -_VJP_TILE = (128, 32, 4) +# H20-tuned output blocks and warp count for chunked dX/dW GEMM. +_DX_DW_BLOCK_M = 128 +_DX_DW_BLOCK_N = 32 +_DX_DW_NUM_WARPS = 4 def projection_gemm_fp32(a, b, *, bias=None): if a.shape[1] <= VIDEO_PATCH_WIDTH: return _triton_gemm_fp32(a, b, bias=bias, aligned_tiles=True) if b.shape[1] == VIDEO_PATCH_WIDTH: - output = _triton_chunked_gemm_fp32(a, b, tile=_VJP_TILE) + output = _triton_chunked_gemm_fp32( + a, b, tile=(_DX_DW_BLOCK_M, _DX_DW_BLOCK_N, _DX_DW_NUM_WARPS) + ) else: output = _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) return output if bias is None else output + bias @@ -56,9 +61,9 @@ def backend_info(self): "_det_gemm_tree_leaf_reduce_rocm_kernel", "triton_version": triton.__version__, "tile": [64, 64, 32], - "vjp_tile": [_VJP_TILE[0], _VJP_TILE[1], _BLOCK_K], - "vjp_num_warps": _VJP_TILE[2], - "vjp_register_allocation": "compiler_default", + "dx_dw_tile": [_DX_DW_BLOCK_M, _DX_DW_BLOCK_N, _BLOCK_K], + "dx_dw_num_warps": _DX_DW_NUM_WARPS, + "dx_dw_register_allocation": "compiler_default", "forward_aligned_tiles": True, "canonical_rows": "live_strict_order_proof_or_stable_sort", "tf32": False, @@ -66,11 +71,11 @@ def backend_info(self): "dot_arithmetic": "ieee_fp32_simt_fma", "forward_bias": "fp32_epilogue", "leaf_loop_bound": "max_logical_leaf_length", - "split_k": "fixed_128_term_vjp_chunks", + "split_k": "fixed_128_term_dx_dw_chunks", "reduction": "fixed_128_chunks_balanced_fp32_merge", - "vjp_chunk_k": 128, - "vjp_chunk_accumulation": "two_64_term_halves", - "vjp_merge_tile": 64, + "dx_dw_chunk_k": 128, + "dx_dw_chunk_accumulation": "two_64_term_halves", + "dx_dw_merge_tile": 64, "bias_reduction": "fixed_midpoint_tree", "reduction_leaf_size": _BLOCK_K, "reduction_backend": "triton", diff --git a/rl_engine/ops/autograd/canonical_backward.py b/rl_engine/ops/autograd/canonical_backward.py index 4ae7819cf..d5b16f804 100644 --- a/rl_engine/ops/autograd/canonical_backward.py +++ b/rl_engine/ops/autograd/canonical_backward.py @@ -14,6 +14,9 @@ import torch WeightReducer = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] +# H20/H3 row-count heuristic for attempting the CUDA order check. +# Small inputs retain stable sorting to avoid the extra check and synchronization. +# Larger, already ordered inputs can skip sorting and gradient gathering. _CANONICAL_ORDER_REUSE_MIN_ROWS = 4096 diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 5cd07ba7a..ce17e256c 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -266,10 +266,10 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "training_mode": "eager", "parameter_gradient_order": "logical_key", "reduction": "fixed_128_chunks_balanced_fp32_merge", - "vjp_tile": (128, 32, 32), - "vjp_num_warps": 4, - "vjp_register_allocation": "compiler_default", - "vjp_chunk_accumulation": "two_64_term_halves", + "dx_dw_tile": (128, 32, 32), + "dx_dw_num_warps": 4, + "dx_dw_register_allocation": "compiler_default", + "dx_dw_chunk_accumulation": "two_64_term_halves", "forward_aligned_tiles": True, "canonical_rows": "live_strict_order_proof_or_stable_sort", }, diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index f97c7c3ce..6cf0419ad 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -17,7 +17,8 @@ def test_video_input_projection_default_and_strict_scope(): ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) - assert desc.determinism_or_alignment_properties["vjp_tile"] == (128, 32, 32) + assert desc.determinism_or_alignment_properties["dx_dw_tile"] == (128, 32, 32) + assert desc.determinism_or_alignment_properties["dx_dw_num_warps"] == 4 assert ( desc.determinism_or_alignment_properties["canonical_rows"] == "live_strict_order_proof_or_stable_sort" From 682ece6fb8cf57bb5d33a6f8cfc24b5ce31ecc33 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Mon, 19 Oct 2026 21:18:29 +0800 Subject: [PATCH 16/26] Restore shared FP32 GEMM compiler fusion defaults Signed-off-by: walker-ai <2398833647@qq.com> --- docs/operators/video-input-projection-fp32.md | 6 ++++-- rl_engine/backends/shared/triton/gemm/det_gemm.py | 1 - .../shared/triton/gemm/video_input_projection_fp32.py | 6 ++++-- rl_engine/runtime/registry.py | 7 +++++-- tests/ops/gemm/test_video_input_projection_fp32.py | 2 +- tests/runtime/test_kernel_registry.py | 2 +- 6 files changed, 15 insertions(+), 9 deletions(-) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 1aace14f5..deb7ac927 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -112,7 +112,7 @@ partials remain FP32. A second grid uses a fixed balanced reduction over the next power-of-two chunk count, with zero padding. The merge uses 64 output elements per program and four warps. VJP GEMMs use fixed 128-by-32 output tiles, 32-term dot tiles, and four warps. The forward retains its 64-by-64 output tile. -These settings define the v10 Triton implementation identity. The K partition +These settings define the v11 Triton implementation identity. The K partition and VJP tile do not depend on M, row placement, or occupancy. The VJP kernels consume positive-stride views directly, including `G.T`. The Triton forward omits K or N load predicates only when the corresponding @@ -131,7 +131,9 @@ The VJP schedule caps its FP32 partial workspace at 512 MiB by chunking output rows. The retained midpoint helper also supports output-column chunking when one complete row exceeds its budget. Existing BF16 callers retain their dtype and schedule. -TF32 and compiler multiply-add fusion remain disabled. +TF32 remains disabled. The shared direct FP32 GEMM, including the H3 forward, +uses the compiler's default multiply-add fusion setting. The chunked VJP and +midpoint reduction kernels explicitly disable compiler multiply-add fusion. The reused Triton `tl.dot` lowers its IEEE FP32 inner products to fixed SIMT `fma.rn.f32` instructions. The compiler fusion flag controls surrounding expressions; it does not remove the intrinsic dot's FMA instructions. This diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index a080831ca..e2b8c8154 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -957,7 +957,6 @@ def _triton_gemm_fp32( HAS_BIAS=bias is not None, EVEN_K=aligned_tiles and K % _BLOCK_K == 0, EVEN_N=aligned_tiles and N % _BLOCK_N == 0, - **({"enable_fp_fusion": False} if torch.version.hip is None else {}), ) return c diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index 9546d7596..a2a513dbc 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -38,7 +38,7 @@ def projection_gemm_fp32(a, b, *, bias=None): class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton" matmul_op = staticmethod(projection_gemm_fp32) - backend_id = "triton-video-input-projection-fp32-v10" + backend_id = "triton-video-input-projection-fp32-v11" def __init__(self): if not _TRITON_AVAILABLE: @@ -67,7 +67,9 @@ def backend_info(self): "forward_aligned_tiles": True, "canonical_rows": "live_strict_order_proof_or_stable_sort", "tf32": False, - "fp_fusion": False, + "direct_gemm_fp_fusion": "compiler_default", + "chunked_gemm_fp_fusion": False, + "tree_gemm_fp_fusion": False, "dot_arithmetic": "ieee_fp32_simt_fma", "forward_bias": "fp32_epilogue", "leaf_loop_bound": "max_logical_leaf_length", diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index ce17e256c..166818f2d 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -241,7 +241,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v10", + backend_id="triton-video-input-projection-fp32-v11", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -256,6 +256,9 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: determinism_or_alignment_properties={ "algorithm": "fixed_fp32_gemm_canonical_parameter_vjp", "forward_bias": "fp32_epilogue", + "direct_gemm_fp_fusion": "compiler_default", + "chunked_gemm_fp_fusion": False, + "tree_gemm_fp_fusion": False, "leaf_loop_bound": "max_logical_leaf_length", "input_width": 96, "output_width": 5376, @@ -278,7 +281,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( "rlkernel-triton-ieee-fp32-dot-chunk128-halves64-tile128x32-merge64-" - "strided-aligned-forward-auto-canonical-proof-h3-v10" + "strided-aligned-forward-default-fusion-auto-canonical-proof-h3-v11" ), ), OperatorBackendDescriptor( diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index ae97c3743..79449276a 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -570,7 +570,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v10" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v11" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op()(x, weight, bias)) diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 6cf0419ad..38634c653 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,7 +13,7 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v10" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v11" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) From 88adb4c2c6360adb857e621f1ff513710353d303 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Mon, 19 Oct 2026 23:42:15 +0800 Subject: [PATCH 17/26] Refactor compact tree workspace planning Replace FP32-specific plan state with the actual workspace size and a compact_workspace request. Build mapped index tensors once and extend existing tests for complete, ragged, and single-leaf plans. Verified on one H20: 231 tests passed, 48 precision cases and 30 BI checks passed. Before/after outputs match bitwise. Full forward/backward median timings change by less than 1%, with the compact workspace preserved. Signed-off-by: walker-ai <2398833647@qq.com> --- .../backends/shared/triton/gemm/det_gemm.py | 56 ++++++++++--------- .../gemm/test_video_input_projection_fp32.py | 47 +++++++++++----- 2 files changed, 63 insertions(+), 40 deletions(-) diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index e2b8c8154..4ac4a33ae 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -197,7 +197,6 @@ class _TreePlan: class _DeviceTreePlan: host: _TreePlan workspace_node_count: int - fp32_fused: bool leaf_starts: torch.Tensor leaf_lengths: torch.Tensor leaf_nodes: torch.Tensor @@ -264,10 +263,10 @@ def visit(begin: int, end: int) -> tuple[int, int]: def _device_tree_plan( - k_size: int, device: torch.device, *, fp32_fuse_leaves: bool = False + k_size: int, device: torch.device, *, compact_workspace: bool = False ) -> _DeviceTreePlan: device_index = device.index if device.index is not None else torch.cuda.current_device() - key = (device_index, k_size, fp32_fuse_leaves) + key = (device_index, k_size, compact_workspace) with _TREE_PLAN_LOCK: cached = _TREE_PLANS.get(key) if cached is not None: @@ -278,10 +277,6 @@ def _device_tree_plan( def indices(values: tuple[int, ...]) -> torch.Tensor: return torch.tensor(values, dtype=torch.int32, device=device) - levels = [] - for operations in host.reduction_levels: - lower_nodes, upper_nodes, output_nodes = zip(*operations, strict=True) - levels.append((indices(lower_nodes), indices(upper_nodes), indices(output_nodes))) # The paired kernels require complete subtrees. Ragged midpoint trees # retain the original per-level schedule and workspace indices. leaf_count = len(host.leaf_nodes) @@ -289,10 +284,10 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: len(level) * (1 << (height + 1)) == leaf_count for height, level in enumerate(host.reduction_levels) ) - fp32_fused = fp32_fuse_leaves and complete + use_compact_plan = compact_workspace and complete node_map = {node: node for node in range(host.node_count)} workspace_node_count = host.node_count - if fp32_fused: + if use_compact_plan: # Alternate disjoint buffers for the retained levels. Each launch # reads its previous buffer completely before the next launch. first_nodes = tuple(op[2] for op in host.reduction_levels[0]) @@ -308,10 +303,21 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: def mapped_indices(values): return indices(tuple(node_map[node] for node in values)) + levels = [] + for operations in host.reduction_levels: + lower_nodes, upper_nodes, output_nodes = zip(*operations, strict=True) + levels.append( + ( + mapped_indices(lower_nodes), + mapped_indices(upper_nodes), + mapped_indices(output_nodes), + ) + ) + fused_pairs = [] leaf_reduction = None fused_pairs_after_leaf = [] - if _device_arch(device_index) == "gfx942" or fp32_fused: + if _device_arch(device_index) == "gfx942" or use_compact_plan: def build_fused_pairs(start_level: int): result = [] @@ -337,7 +343,7 @@ def build_fused_pairs(start_level: int): ) return result - if not fp32_fused: + if not use_compact_plan: fused_pairs = build_fused_pairs(0) fused_pairs_after_leaf = build_fused_pairs(1) if host.reduction_levels: @@ -373,18 +379,10 @@ def build_fused_pairs(start_level: int): result = _DeviceTreePlan( host=host, workspace_node_count=workspace_node_count, - fp32_fused=fp32_fused, leaf_starts=indices(host.leaf_starts), leaf_lengths=indices(host.leaf_lengths), leaf_nodes=mapped_indices(host.leaf_nodes), - reduction_levels=( - tuple( - tuple(mapped_indices(tuple(op[column] for op in level)) for column in range(3)) - for level in host.reduction_levels - ) - if fp32_fused - else tuple(levels) - ), + reduction_levels=tuple(levels), rocm_fused_reduction_pairs=tuple(fused_pairs), rocm_leaf_reduction=leaf_reduction, rocm_fused_reduction_pairs_after_leaf=tuple(fused_pairs_after_leaf), @@ -1014,7 +1012,7 @@ def _triton_tree_gemm( plan = _device_tree_plan( k_size, a.device, - fp32_fuse_leaves=fp32_tree and m_size == 1 and not transpose_output, + compact_workspace=fp32_tree and m_size == 1 and not transpose_output, ) workspace_elements = plan.workspace_node_count * m_size * n_size if workspace_elements > _MAX_TREE_WORKSPACE_ELEMENTS and out is None: @@ -1106,12 +1104,20 @@ def _triton_tree_gemm( else (len(plan.host.leaf_nodes), tiles_m, tiles_n) ) direct_root_output = not transpose_output and (is_gfx942 or fp32_tree) + # Compact slots are shared across levels and require the paired schedule. + use_compact_plan = plan.workspace_node_count < plan.host.node_count fuse_leaf_reduction = ( - (plan.fp32_fused or (_ROCM_FUSE_LEAF_REDUCTION and is_gfx942)) - and direct_root_output - and (use_inference_schedule or plan.fp32_fused) - and (k_size >= 1536 or plan.fp32_fused) + direct_root_output and plan.rocm_leaf_reduction is not None + and ( + use_compact_plan + or ( + _ROCM_FUSE_LEAF_REDUCTION + and is_gfx942 + and use_inference_schedule + and k_size >= 1536 + ) + ) ) if fuse_leaf_reduction: assert plan.rocm_leaf_reduction is not None diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index 79449276a..a91aa2517 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -771,29 +771,46 @@ def gemm(a, b): @pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") -@pytest.mark.parametrize("k", [128, 1024, 13104]) -def test_compact_bias_tree_workspace_preserves_chunked_outputs(k, monkeypatch): +@pytest.mark.parametrize( + "k,compact", + [ + (1, False), + (32, False), + (64, True), + (65, False), + (128, True), + (129, False), + (256, True), + (1024, True), + (13104, True), + ], +) +def test_compact_bias_tree_workspace_preserves_chunked_outputs(k, compact, monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm - from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( - projection_gemm_fp32, - ) a, b = torch.ones(1, k, device="cuda"), torch.randn(k, 5, device="cuda") original_plan = det_gemm._device_tree_plan(k, a.device) - assert not original_plan.fp32_fused assert original_plan.workspace_node_count == original_plan.host.node_count - plan = det_gemm._device_tree_plan(k, a.device, fp32_fuse_leaves=True) - assert plan.fp32_fused - assert plan.workspace_node_count < plan.host.node_count - node_arrays = [plan.rocm_leaf_reduction[-1]] - node_arrays.extend( - nodes for pair in plan.rocm_fused_reduction_pairs_after_leaf for nodes in pair - ) + plan = det_gemm._device_tree_plan(k, a.device, compact_workspace=True) + assert (plan.workspace_node_count < plan.host.node_count) == compact + if compact: + assert plan.rocm_leaf_reduction is not None + node_arrays = [plan.rocm_leaf_reduction[-1]] + node_arrays.extend( + nodes for pair in plan.rocm_fused_reduction_pairs_after_leaf for nodes in pair + ) + if (len(plan.host.reduction_levels) - 1) % 2: + node_arrays.extend(plan.reduction_levels[-1][:2]) + else: + assert plan.workspace_node_count == plan.host.node_count + node_arrays = [plan.leaf_nodes] + node_arrays.extend(nodes for level in plan.reduction_levels for nodes in level) for nodes in node_arrays: assert bool(torch.all((nodes >= 0) & (nodes < plan.workspace_node_count))) - expected = projection_gemm_fp32(a, b) + expected = det_gemm._triton_tree_gemm(a, b, output_dtype=torch.float32) + _assert_bytes(expected.cpu(), _fp32_tree_reference(a, b)) monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.workspace_node_count * 2) - _assert_bytes(projection_gemm_fp32(a, b), expected) + _assert_bytes(det_gemm._triton_tree_gemm(a, b, output_dtype=torch.float32), expected) @pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") From 5d60da6a3732d181adb968f7a0763a3fdae04e2b Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Tue, 20 Oct 2026 00:24:59 +0800 Subject: [PATCH 18/26] Remove compact tree storage after H20 breakdown Keep paired leaf/level scheduling and restore original node slots. Remove compact_workspace, workspace_node_count and slot remapping. Restore the plan signature and cache key; cache generic CUDA descriptors. Update workspace coverage and backend identity. H20: 231 tests passed; operator CLI and public benchmark passed. A/B/C breakdown: paired bias scheduling is 32.1% faster at 13104 rows; compact storage adds no latency benefit or full-call peak-memory reduction. Final implementation matches candidate B outputs and timing. Signed-off-by: walker-ai <2398833647@qq.com> --- docs/operators/video-input-projection-fp32.md | 5 +- .../backends/shared/triton/gemm/det_gemm.py | 91 ++++++------------- .../gemm/video_input_projection_fp32.py | 4 +- rl_engine/runtime/registry.py | 2 +- .../gemm/test_video_input_projection_fp32.py | 30 +++--- tests/runtime/test_kernel_registry.py | 2 +- 6 files changed, 46 insertions(+), 88 deletions(-) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index deb7ac927..0c486c896 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -9,8 +9,9 @@ reference. The model adapter owns the subsequent cast into the BF16 packed strea The implementation reuses RL-Kernel's Triton FP32 GEMM and canonical parameter-gradient session. Input and weight VJPs use fixed 128-term K chunks and one balanced FP32 merge. Long bias VJPs retain the one-row Triton tree. -Only this one-row FP32 tree uses compact paired-leaf scheduling. Other FP32 -tree shapes reuse the existing default tile and per-level schedule. +Only this one-row FP32 tree uses paired-leaf scheduling. Each tree node retains +its original workspace slot. Other FP32 tree shapes reuse the existing default +tile and per-level schedule. ## Entry Point diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 4ac4a33ae..edc5136d3 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -43,7 +43,7 @@ # one shot would otherwise create a multi-terabyte virtual tensor on ROCm. # Chunking over M preserves the exact K-tree and BF16 rounding contract. _MAX_TREE_WORKSPACE_ELEMENTS = 128 * 1024 * 1024 -_TREE_PLANS: dict[tuple[int, int, bool], "_DeviceTreePlan"] = {} +_TREE_PLANS: dict[tuple[int, int], "_DeviceTreePlan"] = {} _TREE_PLAN_LOCK = threading.Lock() _ROCM_TUNE_DECODE_LEAF = True _ROCM_FUSE_LEAF_REDUCTION = True @@ -196,7 +196,6 @@ class _TreePlan: @dataclass(frozen=True) class _DeviceTreePlan: host: _TreePlan - workspace_node_count: int leaf_starts: torch.Tensor leaf_lengths: torch.Tensor leaf_nodes: torch.Tensor @@ -262,11 +261,9 @@ def visit(begin: int, end: int) -> tuple[int, int]: ) -def _device_tree_plan( - k_size: int, device: torch.device, *, compact_workspace: bool = False -) -> _DeviceTreePlan: +def _device_tree_plan(k_size: int, device: torch.device) -> _DeviceTreePlan: device_index = device.index if device.index is not None else torch.cuda.current_device() - key = (device_index, k_size, compact_workspace) + key = (device_index, k_size) with _TREE_PLAN_LOCK: cached = _TREE_PLANS.get(key) if cached is not None: @@ -277,47 +274,21 @@ def _device_tree_plan( def indices(values: tuple[int, ...]) -> torch.Tensor: return torch.tensor(values, dtype=torch.int32, device=device) - # The paired kernels require complete subtrees. Ragged midpoint trees - # retain the original per-level schedule and workspace indices. - leaf_count = len(host.leaf_nodes) - complete = bool(host.reduction_levels) and all( - len(level) * (1 << (height + 1)) == leaf_count - for height, level in enumerate(host.reduction_levels) - ) - use_compact_plan = compact_workspace and complete - node_map = {node: node for node in range(host.node_count)} - workspace_node_count = host.node_count - if use_compact_plan: - # Alternate disjoint buffers for the retained levels. Each launch - # reads its previous buffer completely before the next launch. - first_nodes = tuple(op[2] for op in host.reduction_levels[0]) - next_count = max(1, len(first_nodes) // 4) - workspace_node_count = len(first_nodes) + next_count - for i, node in enumerate(first_nodes): - node_map[node] = i - for step, height in enumerate(range(2, len(host.reduction_levels), 2), start=1): - base = len(first_nodes) if step % 2 else 0 - for i, (_, _, node) in enumerate(host.reduction_levels[height]): - node_map[node] = base + i - - def mapped_indices(values): - return indices(tuple(node_map[node] for node in values)) - levels = [] for operations in host.reduction_levels: lower_nodes, upper_nodes, output_nodes = zip(*operations, strict=True) - levels.append( - ( - mapped_indices(lower_nodes), - mapped_indices(upper_nodes), - mapped_indices(output_nodes), - ) - ) - + levels.append((indices(lower_nodes), indices(upper_nodes), indices(output_nodes))) fused_pairs = [] leaf_reduction = None fused_pairs_after_leaf = [] - if _device_arch(device_index) == "gfx942" or use_compact_plan: + # Complete trees can use the existing paired kernels on CUDA too. + # The execution path decides whether to use these optional descriptors. + leaf_count = len(host.leaf_nodes) + complete = bool(host.reduction_levels) and all( + len(level) * (1 << (height + 1)) == leaf_count + for height, level in enumerate(host.reduction_levels) + ) + if _device_arch(device_index) == "gfx942" or (torch.version.hip is None and complete): def build_fused_pairs(start_level: int): result = [] @@ -334,17 +305,16 @@ def build_fused_pairs(start_level: int): node0, node1, node2, node3 = zip(*grandchildren, strict=True) result.append( ( - mapped_indices(node0), - mapped_indices(node1), - mapped_indices(node2), - mapped_indices(node3), - mapped_indices(tuple(outputs)), + indices(node0), + indices(node1), + indices(node2), + indices(node3), + indices(tuple(outputs)), ) ) return result - if not use_compact_plan: - fused_pairs = build_fused_pairs(0) + fused_pairs = build_fused_pairs(0) fused_pairs_after_leaf = build_fused_pairs(1) if host.reduction_levels: first_level = host.reduction_levels[0] @@ -374,14 +344,13 @@ def build_fused_pairs(start_level: int): indices(tuple(lower_lengths)), indices(tuple(upper_starts)), indices(tuple(upper_lengths)), - mapped_indices(tuple(outputs)), + indices(tuple(outputs)), ) result = _DeviceTreePlan( host=host, - workspace_node_count=workspace_node_count, leaf_starts=indices(host.leaf_starts), leaf_lengths=indices(host.leaf_lengths), - leaf_nodes=mapped_indices(host.leaf_nodes), + leaf_nodes=indices(host.leaf_nodes), reduction_levels=tuple(levels), rocm_fused_reduction_pairs=tuple(fused_pairs), rocm_leaf_reduction=leaf_reduction, @@ -1009,15 +978,11 @@ def _triton_tree_gemm( b = b.contiguous() m_size, k_size = a.shape n_size = b.size(1) - plan = _device_tree_plan( - k_size, - a.device, - compact_workspace=fp32_tree and m_size == 1 and not transpose_output, - ) - workspace_elements = plan.workspace_node_count * m_size * n_size + plan = _device_tree_plan(k_size, a.device) + workspace_elements = plan.host.node_count * m_size * n_size if workspace_elements > _MAX_TREE_WORKSPACE_ELEMENTS and out is None: - if fp32_tree and plan.workspace_node_count * n_size > _MAX_TREE_WORKSPACE_ELEMENTS: - columns_per_chunk = _MAX_TREE_WORKSPACE_ELEMENTS // plan.workspace_node_count + if fp32_tree and plan.host.node_count * n_size > _MAX_TREE_WORKSPACE_ELEMENTS: + columns_per_chunk = _MAX_TREE_WORKSPACE_ELEMENTS // plan.host.node_count if columns_per_chunk < 1: raise RuntimeError("FP32 tree plan exceeds the workspace budget for one element") # A one-row bias VJP cannot be split further over M. Splitting @@ -1036,7 +1001,7 @@ def _triton_tree_gemm( return torch.cat(chunks, dim=0 if transpose_output else 1) rows_per_chunk = max( 1, - _MAX_TREE_WORKSPACE_ELEMENTS // (plan.workspace_node_count * n_size), + _MAX_TREE_WORKSPACE_ELEMENTS // (plan.host.node_count * n_size), ) chunks = [] for start in range(0, m_size, rows_per_chunk): @@ -1070,7 +1035,7 @@ def _triton_tree_gemm( raise ValueError("Triton tree GEMM output buffer must not require gradients") result = out workspace = torch.empty( - (plan.workspace_node_count, m_size, n_size), + (plan.host.node_count, m_size, n_size), dtype=output_dtype, device=a.device, ) @@ -1104,13 +1069,11 @@ def _triton_tree_gemm( else (len(plan.host.leaf_nodes), tiles_m, tiles_n) ) direct_root_output = not transpose_output and (is_gfx942 or fp32_tree) - # Compact slots are shared across levels and require the paired schedule. - use_compact_plan = plan.workspace_node_count < plan.host.node_count fuse_leaf_reduction = ( direct_root_output and plan.rocm_leaf_reduction is not None and ( - use_compact_plan + (fp32_tree and m_size == 1) or ( _ROCM_FUSE_LEAF_REDUCTION and is_gfx942 diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index a2a513dbc..87915e599 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -38,7 +38,7 @@ def projection_gemm_fp32(a, b, *, bias=None): class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton" matmul_op = staticmethod(projection_gemm_fp32) - backend_id = "triton-video-input-projection-fp32-v11" + backend_id = "triton-video-input-projection-fp32-v12" def __init__(self): if not _TRITON_AVAILABLE: @@ -83,5 +83,5 @@ def backend_info(self): "reduction_backend": "triton", "leaf_launch": "paired_leaves_for_one_row_complete_trees", "bias_tree_tile": [1, 128, 32], - "tree_schedule": "paired_levels_compact_workspace", + "tree_schedule": "paired_levels_full_workspace", } diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 166818f2d..8eb7cf886 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -241,7 +241,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v11", + backend_id="triton-video-input-projection-fp32-v12", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index a91aa2517..e6effd938 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -570,7 +570,7 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v11" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v12" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op()(x, weight, bias)) @@ -772,7 +772,7 @@ def gemm(a, b): @pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @pytest.mark.parametrize( - "k,compact", + "k,fused", [ (1, False), (32, False), @@ -785,31 +785,25 @@ def gemm(a, b): (13104, True), ], ) -def test_compact_bias_tree_workspace_preserves_chunked_outputs(k, compact, monkeypatch): +def test_bias_tree_workspace_preserves_chunked_outputs(k, fused, monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm a, b = torch.ones(1, k, device="cuda"), torch.randn(k, 5, device="cuda") - original_plan = det_gemm._device_tree_plan(k, a.device) - assert original_plan.workspace_node_count == original_plan.host.node_count - plan = det_gemm._device_tree_plan(k, a.device, compact_workspace=True) - assert (plan.workspace_node_count < plan.host.node_count) == compact - if compact: - assert plan.rocm_leaf_reduction is not None - node_arrays = [plan.rocm_leaf_reduction[-1]] + plan = det_gemm._device_tree_plan(k, a.device) + assert (plan.rocm_leaf_reduction is not None) == fused + assert plan.leaf_nodes.tolist() == list(plan.host.leaf_nodes) + node_arrays = [plan.leaf_nodes] + node_arrays.extend(nodes for level in plan.reduction_levels for nodes in level) + if fused: + node_arrays.append(plan.rocm_leaf_reduction[-1]) node_arrays.extend( nodes for pair in plan.rocm_fused_reduction_pairs_after_leaf for nodes in pair ) - if (len(plan.host.reduction_levels) - 1) % 2: - node_arrays.extend(plan.reduction_levels[-1][:2]) - else: - assert plan.workspace_node_count == plan.host.node_count - node_arrays = [plan.leaf_nodes] - node_arrays.extend(nodes for level in plan.reduction_levels for nodes in level) for nodes in node_arrays: - assert bool(torch.all((nodes >= 0) & (nodes < plan.workspace_node_count))) + assert bool(torch.all((nodes >= 0) & (nodes < plan.host.node_count))) expected = det_gemm._triton_tree_gemm(a, b, output_dtype=torch.float32) _assert_bytes(expected.cpu(), _fp32_tree_reference(a, b)) - monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.workspace_node_count * 2) + monkeypatch.setattr(det_gemm, "_MAX_TREE_WORKSPACE_ELEMENTS", plan.host.node_count * 2) _assert_bytes(det_gemm._triton_tree_gemm(a, b, output_dtype=torch.float32), expected) diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 38634c653..3af229854 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,7 +13,7 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v11" + "video_input_projection_fp32", "triton-video-input-projection-fp32-v12" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) From 6d621f36a9b6956504fbebb48dcac505396d9d71 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Tue, 20 Oct 2026 14:01:26 +0800 Subject: [PATCH 19/26] Document H3 projection GEMM dispatch and scheduling Signed-off-by: walker-ai <2398833647@qq.com> --- .../backends/shared/triton/gemm/det_gemm.py | 122 ++++++++++++++---- .../gemm/video_input_projection_fp32.py | 26 +++- 2 files changed, 124 insertions(+), 24 deletions(-) diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index edc5136d3..e17264c2f 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -11,9 +11,9 @@ does not match the scalar reference at every BF16 rounding boundary. Autotuning and adaptive split-K are intentionally disabled. Default FP32-output -calls preserve the earlier fixed-order, no-split-K contract. Explicit FP32 VJP -calls use fixed 128-term chunks. The optional CUDA FP32 tree reuses the canonical -plan, leaf kernel, and level reductions with FP32 nodes. +calls preserve the earlier fixed-order, no-split-K contract. Explicit long-K FP32 +input and weight gradients use fixed 128-term chunks. The optional CUDA FP32 tree +reuses the canonical plan, leaf kernel, and level reductions for bias gradients. """ from __future__ import annotations @@ -51,6 +51,8 @@ @dataclass(frozen=True) class _TreeLeafConfig: + """Store leaf tile sizes, warp count, and program-grid ordering.""" + block_m: int block_n: int num_warps: int @@ -125,6 +127,8 @@ def _gfx942_qwen_tree_leaf_config( transpose_output: bool, preserve_a_strides: bool, ) -> _TreeLeafConfig: + """Select a measured gfx942 Qwen tile, or return the default leaf tile.""" + logical_shape = (m_size, k_size, n_size) if ( not transpose_output @@ -155,6 +159,8 @@ def _gfx942_qwen_tree_leaf_config( @functools.lru_cache(maxsize=None) def _device_arch(device_index: int) -> str: + """Return the ROCm architecture name, or an empty string on CUDA.""" + if getattr(torch.version, "hip", None) is None: return "" properties = torch.cuda.get_device_properties(device_index) @@ -171,6 +177,8 @@ def _tree_leaf_config( transpose_output: bool, preserve_a_strides: bool, ) -> _TreeLeafConfig: + """Select the leaf tile for the device, matrix shape, and stride mode.""" + device_index = device.index if device.index is not None else torch.cuda.current_device() if _device_arch(device_index) != "gfx942": return _DEFAULT_TREE_LEAF_CONFIG @@ -185,6 +193,8 @@ def _tree_leaf_config( @dataclass(frozen=True) class _TreePlan: + """Describe midpoint-tree leaves, ordered reductions, and original node IDs.""" + leaf_starts: tuple[int, ...] leaf_lengths: tuple[int, ...] leaf_nodes: tuple[int, ...] @@ -195,6 +205,8 @@ class _TreePlan: @dataclass(frozen=True) class _DeviceTreePlan: + """Store tree indices on the device, including optional fused schedules.""" + host: _TreePlan leaf_starts: torch.Tensor leaf_lengths: torch.Tensor @@ -219,6 +231,8 @@ class _DeviceTreePlan: def _build_tree_plan(k_size: int) -> _TreePlan: + """Split K at each midpoint until every leaf contains at most _BLOCK_K terms.""" + if k_size <= 0: raise ValueError(f"deterministic GEMM K must be positive, got {k_size}") @@ -229,6 +243,8 @@ def _build_tree_plan(k_size: int) -> _TreePlan: next_node = 0 def visit(begin: int, end: int) -> tuple[int, int]: + """Record one subtree and return its root node ID and height.""" + nonlocal next_node if end - begin <= _BLOCK_K: node = next_node @@ -262,6 +278,8 @@ def visit(begin: int, end: int) -> tuple[int, int]: def _device_tree_plan(k_size: int, device: torch.device) -> _DeviceTreePlan: + """Cache device indices for the tree and supported paired-level schedules.""" + device_index = device.index if device.index is not None else torch.cuda.current_device() key = (device_index, k_size) with _TREE_PLAN_LOCK: @@ -272,6 +290,8 @@ def _device_tree_plan(k_size: int, device: torch.device) -> _DeviceTreePlan: host = _build_tree_plan(k_size) def indices(values: tuple[int, ...]) -> torch.Tensor: + """Copy node IDs or leaf bounds to a device index tensor.""" + return torch.tensor(values, dtype=torch.int32, device=device) levels = [] @@ -291,6 +311,8 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: if _device_arch(device_index) == "gfx942" or (torch.version.hip is None and complete): def build_fused_pairs(start_level: int): + """Map adjacent reduction levels to four inputs and one output.""" + result = [] for level_index in range(start_level, len(host.reduction_levels) - 1, 2): first = { @@ -383,6 +405,8 @@ def _det_gemm_tree_leaf_kernel( N_FASTEST: tl.constexpr, LEAF_LOOP_BOUND: tl.constexpr = _BLOCK_K, ): + """Accumulate each leaf in scalar K order and store one tree node.""" + if N_FASTEST: pid_n = tl.program_id(0) pid_m = tl.program_id(1) @@ -400,10 +424,9 @@ def _det_gemm_tree_leaf_kernel( leaf_length = tl.load(leaf_lengths_ptr + leaf).to(tl.int64) leaf_node = tl.load(leaf_nodes_ptr + leaf).to(tl.int64) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - # Keep the leaf's ascending scalar FMA order identical to the native - # gfx942 correctness kernel. FP32 launches trim only the masked tail. - # A tl.dot/MFMA leaf is topology-stable but - # differs from the scalar reference at rare BF16 rounding boundaries. + # Keep each leaf's ascending scalar order. FP32 launches trim only the + # masked tail; BF16 launches retain the native gfx942 FMA behavior. + # A tl.dot/MFMA leaf can change rounding relative to this scalar path. for offset in tl.static_range(0, LEAF_LOOP_BOUND): k_offset = leaf_start + offset active = offset < leaf_length @@ -456,6 +479,8 @@ def _det_gemm_tree_leaf_reduce_rocm_kernel( WRITE_OUTPUT: tl.constexpr, LEAF_LOOP_BOUND: tl.constexpr = _BLOCK_K, ): + """Compute paired leaves and their first reduction in one kernel call.""" + if N_FASTEST: pid_n = tl.program_id(0) pid_m = tl.program_id(1) @@ -486,7 +511,7 @@ def _det_gemm_tree_leaf_reduce_rocm_kernel( other=0.0, ).to(tl.float32) lower_acc += a[:, None] * b[None, :] - # This conversion is the original lower-leaf workspace store boundary. + # Preserve the original lower-leaf rounding to the workspace dtype. lower = lower_acc.to(workspace_ptr.dtype.element_ty) upper_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) @@ -504,8 +529,8 @@ def _det_gemm_tree_leaf_reduce_rocm_kernel( other=0.0, ).to(tl.float32) upper_acc += a[:, None] * b[None, :] - # Match both leaf BF16 stores followed by the original first-level FP32 - # add and BF16 store. Only the intermediate global-memory trip is gone. + # Round both leaves to the workspace dtype before the first-level add. + # Fusing removes the intermediate global-memory writes and reads. result = lower.to(tl.float32) + upper_acc.to(workspace_ptr.dtype.element_ty).to(tl.float32) output_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) offsets = offs_m[:, None].to(tl.int64) * N + offs_n[None, :].to(tl.int64) @@ -533,6 +558,8 @@ def _det_gemm_tree_reduce_kernel( N: tl.constexpr, BLOCK: tl.constexpr, ): + """Add each pair in one tree level and round to the workspace dtype.""" + operation = tl.program_id(0) block = tl.program_id(1) offsets = (block * BLOCK + tl.arange(0, BLOCK)).to(tl.int64) @@ -564,13 +591,13 @@ def _det_gemm_tree_reduce_to_output_rocm_kernel( output_ptr, lower_nodes_ptr, upper_nodes_ptr, - # M is runtime-sized so decode batch changes reuse one compiled - # reduction kernel instead of triggering a Triton JIT per batch size. - # The reduction order and BF16 store boundary are unchanged. M, N: tl.constexpr, BLOCK: tl.constexpr, ): + """Write the final tree addition directly to the output buffer.""" + + # Runtime-sized M lets different decode batches reuse this compiled kernel. offsets = (tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)).to(tl.int64) elements = M * N mask = offsets < elements @@ -587,8 +614,7 @@ def _det_gemm_tree_reduce_to_output_rocm_kernel( other=0.0, ).to(tl.float32) result = lower + upper - # Preserve the canonical root's FP32 add and BF16 store boundary; only - # its destination changes, avoiding a device-to-device copy launch. + # Preserve the root's FP32 add and output rounding while avoiding a copy. tl.store( output_ptr + offsets, result.to(output_ptr.dtype.element_ty), @@ -609,6 +635,8 @@ def _det_gemm_tree_reduce_two_levels_rocm_kernel( BLOCK: tl.constexpr, WRITE_OUTPUT: tl.constexpr, ): + """Combine two tree levels while preserving intermediate rounding.""" + operation = tl.program_id(0) block = tl.program_id(1) offsets = (block * BLOCK + tl.arange(0, BLOCK)).to(tl.int64) @@ -630,8 +658,8 @@ def _det_gemm_tree_reduce_two_levels_rocm_kernel( value3 = tl.load(workspace_ptr + node3 * elements + offsets, mask=mask, other=0.0).to( tl.float32 ) - # Match the two original kernel boundaries exactly: each first-level - # FP32 add is rounded to BF16 before the second-level FP32 add. + # Round each first-level add to the workspace dtype before the second + # level, matching the original two kernel boundaries. lower = (value0 + value1).to(workspace_ptr.dtype.element_ty).to(tl.float32) upper = (value2 + value3).to(workspace_ptr.dtype.element_ty).to(tl.float32) result = lower + upper @@ -657,6 +685,8 @@ def _copy_tree_root_kernel( elements, BLOCK: tl.constexpr, ): + """Copy the completed root node to a contiguous output buffer.""" + offsets = (tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)).to(tl.int64) mask = offsets < elements values = tl.load(workspace_ptr + root * elements + offsets, mask=mask) @@ -672,6 +702,8 @@ def _copy_tree_root_transposed_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, ): + """Copy the completed root from [M, N] to an [N, M] output buffer.""" + offsets_m = (tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)).to(tl.int64) offsets_n = (tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)).to(tl.int64) elements = M * N @@ -680,7 +712,7 @@ def _copy_tree_root_transposed_kernel( workspace_ptr + root * elements + offsets_m[:, None] * N + offsets_n[None, :], mask=mask, ) - # Store the already-rounded BF16 root directly in [N, M] layout. + # Store the already-rounded root directly in [N, M] layout. # This is a pure address permutation: the canonical GEMM leaves, tree, # operand order, and every rounding boundary remain unchanged. tl.store( @@ -699,6 +731,8 @@ def _canonical_key_order_kernel( COLUMNS: tl.constexpr, BLOCK: tl.constexpr, ): + """Write one flag per block for strictly ordered keys without padding.""" + rows = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) first = tl.load(keys_ptr + rows * stride_row, mask=rows < ROWS, other=-1) greater = tl.full((BLOCK,), False, tl.int1) @@ -743,6 +777,13 @@ def _det_gemm_fp32_kernel( EVEN_K: tl.constexpr = False, EVEN_N: tl.constexpr = False, ): + """Compute a full FP32 GEMM tile or one fixed K chunk. + + HAS_BIAS adds one column bias before storing the tile. EVEN_K and + EVEN_N omit load masks only when the caller proves tile divisibility. + Chunked calls store [chunk, M, N] partials and leave HAS_BIAS disabled. + """ + # One program walks the whole K or one declared fixed-width chunk. # The K order does not depend on output rows or their placement. pid_m, pid_n = tl.program_id(0), tl.program_id(1) @@ -798,6 +839,8 @@ def _det_gemm_fp32_chunk_merge_kernel( PADDED_CHUNKS: tl.constexpr, BLOCK_C: tl.constexpr, ): + """Sum [chunk, element] partials with a fixed padded reduction width.""" + offsets = tl.program_id(0) * BLOCK_C + tl.arange(0, BLOCK_C) chunks = tl.arange(0, PADDED_CHUNKS) values = tl.load( @@ -811,7 +854,8 @@ def _det_gemm_fp32_chunk_merge_kernel( def _triton_keys_are_canonical(keys: torch.Tensor) -> bool: - """Check the current key contents before reusing physical row order.""" + """Check live key order and synchronize the block flags to a Python bool.""" + if keys.shape[0] == 0: return True blocks = triton.cdiv(keys.shape[0], 256) @@ -834,7 +878,12 @@ def _triton_chunked_gemm_fp32( *, tile: tuple[int, int, int] = (_BLOCK_M, _BLOCK_N, 4), ) -> torch.Tensor: - """Reuse the FP32 GEMM in fixed K chunks and merge partials in one grid.""" + """Compute FP32 input or weight gradients with fixed 128-term K chunks. + + Merge the partials in one grid. Split independent output rows when needed + to keep the partial-result workspace within its element budget. + """ + m_size, k_size = a.shape n_size = b.shape[1] chunks = triton.cdiv(k_size, _FP32_CHUNK_K) @@ -898,6 +947,12 @@ def _triton_gemm_fp32( bias: torch.Tensor | None = None, aligned_tiles: bool = False, ) -> torch.Tensor: + """Compute a full-K FP32 GEMM with optional column bias. + + aligned_tiles permits omitting K/N load masks after divisibility checks. + The default keeps the existing masked loads and compiler options. + """ + a, b = a.contiguous(), b.contiguous() M, K = a.shape _, N = b.shape @@ -953,6 +1008,13 @@ def _triton_tree_gemm( inference_schedule: bool = False, output_dtype: torch.dtype = torch.bfloat16, ) -> torch.Tensor: + """Evaluate the canonical midpoint tree with BF16 or CUDA FP32 nodes. + + Reuse the original node storage and split independent output dimensions + when the workspace exceeds its budget. Eligible schedules fuse adjacent + steps while preserving the tree's operand order and node rounding. + """ + if not _TRITON_AVAILABLE: raise RuntimeError("Triton is unavailable") if a.dim() != 2 or b.dim() != 2: @@ -985,7 +1047,7 @@ def _triton_tree_gemm( columns_per_chunk = _MAX_TREE_WORKSPACE_ELEMENTS // plan.host.node_count if columns_per_chunk < 1: raise RuntimeError("FP32 tree plan exceeds the workspace budget for one element") - # A one-row bias VJP cannot be split further over M. Splitting + # A one-row bias gradient cannot be split further over M. Splitting # independent output columns retains the same logical K tree. chunks = [ _triton_tree_gemm( @@ -1051,7 +1113,7 @@ def _triton_tree_gemm( preserve_a_strides=preserve_a_strides, ) if fp32_tree and m_size == 1: - # The H3 bias VJP has one output row. Other shapes use the existing tile. + # The H3 bias gradient has one output row. Other shapes use the existing tile. leaf_config = _TreeLeafConfig(1, 128, 4, True) if ( _ROCM_TUNE_DECODE_LEAF @@ -1264,14 +1326,20 @@ def _triton_tree_gemm( class _TritonDetGemmFn(torch.autograd.Function): + """Wrap the existing deterministic GEMM forward and operand gradients.""" + @staticmethod def forward(ctx, a, b, output_fp32=False): + """Save operands and dispatch the requested output precision.""" + ctx.save_for_backward(a, b) ctx.output_fp32 = bool(output_fp32) return _triton_gemm_fp32(a, b) if output_fp32 else _triton_tree_gemm(a, b) @staticmethod def backward(ctx, grad_out): + """Compute requested operand gradients through the BF16 tree path.""" + a, b = ctx.saved_tensors grad_out = grad_out.contiguous() da = ( @@ -1297,11 +1365,15 @@ class TritonDetGemmOp: """Batch-invariant deterministic GEMM, Triton path.""" def __init__(self): + """Require Triton before exposing the deterministic GEMM backend.""" + if not _TRITON_AVAILABLE: raise RuntimeError("Triton not available for TritonDetGemmOp") logger.info("TritonDetGemmOp ready (deterministic, autotune disabled).") def __call__(self, a, b): + """Apply the autograd-enabled BF16 deterministic GEMM.""" + assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16, "BF16 only" assert a.device.type in ("cuda", "musa") assert b.device.type in ("cuda", "musa") @@ -1309,6 +1381,8 @@ def __call__(self, a, b): return _TritonDetGemmFn.apply(a, b, False) def forward_fp32(self, a, b): + """Apply the existing GEMM wrapper with FP32 forward output.""" + if a.dtype not in (torch.bfloat16, torch.float32) or b.dtype not in ( torch.bfloat16, torch.float32, @@ -1322,6 +1396,8 @@ def forward_fp32(self, a, b): forward_accum_fp32 = forward_fp32 def parameter_vjp_contributions_fp32(self, *, a, b, grad_output): + """Return row-local FP32 contributions to the right operand's gradient.""" + del b rows_a = a.float() rows_g = grad_output.float() @@ -1329,4 +1405,6 @@ def parameter_vjp_contributions_fp32(self, *, a, b, grad_output): def deterministic_gemm_triton(a, b): + """Apply the autograd-enabled BF16 deterministic GEMM entry point.""" + return _TritonDetGemmFn.apply(a, b, False) diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index 87915e599..d7c50e1d2 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""MiniMax-H3 projection through RL-Kernel's fixed FP32 Triton GEMM.""" +"""MiniMax-H3 input projection through existing FP32 Triton GEMM paths. + +The shared operator validates inputs and uses the canonical backward session +to collect and order parameter-gradient rows before their final reduction. +""" import torch @@ -17,40 +21,58 @@ VideoInputProjectionFP32Op, ) -# H20-tuned output blocks and warp count for chunked dX/dW GEMM. +# H20-tuned output tile [M, N] and warp count for chunked dX/dW GEMM. +# Output tile dimensions are independent of the fixed K chunk width. _DX_DW_BLOCK_M = 128 _DX_DW_BLOCK_N = 32 _DX_DW_NUM_WARPS = 4 def projection_gemm_fp32(a, b, *, bias=None): + """Select a reused GEMM path from the fixed H3 input-projection dimensions.""" + + # Forward has K=96; short parameter-gradient reductions also use full-K GEMM. if a.shape[1] <= VIDEO_PATCH_WIDTH: return _triton_gemm_fp32(a, b, bias=bias, aligned_tiles=True) + # Long input and weight gradients both have 96 output columns. if b.shape[1] == VIDEO_PATCH_WIDTH: output = _triton_chunked_gemm_fp32( a, b, tile=(_DX_DW_BLOCK_M, _DX_DW_BLOCK_N, _DX_DW_NUM_WARPS) ) else: + # Bias gradients sum canonical token rows through the existing FP32 tree. output = _triton_tree_gemm(a, b, output_dtype=torch.float32, preserve_a_strides=True) return output if bias is None else output + bias class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): + """Bind the shared H3 input-projection interface to CUDA SM90 Triton paths.""" + family = "triton" matmul_op = staticmethod(projection_gemm_fp32) backend_id = "triton-video-input-projection-fp32-v12" def __init__(self): + """Require Triton before constructing this backend.""" + if not _TRITON_AVAILABLE: raise ImportError("video input projection requires Triton") def _check_backend(self, x): + """Enforce CUDA SM90 eligibility for the input tensor.""" + if x.device.type != "cuda" or torch.version.hip is not None: raise ValueError("Triton video input projection is qualified on CUDA only") if torch.cuda.get_device_capability(x.device) != (9, 0): raise ValueError("video input projection requires CUDA SM90") def backend_info(self): + """Describe the implementation configuration and scheduling policies. + + Scheduling fields describe available paths; input shapes determine + which kernels execute for each call. + """ + import triton return { From d3fb6592bb6df58e9728558c3fdc0f9d96cfb1e4 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Tue, 20 Oct 2026 14:41:05 +0800 Subject: [PATCH 20/26] Document H3 projection interface and canonical linear gradients Signed-off-by: walker-ai <2398833647@qq.com> --- rl_engine/ops/gemm/canonical_linear.py | 12 ++++++++++ .../ops/gemm/video_input_projection_fp32.py | 24 +++++++++++++++++-- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/rl_engine/ops/gemm/canonical_linear.py b/rl_engine/ops/gemm/canonical_linear.py index 751820dfe..eeb2ec329 100644 --- a/rl_engine/ops/gemm/canonical_linear.py +++ b/rl_engine/ops/gemm/canonical_linear.py @@ -23,6 +23,8 @@ def _gemm_fp32( bias: torch.Tensor | None = None, matmul_op: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor: + """Dispatch FP32 GEMM through a custom callable or the selected backend.""" + if matmul_op is not None: return matmul_op(a, b) if bias is None else matmul_op(a, b, bias=bias) if bias is not None: @@ -43,6 +45,8 @@ def _gemm_fp32( class _CanonicalLinearFn(torch.autograd.Function): + """Compute linear gradients with canonical parameter-row collection.""" + @staticmethod def forward( ctx, @@ -55,6 +59,8 @@ def forward( bias_parameter_id, matmul_op, ): + """Register trainable parameters and compute the FP32 linear output.""" + session = active_session() parameter_grad = ctx.needs_input_grad[1] or ctx.needs_input_grad[5] if parameter_grad and session is None: @@ -91,6 +97,8 @@ def forward( @staticmethod def backward(ctx, grad_out): + """Compute input gradients and collect rows for parameter gradients.""" + a, weight = ctx.saved_tensors grad32 = grad_out.contiguous().float() da = None @@ -101,6 +109,8 @@ def backward(ctx, grad_out): if ctx.needs_input_grad[1]: def reducer(rows, grads): + """Compute dW from the session's ordered input and gradient rows.""" + lhs = grads.float().t() if ctx.matmul_op is None: lhs = lhs.contiguous() @@ -113,6 +123,8 @@ def reducer(rows, grads): if ctx.needs_input_grad[5]: def bias_reducer(grads): + """Sum the session's ordered gradient rows through FP32 GEMM.""" + if grads.shape[0] == 0: raise RuntimeError(f"no active logical rows for {ctx.bias_parameter_id}") ones = torch.ones((1, grads.shape[0]), device=grads.device, dtype=torch.float32) diff --git a/rl_engine/ops/gemm/video_input_projection_fp32.py b/rl_engine/ops/gemm/video_input_projection_fp32.py index a19a36d7f..1a432d64f 100644 --- a/rl_engine/ops/gemm/video_input_projection_fp32.py +++ b/rl_engine/ops/gemm/video_input_projection_fp32.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""FP32 MiniMax-H3 video input projection with logical-row parameter VJPs.""" +"""FP32 MiniMax-H3 video input projection with logical-row parameter gradients.""" from __future__ import annotations @@ -17,7 +17,8 @@ def validate_projection_inputs(x, weight, bias): - """Validate the public FP32 tensor boundary before backend dispatch.""" + """Check contiguous FP32 projection tensors and fixed H3 dimensions.""" + tensors = {"x": x, "weight": weight} if bias is not None: tensors["bias"] = bias @@ -39,6 +40,12 @@ def validate_projection_inputs(x, weight, bias): def _validate_keys(keys, rows, device, parameter_id, session, weight, bias): + """Require unique active logical keys across each parameter's registered uses. + + Each input row has a two- or three-column key. A negative first component + marks padding, which is excluded from the uniqueness check. + """ + if not isinstance(keys, torch.Tensor): raise TypeError("parameter gradients require logical_keys") if keys.dtype != torch.int64 or keys.device != device or not keys.is_contiguous(): @@ -57,6 +64,7 @@ def _validate_keys(keys, rows, device, parameter_id, session, weight, bias): continue previous = session.uses.get(f"{parameter_id}.{name}", []) signature = tuple(use.contribution for use in previous) + # Matching weight/bias contribution groups allow one key check per group. if all(token is not None for token in signature): if signature in checked: continue @@ -77,11 +85,21 @@ class VideoInputProjectionFP32Op: matmul_op: Callable[..., torch.Tensor] | None = None def _check_backend(self, x): + """Require subclasses to check device and architecture eligibility.""" + raise NotImplementedError def __call__(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None): + """Project [M, 96] or [B, S, 96] to FP32 hidden width 5376. + + When gradients are enabled, trainable weight or bias requires an active + canonical backward session, logical_keys, and parameter_id. Calls that + need only input gradients do not require parameter-gradient registration. + """ + validate_projection_inputs(x, weight, bias) self._check_backend(x) + # Flatten batch and sequence axes without changing logical row order. rows = x.reshape(-1, VIDEO_PATCH_WIDTH) needs_grad = torch.is_grad_enabled() and any( tensor is not None and tensor.requires_grad for tensor in (x, weight, bias) @@ -96,6 +114,7 @@ def __call__(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None _validate_keys( logical_keys, rows.shape[0], x.device, parameter_id, session, weight, bias ) + # Preserve registered row identities if the caller later edits its keys. logical_keys = logical_keys.detach().clone() output = canonical_linear_fp32( rows, @@ -108,6 +127,7 @@ def __call__(self, x, weight, bias=None, *, logical_keys=None, parameter_id=None matmul_op=self.matmul_op, ) else: + # Calls without gradients bypass the backward session. output = _gemm_fp32( rows, weight.t().contiguous(), From f517a377ae74f096f1d38b8b5b0c4086ac9fc1d4 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Tue, 20 Oct 2026 15:33:41 +0800 Subject: [PATCH 21/26] Strengthen H3 projection validation and stabilize backend identity Signed-off-by: walker-ai <2398833647@qq.com> --- docs/operators/video-input-projection-fp32.md | 2 +- .../gemm/video_input_projection_fp32.py | 2 +- rl_engine/runtime/registry.py | 4 +- .../validation/operators/operator_inputs.py | 6 ++ .../validation/operators/operator_specs.py | 2 + .../gemm/test_video_input_projection_fp32.py | 93 ++++++++++++------- tests/runtime/test_kernel_registry.py | 2 +- 7 files changed, 70 insertions(+), 41 deletions(-) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 0c486c896..5f5f8796f 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -113,7 +113,7 @@ partials remain FP32. A second grid uses a fixed balanced reduction over the next power-of-two chunk count, with zero padding. The merge uses 64 output elements per program and four warps. VJP GEMMs use fixed 128-by-32 output tiles, 32-term dot tiles, and four warps. The forward retains its 64-by-64 output tile. -These settings define the v11 Triton implementation identity. The K partition +These settings define the Triton implementation identity. The K partition and VJP tile do not depend on M, row placement, or occupancy. The VJP kernels consume positive-stride views directly, including `G.T`. The Triton forward omits K or N load predicates only when the corresponding diff --git a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py index d7c50e1d2..3f102c9a6 100644 --- a/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -50,7 +50,7 @@ class TritonVideoInputProjectionFP32Op(VideoInputProjectionFP32Op): family = "triton" matmul_op = staticmethod(projection_gemm_fp32) - backend_id = "triton-video-input-projection-fp32-v12" + backend_id = "triton-video-input-projection-fp32" def __init__(self): """Require Triton before constructing this backend.""" diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 8eb7cf886..7dc362b10 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -241,7 +241,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: ), OperatorBackendDescriptor( semantic_op="video_input_projection_fp32", - backend_id="triton-video-input-projection-fp32-v12", + backend_id="triton-video-input-projection-fp32", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cuda"}), supported_dtypes=frozenset({"float32"}), @@ -281,7 +281,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint=( "rlkernel-triton-ieee-fp32-dot-chunk128-halves64-tile128x32-merge64-" - "strided-aligned-forward-default-fusion-auto-canonical-proof-h3-v11" + "strided-aligned-forward-default-fusion-auto-canonical-proof-h3" ), ), OperatorBackendDescriptor( diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index 2e82ca410..1dfdb8e46 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -90,6 +90,12 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: def _make_video_input_projection_fp32_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: + """Build FP32 H3 projection inputs with fixed widths 96 and 5376. + + Logical keys are unique [batch_index, token_index] pairs in the same order + as the flattened input rows, with every row active. + """ + if dtype != torch.float32: raise TypeError("video input projection validation requires FP32") batch, seq = _batch_seq(args) diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 76f4e0a58..b248d5074 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -453,6 +453,8 @@ def make_candidate(args: argparse.Namespace) -> CandidateSpec: @contextmanager def _canonical_validation_context(): + """Keep the session active for forward/backward and check contributions on exit.""" + from rl_engine.ops.autograd.canonical_backward import canonical_backward_session with canonical_backward_session() as session: diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index e6effd938..a4760f7c0 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -82,14 +82,21 @@ def unexpected_sort(*args, **kwargs): assert actual.is_contiguous() -def test_canonical_normalization_reads_changed_keys(monkeypatch): +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param("cuda", marks=pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90")), + ], +) +def test_canonical_normalization_reads_changed_keys(device, monkeypatch): from rl_engine.ops.autograd.canonical_backward import CanonicalBackwardSession monkeypatch.setattr( "rl_engine.ops.autograd.canonical_backward._CANONICAL_ORDER_REUSE_MIN_ROWS", 0 ) - keys = torch.tensor([[0, 0], [0, 1]], dtype=torch.int64) - values = torch.tensor([[1.0], [2.0]]) + keys = torch.tensor([[0, 0], [0, 1]], dtype=torch.int64, device=device) + values = torch.tensor([[1.0], [2.0]], device=device) session = CanonicalBackwardSession() session.register("weight", keys) _, first = session._normalize("weight", [values]) @@ -116,6 +123,7 @@ def unexpected_proof(keys): assert torch.equal(actual, values) +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") def test_small_canonical_uses_retain_stable_normalization(monkeypatch): from rl_engine.ops.autograd import canonical_backward as canonical @@ -124,8 +132,8 @@ def unexpected_proof(keys): monkeypatch.setattr(canonical, "_keys_are_canonical", unexpected_proof) session = canonical.CanonicalBackwardSession() - keys = torch.tensor([[0, 0], [0, 1]], dtype=torch.int64) - values = torch.tensor([[1.0], [2.0]]) + keys = torch.tensor([[0, 0], [0, 1]], dtype=torch.int64, device="cuda") + values = torch.tensor([[1.0], [2.0]], device="cuda") session.register("weight", keys) order, actual = session._normalize("weight", [values]) assert order.tolist() == [0, 1] @@ -350,16 +358,6 @@ def existing_gemm(a, b, *, output_dtype): reset_backward_runtime() -def test_video_projection_declares_backend_and_matmul_separately(): - from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( - TritonVideoInputProjectionFP32Op, - projection_gemm_fp32, - ) - - assert TritonVideoInputProjectionFP32Op.family == "triton" - assert TritonVideoInputProjectionFP32Op.matmul_op is projection_gemm_fp32 - - def test_input_gradient_without_parameter_session(): x = torch.ones(3, 2, requires_grad=True) weight = torch.ones(4, 2) @@ -388,9 +386,20 @@ def test_frozen_parameter_registration(frozen): @pytest.mark.parametrize( - "invalid", ["dtype", "shape", "empty", "bias", "noncontiguous", "keys", "duplicates", "session"] + "invalid,error_type,message", + [ + ("dtype", TypeError, "x must have dtype torch.float32"), + ("shape", ValueError, "x must have non-empty shape"), + ("empty", ValueError, "x must have non-empty shape"), + ("bias", ValueError, "bias must have shape"), + ("noncontiguous", ValueError, "x must be contiguous"), + ("keys", ValueError, "logical_keys must be contiguous int64"), + ("duplicates", ValueError, "active logical keys must be unique"), + ("duplicates_across_uses", ValueError, "active logical keys must be unique"), + ("session", RuntimeError, "parameter gradients require canonical_backward_session"), + ], ) -def test_invalid_inputs(invalid): +def test_invalid_inputs(invalid, error_type, message): x, weight, bias, _ = _data(2, "cpu") keys = _keys(2, "cpu") if invalid == "dtype": @@ -408,12 +417,19 @@ def test_invalid_inputs(invalid): elif invalid == "duplicates": keys.zero_() weight.requires_grad_(True) - with pytest.raises((TypeError, ValueError, RuntimeError)): - if invalid == "session": - _CPUAdapter()(x, weight, bias, logical_keys=keys, parameter_id="p") - else: - with canonical_backward_session(): - _CPUAdapter()(x, weight, bias, logical_keys=keys, parameter_id="p") + op = _CPUAdapter() + if invalid == "session": + with pytest.raises(error_type, match=message): + op(x, weight, bias, logical_keys=keys, parameter_id="p") + return + with canonical_backward_session(): + if invalid == "duplicates_across_uses": + op(x, weight, bias, logical_keys=keys, parameter_id="p") + # Each call has unique keys, but the second overlaps the first. + keys = keys.clone() + keys[:, 1] += 1 + with pytest.raises(error_type, match=message): + op(x, weight, bias, logical_keys=keys, parameter_id="p") @pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @@ -499,9 +515,17 @@ def test_row_placement_and_unrelated_rows(): @pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") -def test_noncontiguous_upstream_and_forward_graph(): +@pytest.mark.parametrize("expanded_upstream", [False, True]) +def test_noncontiguous_upstream_and_forward_graph(expanded_upstream): op = _op() x, weight, bias, grad = _data(17) + # Cover both transposed random gradients and zero-stride expanded gradients. + grad = ( + torch.ones((), device="cuda").expand_as(grad) + if expanded_upstream + else grad.t().contiguous().t() + ) + assert not grad.is_contiguous() with canonical_backward_session() as session: y = op( x.requires_grad_(), @@ -510,9 +534,15 @@ def test_noncontiguous_upstream_and_forward_graph(): logical_keys=_keys(17, "cuda"), parameter_id="p", ) - y.sum().backward() # Expanded scalar upstream has zero strides. + y.backward(grad) session.validate_complete() - _assert_accuracy(x.grad, torch.ones_like(grad).cpu() @ weight.cpu(), "gradient_accuracy") + cx, cw, cb = [t.detach().cpu().requires_grad_() for t in (x, weight, bias)] + gold = NativeVideoInputProjectionFP32Op()(cx, cw, cb) + gold.backward(grad.cpu()) + for actual, expected in zip( + (x.grad, weight.grad, bias.grad), (cx.grad, cw.grad, cb.grad), strict=True + ): + _assert_accuracy(actual, expected, "gradient_accuracy") with torch.no_grad(): stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) @@ -570,21 +600,12 @@ def test_default_cuda_dispatch(): from rl_engine.runtime.registry import KernelRegistry op = KernelRegistry().get_op("video_input_projection_fp32", device="cuda") - assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32-v12" + assert op.backend_info()["backend_id"] == "triton-video-input-projection-fp32" x, weight, bias, _ = _data(3) with torch.no_grad(): _assert_bytes(op(x, weight, bias), _op()(x, weight, bias)) -@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") -def test_reused_fp32_gemm(): - from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 - - x, weight, _, _ = _data(17) - output = _triton_gemm_fp32(x, weight.t().contiguous()) - _assert_accuracy(output, x.cpu() @ weight.cpu().t(), "forward_accuracy") - - def _fp32_tree_reference(a, b): # Evaluate the existing plan with independent CPU FP32 scalar arithmetic. from rl_engine.backends.shared.triton.gemm.det_gemm import _build_tree_plan diff --git a/tests/runtime/test_kernel_registry.py b/tests/runtime/test_kernel_registry.py index 3af229854..abdff31d0 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -13,7 +13,7 @@ def test_video_input_projection_default_and_strict_scope(): OpBackend.TRITON_VIDEO_INPUT_PROJECTION_FP32 ] desc = registry.semantic.backend_descriptor( - "video_input_projection_fp32", "triton-video-input-projection-fp32-v12" + "video_input_projection_fp32", "triton-video-input-projection-fp32" ) assert desc.supported_devices == frozenset({"cuda"}) assert desc.supported_dtypes == frozenset({"float32"}) From c2e20672fee48e0448f1c4002812e9564acdb498 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Tue, 20 Oct 2026 16:15:58 +0800 Subject: [PATCH 22/26] Make H3 projection reproduction use seeded synthetic inputs Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 4 ++- docs/operators/video-input-projection-fp32.md | 15 ++++++---- .../gemm/test_video_input_projection_fp32.py | 30 ++++++++++--------- 3 files changed, 29 insertions(+), 20 deletions(-) diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index 4e5a9b66a..3916bd802 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -125,6 +125,7 @@ def main(): parser.add_argument("--rows", nargs="+", type=int, default=[1, 128, 1024, 4096, 13104]) parser.add_argument("--samples", type=int, default=20) parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--seed", type=int, default=420) parser.add_argument("--weights", type=Path) parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() @@ -137,7 +138,7 @@ def main(): if args.samples < 1 or args.warmup < 0 or min(args.rows) < 1: raise ValueError("rows/samples must be positive and warmup must be non-negative") torch.backends.cuda.matmul.allow_tf32 = False - torch.manual_seed(420) + torch.manual_seed(args.seed) weight = torch.randn(5376, 96, device="cuda") * 0.0625 bias = torch.randn(5376, device="cuda") * 0.125 fixture_hash = None @@ -172,6 +173,7 @@ def main(): report = { "scope": "single H3 FP32 input projection; eager training and no-grad forward graph", "reference_id": NativeVideoInputProjectionFP32Op.reference_id, + "seed": args.seed, "environment": { "gpu": torch.cuda.get_device_name(), "torch": torch.__version__, diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 5f5f8796f..a4568349d 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -163,9 +163,13 @@ Adding a genuine training sample changes the expected parameter gradient. ```bash PYTHONPATH=. python benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py \ --backends pytorch triton --rows 1 128 1024 4096 13104 \ + --seed 420 --warmup 5 --samples 40 \ --output /tmp/h3-input-projection.json ``` +The benchmark generates inputs, weights, bias, and gradients with seed 420 by +default. The JSON records the seed. Use the same row list and dependency versions +to reproduce the generated data. The benchmark first checks independent CPU FP32 forward/backward accuracy and row-permutation VJPs. It measures eager forward, backward, F+B, and no-grad forward CUDA Graph. Timings include layout copies, bias, and the canonical @@ -191,12 +195,13 @@ rows. It checks CPU FP32 accuracy and byte equality across permutations, microbatches, padding, and repeated calls. The benchmark rejects NaN or Inf in either the candidate or reference, including gradients. -Set `H3_PROJECTION_WEIGHTS` to a safetensors fixture containing `proj_in.weight` -and `proj_in.bias` from `MiniMaxAI/MiniMax-H3` at revision -`42ed227ee7df40d41602854ae760620d6eb651fe` to run the latent-pack-projection check. +The latent-pack-projection test uses a synthetic `[1,24,13,48,84]` latent and +random projection weights with seed 420. It runs without a checkpoint. +To use real weights instead, set `H3_PROJECTION_WEIGHTS` to a safetensors fixture +containing `proj_in.weight` and `proj_in.bias` from `MiniMaxAI/MiniMax-H3` at revision +`42ed227ee7df40d41602854ae760620d6eb651fe`. The fixture metadata must include `repository` and `checkpoint_revision` with -these values. The check uses a synthetic `[1,24,13,48,84]` latent with real input -projection weights; it does not load the complete Transformer. +these values. Keep the weights and generated evidence outside tracked source files. ## Known Limitations diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index a4760f7c0..b2ed6a247 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -558,23 +558,25 @@ def test_noncontiguous_upstream_and_forward_graph(expanded_upstream): @pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") -def test_real_weights_patch_path(): - path = os.environ.get("H3_PROJECTION_WEIGHTS") - if not path: - pytest.skip("set H3_PROJECTION_WEIGHTS to pinned H3 safetensors fixture") - from safetensors import safe_open - from safetensors.torch import load_file - +def test_patch_path(): from rl_engine.reference.packing.video_patch_pack_unpack import NativeVideoPatchPackUnpackOp - with safe_open(path, framework="pt") as fixture: - metadata = fixture.metadata() - assert metadata["repository"] == "MiniMaxAI/MiniMax-H3" - assert metadata["checkpoint_revision"] == "42ed227ee7df40d41602854ae760620d6eb651fe" - tensors = load_file(path) - weight = tensors["proj_in.weight"].cuda().float().contiguous() - bias = tensors["proj_in.bias"].cuda().float().contiguous() torch.manual_seed(420) + path = os.environ.get("H3_PROJECTION_WEIGHTS") + if path: + from safetensors import safe_open + from safetensors.torch import load_file + + with safe_open(path, framework="pt") as fixture: + metadata = fixture.metadata() + assert metadata["repository"] == "MiniMaxAI/MiniMax-H3" + assert metadata["checkpoint_revision"] == "42ed227ee7df40d41602854ae760620d6eb651fe" + tensors = load_file(path) + weight = tensors["proj_in.weight"].cuda().float().contiguous() + bias = tensors["proj_in.bias"].cuda().float().contiguous() + else: + weight = torch.randn(5376, 96, device="cuda") * 0.0625 + bias = torch.randn(5376, device="cuda") * 0.125 latent = torch.randn(1, 24, 13, 48, 84, device="cuda", dtype=torch.bfloat16) x = NativeVideoPatchPackUnpackOp().pack(latent).reshape(-1, 96).float() grad = torch.randn(x.shape[0], 5376, device="cuda") * 0.0625 From 9fd297b68c8134675e504703cf6ec09093f632b5 Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Tue, 20 Oct 2026 18:38:34 +0800 Subject: [PATCH 23/26] Record incremental peak memory for H3 projection benchmarks Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 26 ++++++++++++++++++- docs/operators/video-input-projection-fp32.md | 5 ++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index 3916bd802..e916fa16d 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -114,6 +114,23 @@ def _forward_backward(op, x, weight, bias, grad, keys): _backward_call(op, x, weight, bias, grad, keys)() +def _forward_backward_memory(op, x, weight, bias, grad, keys): + """Measure peak allocated bytes above the warmed, gradient-free baseline.""" + x.grad = weight.grad = bias.grad = None + torch.cuda.synchronize() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + _forward_backward(op, x, weight, bias, grad, keys) + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() + x.grad = weight.grad = bias.grad = None + return { + "baseline_allocated_bytes": baseline, + "peak_allocated_bytes": peak, + "incremental_peak_allocated_bytes": peak - baseline, + } + + def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( @@ -193,6 +210,10 @@ def main(): "weight_fixture_sha256": fixture_hash, "weight_fixture_metadata": fixture_metadata, "timing": "CUDA event interval for complete operator calls, including host launch gaps", + "memory": ( + "Warmed eager F+B peak allocated bytes minus the pre-call baseline, " + "excluding preallocated inputs, persistent caches and allocator reservations" + ), "results": [], } for rows in args.rows: @@ -264,6 +285,7 @@ def main(): args.samples, args.warmup, ) + memory = _forward_backward_memory(op, tx, tw, tb, grad, keys) item = { "backend": name, "rows": rows, @@ -282,12 +304,14 @@ def main(): "forward_graph": graph_time, "backward_eager": backward, "forward_backward_eager": combined, + "forward_backward_memory": memory, } report["results"].append(item) print( f"{name} M={rows}: F={forward['median_us']:.1f} us, " f"B={backward['median_us']:.1f} us, " - f"F+B={combined['median_us']:.1f} us" + f"F+B={combined['median_us']:.1f} us, " + f"peak delta={memory['incremental_peak_allocated_bytes'] / 2**20:.2f} MiB" ) args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index a4568349d..31ae8a632 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -175,6 +175,11 @@ row-permutation VJPs. It measures eager forward, backward, F+B, and no-grad forward CUDA Graph. Timings include layout copies, bias, and the canonical parameter-gradient machinery. Raw samples, environment, source hashes, and actual backend metadata are written to the explicitly selected local output. +After timing, a separate eager F+B call measures incremental peak allocated GPU +memory. The measurement subtracts the warmed pre-call allocation after clearing +parameter and input gradients. It includes outputs, gradients and temporary +tensors. Preallocated inputs, persistent caches and allocator reservations are +excluded. PyTorch is a production performance baseline with separately reported BI results. The Triton VJPs use one partial-GEMM launch and one merge launch when the workspace fits. From 520065f9a6e0e54ec84578d32890155815f3effd Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Wed, 21 Oct 2026 14:42:54 +0800 Subject: [PATCH 24/26] Fix FP32 GEMM condition for Triton 3.0 CI Signed-off-by: walker-ai <2398833647@qq.com> --- docs/operators/video-input-projection-fp32.md | 34 +++++++++++++++++-- .../backends/shared/triton/gemm/det_gemm.py | 2 +- tests/ops/gemm/test_det_gemm.py | 23 ++++++++++++- 3 files changed, 54 insertions(+), 5 deletions(-) diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 31ae8a632..1ba6d21b8 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -161,7 +161,7 @@ Adding a genuine training sample changes the expected parameter gradient. ## Performance Notes ```bash -PYTHONPATH=. python benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py \ +PYTHONPATH=. python3.11 benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py \ --backends pytorch triton --rows 1 128 1024 4096 13104 \ --seed 420 --warmup 5 --samples 40 \ --output /tmp/h3-input-projection.json @@ -188,9 +188,37 @@ claim. The benchmark records the reference identity and its source hash. ## Tests +### CI Environment + +Use the image and installation order in `ci/scripts/run_gpu_ci.sh` for development +and acceptance tests. The current CI image is +`runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04`. +The CI installation selects PyTorch 2.4.1+cu124 and Triton 3.0.0. +Run the following commands inside that image, from the repository root on SM90: + +```bash +export FORCE_CUDA=1 MAX_JOBS=8 KERNEL_ALIGN_FORCE_SM90=1 +export TORCH_CUDA_ARCH_LIST="9.0+PTX" +export PYTHONPATH=. OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 + +python3.11 -m pip install --no-cache-dir torch==2.4.1 \ + --index-url https://download.pytorch.org/whl/cu124 +python3.11 -m pip install --no-build-isolation --no-deps -e . +python3.11 -m pip install --no-cache-dir numpy tabulate accelerate \ + "transformers==5.13.1" pytest triton "flashinfer-python>=0.6.0,<0.7" +python3.11 tools/checks/ci_smoke.py +export RL_KERNEL_REQUIRE_EXT=1 +``` + +Keep the image digest and `python3.11 -m pip freeze` output with the test and +benchmark results. Run the performance command above in the same environment. + +### Operator and Shared GEMM Checks + ```bash -python -m pytest tests/ops/gemm/test_video_input_projection_fp32.py -v -python tools/validation/operators/check_operator.py \ +python3.11 -m pytest tests/ops/gemm/test_video_input_projection_fp32.py -v +python3.11 -m pytest tests/ops/gemm/test_det_gemm.py -v +python3.11 tools/validation/operators/check_operator.py \ --op video_input_projection_fp32 --candidate triton --device cuda --dtype fp32 \ --batch 1 --seq 129 --check-grad --grad-mode random ``` diff --git a/rl_engine/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index e17264c2f..bdc60516f 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -814,7 +814,7 @@ def _det_gemm_fp32_kernel( b = b.to(tl.float32) dot = tl.dot(a, b, allow_tf32=False) # Shorten the FP32 accumulation chain inside each fixed chunk. - if CHUNK_K and k >= tl.cdiv(CHUNK_K, 2 * BLOCK_K): + if CHUNK_K != 0 and k >= tl.cdiv(CHUNK_K, 2 * BLOCK_K): second_half += dot else: acc += dot diff --git a/tests/ops/gemm/test_det_gemm.py b/tests/ops/gemm/test_det_gemm.py index 532491338..317261c8d 100644 --- a/tests/ops/gemm/test_det_gemm.py +++ b/tests/ops/gemm/test_det_gemm.py @@ -18,7 +18,7 @@ try: import triton - from rl_engine.backends.shared.triton.gemm import deterministic_gemm_triton + from rl_engine.backends.shared.triton.gemm import TritonDetGemmOp, deterministic_gemm_triton from rl_engine.backends.shared.triton.gemm.det_gemm import ( _DEFAULT_TREE_LEAF_CONFIG, _GFX942_QWEN_FORWARD_LEAF_CONFIGS, @@ -395,6 +395,27 @@ def test_forward_batch_invariance(name, gemm): assert torch.equal(out1[0], outN[0]), f"{name}: forward batch-invariance broken" +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") +@pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float32)) +def test_triton_forward_accum_fp32_accuracy_and_row_invariance(dtype): + # Exercise the shared full-K entry point with ragged tiles and FP32 output. + torch.manual_seed(420) + a = (torch.randn(31, 65, device=DEV) * 0.0625).to(dtype) + b = (torch.randn(65, 70, device=DEV) * 0.0625).to(dtype) + gemm = TritonDetGemmOp().forward_accum_fp32 + + actual = gemm(a, b) + expected = a.cpu().float() @ b.cpu().float() + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual.cpu(), expected, atol=1e-4, rtol=1e-4) + + permutation = torch.randperm(a.shape[0], device=DEV) + permuted = gemm(a[permutation], b) + chunked = torch.cat([gemm(a[:9], b), gemm(a[9:], b)]) + assert torch.equal(permuted.view(torch.int32), actual[permutation].view(torch.int32)) + assert torch.equal(chunked.view(torch.int32), actual.view(torch.int32)) + + @pytest.mark.parametrize("name,gemm", _BACKENDS) def test_forward_chunked_prefill(name, gemm): # Splitting M then concatenating must match the full GEMM bitwise. From 38193cf39ec4ff0e3780725c17685b2edd1eeb7f Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Wed, 21 Oct 2026 16:07:03 +0800 Subject: [PATCH 25/26] Remove projection weight fixtures and simplify reproduction setup Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 16 ------------ docs/operators/video-input-projection-fp32.md | 25 ++++++++----------- .../gemm/test_video_input_projection_fp32.py | 18 ++----------- 3 files changed, 12 insertions(+), 47 deletions(-) diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index e916fa16d..69d9276b3 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -143,7 +143,6 @@ def main(): parser.add_argument("--samples", type=int, default=20) parser.add_argument("--warmup", type=int, default=5) parser.add_argument("--seed", type=int, default=420) - parser.add_argument("--weights", type=Path) parser.add_argument("--output", type=Path, required=True) args = parser.parse_args() if ( @@ -158,19 +157,6 @@ def main(): torch.manual_seed(args.seed) weight = torch.randn(5376, 96, device="cuda") * 0.0625 bias = torch.randn(5376, device="cuda") * 0.125 - fixture_hash = None - fixture_metadata = None - if args.weights: - from safetensors import safe_open - from safetensors.torch import load_file - - with safe_open(str(args.weights), framework="pt") as fixture: - fixture_metadata = fixture.metadata() - tensors = load_file(str(args.weights)) - weight, bias = [ - tensors[f"proj_in.{name}"].cuda().float().contiguous() for name in ("weight", "bias") - ] - fixture_hash = hashlib.sha256(args.weights.read_bytes()).hexdigest() root = Path(__file__).resolve().parents[3] sources = [ "rl_engine/ops/gemm/canonical_linear.py", @@ -207,8 +193,6 @@ def main(): "git_commit": git_commit, "git_dirty": bool(git_status) if git_status is not None else None, "source_sha256": {p: hashlib.sha256((root / p).read_bytes()).hexdigest() for p in sources}, - "weight_fixture_sha256": fixture_hash, - "weight_fixture_metadata": fixture_metadata, "timing": "CUDA event interval for complete operator calls, including host launch gaps", "memory": ( "Warmed eager F+B peak allocated bytes minus the pre-call baseline, " diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 1ba6d21b8..0bdb267c0 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -161,7 +161,7 @@ Adding a genuine training sample changes the expected parameter gradient. ## Performance Notes ```bash -PYTHONPATH=. python3.11 benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py \ +OMP_NUM_THREADS=4 python3.11 -m benchmarks.operators.gemm.benchmark_video_input_projection_fp32 \ --backends pytorch triton --rows 1 128 1024 4096 13104 \ --seed 420 --warmup 5 --samples 40 \ --output /tmp/h3-input-projection.json @@ -170,6 +170,7 @@ PYTHONPATH=. python3.11 benchmarks/operators/gemm/benchmark_video_input_projecti The benchmark generates inputs, weights, bias, and gradients with seed 420 by default. The JSON records the seed. Use the same row list and dependency versions to reproduce the generated data. +`OMP_NUM_THREADS=4` limits CPU reference computations to four threads. The benchmark first checks independent CPU FP32 forward/backward accuracy and row-permutation VJPs. It measures eager forward, backward, F+B, and no-grad forward CUDA Graph. Timings include layout copies, bias, and the canonical @@ -197,28 +198,27 @@ The CI installation selects PyTorch 2.4.1+cu124 and Triton 3.0.0. Run the following commands inside that image, from the repository root on SM90: ```bash -export FORCE_CUDA=1 MAX_JOBS=8 KERNEL_ALIGN_FORCE_SM90=1 -export TORCH_CUDA_ARCH_LIST="9.0+PTX" -export PYTHONPATH=. OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 - python3.11 -m pip install --no-cache-dir torch==2.4.1 \ --index-url https://download.pytorch.org/whl/cu124 -python3.11 -m pip install --no-build-isolation --no-deps -e . +FORCE_CUDA=1 MAX_JOBS=8 TORCH_CUDA_ARCH_LIST=9.0+PTX \ + python3.11 -m pip install --no-build-isolation --no-deps -e . python3.11 -m pip install --no-cache-dir numpy tabulate accelerate \ "transformers==5.13.1" pytest triton "flashinfer-python>=0.6.0,<0.7" -python3.11 tools/checks/ci_smoke.py -export RL_KERNEL_REQUIRE_EXT=1 +python3.11 -m tools.checks.ci_smoke ``` Keep the image digest and `python3.11 -m pip freeze` output with the test and benchmark results. Run the performance command above in the same environment. +The build variables apply only to extension installation, following the CI build +requirements. Run module commands from the repository root. ### Operator and Shared GEMM Checks ```bash +export OMP_NUM_THREADS=4 python3.11 -m pytest tests/ops/gemm/test_video_input_projection_fp32.py -v python3.11 -m pytest tests/ops/gemm/test_det_gemm.py -v -python3.11 tools/validation/operators/check_operator.py \ +python3.11 -m tools.validation.operators.check_operator \ --op video_input_projection_fp32 --candidate triton --device cuda --dtype fp32 \ --batch 1 --seq 129 --check-grad --grad-mode random ``` @@ -230,12 +230,7 @@ either the candidate or reference, including gradients. The latent-pack-projection test uses a synthetic `[1,24,13,48,84]` latent and random projection weights with seed 420. It runs without a checkpoint. -To use real weights instead, set `H3_PROJECTION_WEIGHTS` to a safetensors fixture -containing `proj_in.weight` and `proj_in.bias` from `MiniMaxAI/MiniMax-H3` at revision -`42ed227ee7df40d41602854ae760620d6eb651fe`. -The fixture metadata must include `repository` and `checkpoint_revision` with -these values. -Keep the weights and generated evidence outside tracked source files. +Keep generated evidence outside tracked source files. ## Known Limitations diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index b2ed6a247..b79612b00 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -4,7 +4,6 @@ """H3 FP32 projection accuracy, actual autograd, and logical-row invariance.""" import json -import os import pytest import torch @@ -562,21 +561,8 @@ def test_patch_path(): from rl_engine.reference.packing.video_patch_pack_unpack import NativeVideoPatchPackUnpackOp torch.manual_seed(420) - path = os.environ.get("H3_PROJECTION_WEIGHTS") - if path: - from safetensors import safe_open - from safetensors.torch import load_file - - with safe_open(path, framework="pt") as fixture: - metadata = fixture.metadata() - assert metadata["repository"] == "MiniMaxAI/MiniMax-H3" - assert metadata["checkpoint_revision"] == "42ed227ee7df40d41602854ae760620d6eb651fe" - tensors = load_file(path) - weight = tensors["proj_in.weight"].cuda().float().contiguous() - bias = tensors["proj_in.bias"].cuda().float().contiguous() - else: - weight = torch.randn(5376, 96, device="cuda") * 0.0625 - bias = torch.randn(5376, device="cuda") * 0.125 + weight = torch.randn(5376, 96, device="cuda") * 0.0625 + bias = torch.randn(5376, device="cuda") * 0.125 latent = torch.randn(1, 24, 13, 48, 84, device="cuda", dtype=torch.bfloat16) x = NativeVideoPatchPackUnpackOp().pack(latent).reshape(-1, 96).float() grad = torch.randn(x.shape[0], 5376, device="cuda") * 0.0625 From 8bd28a9828206dcbca8318d8e1931abee9008dec Mon Sep 17 00:00:00 2001 From: walker-ai <2398833647@qq.com> Date: Wed, 21 Oct 2026 19:23:36 +0800 Subject: [PATCH 26/26] Cover H3 video and conditioned task layouts Signed-off-by: walker-ai <2398833647@qq.com> --- .../benchmark_video_input_projection_fp32.py | 29 +++-- docs/operators/video-input-projection-fp32.md | 78 +++++++++++- .../gemm/test_video_input_projection_fp32.py | 118 +++++++++++++++--- 3 files changed, 199 insertions(+), 26 deletions(-) diff --git a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py index 69d9276b3..839a84607 100644 --- a/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -63,14 +63,25 @@ def _accuracy(actual, gold): policy = judgments["forward_accuracy" if i == 0 else "gradient_accuracy"]["by_op_class"][ "reduction" ]["float32"] - a, b = a.detach().cpu(), b.detach().cpu() - error = (a - b).abs() - limit = policy["atol"] + policy["rtol"] * b.abs() - finite = torch.isfinite(a) & torch.isfinite(b) + max_error, violations, non_finite = 0.0, 0, 0 + finite_errors = True + # Bound comparison temporaries for long layouts; check every element. + for start in range(0, a.shape[0], 4096): + ac = a[start : start + 4096].detach().cpu() + bc = b[start : start + 4096].detach().cpu() + error = (ac - bc).abs() + limit = policy["atol"] + policy["rtol"] * bc.abs() + finite = torch.isfinite(ac) & torch.isfinite(bc) + if torch.isfinite(error).all(): + max_error = max(max_error, error.max().item()) + else: + finite_errors = False + violations += int(((error > limit) | ~finite).sum()) + non_finite += int((~finite).sum()) result[name] = { - "max_abs_error": error.max().item() if torch.isfinite(error).all() else None, - "out_of_tolerance": int(((error > limit) | ~finite).sum()), - "non_finite_elements": int((~finite).sum()), + "max_abs_error": max_error if finite_errors else None, + "out_of_tolerance": violations, + "non_finite_elements": non_finite, **policy, } return result @@ -242,6 +253,7 @@ def main(): } if op is not None and any(mismatches.values()): raise RuntimeError(f"{name} failed row permutation BI: {mismatches}") + del actual, permuted with torch.no_grad(): forward = _sample( partial(_forward, op, x, weight, bias, keys), None, args.samples, args.warmup @@ -259,6 +271,7 @@ def main(): if not torch.equal(output.view(torch.int32), expected.view(torch.int32)): raise RuntimeError("forward graph changed output bytes") graph_time = _sample(graph.replay, None, args.samples, args.warmup) + del graph, output, expected tx, tw, tb = [t.detach().clone().requires_grad_() for t in (x, weight, bias)] backward = _sample( None, partial(_backward_call, op, tx, tw, tb, grad, keys), args.samples, args.warmup @@ -270,6 +283,7 @@ def main(): args.warmup, ) memory = _forward_backward_memory(op, tx, tw, tb, grad, keys) + del tx, tw, tb item = { "backend": name, "rows": rows, @@ -297,6 +311,7 @@ def main(): f"F+B={combined['median_us']:.1f} us, " f"peak delta={memory['incremental_peak_allocated_bytes'] / 2**20:.2f} MiB" ) + del gold, cy, cx, cw, cb args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") diff --git a/docs/operators/video-input-projection-fp32.md b/docs/operators/video-input-projection-fp32.md index 0bdb267c0..3839b5dd0 100644 --- a/docs/operators/video-input-projection-fp32.md +++ b/docs/operators/video-input-projection-fp32.md @@ -162,7 +162,8 @@ Adding a genuine training sample changes the expected parameter gradient. ```bash OMP_NUM_THREADS=4 python3.11 -m benchmarks.operators.gemm.benchmark_video_input_projection_fp32 \ - --backends pytorch triton --rows 1 128 1024 4096 13104 \ + --backends pytorch triton \ + --rows 1 128 1024 4096 13104 37296 72576 102816 108576 364608 \ --seed 420 --warmup 5 --samples 40 \ --output /tmp/h3-input-projection.json ``` @@ -228,8 +229,79 @@ rows. It checks CPU FP32 accuracy and byte equality across permutations, microbatches, padding, and repeated calls. The benchmark rejects NaN or Inf in either the candidate or reference, including gradients. -The latent-pack-projection test uses a synthetic `[1,24,13,48,84]` latent and -random projection weights with seed 420. It runs without a checkpoint. +The chunked-GEMM test also covers `K=32768/32769`, where the partial count +crosses 256, and `K=102816`, the weight-gradient reduction length of the longest +tested landscape video. It uses unscaled standard-normal inputs, compares with +the independent CPU FP32 reference, and requires byte equality for strided +inputs and a reduced workspace budget. + +The latent-pack-projection test uses random latents, weights, and gradients with +seed 420. The video layouts follow the pinned +[Diffusers 0.40.0 frame mapping](https://github.com/huggingface/diffusers/blob/v0.40.0/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py#L97): +`F=17n+5` pixel frames map to `T=5n+2` latent frames. At 24 FPS, 345 frames is +the largest aligned frame count within the 15-second duration limit. + +| Case | Pixel frames | Latent shape `[B,C,T,H,W]` | Projection rows | +| --- | --- | --- | ---: | +| Synthetic temporal length, 768p spatial dimensions | N/A | `[1,24,13,48,84]` | 13,104 | +| Default landscape video layout | 124 | `[1,24,37,48,84]` | 37,296 | +| Longer portrait video layout | 243 | `[1,24,72,84,48]` | 72,576 | +| 15-second-limit landscape video layout | 345 | `[1,24,102,48,84]` | 102,816 | + +Each case checks `Y`, `dX`, `dW`, and `dbias` against the independent CPU FP32 +reference, then checks byte equality across row permutations, microbatches, +padding, and repeated calls. Values are synthetic; only the video geometry is +taken from the released pipeline. These cases run without a checkpoint. +Run the single-video and conditioned-layout cases with: + +```bash +OMP_NUM_THREADS=4 python3.11 -m pytest tests/ops/gemm/test_video_input_projection_fp32.py \ + -v -k 'test_patch_path or test_conditioned_patch_path' +``` + +### Maximum Conditioned Task Layouts + +`test_conditioned_patch_path` covers one request (`B=1`) with the released +H3-Base preprocessing defaults. The +[model card](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/42ed227ee7df40d41602854ae760620d6eb651fe/README.md) +allows two keyframes for FL2VA, or nine images and three video references for +Ref2VA. Reference videos must each last 2-15 seconds and total at most 15 seconds; +all reference types together are limited to 12 files. + +The pinned +[canvas rule](https://github.com/huggingface/diffusers/blob/v0.40.0/src/diffusers/modular_pipelines/minimax_h3/modular_pipeline.py#L42) +caps area before rounding both dimensions to multiples of 32. Across the +supported aspect ratios, its largest rounded canvas is `1856x576` (or its +transpose), reached at aspect `23:7`. This produces `58*18=1044` rows per +latent time step. Reference images use a separate 2048-pixel short edge and +aspect ratios up to 4:1, giving `8192x2048` and `256*64=16384` rows per image. +These image and video rules come from +[Ref2VA preprocessing](https://github.com/huggingface/diffusers/blob/v0.40.0/src/diffusers/modular_pipelines/minimax_h3/before_encoder.py). + +| Component | Pixel frames / size | Latent time steps | Projection rows | +| --- | --- | ---: | ---: | +| Target video, both variants | 345 frames, `1856x576` | 102 | 106,488 | +| FL2VA first and last keyframes | Two images, each `1856x576` | 2 | 2,088 | +| Ref2VA image references | Nine images, each `8192x2048` | 9 | 147,456 | +| Ref2VA video references | 56, 56, 243 frames, each `1856x576` | 17+17+72 | 110,664 | +| **FL2VA total** | Target plus two keyframes | | **108,576** | +| **Ref2VA total** | Target plus nine images and three videos | | **364,608** | + +The three reference videos total 355 frames (14.79 seconds). They maximize +the combined latent length: for three aligned videos, `17*sum(n)+15 <= 360`, +so `sum(n) <= 20` and the total latent length is at most `5*20+6=106`. +The test packs each component and concatenates all visual conditions before +the target video, matching +[upstream assembly](https://github.com/huggingface/diffusers/blob/v0.40.0/src/diffusers/modular_pipelines/minimax_h3/before_denoise.py#L1089). +Audio and text use their own projections. + +Both cases check the complete `Y`, `dX`, `dW`, and `dbias` against the existing +CPU FP32 reference, followed by permutation, microbatch, padding, and repeated-call +BI checks. Accuracy comparisons process all output rows in bounded slices; +reference arithmetic and candidate calls still consume the complete input. +Custom canvas settings, larger batches, and the separate 2K regeneration path +have different layout bounds. + Keep generated evidence outside tracked source files. ## Known Limitations diff --git a/tests/ops/gemm/test_video_input_projection_fp32.py b/tests/ops/gemm/test_video_input_projection_fp32.py index b79612b00..20d2dd820 100644 --- a/tests/ops/gemm/test_video_input_projection_fp32.py +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -193,9 +193,10 @@ def _run(op, data, keys, chunks=None, *, train_weight=True, train_bias=True): def _assert_accuracy(actual, expected, kind): threshold = load_contract()["judgments"][kind]["by_op_class"]["reduction"]["float32"] - torch.testing.assert_close( - actual.cpu(), expected.cpu(), atol=threshold["atol"], rtol=threshold["rtol"] - ) + assert actual.shape == expected.shape + # Bound comparison temporaries for full Ref2VA outputs; compare every element. + for a, b in zip(actual.split(4096), expected.split(4096), strict=True): + torch.testing.assert_close(a.cpu(), b.cpu(), atol=threshold["atol"], rtol=threshold["rtol"]) def _assert_bytes(a, b): @@ -237,6 +238,30 @@ def test_benchmark_accuracy_rejects_non_finite( json.dumps(report, allow_nan=False) +def test_benchmark_accuracy_checks_all_comparison_chunks(): + from benchmarks.operators.gemm.benchmark_video_input_projection_fp32 import _accuracy + + actual = [torch.ones(4097, 2) for _ in range(4)] + gold = [torch.ones_like(t) for t in actual] + actual[0][0, 0] = 2.0 + actual[0][-1, 1] = 4.0 + actual[1][-1, 0] = float("nan") + gold[2][-1, 0] = float("inf") + actual[3][-1, 1] += 1e-5 + report = _accuracy(actual, gold) + for name, a, b in zip(("y", "dx", "dw", "dbias"), actual, gold, strict=True): + error = (a - b).abs() + finite = torch.isfinite(a) & torch.isfinite(b) + policy = report[name] + limit = policy["atol"] + policy["rtol"] * b.abs() + assert policy["out_of_tolerance"] == int(((error > limit) | ~finite).sum()) + assert policy["non_finite_elements"] == int((~finite).sum()) + assert policy["max_abs_error"] == ( + error.max().item() if torch.isfinite(error).all() else None + ) + json.dumps(report, allow_nan=False) + + @pytest.mark.parametrize("with_bias", [False, True]) def test_canonical_linear_bias_and_legacy_call(with_bias): torch.manual_seed(420) @@ -477,10 +502,12 @@ def _check_parameter_bi(op, data, keys): inverse = order.argsort() for i in range(4): _assert_bytes(baseline[i], permuted[i].index_select(0, inverse) if i < 2 else permuted[i]) + del permuted chunks = list(torch.tensor_split(order, 3)) chunked = _run(op, data, keys, chunks) for a, b in zip(baseline, chunked, strict=True): _assert_bytes(a, b) + del chunked padded = [ torch.cat((data[0], torch.randn(7, 96, device="cuda"))), data[1], @@ -491,6 +518,7 @@ def _check_parameter_bi(op, data, keys): result = _run(op, padded, padded_keys) for i, (a, b) in enumerate(zip(baseline, result, strict=True)): _assert_bytes(a, b[:rows] if i < 2 else b) + del result, padded repeated = _run(op, data, keys) for a, b in zip(baseline, repeated, strict=True): _assert_bytes(a, b) @@ -557,30 +585,87 @@ def test_noncontiguous_upstream_and_forward_graph(expanded_upstream): @pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") -def test_patch_path(): +@pytest.mark.parametrize( + "num_frames,latent_height,latent_width,expected_rows", + [ + pytest.param(None, 48, 84, 13104, id="synthetic-t13"), + pytest.param(124, 48, 84, 37296, id="124f-1344x768"), + pytest.param(243, 84, 48, 72576, id="243f-768x1344"), + pytest.param(345, 48, 84, 102816, id="345f-1344x768"), + ], +) +def test_patch_path(num_frames, latent_height, latent_width, expected_rows): from rl_engine.reference.packing.video_patch_pack_unpack import NativeVideoPatchPackUnpackOp + # Diffusers 0.40.0 maps F=17n+5 pixel frames to T=5n+2 latent frames. + # T=13 retains the original synthetic case; the other cases use valid frame counts. + latent_frames = 13 + if num_frames is not None: + assert num_frames % 17 == 5 and num_frames <= 15 * 24 + latent_frames = (num_frames - 5) // 17 * 5 + 2 torch.manual_seed(420) weight = torch.randn(5376, 96, device="cuda") * 0.0625 bias = torch.randn(5376, device="cuda") * 0.125 - latent = torch.randn(1, 24, 13, 48, 84, device="cuda", dtype=torch.bfloat16) + latent = torch.randn( + 1, 24, latent_frames, latent_height, latent_width, device="cuda", dtype=torch.bfloat16 + ) x = NativeVideoPatchPackUnpackOp().pack(latent).reshape(-1, 96).float() + assert x.shape == (expected_rows, 96) grad = torch.randn(x.shape[0], 5376, device="cuda") * 0.0625 - actual = _run(_op(), (x, weight, bias, grad), _keys(x.shape[0], "cuda")) + data, keys = (x, weight, bias, grad), _keys(x.shape[0], "cuda") + _check_projection_accuracy_and_bi(data, keys) + + +@pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") +@pytest.mark.parametrize("variant,expected_rows", [("fl2va", 108576), ("ref2va", 364608)]) +def test_conditioned_patch_path(variant, expected_rows): + from rl_engine.reference.packing.video_patch_pack_unpack import NativeVideoPatchPackUnpackOp + + # Diffusers 0.40.0: the default canvas rule reaches 1856x576 at aspect 23:7. + # Both generated and reference videos use this rule; reference images use a + # 2048-pixel short edge and allow 4:1, giving a maximum image of 8192x2048. + # Latents use spatial /16 and each packed row covers a 2x2 latent patch. + if variant == "fl2va": + condition_shapes = [(1, 36, 116)] * 2 + else: + # Nine images plus three videos reach the 12-file limit. The aligned + # video lengths 56+56+243=355 frames fit the 15-second total at 24 FPS. + reference_frames = (56, 56, 243) + assert all(2 * 24 <= f <= 15 * 24 and f % 17 == 5 for f in reference_frames) + assert sum(reference_frames) <= 15 * 24 + condition_shapes = [(1, 128, 512)] * 9 + [ + ((f - 5) // 17 * 5 + 2, 36, 116) for f in reference_frames + ] + target_shape = ((345 - 5) // 17 * 5 + 2, 36, 116) + torch.manual_seed(420) + pack = NativeVideoPatchPackUnpackOp().pack + # Upstream concatenates all visual conditions before the generated video. + x = torch.cat( + [ + pack(torch.randn(1, 24, *shape, device="cuda", dtype=torch.bfloat16)) + .reshape(-1, 96) + .float() + for shape in condition_shapes + [target_shape] + ] + ) + assert x.shape == (expected_rows, 96) + weight = torch.randn(5376, 96, device="cuda") * 0.0625 + bias = torch.randn(5376, device="cuda") * 0.125 + grad = torch.randn(expected_rows, 5376, device="cuda") * 0.0625 + _check_projection_accuracy_and_bi((x, weight, bias, grad), _keys(expected_rows, "cuda")) + + +def _check_projection_accuracy_and_bi(data, keys): + x, weight, bias, grad = data + actual = _run(_op(), data, keys) cx, cw, cb = [t.detach().cpu().requires_grad_() for t in (x, weight, bias)] gold = NativeVideoInputProjectionFP32Op()(cx, cw, cb) gold.backward(grad.cpu()) for i, (a, b) in enumerate(zip(actual, (gold, cx.grad, cw.grad, cb.grad), strict=True)): _assert_accuracy(a, b, "forward_accuracy" if i == 0 else "gradient_accuracy") - order = torch.arange(x.shape[0] - 1, -1, -1, device="cuda") - chunked = _run( - _op(), - (x, weight, bias, grad), - _keys(x.shape[0], "cuda"), - list(torch.tensor_split(order, 3)), - ) - for a, b in zip(actual, chunked, strict=True): - _assert_bytes(a, b) + # Release the accuracy outputs before the separate permutation/chunk/padding checks. + del actual, gold, cx, cw, cb + _check_parameter_bi(_op(), data, keys) @pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") @@ -676,7 +761,7 @@ def test_cpu_fp32_reference_long_dot_accuracy(): @pytest.mark.skipif(not HAS_SM90, reason="requires CUDA SM90") -@pytest.mark.parametrize("k", [127, 128, 129, 5376, 13104]) +@pytest.mark.parametrize("k", [127, 128, 129, 5376, 13104, 32768, 32769, 102816]) def test_fp32_chunked_gemm_strides_and_workspace(k, monkeypatch): from rl_engine.backends.shared.triton.gemm import det_gemm from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( @@ -684,6 +769,7 @@ def test_fp32_chunked_gemm_strides_and_workspace(k, monkeypatch): ) from rl_engine.reference.gemm.video_input_projection_fp32 import _reference_matmul_fp32 + # K=32769 adds the 257th partial; K=102816 covers the longest landscape layout. torch.manual_seed(420) a = torch.randn(k, 7, device="cuda").t() b = torch.randn(k, 96, device="cuda")