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..839a84607 --- /dev/null +++ b/benchmarks/operators/gemm/benchmark_video_input_projection_fp32.py @@ -0,0 +1,320 @@ +# 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 shutil +import statistics +import subprocess +from functools import partial +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 +from rl_engine.reference.gemm.video_input_projection_fp32 import NativeVideoInputProjectionFP32Op + + +def _candidate(name): + if name == "triton": + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + TritonVideoInputProjectionFP32Op, + ) + + return TritonVideoInputProjectionFP32Op() + return None + + +def _tool_version(command): + try: + return subprocess.check_output( + command, text=True, stderr=subprocess.DEVNULL, 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"] + 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": max_error if finite_errors else None, + "out_of_tolerance": violations, + "non_finite_elements": non_finite, + **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 _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( + "--backends", + nargs="+", + 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) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--seed", type=int, default=420) + 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(args.seed) + weight = torch.randn(5376, 96, device="cuda") * 0.0625 + bias = torch.randn(5376, device="cuda") * 0.125 + 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/reference/gemm/video_input_projection_fp32.py", + ] + # 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"]) + if git_commit is not None + else None + ) + 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__, + "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": 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}, + "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: + 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 = NativeVideoInputProjectionFP32Op()(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}") + del actual, permuted + 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) + 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 + ) + combined = _sample( + partial(_forward_backward, op, tx, tw, tb, grad, keys), + None, + args.samples, + args.warmup, + ) + memory = _forward_backward_memory(op, tx, tw, tb, grad, keys) + del tx, tw, tb + 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, + "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"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") + + +if __name__ == "__main__": + main() diff --git a/docs/.nav.yml b/docs/.nav.yml index ed8409fd9..3a383da15 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..3839b5dd0 --- /dev/null +++ b/docs/operators/video-input-projection-fp32.md @@ -0,0 +1,314 @@ +# 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 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 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 + +```python +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: + 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`, `_det_gemm_fp32_chunk_merge_kernel`, bias tree kernels | Default CUDA path; SM90 qualification | +| PyTorch reference | `NativeVideoInputProjectionFP32Op` | CPU `F.linear`; FP32 product/sum VJPs | Accuracy reference | + +## 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. `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. + +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 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 + +### 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: + +- `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 +complete contribution sequences match. Calls with different contribution sets +normalize each parameter independently. + +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. + +### 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. +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 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. +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. +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. 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. Existing BF16 callers retain their +dtype and schedule. +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 +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. +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. + +### 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. +Adding a genuine training sample changes the expected parameter gradient. + +## Performance Notes + +```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 37296 72576 102816 108576 364608 \ + --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. +`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 +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. +Compare complete operator timings, including sorting, before drawing a speedup +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 +python3.11 -m pip install --no-cache-dir torch==2.4.1 \ + --index-url https://download.pytorch.org/whl/cu124 +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 -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 -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 +``` + +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. + +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 + +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/backends/shared/triton/gemm/det_gemm.py b/rl_engine/backends/shared/triton/gemm/det_gemm.py index 7569f815e..bdc60516f 100644 --- a/rl_engine/backends/shared/triton/gemm/det_gemm.py +++ b/rl_engine/backends/shared/triton/gemm/det_gemm.py @@ -10,8 +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 and do not use BF16 tree nodes. +Autotuning and adaptive split-K are intentionally disabled. Default FP32-output +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 @@ -35,6 +37,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. @@ -48,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 @@ -122,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 @@ -152,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) @@ -168,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 @@ -182,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, ...] @@ -192,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 @@ -216,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}") @@ -226,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 @@ -259,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: @@ -269,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 = [] @@ -278,9 +301,18 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: fused_pairs = [] leaf_reduction = None fused_pairs_after_leaf = [] - if _device_arch(device_index) == "gfx942": + # 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): + """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 = { @@ -371,7 +403,10 @@ 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, ): + """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) @@ -389,10 +424,10 @@ 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. 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): + # 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 a = tl.load( @@ -442,7 +477,10 @@ 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, ): + """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) @@ -459,7 +497,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( @@ -473,11 +511,11 @@ 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) - 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( @@ -491,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) @@ -520,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) @@ -551,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 @@ -574,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), @@ -596,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) @@ -617,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 @@ -644,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) @@ -659,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 @@ -667,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( @@ -676,11 +721,44 @@ 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, + ): + """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) + 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, b_ptr, c_ptr, + bias_ptr, M, N, K, @@ -694,35 +772,187 @@ def _det_gemm_fp32_kernel( BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, 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 = one output tile, walks the whole K in fixed order. - # No split-K -> K-accumulation order independent of M -> batch-invariant. + """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) 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 + 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) & (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) 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 != 0 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, :] 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, + ): + """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( + 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_keys_are_canonical(keys: torch.Tensor) -> bool: + """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) + 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), +) -> torch.Tensor: + """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) + 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, + tile=tile, + ) + 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) + 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, + ) + _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=64, + enable_fp_fusion=False, + num_warps=4, + ) + return output + + +def _triton_gemm_fp32( + a: torch.Tensor, + b: torch.Tensor, + *, + 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. + """ -def _triton_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: a, b = a.contiguous(), b.contiguous() M, K = a.shape _, N = b.shape @@ -732,6 +962,7 @@ def _triton_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: a, b, c, + bias, M, N, K, @@ -745,6 +976,9 @@ 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, + EVEN_K=aligned_tiles and K % _BLOCK_K == 0, + EVEN_N=aligned_tiles and N % _BLOCK_N == 0, ) return c @@ -772,15 +1006,29 @@ def _triton_tree_gemm( out: torch.Tensor | None = None, preserve_a_strides: bool = False, 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: 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 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") @@ -795,6 +1043,24 @@ 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 gradient 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, + ) + 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), @@ -808,20 +1074,21 @@ def _triton_tree_gemm( transpose_output=transpose_output, preserve_a_strides=preserve_a_strides, inference_schedule=inference_schedule, + output_dtype=output_dtype, ) 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(): @@ -831,7 +1098,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() @@ -845,6 +1112,9 @@ def _triton_tree_gemm( transpose_output=transpose_output, preserve_a_strides=preserve_a_strides, ) + if fp32_tree and m_size == 1: + # 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 and is_gfx942 @@ -860,13 +1130,19 @@ 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 - and use_inference_schedule - and k_size >= 1536 + direct_root_output and plan.rocm_leaf_reduction is not None + and ( + (fp32_tree and m_size == 1) + 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 @@ -894,6 +1170,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, "LEAF_LOOP_BOUND": max(plan.host.leaf_lengths)} + if fp32_tree + else {} + ), ) else: _det_gemm_kernel[leaf_grid]( @@ -915,6 +1196,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, "LEAF_LOOP_BOUND": max(plan.host.leaf_lengths)} + if fp32_tree + else {} + ), ) reduction_block = 256 if fuse_leaf_reduction: @@ -944,7 +1230,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 @@ -985,8 +1271,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, @@ -1040,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 = ( @@ -1073,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") @@ -1085,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, @@ -1098,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() @@ -1105,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 new file mode 100644 index 000000000..3f102c9a6 --- /dev/null +++ b/rl_engine/backends/shared/triton/gemm/video_input_projection_fp32.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""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 + +from rl_engine.backends.shared.triton.gemm.det_gemm import ( + _BLOCK_K, + _TRITON_AVAILABLE, + _triton_chunked_gemm_fp32, + _triton_gemm_fp32, + _triton_tree_gemm, +) +from rl_engine.ops.gemm.video_input_projection_fp32 import ( + VIDEO_PATCH_WIDTH, + VideoInputProjectionFP32Op, +) + +# 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" + + 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 { + "backend_id": self.backend_id, + "family": "triton", + "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], + "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, + "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", + "split_k": "fixed_128_term_dx_dw_chunks", + "reduction": "fixed_128_chunks_balanced_fp32_merge", + "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", + "leaf_launch": "paired_leaves_for_one_row_complete_trees", + "bias_tree_tile": [1, 128, 32], + "tree_schedule": "paired_levels_full_workspace", + } diff --git a/rl_engine/ops/autograd/canonical_backward.py b/rl_engine/ops/autograd/canonical_backward.py index 75ddd0269..d5b16f804 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 @@ -13,11 +14,35 @@ 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 + + +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 class _Use: keys: torch.Tensor + contribution: object | None = None rows: torch.Tensor | None = None grads: torch.Tensor | None = None @@ -29,16 +54,78 @@ class CanonicalBackwardSession: uses: dict[str, list[_Use]] = field(default_factory=dict) received: dict[str, int] = field(default_factory=dict) backward_started: bool = False + _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) -> int: + 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]) + reuse_order = ( + keys.shape[0] >= _CANONICAL_ORDER_REUSE_MIN_ROWS + and keys.device.type == "cuda" + and torch.version.hip is None + 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 + def submit_linear( self, parameter_id: str, @@ -60,20 +147,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 +174,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 30175b07a..eeb2ec329 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 @@ -13,7 +15,20 @@ 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, + 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: + return _gemm_fp32(a, b, family) + bias if family == "cuda": return _C.det_gemm_rowwise_fwd_fp32(a.contiguous(), b.contiguous()) if family == "triton": @@ -30,36 +45,108 @@ def _gemm_fp32(a: torch.Tensor, b: torch.Tensor, family: str) -> torch.Tensor: class _CanonicalLinearFn(torch.autograd.Function): + """Compute linear gradients with canonical parameter-row collection.""" + @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, + matmul_op, + ): + """Register trainable parameters and compute the FP32 linear output.""" + 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.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 + ctx.slot = ( + 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) + if ctx.needs_input_grad[5] + else None + ) + return _gemm_fp32( + a.float(), + weight.float().t().contiguous(), + ctx.family, + bias=bias.float() if bias is not None else None, + matmul_op=ctx.matmul_op, + ) @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 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): - return _gemm_fp32(grads.float().t().contiguous(), rows.float(), ctx.family).to( + """Compute dW from the session's ordered input and gradient rows.""" + + lhs = grads.float().t() + if ctx.matmul_op is None: + lhs = lhs.contiguous() + 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, grad_out, reducer) - if ctx.family == "triton": + dweight = ctx.session.submit_linear(ctx.parameter_id, ctx.slot, a, grad32, reducer) + dbias = None + 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) + 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.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=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", @@ -73,18 +160,39 @@ def reducer(rows, grads): impl="ascend_det_gemm_canonical_rowfold", family="ascend", ) - return da, dweight, None, None, None + 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, + matmul_op: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor: - return _CanonicalLinearFn.apply(a, weight, logical_keys, parameter_id, family) + """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. + """ + 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", + matmul_op, + ) __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..1a432d64f --- /dev/null +++ b/rl_engine/ops/gemm/video_input_projection_fp32.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""FP32 MiniMax-H3 video input projection with logical-row parameter gradients.""" + +from __future__ import annotations + +from collections.abc import Callable + +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 + + +def validate_projection_inputs(x, weight, bias): + """Check contiguous FP32 projection tensors and fixed H3 dimensions.""" + + 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): + """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(): + 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. + checked = set() + 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}", []) + 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 + 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) 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") + + +class VideoInputProjectionFP32Op: + """Shared semantic adapter; subclasses select one qualified GEMM backend.""" + + family: str + backend_id: str + 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) + ) + 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 + ) + # Preserve registered row identities if the caller later edits its keys. + 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", + matmul_op=self.matmul_op, + ) + else: + # Calls without gradients bypass the backward session. + 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__ + + +__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..98f8f01c6 --- /dev/null +++ b/rl_engine/reference/gemm/video_input_projection_fp32.py @@ -0,0 +1,56 @@ +# 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 +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. + # 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 9bdc0663a..1f1f3d0df 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -114,6 +114,10 @@ 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" + ) 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). @@ -262,6 +266,51 @@ 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="triton-video-input-projection-fp32", + 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", + "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, + "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", + "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", + }, + 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-halves64-tile128x32-merge64-" + "strided-aligned-forward-default-fusion-auto-canonical-proof-h3" + ), + ), OperatorBackendDescriptor( semantic_op="video_patch_pack_unpack", backend_id="pytorch-video-patch-pack-unpack-v1", @@ -677,6 +726,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, @@ -1167,6 +1217,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 b796c39f8..35835fc8d 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -34,6 +34,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, @@ -73,6 +74,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}", @@ -104,6 +106,34 @@ 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: 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) + 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 e6bc19d48..157d16ba6 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,25 @@ 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" + ), + }, + grad_input_names=("x", "weight", "bias"), + ), "rms_norm": OperatorSpec( name="rms_norm", op_class="reduction", @@ -493,6 +513,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"]) @@ -500,3 +525,14 @@ 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(): + """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: + yield session + session.validate_complete() diff --git a/tests/ops/gemm/test_det_gemm.py b/tests/ops/gemm/test_det_gemm.py index 5b6ad2aed..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, @@ -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 @@ -389,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. 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..20d2dd820 --- /dev/null +++ b/tests/ops/gemm/test_video_input_projection_fp32.py @@ -0,0 +1,974 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 FP32 projection accuracy, actual autograd, and logical-row invariance.""" + +import json + +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 +from rl_engine.reference.gemm.video_input_projection_fp32 import NativeVideoInputProjectionFP32Op + +HAS_SM90 = ( + torch.cuda.is_available() + and torch.version.hip is None + and torch.cuda.get_device_capability() == (9, 0) +) + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param("cuda", marks=pytest.mark.skipif(not HAS_SM90, 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) + 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) + 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() + + +@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, device=device) + values = torch.tensor([[1.0], [2.0]], device=device) + session = CanonicalBackwardSession() + session.register("weight", keys) + _, 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_cpu_callers_retain_stable_normalization(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) + + +@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 + + 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, 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] + assert torch.equal(actual, values) + + +def _op(): + from rl_engine.backends.shared.triton.gemm.video_input_projection_fp32 import ( + TritonVideoInputProjectionFP32Op, + ) + + return TritonVideoInputProjectionFP32Op() + + +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", *, 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, + 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"] + 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): + assert torch.equal( + a.detach().contiguous().view(torch.int32), b.detach().contiguous().view(torch.int32) + ) + + +@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) + + +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) + 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)) + + +@pytest.mark.parametrize("with_bias", [False, True]) +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, + ) + 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(parameter_id) + 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, + ) + 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 == (["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" + 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_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,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, error_type, message): + 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) + 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") +@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): + 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()(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(), 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) + 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 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") + _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( + 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]) + 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], + 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) + del result, padded + repeated = _run(op, data, keys) + for a, b in zip(baseline, repeated, strict=True): + _assert_bytes(a, b) + + +@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")) + 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 HAS_SM90, reason="requires CUDA SM90") +@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_(), + weight.requires_grad_(), + bias.requires_grad_(), + logical_keys=_keys(17, "cuda"), + parameter_id="p", + ) + y.backward(grad) + session.validate_complete() + 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()) + 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 HAS_SM90, reason="requires CUDA SM90") +@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, 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 + 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") + # 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") +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" + x, weight, bias, _ = _data(3) + with torch.no_grad(): + _assert_bytes(op(x, weight, bias), _op()(x, weight, bias)) + + +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 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 + + 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 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 ( + projection_gemm_fp32, + ) + + a, b = torch.randn(7, 129, device="cuda"), torch.randn(129, 5, device="cuda") + 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 = 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(projection_gemm_fp32(a, b), expected) + + +@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): + # Unscaled inputs expose cancellation error in the 5376-term input VJP. + 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]] + 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") + _check_parameter_bi(_op(), data, _keys(rows, "cuda")) + + +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 HAS_SM90, reason="requires CUDA SM90") +@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 ( + projection_gemm_fp32, + ) + 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") + 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(projection_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 + + 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 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 + + torch.manual_seed(422) + # 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") + + def gemm(a, b): + return _triton_tree_gemm( + a, + b, + output_dtype=torch.float32, + preserve_a_strides=True, + ) + + 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 HAS_SM90, reason="requires CUDA SM90") +@pytest.mark.parametrize( + "k,fused", + [ + (1, False), + (32, False), + (64, True), + (65, False), + (128, True), + (129, False), + (256, True), + (1024, True), + (13104, True), + ], +) +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") + 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 + ) + for nodes in node_arrays: + 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.host.node_count * 2) + _assert_bytes(det_gemm._triton_tree_gemm(a, b, output_dtype=torch.float32), expected) + + +@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): + 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 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 + + x, weight, bias, _ = _data(rows) + 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()(x, weight, bias), expected) + + +@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() + 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 b2726e67e..abdff31d0 100644 --- a/tests/runtime/test_kernel_registry.py +++ b/tests/runtime/test_kernel_registry.py @@ -7,6 +7,26 @@ 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" + ) + assert desc.supported_devices == frozenset({"cuda"}) + assert desc.supported_dtypes == frozenset({"float32"}) + 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" + ) + 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", [