From 0522865203d6bfaffd9c6cd084c903f4e5864c8e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 6 Oct 2026 23:22:54 +0000 Subject: [PATCH 01/89] [WS1][CUDA][MiniMax-H3] timestep_sinusoid_h3: FP32 timestep features + shared H3 harness RFC #420 step 3, row `timestep_sinusoid_h3`. - CUDA kernel: one thread per (t, k), the provider's exact FP32 op sequence (__fmul_rn, precise expf/sinf/cosf); bitwise equal to diffusers' get_timestep_embedding on CUDA. Analytic row-local backward. Non-1-D, empty, integer, non-finite or out-of-[0, 1] timesteps fail closed (RFC probe H10). - PyTorch reference (provider replay + FP64 golden), registry, gtest. - Shared H3 harness used by the following rows: * h3_manifest.json + prepare_h3_weights.py: pinned revision and sha256, only shard 1 of 14 is downloaded; * h3_provider.py: diffusers replayed op for op; * h3_chain.py: stage-wise chain replay, each stage declaring its promises (provider-bitwise, golden tolerance); used by scripts/h3_chain_replay.py and tests/h3/test_h3_conditioning_e2e.py; * h3_report.py + scripts/h3_evidence.py + scripts/plot_h3_evidence.py: per-op performance/accuracy JSON and figure; * benchmarks/benchmark_h3_conditioning.py. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- benchmarks/benchmark_h3_conditioning.py | 60 ++++++ csrc/cuda/h3/timestep_sinusoid.cu | 83 +++++++++ csrc/ops.cpp | 14 ++ docs/.nav.yml | 2 + docs/operators/README.md | 1 + docs/operators/h3-timestep-sinusoid.md | 102 ++++++++++ rl_engine/_C.pyi | 3 + rl_engine/kernels/gtest/operator_inputs.py | 29 +++ rl_engine/kernels/gtest/operator_specs.py | 16 ++ rl_engine/kernels/ops/cuda/h3/__init__.py | 2 + .../kernels/ops/cuda/h3/timestep_sinusoid.py | 94 ++++++++++ rl_engine/kernels/ops/pytorch/h3/__init__.py | 30 +++ .../kernels/ops/pytorch/h3/fixed_order.py | 46 +++++ .../ops/pytorch/h3/timestep_sinusoid.py | 114 ++++++++++++ rl_engine/kernels/registry.py | 15 ++ rl_engine/testing/h3_cases.py | 35 ++++ rl_engine/testing/h3_chain.py | 170 +++++++++++++++++ rl_engine/testing/h3_manifest.json | 59 ++++++ rl_engine/testing/h3_provider.py | 33 ++++ rl_engine/testing/h3_report.py | 119 ++++++++++++ rl_engine/testing/h3_weights.py | 69 +++++++ rl_engine/tests/test_dispatch.py | 15 ++ scripts/h3_chain_replay.py | 97 ++++++++++ scripts/h3_evidence.py | 64 +++++++ scripts/plot_h3_evidence.py | 176 ++++++++++++++++++ scripts/prepare_h3_weights.py | 95 ++++++++++ setup.py | 2 + tests/h3/conftest.py | 25 +++ tests/h3/test_h3_conditioning_e2e.py | 55 ++++++ tests/h3/test_h3_timestep_sinusoid.py | 171 +++++++++++++++++ tests/test_ws1_gtest_gpu.py | 1 + 31 files changed, 1797 insertions(+) create mode 100644 benchmarks/benchmark_h3_conditioning.py create mode 100644 csrc/cuda/h3/timestep_sinusoid.cu create mode 100644 docs/operators/h3-timestep-sinusoid.md create mode 100644 rl_engine/kernels/ops/cuda/h3/__init__.py create mode 100644 rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/__init__.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/fixed_order.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py create mode 100644 rl_engine/testing/h3_cases.py create mode 100644 rl_engine/testing/h3_chain.py create mode 100644 rl_engine/testing/h3_manifest.json create mode 100644 rl_engine/testing/h3_provider.py create mode 100644 rl_engine/testing/h3_report.py create mode 100644 rl_engine/testing/h3_weights.py create mode 100644 scripts/h3_chain_replay.py create mode 100644 scripts/h3_evidence.py create mode 100644 scripts/plot_h3_evidence.py create mode 100644 scripts/prepare_h3_weights.py create mode 100644 tests/h3/conftest.py create mode 100644 tests/h3/test_h3_conditioning_e2e.py create mode 100644 tests/h3/test_h3_timestep_sinusoid.py diff --git a/benchmarks/benchmark_h3_conditioning.py b/benchmarks/benchmark_h3_conditioning.py new file mode 100644 index 000000000..de1aaceb7 --- /dev/null +++ b/benchmarks/benchmark_h3_conditioning.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark the MiniMax-H3 conditioning ops (RFC #420) against the provider path. + +CUDA-event medians, bandwidth and peak memory per case, plus the backend +the registry dispatched. Cases live in ``rl_engine/testing/h3_report.py``. + + python benchmarks/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 + python benchmarks/benchmark_h3_conditioning.py --op all --json out.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from rl_engine.kernels.registry import KernelRegistry # noqa: E402 +from rl_engine.testing.h3_chain import environment # noqa: E402 +from rl_engine.testing.h3_report import PERF_CASES, TIMED_KEYS, measure # noqa: E402 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", default="all", choices=["all", *PERF_CASES]) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--json", type=Path, default=None) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + torch.backends.cuda.matmul.allow_tf32 = False + + registry = KernelRegistry() + env = environment() + print(f"device={env['gpu']} torch={env['torch']} cuda={env['cuda']}") + results = [] + for name in list(PERF_CASES) if args.op == "all" else [args.op]: + for case in PERF_CASES[name](registry): + row = measure(case, args.warmup, args.iters) + results.append(row) + parts = [ + f"{key}={row[f'{key}_us']:.2f}us ({row[f'{key}_gbps']:.1f} GB/s, " + f"peak {row[f'{key}_peak_mib']:.2f} MiB)" + for key in TIMED_KEYS + if f"{key}_us" in row + ] + print(f"{name} {row['case']} [{row['backend']}]: " + "; ".join(parts)) + if args.json is not None: + args.json.write_text(json.dumps({"environment": env, "results": results}, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/csrc/cuda/h3/timestep_sinusoid.cu b/csrc/cuda/h3/timestep_sinusoid.cu new file mode 100644 index 000000000..778b78965 --- /dev/null +++ b/csrc/cuda/h3/timestep_sinusoid.cu @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 sinusoidal timestep features (RFC #420 `timestep_sinusoid_h3`). +// +// freq[k] = expf((c * k) * (1 / half)) c = (float)(-ln(max_period)) +// arg[t, k] = t[t] * freq[k] +// out[t] = [cosf(arg) | sinf(arg)] (T, 2 * half) FP32 +// +// Every value is produced by the same FP32 operation sequence as diffusers' +// get_timestep_embedding on CUDA (scalar * arange, multiply by the reciprocal +// of the CPU-scalar divisor, exp, multiply, cos/sin), so the output is meant +// to be bitwise equal to that path. Each element depends only on (t, k): +// batch-, position- and repeat-invariant by construction. Build without +// --use_fast_math: expf/sinf/cosf must stay the precise libdevice functions. + +#include +#include +#include +#include + +#include +#include + +namespace { + +__global__ void h3_timestep_sinusoid_kernel( + const float* __restrict__ timestep, + float* __restrict__ out, + int64_t num_timesteps, + int64_t half, + float neg_log_max_period, + float inv_half) { + const int64_t total = num_timesteps * half; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t row = idx / half; + const int64_t k = idx - row * half; + // __fmul_rn keeps each product a separately rounded FP32 multiply (no FMA + // contraction), matching the two elementwise torch kernels it replays. + const float exponent = __fmul_rn(__fmul_rn(neg_log_max_period, static_cast(k)), inv_half); + const float freq = expf(exponent); + const float arg = __fmul_rn(timestep[row], freq); + float* out_row = out + row * 2 * half; + out_row[k] = cosf(arg); + out_row[half + k] = sinf(arg); + } +} + +} // namespace + +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, + double max_period) { + TORCH_CHECK(timestep.is_cuda(), "timestep must be a CUDA tensor"); + TORCH_CHECK(timestep.scalar_type() == at::kFloat, "timestep must be float32, got ", + timestep.scalar_type()); + TORCH_CHECK(timestep.dim() == 1, "timestep must be 1-D (num_timesteps,)"); + TORCH_CHECK(timestep.numel() > 0, "timestep must hold at least one timestep"); + TORCH_CHECK(num_channels > 0 && num_channels % 2 == 0, + "num_channels must be a positive even number, got ", num_channels); + TORCH_CHECK(max_period > 0.0, "max_period must be positive"); + + const c10::cuda::CUDAGuard device_guard(timestep.device()); + auto t = timestep.contiguous(); + const int64_t num_timesteps = t.size(0); + const int64_t half = num_channels / 2; + auto out = torch::empty({num_timesteps, num_channels}, t.options()); + + // torch evaluates `-math.log(max_period) * arange` with the Python double + // rounded to the FP32 opmath type, and `x / half` as x * (1 / half) in FP32. + const float neg_log_max_period = static_cast(-std::log(max_period)); + const float inv_half = 1.0f / static_cast(half); + + const int threads = 256; + const int64_t total = num_timesteps * half; + const int64_t blocks = std::min((total + threads - 1) / threads, 65535); + auto stream = at::cuda::getCurrentCUDAStream(); + h3_timestep_sinusoid_kernel<<(blocks), threads, 0, stream>>>( + t.data_ptr(), out.data_ptr(), num_timesteps, half, neg_log_max_period, + inv_half); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 11fda1159..71b5b5502 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -497,6 +497,13 @@ at::Tensor prefix_shared_attention( #endif #endif +// MiniMax-H3 (RFC #420) conditioning-path declarations. CUDA only. +#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \ + (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, + double max_period); +#endif + // PyBind11 Module Registration PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "RL-Kernel High-Performance Operator Extension Library"; @@ -752,4 +759,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Deterministic GPT-NeoX token-major RoPE apply for ROCm"); #endif #endif + +#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \ + (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) + m.def("h3_timestep_sinusoid_forward", &h3_timestep_sinusoid_forward, + "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", + py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0); +#endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index 7d71230a3..d4c01f6c5 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -29,6 +29,8 @@ nav: - operators/sampling.md - operators/det-gemm.md - operators/embedding.md + - MiniMax-H3: + - operators/h3-timestep-sinusoid.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 00f4cbb45..dee8796e4 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -32,4 +32,5 @@ Every operator page should include: - [Matmul](matmul.md) - [Sampling](sampling.md) - [Token Embedding](embedding.md) +- [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md new file mode 100644 index 000000000..ce45133c3 --- /dev/null +++ b/docs/operators/h3-timestep-sinusoid.md @@ -0,0 +1,102 @@ +# MiniMax-H3 Timestep Sinusoid + +## Summary + +`timestep_sinusoid_h3` produces the FP32 sinusoidal timestep features that feed the +MiniMax-H3 timestep MLP (RFC #420, WS1 step 3). H3 builds +`Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)` and calls it +on the *distinct* timesteps of the packed sequence, unscaled in `[0, 1]` +(`t = 1 - sigma`): + +```text +freq[k] = exp(-ln(10000) * k / 128) k = 0..127 +arg[t, k] = t * freq[k] +out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256), float32 +``` + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227` (`rl_engine/testing/h3_manifest.json`). +Provider reference: `huggingface/diffusers@f53d552`, `get_timestep_embedding`. + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("timestep_sinusoid_h3", device="cuda") +features = op(timestep) # timestep: (T,) in [0, 1] -> (T, 256) float32 +``` + +## Backends + +| Backend | Wrapper | Native symbol | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.kernels.ops.cuda.h3.timestep_sinusoid.H3TimestepSinusoidCudaOp` | `rl_engine._C.h3_timestep_sinusoid_forward` | Bitwise equal to the provider path | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid.NativeH3TimestepSinusoidOp` | n/a | `forward`: provider replay; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `timestep` | `(T,)`, `T >= 1` | fp32 / bf16 / fp16 | Finite, in `[0, 1]`; any stride; CUDA for the CUDA backend | +| `num_channels` | scalar | int | Positive and even; H3 uses 256 | +| output | `(T, num_channels)` | float32 | `[cos | sin]` order | + +Inputs that break the contract fail closed: non-1-D, empty, integer or out-of-range +timesteps raise. `t * 1000` callers (RFC probe H10) therefore fail instead of +silently producing a different embedding. The range check costs one host read-back; +pass `check_range=False` to `forward` when the caller has already validated `t`. + +## Numerics + +- **Operation order.** The kernel evaluates exactly the FP32 operations the provider + runs on CUDA: `(float)(-ln 10000) * k`, multiplied by the FP32 reciprocal of `128` + (PyTorch divides by a CPU scalar this way), `expf`, `t * freq`, `cosf` and `sinf`. + The multiplies use `__fmul_rn`, so they are never fused into an FMA. The extension + is built without `--use_fast_math`, so `expf`, `sinf` and `cosf` are the precise + libdevice functions. +- **Cast points.** bf16 and fp16 timesteps are upcast to FP32 before anything else, + matching `timesteps[:, None].float()`. The output is never rounded below FP32. +- **Invariance.** Each output element depends only on `(t, k)`. Results are + bitwise identical across batch size, position, permutation and repeated runs. +- **Backward.** `d/dt` is the analytic row-local VJP + (`-f * sin(t f)` for the cos half, `f * cos(t f)` for the sin half). It is summed over + the 256 channels with a fixed pairwise tree (`tree_sum_lastdim_fp32`), so it does not + depend on how many timesteps share the launch. + +Measured on a B200 (torch 2.13.0+cu130): + +| Comparison | Result | +| --- | --- | +| CUDA vs provider path, T in {1, 2, 3, 7, 64, 1000, 4097} | bitwise equal | +| CUDA vs FP64 golden | max abs 1.19e-7 (contract: FP32 elementwise atol 1e-5) | +| gradient vs FP64 golden (gtest, T = 7) | max abs 2.4e-7 | + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 +``` + +On a B200 this is a single launch of about 15 µs, independent of `T` for `T <= 64`. With +the range check it takes about 51 µs. The provider path takes about 64 µs, because it is +six eager kernels plus two concatenations. Either way the op is launch-bound: it moves +about 1 KB per timestep. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_timestep_sinusoid.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end, pinned weights +python scripts/check_operator.py --op timestep_sinusoid_h3 --candidate cuda --device cuda \ + --dtype fp32 --batch 7 --check-grad +python scripts/h3_evidence.py --op timestep_sinusoid_h3 \ + --out docs/usage/evidence/h3-timestep-sinusoid-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-timestep-sinusoid-b200/report.json +``` + +## Known Limitations + +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- `max_period` is fixed at 10000, the only value H3 uses. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fb3c4379c..970df8515 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -325,3 +325,6 @@ def deterministic_collective_rocm_ipc_all_gather_input( input: torch.Tensor, output: torch.Tensor, ) -> None: ... +def h3_timestep_sinusoid_forward( + timestep: torch.Tensor, num_channels: int = 256, max_period: float = 10000.0 +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index ca3b7c120..4562f2aad 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -16,6 +16,7 @@ DEFAULT_VOCAB = 151936 DEFAULT_ROPE_THETA = 1.0e6 DEFAULT_RMS_EPS = 1.0e-6 +H3_FREQ_DIM = 256 def make_operator_inputs( @@ -42,6 +43,7 @@ def make_operator_inputs( "embedding": _make_embedding_inputs, "lm_head": _make_lm_head_inputs, "kv_cache_attention": _make_kv_cache_attention_inputs, + "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, } try: return builders[op_name](args, dtype, device) @@ -72,6 +74,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", + "timestep_sinusoid_h3": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}", } try: return names[op_name] @@ -331,6 +334,32 @@ def _make_kv_cache_attention_inputs( } +def _h3_num_timesteps(args: argparse.Namespace) -> int: + # H3 packs a handful of distinct timesteps; reuse --batch as their count. + return _arg_int(args, "num_timesteps", _arg_int(args, "batch", 2)) + + +def _h3_timesteps( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + num = _h3_num_timesteps(args) + mode = _arg_str(args, "input_mode", "random") + if mode == "constant": + t = torch.full((num,), 0.5, device=device) + else: + t = torch.rand((num,), generator=_generator(args, device, offset=7), device=device) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(dtype) + + +def _make_timestep_sinusoid_h3_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + return {"timestep": _h3_timesteps(args, dtype, device)} + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 4b86c2b7e..cc740d0ff 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -245,6 +245,22 @@ def _load_object(path: str) -> Any: }, grad_input_names=("logits",), ), + # MiniMax-H3 (RFC #420) conditioning path. The op is FP32 end to end, so + # its FP32 rows are the declared contract; other dtypes only change the + # timestep input's storage dtype. + "timestep_sinusoid_h3": OperatorSpec( + name="timestep_sinusoid_h3", + op_class="elementwise", + gold_path="rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid.NativeH3TimestepSinusoidOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ), + "cuda": "rl_engine.kernels.ops.cuda.h3.timestep_sinusoid.H3TimestepSinusoidCudaOp", + }, + grad_input_names=("timestep",), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/__init__.py b/rl_engine/kernels/ops/cuda/h3/__init__.py new file mode 100644 index 000000000..86cf4c9da --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py new file mode 100644 index 000000000..a779941b2 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 sinusoidal timestep features (RFC #420 ``timestep_sinusoid_h3``).""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.pytorch.h3 import H3_FREQ_DIM, H3_MAX_PERIOD +from rl_engine.kernels.ops.pytorch.h3.fixed_order import tree_sum_lastdim_fp32 +from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import ( + NativeH3TimestepSinusoidOp, + validate_h3_timesteps, +) + +KERNEL_ID = "rl_engine._C.h3_timestep_sinusoid_forward" +BACKWARD_IMPL = "row_local_fp32_analytic_tree_sum" + + +def h3_sinusoid_cuda_available() -> bool: + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) + + +class _H3TimestepSinusoidCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, timestep: torch.Tensor, num_channels: int): + t32 = timestep.detach().float().contiguous() + out = _C.h3_timestep_sinusoid_forward(t32, int(num_channels), float(H3_MAX_PERIOD)) + ctx.save_for_backward(out) + ctx.num_channels = int(num_channels) + ctx.timestep_dtype = timestep.dtype + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + (out,) = ctx.saved_tensors + half = ctx.num_channels // 2 + freq = NativeH3TimestepSinusoidOp.frequencies_fp32(ctx.num_channels, out.device) + cos_part, sin_part = out[:, :half], out[:, half:] + g = grad_out.float() + # d cos(t f)/dt = -f sin(t f); d sin(t f)/dt = f cos(t f). Row-local, + # summed over channels with a fixed pairwise tree. + contrib = torch.cat( + [-(g[:, :half] * sin_part) * freq, (g[:, half:] * cos_part) * freq], dim=-1 + ) + grad_t = tree_sum_lastdim_fp32(contrib).to(ctx.timestep_dtype) + record_backward( + "timestep_sinusoid_h3", + kernel_id="rl_engine.kernels.ops.pytorch.h3.fixed_order.tree_sum_lastdim_fp32", + impl=BACKWARD_IMPL, + family="cuda", + ) + return grad_t, None + + +class H3TimestepSinusoidCudaOp: + """CUDA candidate: one thread per (timestep, frequency), FP32 output. + + Construction fails when the extension lacks the symbol, so the registry + falls back to the PyTorch reference instead of silently mis-dispatching. + """ + + op_class = "elementwise" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_sinusoid_cuda_available(): + raise RuntimeError( + "rl_engine._C.h3_timestep_sinusoid_forward is unavailable; rebuild the " + "CUDA extension with csrc/cuda/h3/timestep_sinusoid.cu" + ) + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + if not timestep.is_cuda: + raise ValueError("H3TimestepSinusoidCudaOp needs a CUDA timestep tensor") + return _H3TimestepSinusoidCuda.apply(timestep, num_channels) + + def forward_fp32(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + # The op is FP32 end to end; the FP32 path is the op itself. + return self.forward(timestep, num_channels=num_channels) diff --git a/rl_engine/kernels/ops/pytorch/h3/__init__.py b/rl_engine/kernels/ops/pytorch/h3/__init__.py new file mode 100644 index 000000000..56a3e01ef --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 (RFC #420) PyTorch goldens. + +Shape constants are pinned to ``MiniMaxAI/MiniMax-H3@42ed227`` (see +``rl_engine/testing/h3_manifest.json``). The ops accept other sizes so that +tiny synthetic shapes can be tested, but the H3 layout rules (three modality +rows per timestep, six modulation chunks) are fixed. +""" + +from __future__ import annotations + +H3_FREQ_DIM = 256 +H3_MAX_PERIOD = 10000 +H3_TIME_EMBED_HIDDEN_DIM = 5376 +H3_TIME_EMBED_DIM = 2688 +H3_HIDDEN_SIZE = 5376 +# Modality tags of the packed sequence: 0 video, 1 text, 2 audio. +H3_MODALITY_NUM = 3 +# shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp +H3_ADALN_CHUNKS = 6 +H3_ADALN_CHUNK_NAMES = ( + "shift_msa", + "scale_msa", + "gate_msa", + "shift_mlp", + "scale_mlp", + "gate_mlp", +) diff --git a/rl_engine/kernels/ops/pytorch/h3/fixed_order.py b/rl_engine/kernels/ops/pytorch/h3/fixed_order.py new file mode 100644 index 000000000..e95a66054 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/fixed_order.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shape-independent FP32 reductions shared by the H3 goldens and backwards. + +``torch.sum`` picks its reduction tree from the tensor shape, so a row's +result can change with the number of rows next to it. These helpers spell +the order out with elementwise ops only, which keeps every row's bytes a +function of that row alone. +""" + +from __future__ import annotations + +import torch + + +def tree_sum_lastdim_fp32(x: torch.Tensor) -> torch.Tensor: + """Pairwise tree over the last dim: pad to a power of two, halve until 1. + + Level ``l`` adds element ``i`` and ``i + width/2`` for every ``i``. The + tree depends only on the last-dim size. + """ + + x = x.float() + width = x.shape[-1] + if width == 0: + return x.new_zeros(x.shape[:-1]) + padded = 1 << (width - 1).bit_length() + if padded != width: + x = torch.nn.functional.pad(x, (0, padded - width)) + while x.shape[-1] > 1: + half = x.shape[-1] // 2 + x = x[..., :half] + x[..., half:] + return x[..., 0] + + +def fold_rows_fp32(rows: torch.Tensor) -> torch.Tensor: + """Left fold over dim 0 in ascending row order: ((r0 + r1) + r2) + ...""" + + rows = rows.float() + if rows.shape[0] == 0: + return rows.new_zeros(rows.shape[1:]) + acc = rows[0].clone() + for index in range(1, rows.shape[0]): + acc = acc + rows[index] + return acc diff --git a/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py b/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py new file mode 100644 index 000000000..8ca0c0ccb --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 sinusoidal timestep features (RFC #420 row ``timestep_sinusoid_h3``). + +H3 builds ``Timesteps(num_channels=256, flip_sin_to_cos=True, +downscale_freq_shift=0)`` and feeds it the *distinct* timestep values of the +packed sequence, unscaled in ``[0, 1]`` (``t = 1 - sigma``):: + + freq[k] = exp(-ln(10000) * k / 128) k = 0..127, FP32 + arg[t, k] = t * freq[k] FP32 + out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256) FP32 + +The output is always FP32: ``time_embedder`` is a ``_keep_in_fp32_modules`` +module in the checkpoint, so the features never pass through BF16. +""" + +from __future__ import annotations + +import math + +import torch + +from rl_engine.kernels.ops.pytorch.h3 import H3_FREQ_DIM, H3_MAX_PERIOD + +_TIMESTEP_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_timesteps( + timestep: torch.Tensor, num_channels: int, *, check_range: bool = True +) -> None: + """Fail closed on inputs the H3 convention does not define. + + ``check_range`` rejects values outside ``[0, 1]``: H3 consumes + ``t = 1 - sigma`` unscaled, so a ``t * 1000`` caller (RFC probe H10) is a + convention error, not a different valid input. It reads the tensor back + to the host, which is cheap for the handful of distinct timesteps. + """ + + if not isinstance(timestep, torch.Tensor): + raise TypeError("timestep must be a torch.Tensor") + if timestep.dim() != 1: + raise ValueError(f"timestep must be 1-D (num_timesteps,), got {tuple(timestep.shape)}") + if timestep.numel() == 0: + raise ValueError("timestep must hold at least one timestep") + if timestep.dtype not in _TIMESTEP_DTYPES: + raise TypeError(f"timestep must be fp32, bf16 or fp16, got {timestep.dtype}") + if num_channels <= 0 or num_channels % 2 != 0: + raise ValueError(f"num_channels must be a positive even number, got {num_channels}") + if check_range: + t32 = timestep.detach().float() + # One host sync: NaN fails both comparisons, so this also rejects it. + if not bool(((t32 >= 0) & (t32 <= 1)).all()): + raise ValueError( + "timestep must be finite and lie in [0, 1]: H3 consumes t = 1 - sigma " + f"unscaled (got {t32.cpu().tolist()[:8]})" + ) + + +class NativeH3TimestepSinusoidOp: + """PyTorch reference for the H3 sinusoidal timestep features. + + ``forward`` replays diffusers' ``get_timestep_embedding`` op for op on the + input's device (the provider path). ``forward_fp32`` is the independent + golden: the same formula evaluated in FP64 and rounded once to FP32. + """ + + op_class = "elementwise" + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + # Same op sequence and dtypes as diffusers get_timestep_embedding with + # flip_sin_to_cos=True, downscale_freq_shift=0, scale=1. + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / half + freq = torch.exp(exponent) + arg = timestep[:, None].float() * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1) + + def forward_fp32( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + k = torch.arange(half, dtype=torch.float64, device=timestep.device) + freq = torch.exp(-math.log(H3_MAX_PERIOD) * k / half) + arg = timestep.double()[:, None] * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1).float() + + @staticmethod + def frequencies_fp32(num_channels: int, device: torch.device | str) -> torch.Tensor: + """The FP32 frequency table the provider path multiplies by.""" + + half = num_channels // 2 + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=device + ) + return torch.exp(exponent / half) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 9eeb1c42b..533536f57 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -200,6 +200,14 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp" CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" + # MiniMax-H3 (RFC #420) conditioning path + PYTORCH_H3_TIMESTEP_SINUSOID = ( + "rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ) + CUDA_H3_TIMESTEP_SINUSOID = ( + "rl_engine.kernels.ops.cuda.h3.timestep_sinusoid.H3TimestepSinusoidCudaOp" + ) + def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: return ( @@ -609,6 +617,10 @@ def __init__(self): OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE, ], + "timestep_sinusoid_h3": [ + OpBackend.CUDA_H3_TIMESTEP_SINUSOID, + OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID, + ], }, "rocm": { "logp": [ @@ -649,6 +661,7 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -693,6 +706,7 @@ def __init__(self): ], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -718,6 +732,7 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_cases.py b/rl_engine/testing/h3_cases.py new file mode 100644 index 000000000..527a820ec --- /dev/null +++ b/rl_engine/testing/h3_cases.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic MiniMax-H3 (RFC #420) test inputs shared by tests and scripts.""" + +from __future__ import annotations + +import torch + + +def h3_timesteps(num: int, *, seed: int = 0, device: str = "cuda") -> torch.Tensor: + """Distinct-looking timesteps in [0, 1] with both endpoints present.""" + + generator = torch.Generator(device="cpu").manual_seed(seed) + t = torch.rand(num, generator=generator) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(device) + + +def h3_packed_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """(timestep_indices, token_tags) for a packed sequence with every modality. + + Rows are video, text and audio blocks in random order, each block mapped + to a random timestep, like H3's conditioning/target packing. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + token_tags = torch.randint(0, 3, (seq_len,), generator=generator) + timestep_indices = torch.randint(0, num_timesteps, (seq_len,), generator=generator) + token_tags[: min(3, seq_len)] = torch.arange(min(3, seq_len)) + return timestep_indices.to(device), token_tags.to(device) diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py new file mode 100644 index 000000000..fb1f10943 --- /dev/null +++ b/rl_engine/testing/h3_chain.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Stage-by-stage replay of the MiniMax-H3 conditioning chain (RFC #420). + +Each stage runs three ways on the same device: + +* provider: the diffusers op sequence (``rl_engine.testing.h3_provider``); +* candidate: the dispatched RL-Kernel op, fed the previous candidate output + (chained) and, separately, the previous provider output (isolated); +* golden: the FP32/FP64 reference, fed the previous golden output. + +A stage also declares what it promises, so the end-to-end test and the +evidence JSON check the same thing: whether its isolated output is bitwise +equal to the provider, and its tolerance against the golden (rows of +``tolerance_contract.json``). ``first_drift`` is the first stage whose chained +output differs from the provider; ``first_isolated_drift`` the first stage +that differs even on provider inputs. +""" + +from __future__ import annotations + +import platform +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable + +import torch + +from rl_engine.kernels.registry import KernelRegistry +from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps +from rl_engine.testing.h3_provider import provider_time_proj + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +@dataclass(frozen=True) +class Stage: + name: str + op_type: str + # (op, ctx, upstream) -> outputs; ``upstream`` is the previous stage's output. + candidate: Callable[[Any, dict, Any], Any] + provider: Callable[[dict, Any], Any] + golden: Callable[[Any, dict, Any], Any] + # Promises checked by tests/h3/test_h3_conditioning_e2e.py. + provider_bitwise_isolated: bool + golden_atol: float + golden_rtol: float + + +STAGES: list[Stage] = [ + Stage( + name="timestep_sinusoid_h3", + op_type="timestep_sinusoid_h3", + candidate=lambda op, ctx, _up: op(ctx["timestep"]), + provider=lambda ctx, _up: provider_time_proj(ctx["timestep"]), + golden=lambda op, ctx, _up: op.forward_fp32(ctx["timestep"]), + provider_bitwise_isolated=True, + golden_atol=1e-5, # elementwise / float32 + golden_rtol=1e-5, + ), +] + + +def _flat(value: Any) -> list[torch.Tensor]: + if isinstance(value, torch.Tensor): + return [value] + return [tensor for item in value for tensor in _flat(item)] + + +def compare(lhs: Any, rhs: Any, atol: float = 0.0, rtol: float = 0.0) -> dict[str, Any]: + lhs_t, rhs_t = _flat(lhs), _flat(rhs) + bitwise = all(torch.equal(a, b) for a, b in zip(lhs_t, rhs_t, strict=True)) + max_abs = 0.0 + within = True + for a, b in zip(lhs_t, rhs_t, strict=True): + diff = (a.float() - b.float()).abs() + if diff.numel(): + max_abs = max(max_abs, float(diff.max())) + within = within and bool((diff <= atol + rtol * b.float().abs()).all()) + return {"bitwise_equal": bitwise, "max_abs": max_abs, "within_tolerance": within} + + +def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dict[str, Any]: + timestep_indices, token_tags = h3_packed_layout(seq_len, num_timesteps, seed=seed) + return { + "timestep": h3_timesteps(num_timesteps, seed=seed), + "timestep_indices": timestep_indices, + "token_tags": token_tags, + "weights": weights, + } + + +def golden_op(registry: KernelRegistry, op_type: str): + """The PyTorch reference: the last entry of the CPU priority list.""" + + return registry._get_or_create_backend(registry._priority_map["cpu"][op_type][-1]) + + +def run_case( + registry: KernelRegistry, + weights, + *, + num_timesteps: int, + seq_len: int, + seed: int = 0, + stages: list[Stage] | None = None, +) -> dict[str, Any]: + stages = STAGES if stages is None else stages + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + report: dict[str, Any] = { + "num_timesteps": num_timesteps, + "seq_len": seq_len, + "seed": seed, + "stages": [], + } + chained = provider = golden = None + first_drift = first_isolated = None + for stage in stages: + op = registry.get_op(stage.op_type, device="cuda") + gold = golden_op(registry, stage.op_type) + with torch.no_grad(): + isolated = stage.candidate(op, ctx, provider) + chained_input = chained + chained = stage.candidate(op, ctx, chained_input) + repeat = stage.candidate(op, ctx, chained_input) + golden = stage.golden(gold, ctx, golden) + provider = stage.provider(ctx, provider) + entry = { + "stage": stage.name, + "backend": type(op).__name__, + "kernel_id": getattr(op, "kernel_id", type(op).__name__), + "repeat_bitwise_equal": compare(repeat, chained)["bitwise_equal"], + "chained_vs_provider": compare(chained, provider), + "isolated_vs_provider": compare(isolated, provider), + "chained_vs_golden": compare(chained, golden, stage.golden_atol, stage.golden_rtol), + "provider_vs_golden": compare(provider, golden, stage.golden_atol, stage.golden_rtol), + } + if first_drift is None and not entry["chained_vs_provider"]["bitwise_equal"]: + first_drift = stage.name + if first_isolated is None and not entry["isolated_vs_provider"]["bitwise_equal"]: + first_isolated = stage.name + report["stages"].append(entry) + report["first_drift"] = first_drift + report["first_isolated_drift"] = first_isolated + return report + + +def git_state() -> dict[str, Any]: + def run(*args: str) -> str: + return subprocess.check_output(["git", *args], cwd=REPO_ROOT, text=True).strip() + + try: + sha = run("rev-parse", "HEAD") + dirty = bool(run("status", "--porcelain", "--untracked-files=no")) + except Exception: # noqa: BLE001 + sha, dirty = "unknown", True + return {"rl_kernel_commit": sha, "tracked_tree_dirty": dirty} + + +def environment() -> dict[str, Any]: + return { + "gpu": torch.cuda.get_device_name(), + "capability": list(torch.cuda.get_device_capability()), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "python": platform.python_version(), + "tf32": bool(torch.backends.cuda.matmul.allow_tf32), + } diff --git a/rl_engine/testing/h3_manifest.json b/rl_engine/testing/h3_manifest.json new file mode 100644 index 000000000..52b1989b1 --- /dev/null +++ b/rl_engine/testing/h3_manifest.json @@ -0,0 +1,59 @@ +{ + "version": "h3-conditioning-v1", + "rfc": "RL-Align/RL-Kernel#420", + "model_identity": { + "hf_repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "subfolder": "transformer", + "architecture": "MiniMaxH3Transformer3DModel", + "config_sha256": "74c11bff524336576096993cbfcdcdc2ef4fa2fa4409df693bdcbc6c666282ae", + "index_file": "diffusion_pytorch_model.safetensors.index.json", + "index_sha256": "ac30a3b58963f2e735d493475fbb81853a5735ec947619648b3e045acda6783e", + "config_fingerprint": { + "hidden_size": 5376, + "num_attention_heads": 56, + "attention_head_dim": 128, + "num_layers": 50, + "num_refiner_layers": 2, + "ffn_dim": 14336, + "freq_dim": 256, + "time_embed_hidden_dim": 5376, + "time_embed_dim": 2688, + "norm_eps": 1e-05, + "final_norm_eps": 1e-05 + } + }, + "weight_shards": { + "diffusion_pytorch_model-00001-of-00014.safetensors": { + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028", + "size_bytes": 4825958704 + } + }, + "tensors": { + "time_embedder.linear_1.weight": {"dtype": "float32", "shape": [5376, 256]}, + "time_embedder.linear_1.bias": {"dtype": "float32", "shape": [5376]}, + "time_embedder.linear_2.weight": {"dtype": "float32", "shape": [2688, 5376]}, + "time_embedder.linear_2.bias": {"dtype": "float32", "shape": [2688]}, + "transformer_blocks.0.adaln_proj.linear.weight": {"dtype": "bfloat16", "shape": [96768, 2688]}, + "transformer_blocks.0.adaln_proj.linear.bias": {"dtype": "bfloat16", "shape": [96768]} + }, + "reference_implementation": { + "repo": "huggingface/diffusers", + "commit": "f53d552036a0d1bd5570782a39cd40cfabf112bc", + "file": "src/diffusers/models/transformers/transformer_minimax_h3.py", + "file_sha256": "435873d9bc26b8770342961e18cb20324aeeaa4d045e85e0941638d4857d7db3", + "timestep_embedding_file": "src/diffusers/models/embeddings.py" + }, + "conventions": { + "timestep": "distinct timesteps t = 1 - sigma, unscaled in [0, 1], shape (T,)", + "modality_tags": {"video": 0, "text": 1, "audio": 2}, + "adaln_row": "timestep_index * 3 + token_tag", + "adaln_table_rows": "[t0_mod0, t0_mod1, t0_mod2, t1_mod0, ...]", + "adaln_chunk_order": ["shift_msa", "scale_msa", "gate_msa", "shift_mlp", "scale_mlp", "gate_mlp"], + "precision": { + "timestep_sinusoid": "float32", + "timestep_mlp": "float32 weights, float32 activations", + "adaln_projection": "SiLU in float32, cast to bfloat16, bfloat16 weight/bias/output" + } + } +} diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py new file mode 100644 index 000000000..ff4718c00 --- /dev/null +++ b/rl_engine/testing/h3_provider.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The diffusers MiniMax-H3 conditioning path, op for op, without diffusers. + +Each function replays the tensor ops of the pinned reference +(``huggingface/diffusers@f53d552``, see ``h3_manifest.json``) in the same +order and dtypes, so it runs the same kernels the provider does on a given +device. This is the *provider* side of the RFC #420 comparisons; the FP32/FP64 +goldens live in ``rl_engine.kernels.ops.pytorch.h3``. +""" + +from __future__ import annotations + +import math + +import torch + + +def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch.Tensor: + """``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0)``.""" + + half_dim = num_channels // 2 + exponent = -math.log(10000) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / (half_dim - 0) + emb = torch.exp(exponent) + emb = timestep[:, None].float() * emb[None, :] + emb = 1 * emb + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + return emb diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py new file mode 100644 index 000000000..bd5f9c9a6 --- /dev/null +++ b/rl_engine/testing/h3_report.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Per-operator performance and accuracy measurements for the H3 evidence. + +``PERF_CASES[op]`` builds the timed cases that ``benchmarks/benchmark_h3_conditioning.py`` +prints and ``scripts/h3_evidence.py`` stores. ``ACCURACY[op]`` measures the +op against the provider path and its golden. Each RFC #420 row registers both. +""" + +from __future__ import annotations + +import statistics +from typing import Any, Callable + +import torch + +from rl_engine.kernels.registry import KernelRegistry +from rl_engine.testing.h3_cases import h3_timesteps +from rl_engine.testing.h3_provider import provider_time_proj + +# Keys of a perf case that hold a timed callable, in display order. +TIMED_KEYS = ( + "candidate", + "candidate_checked", + "provider", + "candidate_backward", + "provider_backward", +) + + +def time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 200) -> float: + """Median CUDA-event time of ``fn`` in microseconds (includes the Python wrapper).""" + + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1e3) + return statistics.median(samples) + + +def peak_mib(fn: Callable[[], Any]) -> float: + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + fn() + torch.cuda.synchronize() + return (torch.cuda.max_memory_allocated() - base) / 2**20 + + +def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} + for key in TIMED_KEYS: + if key in case: + us = time_us(case[key], warmup, iters) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key]) + return row + + +# --------------------------------------------------------------------------- # +# timestep_sinusoid_h3 +# --------------------------------------------------------------------------- # + + +def _sinusoid_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + cases = [] + for num in (1, 2, 4, 64): + t = h3_timesteps(num) + cases.append( + { + "op": "timestep_sinusoid_h3", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": num * 4 + num * 256 * 4, + "candidate": lambda t=t: op.forward(t, check_range=False), + "candidate_checked": lambda t=t: op.forward(t), + "provider": lambda t=t: provider_time_proj(t), + } + ) + return cases + + +def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["timestep_sinusoid_h3"][-1] + ) + rows = [] + for num in (1, 2, 3, 7, 64, 1000, 4097): + t = h3_timesteps(num, seed=num) + out, ref, gold = op(t), provider_time_proj(t), golden.forward_fp32(t) + rows.append( + { + "num_timesteps": num, + "bitwise_equal_to_provider": bool(torch.equal(out, ref)), + "max_abs_vs_fp64": float((out - gold).abs().max()), + "provider_max_abs_vs_fp64": float((ref - gold).abs().max()), + } + ) + return {"contract_atol": 1e-5, "cases": rows} + + +PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { + "timestep_sinusoid_h3": _sinusoid_perf, +} +ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { + "timestep_sinusoid_h3": _sinusoid_accuracy, +} diff --git a/rl_engine/testing/h3_weights.py b/rl_engine/testing/h3_weights.py new file mode 100644 index 000000000..8090c6e83 --- /dev/null +++ b/rl_engine/testing/h3_weights.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Pinned MiniMax-H3 conditioning weights for the RFC #420 tests. + +Only the tensors listed in ``h3_manifest.json`` are used. They all live in +shard 1 of 14, so ``scripts/prepare_h3_weights.py`` downloads that shard, +checks it against the manifest and extracts the tensors into +``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import torch + +MANIFEST_PATH = Path(__file__).with_name("h3_manifest.json") +WEIGHTS_ENV = "RL_KERNEL_H3_WEIGHTS" +EXTRACTED_FILE = "h3_conditioning.safetensors" + + +def load_h3_manifest() -> dict[str, Any]: + return json.loads(MANIFEST_PATH.read_text()) + + +def h3_weights_dir() -> Path | None: + value = os.environ.get(WEIGHTS_ENV, "").strip() + return Path(value).expanduser() if value else None + + +def sha256_file(path: Path, chunk: int = 1 << 24) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(chunk): + digest.update(block) + return digest.hexdigest() + + +def load_h3_conditioning_weights( + device: torch.device | str = "cpu", names: list[str] | None = None +) -> dict[str, torch.Tensor]: + """Load the extracted pinned tensors; raise if they are missing or off-manifest.""" + + from safetensors.torch import load_file + + root = h3_weights_dir() + if root is None: + raise FileNotFoundError(f"set {WEIGHTS_ENV} to the prepare_h3_weights.py output dir") + path = root / EXTRACTED_FILE + if not path.is_file(): + raise FileNotFoundError(f"{path} missing; run scripts/prepare_h3_weights.py") + manifest = load_h3_manifest() + tensors = load_file(str(path), device=str(device)) + wanted = names if names is not None else list(manifest["tensors"]) + out: dict[str, torch.Tensor] = {} + for name in wanted: + spec = manifest["tensors"][name] + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise ValueError(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise ValueError(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + out[name] = tensor + return out diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index b63a67882..f416295d0 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -4,6 +4,7 @@ import sys from types import ModuleType +import pytest import torch import rl_engine.platforms.device as device_module @@ -225,6 +226,20 @@ def fake_load_backend(backend): ] +H3_OPS = ("timestep_sinusoid_h3",) + + +@pytest.mark.parametrize("op_name", H3_OPS) +def test_h3_ops_prefer_cuda_and_fall_back_to_pytorch_reference(op_name): + registry = KernelRegistry() + cuda_chain = registry._priority_map["cuda"][op_name] + assert cuda_chain[0].name.startswith("CUDA_H3_") + assert cuda_chain[-1].name.startswith("PYTORCH_H3_") + for platform in ("rocm", "musa", "cpu", "npu"): + chain = registry._priority_map[platform][op_name] + assert [backend.name for backend in chain] == [cuda_chain[-1].name] + + def test_npu_available_handles_runtime_failure(monkeypatch): class BrokenNPU: @staticmethod diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py new file mode 100644 index 000000000..e262ca053 --- /dev/null +++ b/scripts/h3_chain_replay.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay the MiniMax-H3 conditioning chain and write JSON evidence (RFC #420). + +See ``rl_engine/testing/h3_chain.py`` for what is compared. + + export RL_KERNEL_H3_WEIGHTS= + python scripts/h3_chain_replay.py --out chain_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from rl_engine.kernels.registry import KernelRegistry # noqa: E402 +from rl_engine.testing import h3_chain # noqa: E402 +from rl_engine.testing.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) + + +def _ints(text: str) -> list[int]: + return [int(v) for v in text.split(",")] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--stages", default=",".join(s.name for s in h3_chain.STAGES)) + parser.add_argument("--timesteps", type=_ints, default=[1, 2, 4], help="comma list of T") + parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("the chain replay needs a CUDA device") + wanted = args.stages.split(",") + unknown = sorted(set(wanted) - {s.name for s in h3_chain.STAGES}) + if unknown: + raise SystemExit(f"unknown stages: {unknown}") + stages = [s for s in h3_chain.STAGES if s.name in wanted] + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + weights = load_h3_conditioning_weights("cuda") + registry = KernelRegistry() + manifest = load_h3_manifest() + + cases = [] + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_case( + registry, + weights, + num_timesteps=num_timesteps, + seq_len=seq_len, + seed=args.seed, + stages=stages, + ) + cases.append(case) + summary = ", ".join( + f"{e['stage']}=" + f"{'eq' if e['chained_vs_provider']['bitwise_equal'] else 'drift'}" + f"(gold {e['chained_vs_golden']['max_abs']:.3e})" + for e in case["stages"] + ) + print(f"T={num_timesteps} S={seq_len}: {summary}; first_drift={case['first_drift']}") + + evidence = { + "kind": "h3_conditioning_chain_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "stages": [s.name for s in stages], + "cases": cases, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/h3_evidence.py b/scripts/h3_evidence.py new file mode 100644 index 000000000..c20778234 --- /dev/null +++ b/scripts/h3_evidence.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write one RFC #420 operator's evidence report (performance + accuracy) as JSON. + +Run it on a clean tree; the report records the commit and whether the tree +was dirty. ``scripts/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python scripts/h3_evidence.py --op timestep_sinusoid_h3 \\ + --out docs/usage/evidence/h3-timestep-sinusoid-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from rl_engine.kernels.registry import KernelRegistry # noqa: E402 +from rl_engine.testing.h3_chain import environment, git_state # noqa: E402 +from rl_engine.testing.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 +from rl_engine.testing.h3_weights import load_h3_manifest # noqa: E402 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + + registry = KernelRegistry() + manifest = load_h3_manifest() + report = { + "kind": "h3_operator_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "reference_commit": manifest["reference_implementation"]["commit"], + **git_state(), + "environment": environment(), + "accuracy": ACCURACY[args.op](registry), + "perf": [measure(case) for case in PERF_CASES[args.op](registry)], + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py new file mode 100644 index 000000000..b0bcc8eb1 --- /dev/null +++ b/scripts/plot_h3_evidence.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render an RFC #420 operator report (``scripts/h3_evidence.py``) as a PNG. + +Needs only the report JSON and matplotlib (not torch or a GPU), so it can run +anywhere: + + python scripts/plot_h3_evidence.py docs/usage/evidence/h3-timestep-sinusoid-b200/report.json + +writes ``figure.png`` next to the report. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any, Callable + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + +# Fixed series identity across every H3 figure (categorical slots 1-3). +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +THIRD = "#1baf7a" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" + + +def _style(ax, title: str, ylabel: str) -> None: + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) + ax.set_ylabel(ylabel, color=MUTED, fontsize=9) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_axisbelow(True) + for side in ("top", "right", "left"): + ax.spines[side].set_visible(False) + ax.spines["bottom"].set_color(GRID) + + +def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt: str) -> None: + """Grouped bars with a 2px-style gap and value labels on top.""" + + n = len(series) + width = 0.8 / n + for i, (name, color, values) in enumerate(series): + xs = [x + (i - (n - 1) / 2) * width for x in range(len(labels))] + bars = ax.bar(xs, values, width * 0.92, color=color, label=name, zorder=2) + for bar, value in zip(bars, values): + ax.annotate( + fmt.format(value), + (bar.get_x() + bar.get_width() / 2, bar.get_height()), + xytext=(0, 2), + textcoords="offset points", + ha="center", + fontsize=7, + color=INK, + ) + ax.set_xticks(range(len(labels)), labels) + top = max(v for _, _, values in series for v in values) + ax.set_ylim(0, top * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=n) + + +# Exact zeros cannot sit on a log axis; draw them on this floor and label them. +LOG_FLOOR = 1e-9 + + +def _log_values(values: list[float]) -> list[float]: + return [v if v > 0 else LOG_FLOOR for v in values] + + +def _figure(report: dict[str, Any], suptitle: str): + fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), facecolor=SURFACE) + env = report["environment"] + fig.suptitle(suptitle, x=0.01, ha="left", fontsize=12, color=INK, fontweight="bold") + fig.text( + 0.01, + 0.01, + f"{env['gpu']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · MiniMax-H3@{report['model_revision'][:7]}", + fontsize=7, + color=MUTED, + ) + return fig, axes + + +def plot_sinusoid(report: dict[str, Any]): + fig, (left, right) = _figure(report, "timestep_sinusoid_h3") + perf = report["perf"] + _style(left, "Latency per call (lower is better)", "µs") + _bars( + left, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + ("CUDA + range check", THIRD, [row["candidate_checked_us"] for row in perf]), + ("diffusers", PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + cases = report["accuracy"]["cases"] + xs = [c["num_timesteps"] for c in cases] + _style(right, "Max |error| vs FP64 golden (log)", "max abs error") + cuda = [c["max_abs_vs_fp64"] for c in cases] + provider = [c["provider_max_abs_vs_fp64"] for c in cases] + right.plot(xs, _log_values(cuda), "o-", color=RL_KERNEL, lw=2, ms=8, label="RL-Kernel CUDA") + right.plot(xs, _log_values(provider), "s--", color=PROVIDER, lw=2, ms=5, label="diffusers") + atol = report["accuracy"]["contract_atol"] + right.axhline(atol, color=MUTED, lw=1, ls=":") + right.annotate( + "contract atol 1e-5", + (xs[-1], atol), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + for x, value in zip(xs, cuda): + if value == 0: + right.annotate( + "exact", + (x, LOG_FLOOR), + xytext=(0, 8), + textcoords="offset points", + fontsize=7, + color=INK, + ha="center", + ) + right.set_xscale("log") + right.set_yscale("log") + right.set_ylim(LOG_FLOOR / 3, atol * 30) + right.set_xlabel("number of timesteps T", color=MUTED, fontsize=9) + equal = sum(c["bitwise_equal_to_provider"] for c in cases) + right.legend( + frameon=False, + fontsize=8, + labelcolor=INK, + loc="upper left", + ncol=2, + title=f"bitwise equal to diffusers at {equal}/{len(cases)} sizes", + title_fontsize=8, + alignment="left", + ) + return fig + + +PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { + "timestep_sinusoid_h3": plot_sinusoid, +} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + fig = PLOTS[report["op"]](report) + fig.tight_layout(rect=(0, 0.04, 1, 0.95)) + out = args.out or args.report.with_name("figure.png") + fig.savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_h3_weights.py b/scripts/prepare_h3_weights.py new file mode 100644 index 000000000..dc2ed209e --- /dev/null +++ b/scripts/prepare_h3_weights.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fetch and pin the MiniMax-H3 conditioning tensors used by the RFC #420 tests. + +Downloads only the config, the shard index and the shards that hold the +manifest tensors (shard 1 of 14, ~4.8 GB) at the pinned revision, verifies +every sha256 against ``rl_engine/testing/h3_manifest.json`` and writes the +tensors to ``/h3_conditioning.safetensors``. Point the tests at it with +``export RL_KERNEL_H3_WEIGHTS=``. + + python scripts/prepare_h3_weights.py --out ~/.cache/rl-kernel/h3 +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.testing.h3_weights import EXTRACTED_FILE, load_h3_manifest, sha256_file # noqa: E402 + + +def _check(path: Path, expected: str) -> None: + actual = sha256_file(path) + if actual != expected: + raise SystemExit(f"sha256 mismatch for {path.name}: {actual} != {expected}") + print(f"ok {path.name} {actual}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, required=True, help="output directory") + parser.add_argument( + "--download-dir", + type=Path, + default=None, + help="where the raw shard goes (default: /hf)", + ) + args = parser.parse_args() + + from huggingface_hub import hf_hub_download + from safetensors import safe_open + from safetensors.torch import save_file + + manifest = load_h3_manifest() + identity = manifest["model_identity"] + out_dir: Path = args.out.expanduser() + download_dir = (args.download_dir or out_dir / "hf").expanduser() + out_dir.mkdir(parents=True, exist_ok=True) + + def fetch(filename: str) -> Path: + return Path( + hf_hub_download( + identity["hf_repo"], + f"{identity['subfolder']}/{filename}", + revision=identity["revision"], + local_dir=str(download_dir), + ) + ) + + _check(fetch("config.json"), identity["config_sha256"]) + index_path = fetch(identity["index_file"]) + _check(index_path, identity["index_sha256"]) + weight_map = json.loads(index_path.read_text())["weight_map"] + + tensors = {} + for shard in sorted({weight_map[name] for name in manifest["tensors"]}): + if shard not in manifest["weight_shards"]: + raise SystemExit(f"shard {shard} is not pinned in the manifest") + shard_path = fetch(shard) + _check(shard_path, manifest["weight_shards"][shard]["sha256"]) + with safe_open(str(shard_path), framework="pt") as handle: + for name in manifest["tensors"]: + if weight_map[name] == shard: + tensors[name] = handle.get_tensor(name).contiguous() + + for name, spec in manifest["tensors"].items(): + tensor = tensors[name] + assert str(tensor.dtype).removeprefix("torch.") == spec["dtype"], name + assert list(tensor.shape) == spec["shape"], name + + target = out_dir / EXTRACTED_FILE + save_file(tensors, str(target), metadata={"revision": identity["revision"]}) + print(f"wrote {len(tensors)} tensors to {target}") + print(f"export RL_KERNEL_H3_WEIGHTS={out_dir}") + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 9831c7022..29f870cc9 100644 --- a/setup.py +++ b/setup.py @@ -156,6 +156,8 @@ def get_extensions(): # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. cuda_sources.append("csrc/cuda/distributed/deterministic_collective.cu") + # MiniMax-H3 (RFC #420) conditioning path; plain-arch, runs on SM90 and SM100. + cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/conftest.py b/tests/h3/conftest.py new file mode 100644 index 000000000..44c8fe117 --- /dev/null +++ b/tests/h3/conftest.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared fixtures for the MiniMax-H3 (RFC #420) operator tests. + +Real-weight tests read the pinned tensors from ``$RL_KERNEL_H3_WEIGHTS`` +(written by ``scripts/prepare_h3_weights.py``) and skip when it is unset. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.testing.h3_weights import WEIGHTS_ENV, h3_weights_dir, load_h3_conditioning_weights + + +@pytest.fixture(scope="session") +def h3_weights_cpu() -> dict[str, torch.Tensor]: + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run scripts/prepare_h3_weights.py") + try: + return load_h3_conditioning_weights("cpu") + except FileNotFoundError as exc: + pytest.skip(str(exc)) diff --git a/tests/h3/test_h3_conditioning_e2e.py b/tests/h3/test_h3_conditioning_e2e.py new file mode 100644 index 000000000..edd75a3b2 --- /dev/null +++ b/tests/h3/test_h3_conditioning_e2e.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""End to end: the H3 conditioning chain on the pinned checkpoint, through the registry. + +Runs every stage registered in ``rl_engine.testing.h3_chain.STAGES`` (each +RFC #420 row adds its own), chained exactly as the model calls them, and +checks every stage's promises: + +* the registry dispatched the CUDA backend; +* a repeated run is bitwise equal; +* the chained output is within the stage's contract tolerance of the golden; +* stages that promise it are bitwise equal to diffusers on diffusers' inputs, + so the first isolated drift can only be a stage that declares a reduction. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.registry import KernelRegistry +from rl_engine.testing import h3_chain + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + +# (T distinct timesteps, S packed rows): single row, short packed, realistic-size packed. +CASES = [(1, 3), (3, 257), (4, 4097)] + + +@pytest.fixture(scope="module") +def chain_weights(h3_weights_cpu): + return {name: tensor.cuda() for name, tensor in h3_weights_cpu.items()} + + +@pytest.fixture(scope="module") +def registry(): + return KernelRegistry() + + +@pytest.mark.parametrize("num_timesteps, seq_len", CASES) +def test_chain_forward(registry, chain_weights, num_timesteps, seq_len): + report = h3_chain.run_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + stages = {stage.name: stage for stage in h3_chain.STAGES} + for entry in report["stages"]: + stage = stages[entry["stage"]] + assert entry["backend"].endswith("CudaOp"), entry + assert entry["repeat_bitwise_equal"], entry["stage"] + assert entry["chained_vs_golden"]["within_tolerance"], entry + if stage.provider_bitwise_isolated: + assert entry["isolated_vs_provider"]["bitwise_equal"], entry["stage"] + drift = report["first_isolated_drift"] + assert drift is None or not stages[drift].provider_bitwise_isolated diff --git a/tests/h3/test_h3_timestep_sinusoid.py b/tests/h3/test_h3_timestep_sinusoid.py new file mode 100644 index 000000000..8575041a3 --- /dev/null +++ b/tests/h3/test_h3_timestep_sinusoid.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_sinusoid_h3``: FP32 [cos | sin] timestep features. + +Three properties, each checked separately: + +* provider parity: the CUDA kernel is bitwise equal to the diffusers + ``get_timestep_embedding`` replay on the same device (an elementwise op + with the same FP32 operation sequence); +* accuracy: within the FP32 elementwise contract of an FP64 golden; +* invariance: a row's bytes do not depend on how many timesteps share the + launch, their order, or repeated runs. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.testing.h3_cases import h3_timesteps +from rl_engine.testing.h3_provider import provider_time_proj + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.timestep_sinusoid import ( + H3TimestepSinusoidCudaOp, + h3_sinusoid_cuda_available, + ) + + if not h3_sinusoid_cuda_available(): + pytest.skip("rl_engine._C lacks h3_timestep_sinusoid_forward") + return H3TimestepSinusoidCudaOp() + + +class TestReference: + def test_layout_is_cos_then_sin(self): + t = torch.tensor([0.0, 1.0]) + out = NativeH3TimestepSinusoidOp().forward(t) + assert out.shape == (2, 256) and out.dtype == torch.float32 + assert torch.equal(out[0, :128], torch.ones(128)) + assert torch.equal(out[0, 128:], torch.zeros(128)) + # k = 0 has frequency 1: cos(1), sin(1) at t = 1. + assert out[1, 0].item() == pytest.approx(math.cos(1.0), abs=1e-7) + assert out[1, 128].item() == pytest.approx(math.sin(1.0), abs=1e-7) + + def test_provider_path_matches_diffusers_cpu(self): + t = h3_timesteps(9, device="cpu") + ours = NativeH3TimestepSinusoidOp().forward(t) + assert torch.equal(ours, provider_time_proj(t, 256)) + + def test_golden_close_to_provider_path(self): + t = h3_timesteps(33, device="cpu") + op = NativeH3TimestepSinusoidOp() + torch.testing.assert_close(op.forward(t), op.forward_fp32(t), atol=1e-6, rtol=0) + + @pytest.mark.parametrize( + "bad, error", + [ + (torch.tensor([[0.5]]), ValueError), # not 1-D + (torch.empty(0), ValueError), # empty + (torch.tensor([1, 0]), TypeError), # integer + (torch.tensor([0.5, 1.5]), ValueError), # H10: t outside [0, 1] + (torch.tensor([500.0]), ValueError), # H10: t * 1000 convention + (torch.tensor([float("nan")]), ValueError), + ], + ) + def test_rejects_invalid_timesteps(self, bad, error): + with pytest.raises(error): + NativeH3TimestepSinusoidOp().forward(bad) + + def test_rejects_odd_channels(self): + with pytest.raises(ValueError): + NativeH3TimestepSinusoidOp().forward(torch.tensor([0.5]), num_channels=255) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("num", [1, 2, 3, 7, 64, 1000, 4097]) + def test_bitwise_equal_to_diffusers_cuda_path(self, num): + t = h3_timesteps(num, seed=num) + out = _cuda_op()(t) + assert out.dtype == torch.float32 and out.shape == (num, 256) + assert torch.equal(out, provider_time_proj(t, 256)) + assert torch.equal(out, NativeH3TimestepSinusoidOp().forward(t)) + + @pytest.mark.parametrize("num_channels", [2, 8, 96, 256, 320]) + def test_other_channel_counts_match_reference(self, num_channels): + t = h3_timesteps(5) + out = _cuda_op()(t, num_channels=num_channels) + ref = NativeH3TimestepSinusoidOp().forward(t, num_channels=num_channels) + assert torch.equal(out, ref) + + def test_fp32_contract_against_fp64_golden(self): + t = h3_timesteps(257) + out = _cuda_op()(t) + gold = NativeH3TimestepSinusoidOp().forward_fp32(t) + # tolerance_contract.json forward_accuracy / elementwise / float32. + torch.testing.assert_close(out, gold, atol=1e-5, rtol=1e-5) + assert (out - gold).abs().max().item() <= 2.0**-23 + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + def test_low_precision_timesteps_are_upcast_first(self, dtype): + t = h3_timesteps(6).to(dtype) + assert torch.equal(_cuda_op()(t), provider_time_proj(t, 256)) + + def test_batch_position_and_repeat_invariance(self): + op = _cuda_op() + base = torch.tensor(0.3141592, device="cuda") + single = op(base.view(1)) + for num in (2, 5, 64, 1023): + for pos in (0, num // 2, num - 1): + t = h3_timesteps(num, seed=num) + t[pos] = base + assert torch.equal(op(t)[pos], single[0]), (num, pos) + t = h3_timesteps(129, seed=11) + first = op(t) + for _ in range(3): + assert torch.equal(op(t), first) + perm = torch.randperm(129, generator=torch.Generator().manual_seed(0)).cuda() + assert torch.equal(op(t[perm]), first[perm]) + + def test_non_contiguous_input(self): + t = h3_timesteps(20) + strided = t[::2] + assert not strided.is_contiguous() + assert torch.equal(_cuda_op()(strided), _cuda_op()(strided.contiguous())) + + def test_rejects_cpu_input(self): + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.5])) + + def test_rejects_out_of_range(self): + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.0, 999.0], device="cuda")) + + def test_backward_matches_autograd_reference(self): + t = h3_timesteps(17).requires_grad_(True) + t_ref = t.detach().clone().requires_grad_(True) + grad = torch.randn(17, 256, device="cuda", generator=torch.Generator("cuda").manual_seed(1)) + _cuda_op()(t).backward(grad) + NativeH3TimestepSinusoidOp().forward_fp32(t_ref).backward(grad) + torch.testing.assert_close(t.grad, t_ref.grad, atol=1e-4, rtol=1e-4) + + def test_backward_is_row_invariant(self): + op = _cuda_op() + grad_row = torch.randn(256, device="cuda", generator=torch.Generator("cuda").manual_seed(2)) + + def grad_of(t, pos): + t = t.clone().requires_grad_(True) + grad = torch.zeros(t.shape[0], 256, device="cuda") + grad[pos] = grad_row + op(t).backward(grad) + return t.grad[pos] + + single = grad_of(torch.tensor([0.7], device="cuda"), 0) + t = h3_timesteps(40) + t[13] = 0.7 + assert torch.equal(grad_of(t, 13), single) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda") + assert type(op).__name__ == "H3TimestepSinusoidCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 2f4ec8648..c688cd3d2 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -49,6 +49,7 @@ def test_all_ws1_single_ops_are_registered(): "pack", "linear_logp", "prefix_shared_attention", + "timestep_sinusoid_h3", } <= names From 23a530e267c9be17d4a27d3ea96db921e84dd04e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:28:41 +0000 Subject: [PATCH 02/89] [WS1][MiniMax-H3] timestep_sinusoid_h3: B200 evidence report and figure Produced by scripts/h3_evidence.py at 0522865 on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-sinusoid.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index ce45133c3..514ee651f 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -83,6 +83,14 @@ the range check it takes about 51 µs. The provider path takes about 64 µs, bec six eager kernels plus two concatenations. Either way the op is launch-bound: it moves about 1 KB per timestep. +## Evidence + +![timestep_sinusoid_h3 on B200: latency and error vs FP64](../usage/evidence/h3-timestep-sinusoid-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-timestep-sinusoid-b200/report.json). +`scripts/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report +records that commit and the environment. + ## Tests ```bash From 418998c1a8c76731fd9e70efe7625e5912708608 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:32:38 +0000 Subject: [PATCH 03/89] [WS1][CUDA][MiniMax-H3] timestep_mlp_fp32: deterministic FP32 timestep MLP RFC #420 step 3, row `timestep_mlp_fp32` (stacked on timestep_sinusoid_h3). - New CUDA module csrc/cuda/h3/det_linear.cu, contract h3-det-linear-v1: warp per output-column group, lane-strided 16-byte K chunks, ascending fmaf chains, xor butterfly, bias then fused SiLU. The launch shape never changes a column's summation order, so rows are batch/position invariant and runs are repeat-bitwise. - Backward without cuBLAS or atomics: dinput via fixed 64-row N chunks folded in order (row-local), dweight/dbias as ascending-row FP32 folds. - PyTorch reference (provider path + FP64 golden); BF16 anywhere is rejected (RFC probe H7: the time_embedder is a declared FP32 module). - On the pinned weights both CUDA and cuBLAS are ~1e-6 from the FP64 golden (100x inside the contract); the op is 2x faster for T >= 2. - Chain stage (covered by the end-to-end test), report/figure hooks, registry, gtest, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 476 ++++++++++++++++++ csrc/ops.cpp | 18 + docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-timestep-mlp.md | 128 +++++ rl_engine/_C.pyi | 13 + rl_engine/kernels/gtest/operator_inputs.py | 25 + rl_engine/kernels/gtest/operator_specs.py | 11 + rl_engine/kernels/ops/cuda/h3/det_linear.py | 79 +++ rl_engine/kernels/ops/cuda/h3/timestep_mlp.py | 99 ++++ .../kernels/ops/pytorch/h3/timestep_mlp.py | 77 +++ rl_engine/kernels/registry.py | 9 + rl_engine/testing/h3_chain.py | 17 +- rl_engine/testing/h3_provider.py | 15 + rl_engine/testing/h3_report.py | 71 ++- rl_engine/tests/test_dispatch.py | 2 +- scripts/plot_h3_evidence.py | 71 +++ setup.py | 1 + tests/h3/test_h3_timestep_mlp.py | 193 +++++++ tests/test_ws1_gtest_gpu.py | 1 + 20 files changed, 1305 insertions(+), 3 deletions(-) create mode 100644 csrc/cuda/h3/det_linear.cu create mode 100644 docs/operators/h3-timestep-mlp.md create mode 100644 rl_engine/kernels/ops/cuda/h3/det_linear.py create mode 100644 rl_engine/kernels/ops/cuda/h3/timestep_mlp.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py create mode 100644 tests/h3/test_h3_timestep_mlp.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu new file mode 100644 index 000000000..cc82acafa --- /dev/null +++ b/csrc/cuda/h3/det_linear.cu @@ -0,0 +1,476 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Deterministic, batch-invariant row-wise linear layers for the MiniMax-H3 +// conditioning path (RFC #420: `timestep_mlp_fp32`, `adaln_projection_3mod`). +// +// These layers run on a handful of rows (one per distinct timestep) against +// large weights, so they are GEMVs bound by weight bandwidth. The reduction +// order is fixed and depends only on K (and N for the input gradient), never +// on the number of rows or their position. +// +// Reduction contract h3-det-linear-v1 +// ----------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * one warp per output column n; lane l owns the 16-byte K chunks +// c = l, l + 32, l + 64, ... (VEC = 16 / sizeof(w) elements each); +// * each lane accumulates its chunks in ascending order and the elements of +// a chunk in ascending order with fmaf into an FP32 accumulator from 0; +// * the 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1), +// whose result is identical on every lane (FP add is commutative); +// * the bias is added once after the butterfly; the activation runs in FP32 +// (SiLU as v / (1 + expf(-v)), PyTorch's formula); one cast at the store. +// d_input dx[t, k] = sum_n g[t, n] * W[n, k] +// * N is split into fixed chunks of 64 rows; each chunk is an ascending +// fmaf chain into FP32 from 0; the chunk partials are then left-folded in +// ascending chunk order and cast once. +// d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] +// * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one +// cross-row reduction; its order is the logical row order. + +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +#ifndef H3_FWD_F32_COLS +#define H3_FWD_F32_COLS 2 +#define H3_FWD_F32_AHEAD 6 +#define H3_FWD_BF16_COLS 4 +#define H3_FWD_BF16_AHEAD 2 +#endif +#ifndef H3_FWD_CONFIGS +#define H3_FWD_CONFIGS H3_FWD_CASE(H3_FWD_F32_COLS, H3_FWD_F32_AHEAD) \ + H3_FWD_CASE(H3_FWD_BF16_COLS, H3_FWD_BF16_AHEAD) +#endif + +constexpr int kWarp = 32; +constexpr int kWarpsPerBlock = 8; +constexpr int kRowTile = 4; +constexpr int kDInputChunk = 64; + +enum Activation : int64_t { kActNone = 0, kActSilu = 1 }; + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ float from_float(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); // round to nearest even, like at::BFloat16 +} + +// 16-byte vector load converted to FP32. +template +struct Vec16 { + static constexpr int kN = 16 / sizeof(T); + __device__ __forceinline__ static void load(const T* ptr, float (&out)[kN]) { + unpack(load_raw(ptr), out); + } + __device__ __forceinline__ static uint4 load_raw(const T* ptr) { + return *reinterpret_cast(ptr); + } + // Streaming weights are read once: bypass L1 allocation. + __device__ __forceinline__ static uint4 load_stream(const T* ptr) { + return __ldcs(reinterpret_cast(ptr)); + } + __device__ __forceinline__ static void unpack(const uint4& raw, float (&out)[kN]) { + const T* vals = reinterpret_cast(&raw); +#pragma unroll + for (int i = 0; i < kN; ++i) out[i] = to_float(vals[i]); + } +}; + +__device__ __forceinline__ float warp_butterfly_sum(float v) { +#pragma unroll + for (int offset = kWarp / 2; offset > 0; offset >>= 1) { + v += __shfl_xor_sync(0xffffffffu, v, offset); + } + return v; +} + +// Each warp owns kCols adjacent output columns so that kLoadAhead x kCols +// 16-byte weight loads are in flight per lane. Grouping columns and chunks +// changes only when loads are issued: every column still accumulates its own +// chunks in ascending order, exactly as a one-column-per-warp kernel would. +template +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_kernel(const x_t* __restrict__ x, const w_t* __restrict__ w, + const w_t* __restrict__ bias, out_t* __restrict__ out, + float* __restrict__ pre_act, int64_t rows, int64_t n_out, + int64_t k_in, int64_t activation) { + using WV = Vec16; + using XV = Vec16; + constexpr int kVec = WV::kN; + static_assert(XV::kN == kVec, "x and weight share a dtype"); + constexpr int64_t kStride = static_cast(kWarp) * kVec; + + const int lane = threadIdx.x % kWarp; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + + for (int64_t n0 = warp_global * kCols; n0 < n_out; n0 += warp_stride * kCols) { + for (int64_t t0 = 0; t0 < rows; t0 += kRows) { + float acc[kCols][kRows]; +#pragma unroll + for (int c = 0; c < kCols; ++c) { +#pragma unroll + for (int r = 0; r < kRows; ++r) acc[c][r] = 0.0f; + } + for (int64_t g0 = static_cast(lane) * kVec; g0 < k_in; + g0 += kStride * kLoadAhead) { + uint4 raw[kLoadAhead][kCols]; +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; +#pragma unroll + for (int c = 0; c < kCols; ++c) { + if (k0 < k_in && n0 + c < n_out) raw[u][c] = WV::load_stream(w + (n0 + c) * k_in + k0); + } + } +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; + if (k0 >= k_in) break; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + if (t0 + r < rows) { + float xv[kVec]; + XV::load(x + (t0 + r) * k_in + k0, xv); +#pragma unroll + for (int c = 0; c < kCols; ++c) { + float wv[kVec]; + WV::unpack(raw[u][c], wv); +#pragma unroll + for (int j = 0; j < kVec; ++j) acc[c][r] = fmaf(xv[j], wv[j], acc[c][r]); + } + } + } + } + } +#pragma unroll + for (int c = 0; c < kCols; ++c) { + const int64_t n = n0 + c; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + const int64_t t = t0 + r; + const float sum = warp_butterfly_sum(acc[c][r]); + if (lane == c * kRows + r && t < rows && n < n_out) { + float v = sum + (bias != nullptr ? to_float(bias[n]) : 0.0f); + if (pre_act != nullptr) pre_act[t * n_out + n] = v; + if (activation == kActSilu) v = v / (1.0f + expf(-v)); + out[t * n_out + n] = from_float(v); + } + } + } + } + } +} + +// partial[c, t, k] = sum_{n in chunk c} g[t, n] * w[n, k] +template +__global__ void det_linear_dinput_partial_kernel(const g_t* __restrict__ grad, + const w_t* __restrict__ w, + float* __restrict__ partial, int64_t rows, + int64_t n_out, int64_t k_in) { + const int64_t k = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t chunk = blockIdx.y; + if (k >= k_in) return; + const int64_t n_begin = chunk * kDInputChunk; + const int64_t n_end = min(n_begin + kDInputChunk, n_out); + for (int64_t t0 = 0; t0 < rows; t0 += kRowTile) { + float acc[kRowTile]; +#pragma unroll + for (int r = 0; r < kRowTile; ++r) acc[r] = 0.0f; + for (int64_t n = n_begin; n < n_end; ++n) { + const float wv = to_float(w[n * k_in + k]); +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) acc[r] = fmaf(to_float(grad[t * n_out + n]), wv, acc[r]); + } + } +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) partial[(chunk * rows + t) * k_in + k] = acc[r]; + } + } +} + +template +__global__ void det_fold_chunks_kernel(const float* __restrict__ partial, out_t* __restrict__ out, + int64_t chunks, int64_t elems) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= elems) return; + float acc = partial[i]; + for (int64_t c = 1; c < chunks; ++c) acc += partial[c * elems + i]; + out[i] = from_float(acc); +} + +// dW[n, k] = sum_t g[t, n] * x[t, k] ; ascending t, FP32 from 0, one cast. +template +__global__ void det_linear_dweight_kernel(const g_t* __restrict__ grad, const x_t* __restrict__ x, + w_t* __restrict__ dw, int64_t rows, int64_t n_out, + int64_t k_in) { + const int64_t total = n_out * k_in; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < total; + i += static_cast(gridDim.x) * blockDim.x) { + const int64_t n = i / k_in; + const int64_t k = i - n * k_in; + float acc = 0.0f; + for (int64_t t = 0; t < rows; ++t) { + acc = fmaf(to_float(grad[t * n_out + n]), to_float(x[t * k_in + k]), acc); + } + dw[i] = from_float(acc); + } +} + +template +__global__ void det_linear_dbias_kernel(const g_t* __restrict__ grad, w_t* __restrict__ db, + int64_t rows, int64_t n_out) { + const int64_t n = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (n >= n_out) return; + float acc = 0.0f; + for (int64_t t = 0; t < rows; ++t) acc += to_float(grad[t * n_out + n]); + db[n] = from_float(acc); +} + +template +struct CudaType { + using type = T; +}; +template <> +struct CudaType { + using type = __nv_bfloat16; +}; + +template +const typename CudaType::type* cptr(const torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} +template +typename CudaType::type* mptr(torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} + +void check_matrix(const torch::Tensor& t, const char* name) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.dim() == 2, name, " must be 2-D"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(t.scalar_type() == at::kFloat || t.scalar_type() == at::kBFloat16, name, + " must be float32 or bfloat16, got ", t.scalar_type()); + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, name, + " must be 16-byte aligned"); +} + +int64_t grid_for(int64_t work, int64_t per_block) { + const int64_t blocks = (work + per_block - 1) / per_block; + return std::max(1, std::min(blocks, 1 << 20)); +} + +// Launch shape only: the row tile, the columns per warp and the load depth +// change which loads are in flight, never a column's accumulation order. +struct ForwardConfig { + int rows; + int cols; + int ahead; +}; + +ForwardConfig pick_forward_config(int64_t rows, int64_t elem_size) { + const int row_tile = rows <= 1 ? 1 : (rows <= 2 ? 2 : 4); + if (elem_size == 2) return {row_tile, H3_FWD_BF16_COLS, H3_FWD_BF16_AHEAD}; + return {row_tile, H3_FWD_F32_COLS, H3_FWD_F32_AHEAD}; +} + +template +void launch_forward_impl(const T* x, const T* w, const T* bias, T* out, float* pre, + int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + const int64_t blocks = grid_for((n_out + kCols - 1) / kCols, kWarpsPerBlock); + det_linear_forward_kernel + <<(blocks), kWarpsPerBlock * kWarp, 0, stream>>>( + x, w, bias, out, pre, rows, n_out, k_in, activation); +} + +template +void launch_forward_rows(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation, cudaStream_t stream) { +#define H3_FWD_CASE(C, A) \ + if (cfg.cols == C && cfg.ahead == A) { \ + launch_forward_impl(x, w, bias, out, pre, rows, n_out, k_in, activation, \ + stream); \ + return; \ + } + H3_FWD_CONFIGS +#undef H3_FWD_CASE + TORCH_CHECK(false, "unsupported det_linear forward config cols=", cfg.cols, + " ahead=", cfg.ahead); +} + +template +void launch_forward(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + if (cfg.rows == 1) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else if (cfg.rows == 2) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } +} + +} // namespace + +std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor weight, + c10::optional bias, + int64_t activation, bool save_pre_activation) { + check_matrix(x, "x"); + check_matrix(weight, "weight"); + TORCH_CHECK(x.device() == weight.device(), "x and weight must be on the same device"); + TORCH_CHECK(x.scalar_type() == weight.scalar_type(), + "x and weight must share a dtype (cast at the declared boundary first), got ", + x.scalar_type(), " and ", weight.scalar_type()); + TORCH_CHECK(activation == kActNone || activation == kActSilu, "unknown activation ", + activation); + const int64_t rows = x.size(0); + const int64_t k_in = x.size(1); + const int64_t n_out = weight.size(0); + TORCH_CHECK(rows > 0, "x must have at least one row"); + TORCH_CHECK(weight.size(1) == k_in, "weight is [", n_out, ", ", weight.size(1), + "] but x has K=", k_in); + const int64_t vec = 16 / x.element_size(); + TORCH_CHECK(k_in % vec == 0, "K=", k_in, " must be a multiple of ", vec, + " for 16-byte vector loads"); + if (bias.has_value()) { + TORCH_CHECK(bias->is_cuda() && bias->dim() == 1 && bias->size(0) == n_out && + bias->is_contiguous() && bias->scalar_type() == weight.scalar_type(), + "bias must be a contiguous [N] CUDA tensor with the weight dtype"); + } + + const c10::cuda::CUDAGuard device_guard(x.device()); + auto out = torch::empty({rows, n_out}, x.options()); + torch::Tensor pre; + if (save_pre_activation) pre = torch::empty({rows, n_out}, x.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + const ForwardConfig cfg = pick_forward_config(rows, x.element_size()); + if (x.scalar_type() == at::kFloat) { + launch_forward(cfg, cptr(x), cptr(weight), + bias.has_value() ? cptr(*bias) : nullptr, mptr(out), + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, + k_in, activation, stream); + } else { + launch_forward<__nv_bfloat16>( + cfg, cptr(x), cptr(weight), + bias.has_value() ? cptr(*bias) : nullptr, mptr(out), + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation, + stream); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (save_pre_activation) return {out, pre}; + return {out}; +} + +// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype. +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(weight, "weight"); + TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N"); + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + const int64_t rows = grad.size(0); + const int64_t n_out = weight.size(0); + const int64_t k_in = weight.size(1); + const c10::cuda::CUDAGuard device_guard(grad.device()); + const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; + auto partial = torch::empty({chunks, rows, k_in}, grad.options()); + auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); + if (weight.scalar_type() == at::kFloat) { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, n_out, + k_in); + } else { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, + n_out, k_in); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const int64_t elems = rows * k_in; + const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); + if (out_dtype == at::kFloat) { + det_fold_chunks_kernel<<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } else { + det_fold_chunks_kernel<__nv_bfloat16><<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// grad [T, N] float32, x [T, K] (float32 or bfloat16) -> (dW [N, K], dbias [N]) in w_dtype. +std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, + c10::ScalarType w_dtype, + bool with_bias) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(x, "x"); + TORCH_CHECK(grad.size(0) == x.size(0), "grad rows != x rows"); + TORCH_CHECK(w_dtype == at::kFloat || w_dtype == at::kBFloat16, + "w_dtype must be float32 or bfloat16"); + const int64_t rows = grad.size(0); + const int64_t n_out = grad.size(1); + const int64_t k_in = x.size(1); + const c10::cuda::CUDAGuard device_guard(grad.device()); + auto dw = torch::empty({n_out, k_in}, grad.options().dtype(w_dtype)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const int64_t total = n_out * k_in; + const unsigned blocks = static_cast(grid_for(total, threads)); + const float* g = grad.data_ptr(); +#define H3_DW_LAUNCH(XT, XC, WT, WC) \ + det_linear_dweight_kernel<<>>( \ + g, cptr(x), mptr(dw), rows, n_out, k_in) + if (x.scalar_type() == at::kFloat && w_dtype == at::kFloat) { + H3_DW_LAUNCH(float, float, float, float); + } else if (x.scalar_type() == at::kFloat) { + H3_DW_LAUNCH(float, float, at::BFloat16, __nv_bfloat16); + } else if (w_dtype == at::kFloat) { + H3_DW_LAUNCH(at::BFloat16, __nv_bfloat16, float, float); + } else { + H3_DW_LAUNCH(at::BFloat16, __nv_bfloat16, at::BFloat16, __nv_bfloat16); + } +#undef H3_DW_LAUNCH + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!with_bias) return {dw}; + auto db = torch::empty({n_out}, grad.options().dtype(w_dtype)); + const unsigned bias_blocks = static_cast((n_out + threads - 1) / threads); + if (w_dtype == at::kFloat) { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } else { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dw, db}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 71b5b5502..f245f8e15 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -502,6 +502,14 @@ at::Tensor prefix_shared_attention( (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, double max_period); +std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor weight, + c10::optional bias, + int64_t activation, bool save_pre_activation); +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype); +std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, + c10::ScalarType w_dtype, + bool with_bias); #endif // PyBind11 Module Registration @@ -765,5 +773,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("h3_timestep_sinusoid_forward", &h3_timestep_sinusoid_forward, "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0); + m.def("h3_det_linear_forward", &h3_det_linear_forward, + "Batch-invariant warp-per-column linear (contract h3-det-linear-v1), optional SiLU", + py::arg("x"), py::arg("weight"), py::arg("bias") = py::none(), + py::arg("activation") = 0, py::arg("save_pre_activation") = false); + m.def("h3_det_linear_backward_input", &h3_det_linear_backward_input, + "Deterministic grad @ weight with fixed 64-row N chunks folded in order", + py::arg("grad"), py::arg("weight"), py::arg("out_dtype")); + m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight, + "Deterministic dW/dbias as ascending-row FP32 folds", + py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index d4c01f6c5..0dbcaec7d 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -31,6 +31,7 @@ nav: - operators/embedding.md - MiniMax-H3: - operators/h3-timestep-sinusoid.md + - operators/h3-timestep-mlp.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index dee8796e4..ac127f3c4 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -33,4 +33,5 @@ Every operator page should include: - [Sampling](sampling.md) - [Token Embedding](embedding.md) - [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) +- [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md new file mode 100644 index 000000000..d5077f3f9 --- /dev/null +++ b/docs/operators/h3-timestep-mlp.md @@ -0,0 +1,128 @@ +# MiniMax-H3 FP32 Timestep MLP + +## Summary + +`timestep_mlp_fp32` is H3's `TimestepEmbedding(in_channels=256, time_embed_dim=5376, +out_dim=2688)`. It runs on the sinusoidal features of the distinct timesteps +(RFC #420, WS1 step 3): + +```text +temb = linear_2(silu(linear_1(features))) (T, 256) -> (T, 5376) -> (T, 2688) +``` + +`time_embedder` is a `_keep_in_fp32_modules` module. Its weights, biases, +activations and output are all FP32. `temb` stays FP32 because every AdaLN +projection applies its own SiLU before casting to BF16. + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227`. The tensors are +`time_embedder.linear_{1,2}.{weight,bias}` (see `rl_engine/testing/h3_manifest.json`). + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("timestep_mlp_fp32", device="cuda") +temb = op(features, w1, b1, w2, b2) # all float32 +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.kernels.ops.cuda.h3.timestep_mlp.H3TimestepMLPCudaOp` | `rl_engine._C.h3_det_linear_{forward,backward_input,backward_weight}` | Contract `h3-det-linear-v1` | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp` | n/a | `forward`: provider path (`F.linear`/`F.silu`); `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `x` | `(T, K)`, `T >= 1` | float32 | H3: `K = 256` | +| `w1`, `b1` | `(H, K)`, `(H,)` | float32 | H3: `H = 5376` | +| `w2`, `b2` | `(D, H)`, `(D,)` | float32 | H3: `D = 2688` | +| output | `(T, D)` | float32 | | + +The CUDA kernel needs `K` and `H` to be multiples of 4, for 16-byte vector loads. +These inputs fail closed: BF16 in any argument (RFC probe H7), mismatched shapes, +an empty `x`, mixed devices, and unaligned `K`. + +## Numerics: contract `h3-det-linear-v1` + +Forward, for each output `y[t, n]`: + +- One warp owns a group of output columns. Lane `l` reads the 16-byte K chunks + `l, l + 32, l + 64, ...`. +- Each lane accumulates its chunks in ascending order, and the elements inside a chunk + in ascending order, using `fmaf` into an FP32 accumulator that starts at 0. +- The 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1). +- The bias is added once, then SiLU runs in FP32 as `v / (1 + expf(-v))` + (PyTorch's formula). The result is stored as FP32. + +How many columns a warp owns, the row tile (1, 2 or 4), and how many chunks are +loaded ahead change only when loads are issued, never the order in which a column +is summed. A timestep's `temb` is therefore bitwise identical however many other +timesteps share the call, and wherever it sits among them. + +Backward is deterministic. There is no cuBLAS and there are no atomics: + +- `d_hidden = g @ W2` and `dx = d_pre @ W1`: N is split into fixed 64-row chunks. + Each chunk is an ascending `fmaf` chain, and the chunks are then left-folded in + ascending order. These gradients are row-local and batch-invariant. +- `d_pre = d_hidden * s * (1 + z * (1 - s))`, elementwise in FP32, using the saved + pre-activation `z`. +- `dW[n, k] = sum_t g[t, n] * x[t, k]` and `db[n] = sum_t g[t, n]`: an ascending-`t` + FP32 fold that starts from 0, in logical row order. + +Accuracy, measured on a B200 with the pinned checkpoint weights, TF32 disabled and 200 +draws of 4 timesteps (each draw includes t = 0 and t = 1): + +| Comparison | CUDA | provider (cuBLAS) | +| --- | --- | --- | +| max abs error vs FP64 golden, median over draws | 1.1e-6 | 5.4e-7 | +| gradients vs FP64 autograd (`x`, `w1`, `b1`, `w2`, `b2`) | <= 1.2e-5 | | + +Both outputs are about 100× inside the contract (`reduction` / `float32`, atol and rtol +1e-4). Neither is more accurate in general: cuBLAS's tree happens to do slightly better +on these weights. What the CUDA kernel adds is a fixed summation order, which makes each +row batch- and position-invariant and every run repeat-bitwise, and it is 2× faster +for T >= 2. The two outputs are not bitwise equal, so `scripts/h3_chain_replay.py` +reports this stage as the chain's `first_drift`, which is expected for a reduction. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op timestep_mlp_fp32 +``` + +B200, pinned weights (63.3 MB FP32). End-to-end op times include the Python wrapper: + +| T | CUDA op | provider | kernels only (CUDA) | +| --- | --- | --- | --- | +| 1 | 32.8 µs | 36.5 µs | 3.8 + 11.6 µs | +| 2 | 38.8 µs | 79.9 µs | 3.9 + 13.7 µs | +| 4 | 42.7 µs | 83.1 µs | 5.0 + 15.8 µs | + +The 5376→2688 layer reads its weights at about 5.0 TB/s at T = 1. For T >= 2 the +provider switches from cuBLAS GEMV to an SGEMM path, which takes about 52 µs of +kernel time. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_timestep_mlp.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end: sinusoid -> MLP +python scripts/check_operator.py --op timestep_mlp_fp32 --candidate cuda --device cuda \ + --dtype fp32 --batch 3 --check-grad +python scripts/h3_evidence.py --op timestep_mlp_fp32 \ + --out docs/usage/evidence/h3-timestep-mlp-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-timestep-mlp-b200/report.json +``` + +## Known Limitations + +- FP32 only, by contract. The BF16 and FP16 gtest rows do not apply to this op. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- Weight gradients are reductions across timesteps, so they follow logical row + order. They are deterministic, but not batch-invariant (no cross-row reduction can be). diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 970df8515..b8cf1d58a 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -328,3 +328,16 @@ def deterministic_collective_rocm_ipc_all_gather_input( def h3_timestep_sinusoid_forward( timestep: torch.Tensor, num_channels: int = 256, max_period: float = 10000.0 ) -> torch.Tensor: ... +def h3_det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: int = 0, + save_pre_activation: bool = False, +) -> list[torch.Tensor]: ... +def h3_det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: ... +def h3_det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 4562f2aad..435d42023 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -17,6 +17,8 @@ DEFAULT_ROPE_THETA = 1.0e6 DEFAULT_RMS_EPS = 1.0e-6 H3_FREQ_DIM = 256 +H3_TIME_HIDDEN = 5376 +H3_TIME_EMBED = 2688 def make_operator_inputs( @@ -44,6 +46,7 @@ def make_operator_inputs( "lm_head": _make_lm_head_inputs, "kv_cache_attention": _make_kv_cache_attention_inputs, "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, + "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, } try: return builders[op_name](args, dtype, device) @@ -75,6 +78,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", "timestep_sinusoid_h3": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}", + "timestep_mlp_fp32": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}x{H3_TIME_HIDDEN}" + f"x{H3_TIME_EMBED}", } try: return names[op_name] @@ -360,6 +365,26 @@ def _make_timestep_sinusoid_h3_inputs( return {"timestep": _h3_timesteps(args, dtype, device)} +def _make_timestep_mlp_fp32_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # The H3 time_embedder is declared FP32; ``dtype`` does not apply to it. + del dtype + from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import NativeH3TimestepSinusoidOp + + features = NativeH3TimestepSinusoidOp().forward(_h3_timesteps(args, torch.float32, device)) + scale1, scale2 = H3_FREQ_DIM**-0.5, H3_TIME_HIDDEN**-0.5 + return { + "x": features, + "w1": _floating_tensor((H3_TIME_HIDDEN, H3_FREQ_DIM), args, torch.float32, device, 1) + * scale1, + "b1": _floating_tensor((H3_TIME_HIDDEN,), args, torch.float32, device, 2) * 0.1, + "w2": _floating_tensor((H3_TIME_EMBED, H3_TIME_HIDDEN), args, torch.float32, device, 3) + * scale2, + "b2": _floating_tensor((H3_TIME_EMBED,), args, torch.float32, device, 4) * 0.1, + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index cc740d0ff..cefebd46f 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -261,6 +261,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("timestep",), ), + "timestep_mlp_fp32": OperatorSpec( + name="timestep_mlp_fp32", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp", + "cuda": "rl_engine.kernels.ops.cuda.h3.timestep_mlp.H3TimestepMLPCudaOp", + }, + grad_input_names=("x", "w1", "b1", "w2", "b2"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/det_linear.py b/rl_engine/kernels/ops/cuda/h3/det_linear.py new file mode 100644 index 000000000..f5dab854e --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/det_linear.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Python surface of the H3 deterministic row-wise linear (contract h3-det-linear-v1). + +See ``csrc/cuda/h3/det_linear.cu`` for the reduction order. Every helper +raises instead of falling back to cuBLAS when the extension is missing. +""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + +ACT_NONE = 0 +ACT_SILU = 1 +CONTRACT = "h3-det-linear-v1" +_SYMBOLS = ( + "h3_det_linear_forward", + "h3_det_linear_backward_input", + "h3_det_linear_backward_weight", +) + + +def det_linear_available() -> bool: + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +def _require() -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks the h3_det_linear_* symbols; rebuild the CUDA extension " + "with csrc/cuda/h3/det_linear.cu (no cuBLAS fallback is allowed)" + ) + + +def det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + *, + activation: int = ACT_NONE, + save_pre_activation: bool = False, +) -> list[torch.Tensor]: + _require() + return _C.h3_det_linear_forward( + x.contiguous(), + weight.contiguous(), + None if bias is None else bias.contiguous(), + int(activation), + bool(save_pre_activation), + ) + + +def det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: + _require() + return _C.h3_det_linear_backward_input( + grad.float().contiguous(), weight.contiguous(), out_dtype + ) + + +def det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True +) -> list[torch.Tensor]: + _require() + return _C.h3_det_linear_backward_weight( + grad.float().contiguous(), x.contiguous(), w_dtype, bool(with_bias) + ) + + +def silu_backward_fp32(grad: torch.Tensor, pre_activation: torch.Tensor) -> torch.Tensor: + """Elementwise SiLU VJP in FP32: g * s * (1 + z * (1 - s)), s = sigmoid(z).""" + + z = pre_activation.float() + s = torch.sigmoid(z) + return grad.float() * (s * (1.0 + z * (1.0 - s))) diff --git a/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py b/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py new file mode 100644 index 000000000..da85e87d8 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 FP32 timestep MLP (RFC #420 ``timestep_mlp_fp32``).""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + ACT_NONE, + ACT_SILU, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.pytorch.h3.timestep_mlp import validate_h3_timestep_mlp + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[silu]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.kernels.ops.cuda.h3.det_linear.silu_backward_fp32" +) + + +class _H3TimestepMLPCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w1, b1, w2, b2): + x = x.contiguous() + hidden, pre = det_linear_forward(x, w1, b1, activation=ACT_SILU, save_pre_activation=True) + (out,) = det_linear_forward(hidden, w2, b2, activation=ACT_NONE) + ctx.save_for_backward(x, w1, w2, pre, hidden) + return out + + @staticmethod + def backward(ctx, grad_out): + x, w1, w2, pre, hidden = ctx.saved_tensors + need_x, need_w1, need_b1, need_w2, need_b2 = ctx.needs_input_grad + grad_out = grad_out.float().contiguous() + dw2 = db2 = dw1 = db1 = dx = None + if need_w2 or need_b2: + dw2, db2 = det_linear_backward_weight(grad_out, hidden, torch.float32) + d_pre = None + if need_x or need_w1 or need_b1: + d_hidden = det_linear_backward_input(grad_out, w2, torch.float32) + d_pre = silu_backward_fp32(d_hidden, pre).contiguous() + if need_w1 or need_b1: + dw1, db1 = det_linear_backward_weight(d_pre, x, torch.float32) + if need_x: + dx = det_linear_backward_input(d_pre, w1, torch.float32) + record_backward( + "timestep_mlp_fp32", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return ( + dx, + dw1 if need_w1 else None, + db1 if need_b1 else None, + dw2 if need_w2 else None, + db2 if need_b2 else None, + ) + + +class H3TimestepMLPCudaOp: + """CUDA candidate: two warp-per-column GEMVs, SiLU fused into the first. + + The weights dominate the traffic (5.5 MB + 57.8 MB FP32), the rows are the + handful of distinct timesteps, and every output element follows contract + ``h3-det-linear-v1``, so a timestep's ``temb`` bytes do not depend on which + other timesteps share the call. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, x, w1, b1, w2, b2): + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + if not x.is_cuda: + raise ValueError("H3TimestepMLPCudaOp needs CUDA tensors") + return _H3TimestepMLPCuda.apply(x, w1, b1, w2, b2) + + def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + return self.forward(x, w1, b1, w2, b2) diff --git a/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py b/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py new file mode 100644 index 000000000..ac54d295f --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 FP32 timestep MLP (RFC #420 row ``timestep_mlp_fp32``). + +``TimestepEmbedding(in_channels=256, time_embed_dim=5376, out_dim=2688)``:: + + temb = linear_2(silu(linear_1(features))) 256 -> 5376 -> 2688 + +``time_embedder`` is in ``_keep_in_fp32_modules``: weights, biases, +activations and the output are FP32, and ``temb`` stays FP32 because every +AdaLN projection applies its own SiLU to it before casting (RFC #420 §4). +A BF16 argument is a contract violation (RFC probe H7), not a lower-precision +variant, so it is rejected. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def validate_h3_timestep_mlp( + x: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> None: + named = {"x": x, "w1": w1, "b1": b1, "w2": w2, "b2": b2} + for name, tensor in named.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 be float32: the H3 time_embedder is an FP32 module " + f"(got {tensor.dtype})" + ) + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if x.dim() != 2 or x.shape[0] == 0: + raise ValueError(f"x must be a non-empty (T, K) matrix, got {tuple(x.shape)}") + hidden, k_in = w1.shape if w1.dim() == 2 else (None, None) + if k_in != x.shape[1]: + raise ValueError(f"w1 must be (H, {x.shape[1]}), got {tuple(w1.shape)}") + if b1.shape != (hidden,): + raise ValueError(f"b1 must be ({hidden},), got {tuple(b1.shape)}") + if w2.dim() != 2 or w2.shape[1] != hidden: + raise ValueError(f"w2 must be (D, {hidden}), got {tuple(w2.shape)}") + if b2.shape != (w2.shape[0],): + raise ValueError(f"b2 must be ({w2.shape[0]},), got {tuple(b2.shape)}") + + +class NativeH3TimestepMLPOp: + """PyTorch reference for the H3 timestep MLP. + + ``forward`` is the provider path (``nn.Linear`` / ``F.silu`` in FP32; run + it with TF32 disabled, as the RL-Kernel contract requires). + ``forward_fp32`` is the independent golden: the same graph in FP64, + rounded once to FP32, so its own rounding error (~1e-16 relative) is far + below the FP32 contract and its reduction order does not matter. + """ + + op_class = "reduction" + + def __call__(self, x, w1, b1, w2, b2): + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + return F.linear(F.silu(F.linear(x, w1, b1)), w2, b2) + + def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + z = F.linear(x.double(), w1.double(), b1.double()) + h = z * torch.sigmoid(z) + return F.linear(h, w2.double(), b2.double()).float() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 533536f57..e3b3765d5 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -207,6 +207,8 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_TIMESTEP_SINUSOID = ( "rl_engine.kernels.ops.cuda.h3.timestep_sinusoid.H3TimestepSinusoidCudaOp" ) + PYTORCH_H3_TIMESTEP_MLP = "rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp" + CUDA_H3_TIMESTEP_MLP = "rl_engine.kernels.ops.cuda.h3.timestep_mlp.H3TimestepMLPCudaOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -621,6 +623,10 @@ def __init__(self): OpBackend.CUDA_H3_TIMESTEP_SINUSOID, OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID, ], + "timestep_mlp_fp32": [ + OpBackend.CUDA_H3_TIMESTEP_MLP, + OpBackend.PYTORCH_H3_TIMESTEP_MLP, + ], }, "rocm": { "logp": [ @@ -662,6 +668,7 @@ def __init__(self): "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -707,6 +714,7 @@ def __init__(self): "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -733,6 +741,7 @@ def __init__(self): "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index fb1f10943..a2c6a80c5 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -30,7 +30,7 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps -from rl_engine.testing.h3_provider import provider_time_proj +from rl_engine.testing.h3_provider import provider_time_embedder, provider_time_proj REPO_ROOT = Path(__file__).resolve().parents[2] @@ -49,6 +49,11 @@ class Stage: golden_rtol: float +def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: + weights = ctx["weights"] + return [weights[f"time_embedder.linear_{i}.{p}"] for i in (1, 2) for p in ("weight", "bias")] + + STAGES: list[Stage] = [ Stage( name="timestep_sinusoid_h3", @@ -60,6 +65,16 @@ class Stage: golden_atol=1e-5, # elementwise / float32 golden_rtol=1e-5, ), + Stage( + name="timestep_mlp_fp32", + op_type="timestep_mlp_fp32", + candidate=lambda op, ctx, up: op(up, *mlp_params(ctx)), + provider=lambda ctx, up: provider_time_embedder(up, *mlp_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *mlp_params(ctx)), + provider_bitwise_isolated=False, # reduction: different tree from cuBLAS + golden_atol=1e-4, # reduction / float32 + golden_rtol=1e-4, + ), ] diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index ff4718c00..6ae2c5790 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -15,6 +15,7 @@ import math import torch +import torch.nn.functional as F def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch.Tensor: @@ -31,3 +32,17 @@ def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) return emb + + +def provider_time_embedder( + features: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> torch.Tensor: + """``TimestepEmbedding(256, 5376, out_dim=2688)``: linear_1 -> SiLU -> linear_2, FP32.""" + + sample = F.linear(features.to(w1.dtype), w1, b1) + sample = F.silu(sample) + return F.linear(sample, w2, b2) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index bd5f9c9a6..175547a21 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -17,7 +17,8 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing.h3_cases import h3_timesteps -from rl_engine.testing.h3_provider import provider_time_proj +from rl_engine.testing.h3_provider import provider_time_embedder, provider_time_proj +from rl_engine.testing.h3_weights import h3_weights_dir, load_h3_conditioning_weights # Keys of a perf case that hold a timed callable, in display order. TIMED_KEYS = ( @@ -111,9 +112,77 @@ def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: return {"contract_atol": 1e-5, "cases": rows} +def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Tensor]: + """Pinned checkpoint tensors when available, otherwise same-shape random ones.""" + + if h3_weights_dir() is not None: + weights = load_h3_conditioning_weights("cuda", names) + return [weights[name] for name in names] + generator = torch.Generator(device="cuda").manual_seed(0) + return [ + torch.randn(shape, device="cuda", generator=generator) * shape[-1] ** -0.5 + for shape in shapes + ] + + +# --------------------------------------------------------------------------- # +# timestep_mlp_fp32 +# --------------------------------------------------------------------------- # + +MLP_NAMES = [f"time_embedder.linear_{i}.{p}" for i in (1, 2) for p in ("weight", "bias")] +MLP_SHAPES = [(5376, 256), (5376,), (2688, 5376), (2688,)] + + +def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("timestep_mlp_fp32", device="cuda") + params = h3_params(MLP_NAMES, MLP_SHAPES) + weight_bytes = sum(p.numel() * p.element_size() for p in params) + cases = [] + for num in (1, 2, 3, 4): + x = torch.rand(num, 256, device="cuda") * 2 - 1 + cases.append( + { + "op": "timestep_mlp_fp32", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight_bytes, + "candidate": lambda x=x: op(x, *params), + "provider": lambda x=x: provider_time_embedder(x, *params), + } + ) + return cases + + +def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: + op = registry.get_op("timestep_mlp_fp32", device="cuda") + golden = registry._get_or_create_backend(registry._priority_map["cpu"]["timestep_mlp_fp32"][-1]) + sinusoid = registry.get_op("timestep_sinusoid_h3", device="cuda") + params = h3_params(MLP_NAMES, MLP_SHAPES) + cuda_err, provider_err = [], [] + for seed in range(draws): + x = sinusoid(h3_timesteps(4, seed=seed)) + gold = golden.forward_fp32(x, *params) + cuda_err.append(float((op(x, *params) - gold).abs().max())) + provider_err.append(float((provider_time_embedder(x, *params) - gold).abs().max())) + # Row invariance: each timestep alone equals its row in the batch of 9. + x = sinusoid(h3_timesteps(9, seed=1)) + full = op(x, *params) + invariant = all(torch.equal(op(x[i : i + 1], *params)[0], full[i]) for i in range(9)) + return { + "contract_atol": 1e-4, + "draws": draws, + "num_timesteps": 4, + "cuda_max_abs_vs_fp64": cuda_err, + "provider_max_abs_vs_fp64": provider_err, + "rows_batch_invariant": invariant, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, + "timestep_mlp_fp32": _mlp_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, + "timestep_mlp_fp32": _mlp_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index f416295d0..b436e3f2d 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -226,7 +226,7 @@ def fake_load_backend(backend): ] -H3_OPS = ("timestep_sinusoid_h3",) +H3_OPS = ("timestep_sinusoid_h3", "timestep_mlp_fp32") @pytest.mark.parametrize("op_name", H3_OPS) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index b0bcc8eb1..2baa21bb6 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -154,8 +154,79 @@ def plot_sinusoid(report: dict[str, Any]): return fig +def _latency_panel(ax, perf: list[dict[str, Any]], provider_name: str) -> None: + _style(ax, "Latency per call (lower is better)", "µs") + _bars( + ax, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + (provider_name, PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + + +def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], contract: float): + """Per-draw max |error| as boxes with every draw overlaid, log scale.""" + + _style(ax, title, "max abs error per draw") + for i, (name, color, values) in enumerate(series): + ax.boxplot( + [values], + positions=[i], + widths=0.45, + showfliers=False, + patch_artist=True, + boxprops={"facecolor": color, "alpha": 0.25, "edgecolor": color}, + medianprops={"color": color, "linewidth": 2}, + whiskerprops={"color": color}, + capprops={"color": color}, + ) + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax.scatter(jitter, values, s=8, color=color, alpha=0.6, linewidths=0, zorder=3) + median = sorted(values)[len(values) // 2] + ax.annotate( + f"median {median:.1e}", + (i + 0.28, median), + fontsize=8, + color=INK, + va="center", + ) + ax.axhline(contract, color=MUTED, lw=1, ls=":") + ax.annotate( + f"contract atol {contract:g}", + (len(series) - 0.5, contract), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + ax.set_yscale("log") + ax.set_xticks(range(len(series)), [name for name, _, _ in series]) + ax.set_xlim(-0.6, len(series) - 0.1) + + +def plot_mlp(report: dict[str, Any]): + fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _error_boxes( + right, + f"Error vs FP64 golden, {acc['draws']} draws of T={acc['num_timesteps']}", + [ + ("RL-Kernel CUDA", RL_KERNEL, acc["cuda_max_abs_vs_fp64"]), + ("diffusers (cuBLAS)", PROVIDER, acc["provider_max_abs_vs_fp64"]), + ], + acc["contract_atol"], + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, + "timestep_mlp_fp32": plot_mlp, } diff --git a/setup.py b/setup.py index 29f870cc9..b388bdd29 100644 --- a/setup.py +++ b/setup.py @@ -158,6 +158,7 @@ def get_extensions(): cuda_sources.append("csrc/cuda/distributed/deterministic_collective.cu") # MiniMax-H3 (RFC #420) conditioning path; plain-arch, runs on SM90 and SM100. cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") + cuda_sources.append("csrc/cuda/h3/det_linear.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py new file mode 100644 index 000000000..7e09da4ba --- /dev/null +++ b/tests/h3/test_h3_timestep_mlp.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_mlp_fp32``: FP32 256 -> 5376 -> 2688 timestep MLP. + +* accuracy: forward and every gradient against an FP64 golden, on the pinned + checkpoint weights and on synthetic non-H3 sizes; +* invariance: a timestep's ``temb`` bytes (and its row-local ``dx``) do not + depend on how many timesteps share the call or where it sits; +* determinism: repeated forward/backward runs are bitwise equal; +* contract: BF16 inputs (RFC probe H7) and malformed shapes fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3.timestep_mlp import NativeH3TimestepMLPOp +from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.testing.h3_cases import h3_timesteps + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NAMES = ("x", "w1", "b1", "w2", "b2") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available + from rl_engine.kernels.ops.cuda.h3.timestep_mlp import H3TimestepMLPCudaOp + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3TimestepMLPCudaOp() + + +def _features(num, device="cuda", seed=0): + return NativeH3TimestepSinusoidOp().forward(h3_timesteps(num, seed=seed, device=device)) + + +def _synthetic_params(k_in, hidden, out, device="cuda", seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + w1 = torch.randn(hidden, k_in, generator=g) / k_in**0.5 + b1 = torch.randn(hidden, generator=g) * 0.1 + w2 = torch.randn(out, hidden, generator=g) / hidden**0.5 + b2 = torch.randn(out, generator=g) * 0.1 + return [t.to(device) for t in (w1, b1, w2, b2)] + + +def _real_params(weights): + return [ + weights[f"time_embedder.linear_{i}.{p}"].cuda() for i in (1, 2) for p in ("weight", "bias") + ] + + +def _fp64_grads(x, params, grad_out): + leaves = [t.detach().double().requires_grad_(True) for t in (x, *params)] + z = F.linear(leaves[0], leaves[1], leaves[2]) + F.linear(z * torch.sigmoid(z), leaves[3], leaves[4]).backward(grad_out.double()) + return [leaf.grad for leaf in leaves] + + +class TestReference: + def test_provider_and_golden_agree(self): + x = _features(5, device="cpu") + params = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + torch.testing.assert_close( + op.forward(x, *params), op.forward_fp32(x, *params), atol=1e-5, rtol=1e-5 + ) + + @pytest.mark.parametrize("which", range(5)) + def test_rejects_bf16_anywhere(self, which): + # RFC probe H7: casting a declared FP32 path to BF16 early. + args = [_features(2, device="cpu"), *_synthetic_params(256, 64, 32, device="cpu")] + args[which] = args[which].bfloat16() + with pytest.raises(TypeError, match="float32"): + NativeH3TimestepMLPOp().forward(*args) + + def test_rejects_shape_mismatches(self): + x = _features(2, device="cpu") + w1, b1, w2, b2 = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + with pytest.raises(ValueError): + op.forward(x[:, :128], w1, b1, w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1[:10], w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1, w2[:, :10], b2) + with pytest.raises(ValueError): + op.forward(x[:0], w1, b1, w2, b2) + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4, 7]) + def test_forward_against_fp64_golden(self, h3_weights_cpu, num): + params = _real_params(h3_weights_cpu) + x = _features(num, seed=num) + out = _cuda_op()(x, *params) + gold = NativeH3TimestepMLPOp().forward_fp32(x, *params) + assert out.dtype == torch.float32 and out.shape == (num, 2688) + # tolerance_contract.json forward_accuracy / reduction / float32. + torch.testing.assert_close(out, gold, atol=1e-4, rtol=1e-4) + assert (out - gold).abs().max().item() < 1e-5 + + def test_rows_are_batch_and_position_invariant(self, h3_weights_cpu): + op = _cuda_op() + params = _real_params(h3_weights_cpu) + x_all = _features(9, seed=1) + full = op(x_all, *params) + for i in range(9): + assert torch.equal(op(x_all[i : i + 1], *params)[0], full[i]) + for num in (2, 3, 4, 5, 8): + for pos in (0, num - 1): + x = _features(num, seed=10 + num) + x[pos] = x_all[4] + assert torch.equal(op(x, *params)[pos], full[4]), (num, pos) + + def test_backward_against_fp64_golden(self, h3_weights_cpu): + params = _real_params(h3_weights_cpu) + x = _features(3, seed=2) + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + grad_out = torch.randn( + 3, 2688, device="cuda", generator=torch.Generator("cuda").manual_seed(0) + ) + _cuda_op()(*leaves).backward(grad_out) + for name, leaf, ref in zip(NAMES, leaves, _fp64_grads(x, params, grad_out)): + # tolerance_contract.json gradient_accuracy / reduction / float32. + torch.testing.assert_close(leaf.grad.double(), ref, atol=1e-4, rtol=1e-4, msg=name) + + +@requires_cuda +class TestCudaSynthetic: + @pytest.mark.parametrize("k_in, hidden, out", [(8, 24, 16), (12, 40, 20), (256, 5376, 2688)]) + def test_other_sizes(self, k_in, hidden, out): + params = _synthetic_params(k_in, hidden, out, seed=k_in) + x = torch.randn(5, k_in, device="cuda") + torch.testing.assert_close( + _cuda_op()(x, *params), + NativeH3TimestepMLPOp().forward_fp32(x, *params), + atol=1e-4, + rtol=1e-4, + ) + + def test_rejects_unaligned_k(self): + params = _synthetic_params(6, 8, 8) + with pytest.raises(RuntimeError, match="multiple of 4"): + _cuda_op()(torch.randn(2, 6, device="cuda"), *params) + + def test_rejects_cpu_and_mixed_devices(self): + params = _synthetic_params(8, 24, 16) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *[p.cpu() for p in params]) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *params) + + def test_repeat_and_backward_are_deterministic(self): + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=3) + x = _features(4, seed=3) + grad_out = torch.randn(4, 2688, device="cuda") + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + out = op(*leaves) + out.backward(grad_out) + runs.append([out.detach(), *[leaf.grad for leaf in leaves]]) + for later in runs[1:]: + for a, b in zip(runs[0], later): + assert torch.equal(a, b) + + def test_dx_rows_are_batch_invariant(self): + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=4) + x = _features(6, seed=4) + grad_out = torch.randn(6, 2688, device="cuda") + + def dx(rows): + leaf = x[rows].detach().clone().requires_grad_(True) + op(leaf, *params).backward(grad_out[rows]) + return leaf.grad + + full = dx(slice(0, 6)) + for i in range(6): + assert torch.equal(dx(slice(i, i + 1))[0], full[i]) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_mlp_fp32", device="cuda") + assert type(op).__name__ == "H3TimestepMLPCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index c688cd3d2..6dde75c53 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -50,6 +50,7 @@ def test_all_ws1_single_ops_are_registered(): "linear_logp", "prefix_shared_attention", "timestep_sinusoid_h3", + "timestep_mlp_fp32", } <= names From e32bb4ab8c946ff27ec40c7f37b5751d7e313b61 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:33:06 +0000 Subject: [PATCH 04/89] [WS1][MiniMax-H3] timestep_mlp_fp32: B200 evidence report and figure Produced by scripts/h3_evidence.py at 65ef7f6 on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index d5077f3f9..1da2434ac 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -107,6 +107,14 @@ The 5376→2688 layer reads its weights at about 5.0 TB/s at T = 1. For T >= 2 t provider switches from cuBLAS GEMV to an SGEMM path, which takes about 52 µs of kernel time. +## Evidence + +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../usage/evidence/h3-timestep-mlp-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-timestep-mlp-b200/report.json), written +by `scripts/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +timestep's row is bitwise identical whether it runs alone or in a batch of 9. + ## Tests ```bash From 503305ad83907c4a5724f292eb2dc599e9d34273 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:36:25 +0000 Subject: [PATCH 05/89] [WS1][CUDA][MiniMax-H3] adaln_projection_3mod: three-modality AdaLN projection RFC #420 step 3, row `adaln_projection_3mod` (stacked on timestep_mlp_fp32). - SiLU at temb's FP32 precision and exactly one cast to the projection dtype (bitwise equal to the provider activation); a BF16 temb is rejected (RFC probe H7). - 2688 -> 96768 projection on a BF16 tensor-core path in det_linear.cu, contract h3-det-linear-bf16-mma-v1: warp = 16 output columns x 8 input rows (zero-padded), K in ascending 16-wide groups, each mma.sync m16n8k16 from a zero accumulator added into an FP32 running sum. Same instruction sequence for every T <= 8, so rows are bitwise batch/order/tile invariant; 99.98% of outputs correctly rounded vs 99.84% for cuBLAS; kernel time flat at ~78 us for T = 1..8. SM80+. - Six (3T, 5376) views in diffusers' view(-1, 6H).chunk(6) layout, row t*3+m; a channel-tagging test pins chunk/modality placement (probe H4). - Backward: FP32 d_temb through the identity-VJP cast, ascending-row dW/db folds, no cuBLAS or atomics. The FP64 golden applies the cast straight-through so its own gradient is not rounded to BF16. - Chain stage (end-to-end test), report/figure hooks, registry, gtest, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 124 ++++++++- docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-adaln-projection.md | 142 ++++++++++ rl_engine/kernels/gtest/operator_inputs.py | 26 ++ rl_engine/kernels/gtest/operator_specs.py | 13 + .../kernels/ops/cuda/h3/adaln_projection.py | 106 ++++++++ .../ops/pytorch/h3/adaln_projection.py | 98 +++++++ rl_engine/kernels/registry.py | 13 + rl_engine/testing/h3_chain.py | 21 +- rl_engine/testing/h3_provider.py | 10 + rl_engine/testing/h3_report.py | 84 +++++- rl_engine/tests/test_dispatch.py | 2 +- scripts/plot_h3_evidence.py | 38 +++ tests/h3/test_h3_adaln_projection.py | 249 ++++++++++++++++++ tests/test_ws1_gtest_gpu.py | 1 + 16 files changed, 915 insertions(+), 14 deletions(-) create mode 100644 docs/operators/h3-adaln-projection.md create mode 100644 rl_engine/kernels/ops/cuda/h3/adaln_projection.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/adaln_projection.py create mode 100644 tests/h3/test_h3_adaln_projection.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index cc82acafa..04f7060fb 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -20,6 +20,8 @@ // whose result is identical on every lane (FP add is commutative); // * the bias is added once after the butterfly; the activation runs in FP32 // (SiLU as v / (1 + expf(-v)), PyTorch's formula); one cast at the store. +// * this FP32 forward is used for FP32 weights; BF16 weights take the +// tensor-core forward of contract h3-det-linear-bf16-mma-v1 below. // d_input dx[t, k] = sum_n g[t, n] * W[n, k] // * N is split into fixed chunks of 64 rows; each chunk is an ascending // fmaf chain into FP32 from 0; the chunk partials are then left-folded in @@ -27,6 +29,20 @@ // d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] // * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one // cross-row reduction; its order is the logical row order. +// +// Reduction contract h3-det-linear-bf16-mma-v1 (BF16 x and W, FP32 accumulate) +// ----------------------------------------------------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * a warp owns 16 output columns (MMA rows) and 8 input rows (MMA columns; +// rows past T are zero), so every launch runs the same instruction +// sequence for every T <= 8, and an MMA column never sees another's data; +// * K is visited in groups of 16 in ascending order. Within a 32-wide step, +// quad q supplies k = 32s + 8q + {0..3} to the first mma.sync m16n8k16 and +// k = 32s + 8q + {4..7} to the second (the same k map for W and x); +// * every mma.sync starts from a zero accumulator, so the tensor core only +// sums 16 products; its FP32 result is added to the running FP32 sum with +// an IEEE add, group by group in ascending k; +// * then bias, FP32 activation and one cast, as above. #include #include @@ -42,12 +58,9 @@ namespace { #ifndef H3_FWD_F32_COLS #define H3_FWD_F32_COLS 2 #define H3_FWD_F32_AHEAD 6 -#define H3_FWD_BF16_COLS 4 -#define H3_FWD_BF16_AHEAD 2 #endif #ifndef H3_FWD_CONFIGS -#define H3_FWD_CONFIGS H3_FWD_CASE(H3_FWD_F32_COLS, H3_FWD_F32_AHEAD) \ - H3_FWD_CASE(H3_FWD_BF16_COLS, H3_FWD_BF16_AHEAD) +#define H3_FWD_CONFIGS H3_FWD_CASE(H3_FWD_F32_COLS, H3_FWD_F32_AHEAD) #endif constexpr int kWarp = 32; @@ -177,6 +190,93 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) } } +__device__ __forceinline__ void mma_m16n8k16_bf16(float (&c)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + // BF16 mma.sync needs SM80+; the host refuses to launch below that. + __trap(); +#else + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +#endif +} + +constexpr int kMmaCols = 16; // output columns per warp (MMA M) +constexpr int kMmaRows = 8; // input rows per pass (MMA N) +constexpr int kMmaAhead = 4; // 32-wide K steps loaded before use + +__device__ __forceinline__ void mma_group_accumulate(float (&acc)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { + float part[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + mma_m16n8k16_bf16(part, a, b0, b1); +#pragma unroll + for (int i = 0; i < 4; ++i) acc[i] += part[i]; +} + +// Contract h3-det-linear-bf16-mma-v1. Lane = 4 * g + q: g picks the MMA row +// (output columns n0 + g and n0 + g + 8) and the MMA column (input row t0 + g). +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_bf16_mma_kernel(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, float* __restrict__ pre_act, + int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation) { + const int lane = threadIdx.x % kWarp; + const int g = lane >> 2; + const int q = lane & 3; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + + for (int64_t n0 = warp_global * kMmaCols; n0 < n_out; n0 += warp_stride * kMmaCols) { + const bool lo_ok = n0 + g < n_out; + const bool hi_ok = n0 + g + 8 < n_out; + const __nv_bfloat16* w_lo = w + (n0 + g) * k_in; + const __nv_bfloat16* w_hi = w + (n0 + g + 8) * k_in; + for (int64_t t0 = 0; t0 < rows; t0 += kMmaRows) { + const bool x_ok = t0 + g < rows; + const __nv_bfloat16* x_row = x + (t0 + g) * k_in; + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int64_t kb = 0; kb < k_in; kb += 32 * kMmaAhead) { + uint4 a_lo[kMmaAhead], a_hi[kMmaAhead], bx[kMmaAhead]; +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + const int64_t k = kb + 32 * u + 8 * q; + const bool k_ok = k < k_in; + a_lo[u] = (k_ok && lo_ok) ? __ldcs(reinterpret_cast(w_lo + k)) : zero; + a_hi[u] = (k_ok && hi_ok) ? __ldcs(reinterpret_cast(w_hi + k)) : zero; + bx[u] = (k_ok && x_ok) ? *reinterpret_cast(x_row + k) : zero; + } +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + if (kb + 32 * u >= k_in) break; + const uint32_t first[4] = {a_lo[u].x, a_hi[u].x, a_lo[u].y, a_hi[u].y}; + const uint32_t second[4] = {a_lo[u].z, a_hi[u].z, a_lo[u].w, a_hi[u].w}; + mma_group_accumulate(acc, first, bx[u].x, bx[u].y); + mma_group_accumulate(acc, second, bx[u].z, bx[u].w); + } + } + // acc = {(n0+g, t0+2q), (n0+g, t0+2q+1), (n0+g+8, t0+2q), (n0+g+8, t0+2q+1)} +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int64_t n = n0 + g + (i >= 2 ? 8 : 0); + const int64_t t = t0 + 2 * q + (i & 1); + if (n < n_out && t < rows) { + float v = acc[i] + (bias != nullptr ? __bfloat162float(bias[n]) : 0.0f); + if (pre_act != nullptr) pre_act[t * n_out + n] = v; + if (activation == kActSilu) v = v / (1.0f + expf(-v)); + out[t * n_out + n] = __float2bfloat16(v); + } + } + } + } +} + // partial[c, t, k] = sum_{n in chunk c} g[t, n] * w[n, k] template __global__ void det_linear_dinput_partial_kernel(const g_t* __restrict__ grad, @@ -287,9 +387,8 @@ struct ForwardConfig { int ahead; }; -ForwardConfig pick_forward_config(int64_t rows, int64_t elem_size) { +ForwardConfig pick_forward_config(int64_t rows) { const int row_tile = rows <= 1 ? 1 : (rows <= 2 ? 2 : 4); - if (elem_size == 2) return {row_tile, H3_FWD_BF16_COLS, H3_FWD_BF16_AHEAD}; return {row_tile, H3_FWD_F32_COLS, H3_FWD_F32_AHEAD}; } @@ -365,18 +464,21 @@ std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor torch::Tensor pre; if (save_pre_activation) pre = torch::empty({rows, n_out}, x.options().dtype(at::kFloat)); auto stream = at::cuda::getCurrentCUDAStream(); - const ForwardConfig cfg = pick_forward_config(rows, x.element_size()); if (x.scalar_type() == at::kFloat) { + const ForwardConfig cfg = pick_forward_config(rows); launch_forward(cfg, cptr(x), cptr(weight), bias.has_value() ? cptr(*bias) : nullptr, mptr(out), save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation, stream); } else { - launch_forward<__nv_bfloat16>( - cfg, cptr(x), cptr(weight), + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 8, + "the BF16 h3 det_linear forward uses mma.sync and needs SM80 or newer"); + const int64_t blocks = grid_for((n_out + kMmaCols - 1) / kMmaCols, kWarpsPerBlock); + det_linear_forward_bf16_mma_kernel<<(blocks), kWarpsPerBlock * kWarp, + 0, stream>>>( + cptr(x), cptr(weight), bias.has_value() ? cptr(*bias) : nullptr, mptr(out), - save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation, - stream); + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation); } C10_CUDA_KERNEL_LAUNCH_CHECK(); if (save_pre_activation) return {out, pre}; diff --git a/docs/.nav.yml b/docs/.nav.yml index 0dbcaec7d..b5a3c718b 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -32,6 +32,7 @@ nav: - MiniMax-H3: - operators/h3-timestep-sinusoid.md - operators/h3-timestep-mlp.md + - operators/h3-adaln-projection.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index ac127f3c4..d713f52d1 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -34,4 +34,5 @@ Every operator page should include: - [Token Embedding](embedding.md) - [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) - [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) +- [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md new file mode 100644 index 000000000..8e91cf9bf --- /dev/null +++ b/docs/operators/h3-adaln-projection.md @@ -0,0 +1,142 @@ +# MiniMax-H3 Three-Modality AdaLN Projection + +## Summary + +`adaln_projection_3mod` is the per-block `MiniMaxH3AdaLayerNormModulation` +(RFC #420, WS1 step 3). It turns the shared FP32 timestep embedding into the six +modulation tensors of one transformer block, for all three modalities: + +```text +act = silu(temb).to(bf16) SiLU in FP32, one declared cast +table = act @ W.T + b (T, 96768) BF16, W: (96768, 2688) +rows = table.view(3T, 32256) row t * 3 + m, m: 0 video, 1 text, 2 audio +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) # (3T, 5376) each +``` + +Projection output channel `o = m * 32256 + c * 5376 + h` is chunk `c`, hidden index `h` +of modality `m`. The table rows are laid out `[t0m0, t0m1, t0m2, t1m0, ...]`. That is the +layout `adaln_row_gather` addresses with `timestep_indices * 3 + token_tags`. + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227`, `transformer_blocks.0.adaln_proj.linear` +(BF16 weight and bias). All 50 blocks have the same shape. + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("adaln_projection_3mod", device="cuda") +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, weight, bias) +table = op.forward_table(temb, weight, bias) # raw (T, 96768) for a fused gather +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+, validated on SM100) | `rl_engine.kernels.ops.cuda.h3.adaln_projection.H3AdaLNProjectionCudaOp` | `rl_engine._C.h3_det_linear_*` | BF16 weights: contract `h3-det-linear-bf16-mma-v1`; FP32 weights: `h3-det-linear-v1` | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.adaln_projection.NativeH3AdaLNProjectionOp` | n/a | `forward`: provider path; `forward_fp32`: declared-cast FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `temb` | `(T, D)`, `T >= 1` | float32 | H3: `D = 2688`. BF16 is rejected (RFC probe H7) | +| `weight` | `(6 * H * 3, D)` | bfloat16 (checkpoint) or float32 | First dim must be a multiple of 18. H3: `H = 5376` | +| `bias` | `(6 * H * 3,)` | same as `weight` | | +| outputs | six `(3T, H)` views | weight dtype | Views of one `(T, 18H)` table, as in diffusers | + +The CUDA kernel needs `D` to be a multiple of 8 for BF16 or 4 for FP32, and BF16 needs SM80+. These inputs fail +closed: a BF16 `temb`, mismatched dtypes or shapes, an empty `temb`, and mixed devices. + +## Numerics + +- **Mixed-precision boundary.** SiLU runs in FP32 at `temb`'s precision, and the result is + rounded to BF16 exactly once. These are the provider's own elementwise ops, so the + activation is bitwise equal to diffusers. Casting before the SiLU (probe H7) is + rejected at the API. It is also numerically visible: on block 0 it changes about 52% of + the BF16 outputs. +- **Projection (BF16 weights).** It follows contract `h3-det-linear-bf16-mma-v1` in + `csrc/cuda/h3/det_linear.cu`: + - Each warp owns 16 output columns and 8 input rows. Rows past T are zero, so every + T <= 8 runs the same instruction sequence. A tensor-core column never sees another + column's data. + - K is visited in groups of 16 in ascending order. Each `mma.sync m16n8k16` (BF16 in, + FP32 out) starts from zero, so the tensor core only sums 16 products. Its result is + added to an FP32 running sum with an IEEE add. + - Then the bias is added and the result is rounded to BF16 once. + + A timestep's 18 modulation rows are bitwise independent of the other timesteps in the + call and of their order. Restarting the accumulator for every 16-wide group is what keeps + accuracy at the level of an FP32 FMA chain. Letting the tensor core accumulate across + the whole of K, as cuBLAS does, loses about 0.15% of correctly rounded outputs. FP32 + weights use the FMA path of `h3-det-linear-v1` instead, because FP32 on tensor cores + would mean TF32, which the contract forbids. +- **Golden.** `forward_fp32` computes the SiLU in FP64, rounds it to BF16 at the declared + boundary (the cast is model semantics), and runs the projection in FP64. It returns + FP32 without the final rounding. Its gradient stays in FP64 end to end, because the + cast is applied straight-through (identity VJP). Writing the cast as + `.to(bf16).double()` would make autograd round the golden's own `d_temb` to BF16. +- **Backward.** No cuBLAS and no atomics. `d_act` is computed with the chunked + deterministic `grad @ W`, kept in FP32 through the cast (whose VJP is the identity), + and passed through the FP32 SiLU VJP; it is row-local and batch-invariant. + `dW` and `db` are ascending-row FP32 folds, rounded to BF16 once. Unlike the provider, + `d_temb` is not rounded to BF16 on its way through the cast. + +Measured on a B200 with the pinned block-0 weights and torch 2.13.0+cu130. Fractions are +medians over 20 random draws of `temb`, with T = 4: + +| Comparison | CUDA | provider (cuBLAS) | +| --- | --- | --- | +| outputs equal to the correctly rounded golden | 99.982% (worst draw 99.976%) | 99.842% (worst draw 99.787%) | +| CUDA equal to provider | 99.84% | | +| early-cast golden (probe H7) equal to CUDA | 43% | | +| `d_temb` vs FP64 golden (gtest, T = 3) | max abs 2.3e-5 | 6.1e-2 (`d_act` rounded to BF16) | +| `dW`, `db` vs FP64 golden rounded to BF16 | > 99.9% bitwise equal | | + +The contract tolerance is `reduction` / `bfloat16`: atol 5e-2, rtol 2e-2, and atol 1e-1 +for gradients. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op adaln_projection_3mod +``` + +B200, pinned BF16 weights (520 MB per block, streamed once per call): + +| T | CUDA op | provider (SiLU + cast + cuBLAS) | CUDA GEMV kernel | cuBLAS kernel | +| --- | --- | --- | --- | --- | +| 1 | 102 µs | 112 µs | 78.3 µs (6.6 TB/s) | 93 µs | +| 2 | 102 µs | 101 µs | 78.4 µs | 79 µs | +| 4 | 102 µs | 101 µs | 79.0 µs | 79 µs | + +The kernel takes the same time for every T <= 8, because it always computes 8 rows. It +matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 24 µs of the op time is +the SiLU and cast kernels plus the Python wrapper. + +An earlier FMA-only BF16 kernel ran at 81, 98 and 166 µs for T = 1, 2 and 4. Its time +grew with T because every extra row adds FMAs per weight element. Packed FP32 FMA +(`__ffma2_rn`) gave identical bits but no speedup. Moving to tensor cores removed the +dependence on T. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_adaln_projection.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end: sinusoid -> MLP -> projection +python scripts/check_operator.py --op adaln_projection_3mod --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op adaln_projection_3mod \ + --out docs/usage/evidence/h3-adaln-projection-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-adaln-projection-b200/report.json +``` + +## Known Limitations + +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- The BF16 path needs SM80 or newer (`mma.sync`). It is rejected at launch on older GPUs. +- T > 8 runs one more weight pass per 8 rows. +- Sharding the 96768-wide projection (`tp_adaln_3mod`) is a separate WS2 row. diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 435d42023..255a48c0c 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -19,6 +19,7 @@ H3_FREQ_DIM = 256 H3_TIME_HIDDEN = 5376 H3_TIME_EMBED = 2688 +H3_HIDDEN = 5376 def make_operator_inputs( @@ -47,6 +48,7 @@ def make_operator_inputs( "kv_cache_attention": _make_kv_cache_attention_inputs, "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, + "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, } try: return builders[op_name](args, dtype, device) @@ -80,6 +82,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "timestep_sinusoid_h3": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}", "timestep_mlp_fp32": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}x{H3_TIME_HIDDEN}" f"x{H3_TIME_EMBED}", + "adaln_projection_3mod": f"{_h3_num_timesteps(args)}x{H3_TIME_EMBED}" + f"x{6 * 3 * _h3_hidden(args)}", } try: return names[op_name] @@ -385,6 +389,28 @@ def _make_timestep_mlp_fp32_inputs( } +def _h3_hidden(args: argparse.Namespace) -> int: + return _arg_int(args, "normalized_dim", H3_HIDDEN) + + +def _make_adaln_projection_3mod_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # temb is FP32 by contract (the SiLU runs before the cast); ``dtype`` is the + # projection's weight dtype, BF16 in the checkpoint. + n_out = 6 * 3 * _h3_hidden(args) + temb = _floating_tensor( + (_h3_num_timesteps(args), H3_TIME_EMBED), args, torch.float32, device, 0 + ) + weight = _floating_tensor((n_out, H3_TIME_EMBED), args, torch.float32, device, 1) + bias = _floating_tensor((n_out,), args, torch.float32, device, 2) + return { + "temb": temb * 2.0, + "weight": (weight * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (bias * 0.1).to(dtype), + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index cefebd46f..91d877b62 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -272,6 +272,19 @@ def _load_object(path: str) -> Any: }, grad_input_names=("x", "w1", "b1", "w2", "b2"), ), + "adaln_projection_3mod": OperatorSpec( + name="adaln_projection_3mod", + op_class="reduction", + gold_path=("rl_engine.kernels.ops.pytorch.h3.adaln_projection.NativeH3AdaLNProjectionOp"), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.kernels.ops.pytorch.h3.adaln_projection.NativeH3AdaLNProjectionOp" + ), + "cuda": "rl_engine.kernels.ops.cuda.h3.adaln_projection.H3AdaLNProjectionCudaOp", + }, + grad_input_names=("temb", "weight", "bias"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py new file mode 100644 index 000000000..c8a9b42cc --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 three-modality AdaLN projection (RFC #420 ``adaln_projection_3mod``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + ACT_NONE, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.pytorch.h3.adaln_projection import ( + split_adaln_table, + validate_h3_adaln_projection, +) + +KERNEL_ID = "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.kernels.ops.cuda.h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNProjectionCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias): + temb = temb.contiguous() + # Declared mixed-precision boundary: SiLU at temb's FP32 precision, + # then exactly one rounding to the projection dtype. These are the + # provider's own elementwise ops, so the activation is bitwise equal. + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias, activation=ACT_NONE) + ctx.save_for_backward(temb, act, weight) + return table + + @staticmethod + def backward(ctx, grad_table): + temb, act, weight = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad + grad_table = grad_table.float().contiguous() + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(grad_table, act, weight.dtype) + if need_temb: + # The cast's VJP is the identity; the SiLU VJP runs in FP32. + d_act = det_linear_backward_input(grad_table, weight, torch.float32) + d_temb = silu_backward_fp32(d_act, temb) + record_backward( + "adaln_projection_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return d_temb, dw if need_weight else None, db if need_bias else None + + +class H3AdaLNProjectionCudaOp: + """CUDA candidate: FP32 SiLU, one BF16 cast, deterministic GEMV, six views. + + The projection weight (96768 x 2688 BF16, 520 MB per block) is streamed + once per call; every output element follows contract ``h3-det-linear-v1``, + so a timestep's 3 x 6 modulation rows do not depend on the other + timesteps in the call. The six outputs are views of one ``(T, 6H*3)`` + table, laid out exactly like diffusers' ``view(-1, 6H).chunk(6)``. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, temb, weight, bias): + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") + table = _H3AdaLNProjectionCuda.apply(temb, weight, bias) + return split_adaln_table(table, hidden) + + def forward_table(self, temb, weight, bias) -> torch.Tensor: + """The raw ``(T, 6H*3)`` projection, for consumers that gather it directly.""" + + validate_h3_adaln_projection(temb, weight, bias) + return _H3AdaLNProjectionCuda.apply(temb, weight, bias) + + def forward_fp32(self, temb, weight, bias): + return self.forward(temb, weight, bias) diff --git a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py new file mode 100644 index 000000000..75ec2b8bd --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 three-modality AdaLN projection (RFC #420 row ``adaln_projection_3mod``). + +One ``MiniMaxH3AdaLayerNormModulation`` per transformer block:: + + act = silu(temb).to(weight.dtype) SiLU in FP32, one cast to BF16 + table = act @ W.T + b (T, 6 * H * 3), BF16 + rows = table.view(3 * T, 6 * H) row t * 3 + m, m in {video, text, audio} + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) + +Output channel ``o = m * 6H + c * H + h`` of the projection is chunk ``c``, hidden +index ``h`` of modality ``m``. Each of the six outputs has shape ``(3T, H)``, +and row ``t * 3 + m`` is what ``timestep_indices * 3 + token_tags`` addresses. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_WEIGHT_DTYPES = (torch.bfloat16, torch.float32) + + +def adaln_hidden_size(weight: torch.Tensor) -> int: + rows_per_hidden = H3_ADALN_CHUNKS * H3_MODALITY_NUM + if weight.dim() != 2 or weight.shape[0] % rows_per_hidden != 0: + raise ValueError( + f"AdaLN weight must be (6 * H * 3, D); got {tuple(weight.shape)}, whose first " + f"dim is not a multiple of {rows_per_hidden}" + ) + return weight.shape[0] // rows_per_hidden + + +def validate_h3_adaln_projection(temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.device != temb.device: + raise ValueError(f"{name} is on {tensor.device}, temb is on {temb.device}") + if temb.dtype != torch.float32: + # RFC probe H7: the activation must run at the timestep-embedding + # precision; a BF16 temb means the cast happened before the SiLU. + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in _WEIGHT_DTYPES or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share a dtype in {_WEIGHT_DTYPES}, got {weight.dtype}/{bias.dtype}" + ) + hidden = adaln_hidden_size(weight) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + if bias.shape != (weight.shape[0],): + raise ValueError(f"bias must be ({weight.shape[0]},), got {tuple(bias.shape)}") + return hidden + + +def split_adaln_table(table: torch.Tensor, hidden: int) -> tuple[torch.Tensor, ...]: + """(T, 6H*3) projection output -> six (3T, H) views, diffusers order.""" + + return table.view(-1, H3_ADALN_CHUNKS * hidden).chunk(H3_ADALN_CHUNKS, dim=-1) + + +class NativeH3AdaLNProjectionOp: + """PyTorch reference for the H3 AdaLN projection. + + ``forward`` is the provider path (``F.silu`` in FP32, cast, ``F.linear``). + ``forward_fp32`` is the golden: the SiLU in FP64, rounded once to the + weight dtype at the declared boundary (the cast is model semantics, not + an implementation detail), then the projection in FP64. The six outputs are + returned in FP32 without the final BF16 rounding. Its gradient is FP64 + end to end: the cast is applied straight-through (identity VJP). + """ + + op_class = "reduction" + + def __call__(self, temb, weight, bias): + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + hidden = validate_h3_adaln_projection(temb, weight, bias) + table = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + return split_adaln_table(table, hidden) + + def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + hidden = validate_h3_adaln_projection(temb, weight, bias) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + # The declared cast rounds the value; its VJP is the identity. Written + # straight-through so autograd does not round the FP64 gradient to the + # weight dtype on its way back (``.to(bf16)`` would). + act = act + (act.to(weight.dtype).double() - act).detach() + table = F.linear(act, weight.double(), bias.double()) + return tuple(chunk.float() for chunk in split_adaln_table(table, hidden)) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index e3b3765d5..22bbab259 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -209,6 +209,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ) PYTORCH_H3_TIMESTEP_MLP = "rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp" CUDA_H3_TIMESTEP_MLP = "rl_engine.kernels.ops.cuda.h3.timestep_mlp.H3TimestepMLPCudaOp" + PYTORCH_H3_ADALN_PROJECTION = ( + "rl_engine.kernels.ops.pytorch.h3.adaln_projection.NativeH3AdaLNProjectionOp" + ) + CUDA_H3_ADALN_PROJECTION = ( + "rl_engine.kernels.ops.cuda.h3.adaln_projection.H3AdaLNProjectionCudaOp" + ) def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -627,6 +633,10 @@ def __init__(self): OpBackend.CUDA_H3_TIMESTEP_MLP, OpBackend.PYTORCH_H3_TIMESTEP_MLP, ], + "adaln_projection_3mod": [ + OpBackend.CUDA_H3_ADALN_PROJECTION, + OpBackend.PYTORCH_H3_ADALN_PROJECTION, + ], }, "rocm": { "logp": [ @@ -669,6 +679,7 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -715,6 +726,7 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -742,6 +754,7 @@ def __init__(self): "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index a2c6a80c5..8671ef141 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -30,7 +30,11 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps -from rl_engine.testing.h3_provider import provider_time_embedder, provider_time_proj +from rl_engine.testing.h3_provider import ( + provider_adaln_modulation, + provider_time_embedder, + provider_time_proj, +) REPO_ROOT = Path(__file__).resolve().parents[2] @@ -54,6 +58,11 @@ def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: return [weights[f"time_embedder.linear_{i}.{p}"] for i in (1, 2) for p in ("weight", "bias")] +def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + prefix = "transformer_blocks.0.adaln_proj.linear" + return ctx["weights"][f"{prefix}.weight"], ctx["weights"][f"{prefix}.bias"] + + STAGES: list[Stage] = [ Stage( name="timestep_sinusoid_h3", @@ -75,6 +84,16 @@ def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: golden_atol=1e-4, # reduction / float32 golden_rtol=1e-4, ), + Stage( + name="adaln_projection_3mod", + op_type="adaln_projection_3mod", + candidate=lambda op, ctx, up: op(up, *adaln_params(ctx)), + provider=lambda ctx, up: provider_adaln_modulation(up, *adaln_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *adaln_params(ctx)), + provider_bitwise_isolated=False, # reduction: tensor-core tree differs from cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ), ] diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index 6ae2c5790..f64f7174f 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -46,3 +46,13 @@ def provider_time_embedder( sample = F.linear(features.to(w1.dtype), w1, b1) sample = F.silu(sample) return F.linear(sample, w2, b2) + + +def provider_adaln_modulation( + temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, hidden_size: int = 5376 +) -> tuple[torch.Tensor, ...]: + """``MiniMaxH3AdaLayerNormModulation.forward``: six (3T, H) views.""" + + temb = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + temb = temb.view(-1, 6 * hidden_size) + return temb.chunk(6, dim=-1) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 175547a21..3f84eb520 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -17,7 +17,11 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing.h3_cases import h3_timesteps -from rl_engine.testing.h3_provider import provider_time_embedder, provider_time_proj +from rl_engine.testing.h3_provider import ( + provider_adaln_modulation, + provider_time_embedder, + provider_time_proj, +) from rl_engine.testing.h3_weights import h3_weights_dir, load_h3_conditioning_weights # Keys of a perf case that hold a timed callable, in display order. @@ -178,11 +182,89 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: } +# --------------------------------------------------------------------------- # +# adaln_projection_3mod +# --------------------------------------------------------------------------- # + +ADALN_NAMES = [ + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +] + + +def _adaln_params() -> tuple[torch.Tensor, torch.Tensor]: + weight, bias = h3_params(ADALN_NAMES, [(96768, 2688), (96768,)]) + return weight.bfloat16(), bias.bfloat16() + + +def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_projection_3mod", device="cuda") + weight, bias = _adaln_params() + cases = [] + for num in (1, 2, 3, 4): + temb = torch.randn(num, 2688, device="cuda") + cases.append( + { + "op": "adaln_projection_3mod", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight.numel() * weight.element_size(), + "candidate": lambda temb=temb: op(temb, weight, bias), + "provider": lambda temb=temb: provider_adaln_modulation(temb, weight, bias), + } + ) + return cases + + +def _flat_cat(outputs) -> torch.Tensor: + return torch.cat([out.reshape(-1) for out in outputs]) + + +def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, Any]: + """Fraction of BF16 outputs equal to the correctly rounded FP64 golden, per draw.""" + + op = registry.get_op("adaln_projection_3mod", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["adaln_projection_3mod"][-1] + ) + weight, bias = _adaln_params() + cuda_frac, provider_frac, early_frac = [], [], [] + for seed in range(draws): + g = torch.Generator(device="cuda").manual_seed(seed) + temb = torch.randn(4, 2688, device="cuda", generator=g) * 2 + gold = _flat_cat(golden.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat_cat(golden.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + ours = _flat_cat(op(temb, weight, bias)) + theirs = _flat_cat(provider_adaln_modulation(temb, weight, bias)) + cuda_frac.append(float((ours == gold).float().mean())) + provider_frac.append(float((theirs == gold).float().mean())) + early_frac.append(float((ours == early).float().mean())) + temb = torch.randn(9, 2688, device="cuda") + full = op(temb, weight, bias) + invariant = all( + all( + torch.equal(s, f[3 * i : 3 * i + 3]) + for s, f in zip(op(temb[i : i + 1], weight, bias), full) + ) + for i in range(9) + ) + return { + "draws": draws, + "num_timesteps": 4, + "cuda_correctly_rounded": cuda_frac, + "provider_correctly_rounded": provider_frac, + "early_cast_golden_match": early_frac, + "rows_batch_invariant": invariant, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, + "adaln_projection_3mod": _projection_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, "timestep_mlp_fp32": _mlp_accuracy, + "adaln_projection_3mod": _projection_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index b436e3f2d..c20e157e4 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -226,7 +226,7 @@ def fake_load_backend(backend): ] -H3_OPS = ("timestep_sinusoid_h3", "timestep_mlp_fp32") +H3_OPS = ("timestep_sinusoid_h3", "timestep_mlp_fp32", "adaln_projection_3mod") @pytest.mark.parametrize("op_name", H3_OPS) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 2baa21bb6..6385b4aa4 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -224,9 +224,47 @@ def plot_mlp(report: dict[str, Any]): return fig +def plot_projection(report: dict[str, Any]): + fig, (left, right) = _figure(report, "adaln_projection_3mod (pinned block-0 AdaLN weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _style( + right, + f"Outputs equal to the correctly rounded FP64 golden, {acc['draws']} draws", + "% of BF16 outputs", + ) + series = [ + ("RL-Kernel CUDA", RL_KERNEL, [100 * v for v in acc["cuda_correctly_rounded"]]), + ("diffusers (cuBLAS)", PROVIDER, [100 * v for v in acc["provider_correctly_rounded"]]), + ] + for i, (name, color, values) in enumerate(series): + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax_vals = sorted(values) + median = ax_vals[len(ax_vals) // 2] + right.scatter(jitter, values, s=36, color=color, alpha=0.8, linewidths=0, zorder=3) + right.hlines(median, i - 0.25, i + 0.25, color=color, lw=2.5, zorder=4) + right.annotate( + f"median {median:.3f}%", (i + 0.28, median), fontsize=8, color=INK, va="center" + ) + right.set_xticks(range(len(series)), [name for name, _, _ in series]) + right.set_xlim(-0.6, len(series) - 0.1) + early = sorted(acc["early_cast_golden_match"])[len(acc["early_cast_golden_match"]) // 2] + right.text( + 0.98, + 0.62, + f"an early BF16 cast (probe H7)\nwould match only {100 * early:.0f}%", + transform=right.transAxes, + fontsize=8, + color=MUTED, + ha="right", + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, + "adaln_projection_3mod": plot_projection, } diff --git a/tests/h3/test_h3_adaln_projection.py b/tests/h3/test_h3_adaln_projection.py new file mode 100644 index 000000000..c9aac92d1 --- /dev/null +++ b/tests/h3/test_h3_adaln_projection.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_projection_3mod``: SiLU + 2688 -> 96768 projection, 6 x 3 table. + +* layout: the six outputs and the three modality rows per timestep land + exactly where diffusers puts them (RFC probe H4); +* mixed precision: SiLU in FP32, one cast to BF16; an early cast (probe H7) + is rejected at the API and is measurably different numerically; +* accuracy: forward and gradients against the declared-cast golden on the + pinned block-0 weights; +* invariance and determinism: a timestep's 18 modulation rows do not depend + on the other timesteps, and repeated runs are bitwise equal. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.kernels.ops.pytorch.h3.adaln_projection import NativeH3AdaLNProjectionOp +from rl_engine.testing.h3_provider import provider_adaln_modulation + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3AdaLNProjectionCudaOp() + + +def _temb(num, dim=2688, seed=0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(num, dim, generator=g) * 2).to(device) + + +def _synthetic(hidden, dim, dtype=torch.bfloat16, seed=0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).to(dtype) + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype) + return weight.to(device), bias.to(device) + + +def _flat(outputs): + return torch.cat([out.reshape(-1) for out in outputs]) + + +class TestReference: + def test_provider_layout_matches_diffusers(self): + temb = _temb(2, dim=16, device="cpu") + weight, bias = _synthetic(8, 16, device="cpu") + ours = NativeH3AdaLNProjectionOp().forward(temb, weight, bias) + theirs = provider_adaln_modulation(temb, weight, bias, hidden_size=8) + assert len(ours) == 6 + for a, b in zip(ours, theirs): + assert a.shape == (6, 8) and torch.equal(a, b) + + def test_rejects_bf16_temb(self): + # RFC probe H7: casting temb before the SiLU. + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(TypeError, match="float32"): + NativeH3AdaLNProjectionOp().forward(_temb(2, 16, device="cpu").bfloat16(), weight, bias) + + def test_rejects_bad_shapes(self): + weight, bias = _synthetic(8, 16, device="cpu") + op = NativeH3AdaLNProjectionOp() + with pytest.raises(ValueError): # not a multiple of 6 * 3 rows + op.forward(_temb(2, 16, device="cpu"), weight[:-1], bias[:-1]) + with pytest.raises(ValueError): + op.forward(_temb(2, 15, device="cpu"), weight, bias) + with pytest.raises(ValueError): + op.forward(_temb(0, 16, device="cpu"), weight, bias) + with pytest.raises(TypeError): + op.forward(_temb(2, 16, device="cpu"), weight, bias.float()) + + +@requires_cuda +class TestLayout: + def test_channel_identity_lands_in_its_chunk_and_modality_row(self): + """Probe H4: bias o = m*6H + c*H + h tagged (m, c); zero weight passes it through.""" + + hidden, dim, num = 16, 32, 3 + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + o = torch.arange(n_out) + modality, chunk = o // (H3_ADALN_CHUNKS * hidden), (o // hidden) % H3_ADALN_CHUNKS + bias = (modality * 10 + chunk).to(torch.bfloat16).cuda() + weight = torch.zeros(n_out, dim, dtype=torch.bfloat16, device="cuda") + outputs = _cuda_op()(_temb(num, dim), weight, bias) + assert len(outputs) == 6 + for c, out in enumerate(outputs): + assert out.shape == (num * 3, hidden) + for t in range(num): + for m in range(3): + expected = torch.full((hidden,), float(m * 10 + c), device="cuda") + assert torch.equal(out[t * 3 + m].float(), expected), (c, t, m) + + def test_outputs_are_views_of_one_table_like_diffusers(self): + weight, bias = _synthetic(16, 32) + outputs = _cuda_op()(_temb(2, 32), weight, bias) + base = outputs[0].untyped_storage().data_ptr() + assert all(out.untyped_storage().data_ptr() == base for out in outputs) + assert outputs[1].data_ptr() - outputs[0].data_ptr() == 16 * outputs[0].element_size() + + +@pytest.fixture(scope="module") +def block0(h3_weights_cpu): + return h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4]) + def test_forward_is_correctly_rounded_declared_golden(self, block0, num): + weight, bias = block0 + temb = _temb(num, seed=num) + outputs = _cuda_op()(temb, weight, bias) + golden = NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias) + for out, gold in zip(outputs, golden): + assert out.dtype == torch.bfloat16 and out.shape == (3 * num, 5376) + # tolerance_contract.json forward_accuracy / reduction / bfloat16. + torch.testing.assert_close(out.float(), gold, atol=5e-2, rtol=2e-2) + # Beyond the contract: almost every element is the correctly rounded + # golden value; the rest are 1-ULP ties of the FP32 accumulation. + match = (_flat(outputs) == _flat(golden).bfloat16()).float().mean().item() + assert match > 0.999 + + def test_early_cast_is_numerically_detectable(self, block0): + """Probe H7: a BF16 cast before the SiLU moves about half the outputs.""" + + weight, bias = block0 + temb = _temb(3, seed=7) + op = NativeH3AdaLNProjectionOp() + ours = _flat(_cuda_op()(temb, weight, bias)) + declared = _flat(op.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat(op.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + assert (ours == declared).float().mean() > 0.999 + assert (ours == early).float().mean() < 0.8 + + def test_rows_are_batch_and_position_invariant(self, block0): + weight, bias = block0 + op = _cuda_op() + temb = _temb(5, seed=3) + full = op(temb, weight, bias) + for i in range(5): + single = op(temb[i : i + 1], weight, bias) + for s, f in zip(single, full): + assert torch.equal(s, f[3 * i : 3 * i + 3]) + swapped = op(temb.flip(0), weight, bias) + for s, f in zip(swapped, full): + assert torch.equal(s.view(5, 3, -1).flip(0), f.view(5, 3, -1)) + + def test_backward_against_golden(self, block0): + weight, bias = block0 + temb = _temb(2, seed=5) + g = torch.Generator(device="cuda").manual_seed(0) + grads = [torch.randn(6, 5376, device="cuda", generator=g).bfloat16() for _ in range(6)] + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(_cuda_op()(*leaves)), grads) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + act = act + (act.to(torch.bfloat16).double() - act).detach() # declared cast, identity VJP + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376).chunk(6, dim=-1) + torch.autograd.backward(list(table), [gr.double() for gr in grads]) + # tolerance_contract.json gradient_accuracy / reduction: FP32 temb grad + # and BF16 weight/bias grads (rounded once from an FP32 accumulation). + torch.testing.assert_close(leaves[0].grad.double(), ref[0].grad, atol=1e-4, rtol=1e-4) + for leaf, r in zip(leaves[1:], ref[1:]): + torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) + assert (leaf.grad == r.grad.to(torch.bfloat16)).float().mean() > 0.999 + + +@requires_cuda +class TestCudaSynthetic: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) + @pytest.mark.parametrize("hidden, dim", [(24, 16), (40, 72), (8, 2688)]) + def test_other_sizes(self, dtype, hidden, dim): + weight, bias = _synthetic(hidden, dim, dtype=dtype, seed=hidden) + temb = _temb(3, dim, seed=hidden) + tol = dict(atol=5e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-4, rtol=1e-4) + for out, gold in zip( + _cuda_op()(temb, weight, bias), + NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias), + ): + torch.testing.assert_close(out.float(), gold, **tol) + + def test_repeat_forward_backward_bitwise(self): + weight, bias = _synthetic(64, 2688, seed=9) + temb = _temb(3, seed=9) + grads = [torch.randn(9, 64, device="cuda").bfloat16() for _ in range(6)] + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + outs = _cuda_op()(*leaves) + torch.autograd.backward(list(outs), grads) + runs.append([*[o.detach() for o in outs], *[leaf.grad for leaf in leaves]]) + for later in runs[1:]: + for a, b in zip(runs[0], later): + assert torch.equal(a, b) + + def test_temb_grad_rows_are_batch_invariant(self): + weight, bias = _synthetic(64, 2688, seed=11) + temb = _temb(4, seed=11) + grads = [torch.randn(12, 64, device="cuda").bfloat16() for _ in range(6)] + + def d_temb(rows): + leaf = temb[rows].detach().clone().requires_grad_(True) + outs = _cuda_op()(leaf, weight, bias) + torch.autograd.backward( + list(outs), [g.view(4, 3, -1)[rows].reshape(-1, 64) for g in grads] + ) + return leaf.grad + + full = d_temb(slice(0, 4)) + for i in range(4): + assert torch.equal(d_temb(slice(i, i + 1))[0], full[i]) + + @pytest.mark.parametrize("num", [8, 9, 17]) + def test_rows_invariant_across_tensor_core_row_tiles(self, num): + """The BF16 path computes rows in tiles of 8; a row's bits must not depend on its tile.""" + + weight, bias = _synthetic(64, 2688, seed=13) + temb = _temb(num, seed=num) + full = _cuda_op()(temb, weight, bias) + for i in (0, 7, num - 1): + single = _cuda_op()(temb[i : i + 1], weight, bias) + for s_out, f_out in zip(single, full): + assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) + + def test_rejects_cpu(self): + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(ValueError): + _cuda_op()(_temb(2, 16, device="cpu"), weight, bias) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_projection_3mod", device="cuda") + assert type(op).__name__ == "H3AdaLNProjectionCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 6dde75c53..fdef637f7 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -51,6 +51,7 @@ def test_all_ws1_single_ops_are_registered(): "prefix_shared_attention", "timestep_sinusoid_h3", "timestep_mlp_fp32", + "adaln_projection_3mod", } <= names From 404e766bb1d5131776b03a44edc12bd93c4878a6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:36:46 +0000 Subject: [PATCH 06/89] [WS1][MiniMax-H3] adaln_projection_3mod: B200 evidence report and figure Produced by scripts/h3_evidence.py at d06e120 on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 8e91cf9bf..0e27a5761 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -108,19 +108,29 @@ B200, pinned BF16 weights (520 MB per block, streamed once per call): | T | CUDA op | provider (SiLU + cast + cuBLAS) | CUDA GEMV kernel | cuBLAS kernel | | --- | --- | --- | --- | --- | -| 1 | 102 µs | 112 µs | 78.3 µs (6.6 TB/s) | 93 µs | -| 2 | 102 µs | 101 µs | 78.4 µs | 79 µs | -| 4 | 102 µs | 101 µs | 79.0 µs | 79 µs | +| 1 | 105 µs | 115 µs | 78.3 µs (6.6 TB/s) | 93 µs | +| 2 | 105 µs | 103 µs | 78.4 µs | 79 µs | +| 4 | 106 µs | 104 µs | 79.0 µs | 79 µs | The kernel takes the same time for every T <= 8, because it always computes 8 rows. It matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 24 µs of the op time is -the SiLU and cast kernels plus the Python wrapper. +the SiLU and cast kernels plus the Python wrapper. With those included, the op is within +about 2% of the provider for T >= 2. An earlier FMA-only BF16 kernel ran at 81, 98 and 166 µs for T = 1, 2 and 4. Its time grew with T because every extra row adds FMAs per weight element. Packed FP32 FMA (`__ffma2_rn`) gave identical bits but no speedup. Moving to tensor cores removed the dependence on T. +## Evidence + +![adaln_projection_3mod on B200: latency and correctly rounded outputs](../usage/evidence/h3-adaln-projection-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-adaln-projection-b200/report.json), +written by `scripts/h3_evidence.py` from a clean tree at commit `d06e120`. The report also +records that a timestep's 18 modulation rows are bitwise identical whether it runs alone +or in a batch of 9. + ## Tests ```bash From d275a06f96169f9d4b3b79c7084182b95668031d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:40:27 +0000 Subject: [PATCH 07/89] [WS1][CUDA][MiniMax-H3] adaln_row_gather: fused six-way AdaLN row gather RFC #420 step 3, row `adaln_row_gather` (stacked on adaln_projection_3mod). - Forward: one launch gathers all six (S, H) modulation tensors by timestep_index * 3 + token_tag with 16-byte copies; bitwise equal to diffusers' six index_select calls. Out-of-range tags or timesteps fail closed (RFC probes H2/H3), never clamped. - Backward: deterministic FP32 segmented sum (stable sort by row, fixed 256-position tiles, ordered fold, one cast; no atomics) instead of index_select's atomic BF16 scatter-add. The PyTorch reference gets an equally deterministic backward. - H3AdaLNModulationCudaOp: projection + gather as one autograd node, bitwise-equal forward, FP32 table gradient into the projection backward. - Chain replay gains the gather stage and a backward mode (separate, fused and diffusers chains vs an FP64 golden); the end-to-end test now covers the whole chain forward and backward. - Report/figure hooks, registry, gtest (reduction class: the VJP is a segmented sum), doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 243 ++++++++++++++++++ csrc/ops.cpp | 15 ++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-adaln-row-gather.md | 141 ++++++++++ rl_engine/_C.pyi | 15 ++ rl_engine/kernels/gtest/operator_inputs.py | 20 ++ rl_engine/kernels/gtest/operator_specs.py | 15 ++ .../kernels/ops/cuda/h3/adaln_modulation.py | 133 ++++++++++ .../kernels/ops/cuda/h3/adaln_row_gather.py | 113 ++++++++ .../ops/pytorch/h3/adaln_row_gather.py | 141 ++++++++++ rl_engine/kernels/registry.py | 13 + rl_engine/testing/h3_chain.py | 99 +++++++ rl_engine/testing/h3_provider.py | 9 + rl_engine/testing/h3_report.py | 95 ++++++- rl_engine/tests/test_dispatch.py | 7 +- scripts/h3_chain_replay.py | 24 ++ scripts/plot_h3_evidence.py | 99 ++++++- setup.py | 1 + tests/h3/test_h3_adaln_modulation.py | 117 +++++++++ tests/h3/test_h3_adaln_row_gather.py | 207 +++++++++++++++ tests/h3/test_h3_conditioning_e2e.py | 19 ++ tests/test_ws1_gtest_gpu.py | 1 + 23 files changed, 1520 insertions(+), 9 deletions(-) create mode 100644 csrc/cuda/h3/adaln_row_gather.cu create mode 100644 docs/operators/h3-adaln-row-gather.md create mode 100644 rl_engine/kernels/ops/cuda/h3/adaln_modulation.py create mode 100644 rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py create mode 100644 tests/h3/test_h3_adaln_modulation.py create mode 100644 tests/h3/test_h3_adaln_row_gather.py diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu new file mode 100644 index 000000000..192faa187 --- /dev/null +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -0,0 +1,243 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 AdaLN row gather (RFC #420 `adaln_row_gather`). +// +// rows: (3T, C * H) modulation rows, row r = timestep * 3 + modality and +// column block c = shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, +// gate_mlp. For every packed-sequence position s: +// +// r[s] = timestep_indices[s] * modality_num + token_tags[s] +// out[c, s, :] = rows[r[s], c * H : (c + 1) * H] +// +// Forward is a pure copy: bitwise equal to C index_select calls, one launch, +// independent of S, packing order and repeats. +// +// Backward d_rows[r, j] = sum_{s : r[s] = r} grad[c(j), s, h(j)] is the one +// cross-row reduction. Positions are visited in a fixed order: the caller +// sorts them by (r, s) with a stable sort and cuts every segment into tiles of +// at most kTile positions. Each tile is an ascending FP32 chain from 0; the +// tiles of a segment are then left-folded in ascending order and cast once. +// No atomics: the result is a function of (r[s], grad) only. + +#include +#include +#include +#include +#include + +#include + +namespace { + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_float(__half v) { return __half2float(v); } +__device__ __forceinline__ float to_float(double v) { return static_cast(v); } + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ float from_float(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); +} +template <> +__device__ __forceinline__ __half from_float<__half>(float v) { return __float2half(v); } +template <> +__device__ __forceinline__ double from_float(float v) { return v; } + +template +struct CudaType { + using type = T; +}; +template <> +struct CudaType { + using type = __nv_bfloat16; +}; +template <> +struct CudaType { + using type = __half; +}; + +template +__device__ __forceinline__ int64_t adaln_row(const idx_t* ti, const idx_t* tags, int64_t s, + int64_t modality_num) { + return static_cast(ti[s]) * modality_num + static_cast(tags[s]); +} + +// One block per (s, c); 16-byte copies when every row start is 16-byte aligned. +template +__global__ void adaln_row_gather_kernel(const T* __restrict__ rows, int64_t row_stride, + const idx_t* __restrict__ ti, + const idx_t* __restrict__ tags, T* __restrict__ out, + int64_t seq, int64_t hidden, int64_t chunks, + int64_t modality_num) { + for (int64_t s = blockIdx.x; s < seq; s += gridDim.x) { + const int64_t c = blockIdx.y; + const int64_t r = adaln_row(ti, tags, s, modality_num); + const T* src = rows + r * row_stride + c * hidden; + T* dst = out + (c * seq + s) * hidden; + if constexpr (kVector) { + constexpr int kVec = 16 / sizeof(T); + const uint4* src4 = reinterpret_cast(src); + uint4* dst4 = reinterpret_cast(dst); + for (int64_t i = threadIdx.x; i < hidden / kVec; i += blockDim.x) dst4[i] = src4[i]; + } else { + for (int64_t i = threadIdx.x; i < hidden; i += blockDim.x) dst[i] = src[i]; + } + } +} + +// partial[tile, j] = sum over sorted positions p in tile (ascending) of grad[c, s_p, h] +template +__global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, + const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int64_t seq, + int64_t hidden, int64_t width) { + const int64_t tile = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= width) return; + const int64_t c = j / hidden; + const int64_t h = j - c * hidden; + const g_t* g = grad + c * seq * hidden + h; + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + acc += to_float(g[sorted_pos[p] * hidden]); + } + partial[tile * width + j] = acc; +} + +template +__global__ void adaln_row_gather_fold_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + out_t* __restrict__ out, int64_t width) { + const int64_t r = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= width) return; + float acc = 0.0f; + for (int64_t tile = seg_first_tile[r]; tile < seg_first_tile[r + 1]; ++tile) { + acc += partial[tile * width + j]; + } + out[r * width + j] = from_float(acc); +} + +void check_index(const torch::Tensor& t, const char* name, int64_t seq, + const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device(), name, " must be on ", like.device()); + TORCH_CHECK(t.dim() == 1 && t.size(0) == seq, name, " must be (", seq, ",)"); + TORCH_CHECK(t.scalar_type() == at::kLong || t.scalar_type() == at::kInt, name, + " must be int64 or int32"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +// rows (R, C * H) with unit column stride -> out (C, S, H) contiguous. +torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor timestep_indices, + torch::Tensor token_tags, int64_t chunks, + int64_t modality_num) { + TORCH_CHECK(rows.is_cuda() && rows.dim() == 2 && rows.stride(1) == 1, + "rows must be a 2-D CUDA tensor with unit column stride"); + TORCH_CHECK(chunks > 0 && rows.size(1) % chunks == 0, "rows width ", rows.size(1), + " is not a multiple of ", chunks); + TORCH_CHECK(modality_num > 0 && rows.size(0) % modality_num == 0, "rows count ", + rows.size(0), " is not a multiple of ", modality_num); + const int64_t seq = timestep_indices.numel(); + TORCH_CHECK(seq > 0, "the packed sequence must not be empty"); + check_index(timestep_indices, "timestep_indices", seq, rows); + check_index(token_tags, "token_tags", seq, rows); + TORCH_CHECK(timestep_indices.scalar_type() == token_tags.scalar_type(), + "timestep_indices and token_tags must share an integer dtype"); + const int64_t hidden = rows.size(1) / chunks; + const c10::cuda::CUDAGuard device_guard(rows.device()); + auto out = torch::empty({chunks, seq, hidden}, rows.options()); + auto stream = at::cuda::getCurrentCUDAStream(); + const dim3 grid(static_cast(std::min(seq, 1 << 20)), + static_cast(chunks)); + const int threads = 128; + const int64_t elem = rows.element_size(); + const bool vector = (hidden * elem) % 16 == 0 && (rows.stride(0) * elem) % 16 == 0 && + reinterpret_cast(rows.data_ptr()) % 16 == 0; + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, rows.scalar_type(), "h3_adaln_row_gather_forward", [&] { + using T = typename CudaType::type; + const T* src = reinterpret_cast(rows.data_ptr()); + T* dst = reinterpret_cast(out.data_ptr()); + auto launch = [&](auto idx_tag) { + using idx_t = decltype(idx_tag); + const idx_t* ti = timestep_indices.data_ptr(); + const idx_t* tags = token_tags.data_ptr(); + if (vector) { + adaln_row_gather_kernel<<>>( + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num); + } else { + adaln_row_gather_kernel<<>>( + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num); + } + }; + if (timestep_indices.scalar_type() == at::kLong) { + launch(int64_t{0}); + } else { + launch(int32_t{0}); + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// grad (C, S, H) contiguous; sorted_pos (S,) positions sorted stably by row; +// tiles [tile_begin, tile_end) in sorted order; seg_first_tile (R + 1,). +torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sorted_pos, + torch::Tensor tile_begin, torch::Tensor tile_end, + torch::Tensor seg_first_tile, + c10::ScalarType out_dtype) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 3 && grad.is_contiguous(), + "grad must be a contiguous (C, S, H) CUDA tensor"); + for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { + TORCH_CHECK(t->is_cuda() && t->dim() == 1 && t->is_contiguous() && + t->scalar_type() == at::kLong, + "tile metadata must be contiguous 1-D int64 CUDA tensors"); + } + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); + const int64_t chunks = grad.size(0); + const int64_t seq = grad.size(1); + const int64_t hidden = grad.size(2); + const int64_t width = chunks * hidden; + const int64_t num_rows = seg_first_tile.numel() - 1; + const int64_t tiles = tile_begin.numel(); + TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries"); + const c10::cuda::CUDAGuard device_guard(grad.device()); + auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype)); + auto partial = torch::empty({std::max(tiles, 1), width}, + grad.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((width + threads - 1) / threads); + if (tiles > 0) { + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, grad.scalar_type(), "h3_adaln_row_gather_partial", [&] { + using G = typename CudaType::type; + adaln_row_gather_partial_kernel + <<(tiles), col_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), + sorted_pos.data_ptr(), tile_begin.data_ptr(), + tile_end.data_ptr(), partial.data_ptr(), seq, hidden, width); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, out_dtype, "h3_adaln_row_gather_fold", [&] { + using O = typename CudaType::type; + adaln_row_gather_fold_kernel + <<(num_rows), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(out.data_ptr()), width); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index f245f8e15..27159d5e9 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -510,6 +510,13 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, c10::ScalarType w_dtype, bool with_bias); +torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor timestep_indices, + torch::Tensor token_tags, int64_t chunks, + int64_t modality_num); +torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sorted_pos, + torch::Tensor tile_begin, torch::Tensor tile_end, + torch::Tensor seg_first_tile, + c10::ScalarType out_dtype); #endif // PyBind11 Module Registration @@ -783,5 +790,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight, "Deterministic dW/dbias as ascending-row FP32 folds", py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true); + m.def("h3_adaln_row_gather_forward", &h3_adaln_row_gather_forward, + "Fused six-way AdaLN row gather by timestep_index * 3 + token_tag (pure copy)", + py::arg("rows"), py::arg("timestep_indices"), py::arg("token_tags"), + py::arg("chunks") = 6, py::arg("modality_num") = 3); + m.def("h3_adaln_row_gather_backward", &h3_adaln_row_gather_backward, + "Deterministic segmented sum (sorted tiles folded in order) for the row gather", + py::arg("grad"), py::arg("sorted_pos"), py::arg("tile_begin"), py::arg("tile_end"), + py::arg("seg_first_tile"), py::arg("out_dtype")); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index b5a3c718b..ff55306a7 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -33,6 +33,7 @@ nav: - operators/h3-timestep-sinusoid.md - operators/h3-timestep-mlp.md - operators/h3-adaln-projection.md + - operators/h3-adaln-row-gather.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index d713f52d1..33e405acd 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -35,4 +35,5 @@ Every operator page should include: - [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) - [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) - [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) +- [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md new file mode 100644 index 000000000..54636dd09 --- /dev/null +++ b/docs/operators/h3-adaln-row-gather.md @@ -0,0 +1,141 @@ +# MiniMax-H3 AdaLN Row Gather + +## Summary + +`adaln_row_gather` selects every packed-sequence position's six modulation vectors from +one block's AdaLN table (RFC #420, WS1 step 3): + +```text +adaln_indices = timestep_indices * 3 + token_tags (S,) 0 video, 1 text, 2 audio +shift_msa[adaln_indices], scale_msa[...], gate_msa[...], +shift_mlp[...], scale_mlp[...], gate_mlp[...] six (S, 5376) +``` + +`rows` is the `(3T, 6H)` view of the [AdaLN projection](h3-adaln-projection.md) table, and +its six column blocks are the six modulation tensors. `timestep_indices` and `token_tags` +are semantic inputs (RFC #420 §4): they are validated, never clamped. + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("adaln_row_gather", device="cuda") +outs = op(table.view(-1, 6 * 5376), timestep_indices, token_tags) # six (S, 5376) +outs = op.gather_chunks((shift_msa, ..., gate_mlp), timestep_indices, token_tags) # drop-in +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.kernels.ops.cuda.h3.adaln_row_gather.H3AdaLNRowGatherCudaOp` | `rl_engine._C.h3_adaln_row_gather_{forward,backward}` | Forward bitwise equal to `index_select`; deterministic backward | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.adaln_row_gather.NativeH3AdaLNRowGatherOp` | n/a | `index_select` forward; deterministic per-row FP32 backward | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +The raw diffusers path, including its atomic backward, is kept as +`rl_engine.testing.h3_provider.provider_adaln_row_gather` for evidence. + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `rows` | `(3T, 6H)` | bf16 / fp16 / fp32 | Unit column stride. Any row stride (16-byte copies when aligned) | +| `timestep_indices` | `(S,)`, `S >= 1` | int64 or int32 | In `[0, T)` | +| `token_tags` | `(S,)` | same as `timestep_indices` | In `{0, 1, 2}` | +| outputs | six `(S, H)` | `rows.dtype` | Contiguous; slices of one `(6, S, H)` buffer | + +These inputs fail closed: out-of-range tags (a fourth modality, probe H2) or timesteps +(an offset past `T`, probe H3), length or dtype mismatches, empty sequences, and tables +that are not 3 rows per timestep or 6 blocks wide. The range check is one host +read-back; pass `check_range=False` when the indices are already validated. + +## Numerics + +- **Forward** is a byte copy: one launch, one block per `(position, chunk)`, using + 16-byte vector copies. It is bitwise equal to diffusers' six `index_select` calls for + every packing, length and index dtype tested (S up to 131072). +- **Backward**: `d_rows[r] = sum over {s : r[s] = r} of grad[s]` is the op's only + reduction. Positions are sorted stably by `(row, position)` and cut into tiles of 256. + Each tile is an ascending FP32 chain from 0, and a row's tiles are left-folded in order + and rounded once. There are no atomics, so the result depends only on the indices and + the gradient. `index_select`'s own backward is an atomic scatter-add in the gradient + dtype (BF16), which is neither deterministic nor accurate. +- **Tolerance class.** The VJP is a segmented reduction, so gtest judges the op as + `reduction`. The forward is asserted bitwise in `tests/h3`. + +Measured on a B200 (torch 2.13.0+cu130): + +| Comparison (T = 3, S = 4097, H = 5376, BF16) | CUDA | provider (`index_select`) | +| --- | --- | --- | +| Forward vs `index_select` | bitwise equal | | +| `d_rows` repeat-bitwise | yes | no | +| `d_rows` correctly rounded from the FP64 sum | 99.99999% | 7.8% | +| `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | + +### Fused modulation: `H3AdaLNModulationCudaOp` + +`rl_engine.kernels.ops.cuda.h3.adaln_modulation.H3AdaLNModulationCudaOp(temb, weight, bias, +timestep_indices, token_tags)` runs [`adaln_projection_3mod`](h3-adaln-projection.md) and +this gather as one autograd node. Its forward uses the same kernels and is bitwise equal +to calling the two ops in sequence. + +The difference is in the backward. When the ops are called separately, autograd returns +the gather's FP32 segment sum to the projection in the table's BF16 dtype, which rounds it +once more. The fused node feeds the FP32 table gradient straight into the projection +backward. Use it wherever both ops run back to back, as they do in an H3 block. + +### Whole-chain backward + +`scripts/h3_chain_replay.py --backward` runs on the pinned weights for +T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. It checks six parameter gradients in each +of the 16 cases (the time embedder's `linear_{1,2}.{weight,bias}` and block 0's AdaLN +weight and bias) against an FP64 golden in which every declared cast is straight-through: + +| S >= 257 (72 gradients) | RL-Kernel, separate ops | RL-Kernel, fused | diffusers | +| --- | --- | --- | --- | +| repeat-bitwise | 72 / 72 | 72 / 72 | 0 / 72 | +| time embedder: max abs error / golden max | 5.3e-4 to 2.4e-3 | 6.5e-7 to 4.9e-6 | 5.6e-3 to 1.2e-1 | +| AdaLN weight/bias: correctly rounded BF16, worst case | 59.8% | 99.57% | 1.6% | + +At S = 3 all three chains are deterministic. Both RL-Kernel variants reach 1.5e-6 on the +time embedder, against 3.2e-3 for diffusers. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op adaln_row_gather +``` + +B200, T = 3, H = 5376, BF16. "write BW" counts the 6·S·H outputs; the table stays in L2. + +| S | CUDA forward | provider forward | CUDA backward | provider backward | +| --- | --- | --- | --- | --- | +| 4097 | 62 µs | 92 µs | 1.10 ms | 1.44 ms | +| 32768 | 315 µs | 587 µs | 2.10 ms | 8.66 ms | +| 131072 | 1.24 ms (6.8 TB/s write) | 2.37 ms | 6.84 ms | 34.9 ms | + +The backward's peak memory is about twice the provider's, because autograd stacks the +six output gradients and the per-tile FP32 partials are materialised. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_adaln_row_gather.py tests/h3/test_h3_adaln_modulation.py -v +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # whole chain, forward + backward +python scripts/check_operator.py --op adaln_row_gather --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op adaln_row_gather \ + --out docs/usage/evidence/h3-adaln-row-gather-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-adaln-row-gather-b200/report.json +python scripts/h3_chain_replay.py --timesteps 1,2,3,4 --seq-lens 3,257,4097,32768 \ + --backward --out docs/usage/evidence/h3-adaln-row-gather-b200/chain_replay.json +``` + +## Known Limitations + +- Like diffusers, the op materialises six `(S, H)` tensors, which is 8 GB at + S = 131072. Fusing the gather into the norm/modulate and gate-residual consumers + (`h3_rmsnorm`, `adaln_gate_residual`) avoids this. Those are separate rows. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index b8cf1d58a..144b06e27 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -341,3 +341,18 @@ def h3_det_linear_backward_input( def h3_det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True ) -> list[torch.Tensor]: ... +def h3_adaln_row_gather_forward( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + chunks: int = 6, + modality_num: int = 3, +) -> torch.Tensor: ... +def h3_adaln_row_gather_backward( + grad: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 255a48c0c..ef77fe7bc 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -49,6 +49,7 @@ def make_operator_inputs( "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, + "adaln_row_gather": _make_adaln_row_gather_inputs, } try: return builders[op_name](args, dtype, device) @@ -84,6 +85,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: f"x{H3_TIME_EMBED}", "adaln_projection_3mod": f"{_h3_num_timesteps(args)}x{H3_TIME_EMBED}" f"x{6 * 3 * _h3_hidden(args)}", + "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" + f"->{batch * seq}", } try: return names[op_name] @@ -411,6 +414,23 @@ def _make_adaln_projection_3mod_inputs( } +def _make_adaln_row_gather_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # T distinct timesteps (--batch) and a packed sequence of batch * seq rows + # mixing all three modalities. + batch, seq = _batch_seq(args) + num_timesteps = _h3_num_timesteps(args) + generator = _generator(args, device, offset=29) + packed = batch * seq + token_tags = torch.randint(0, 3, (packed,), generator=generator, device=device) + timestep_indices = torch.randint( + 0, num_timesteps, (packed,), generator=generator, device=device + ) + rows = _floating_tensor((3 * num_timesteps, 6 * _h3_hidden(args)), args, dtype, device, 0) + return {"rows": rows, "timestep_indices": timestep_indices, "token_tags": token_tags} + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 91d877b62..11ce047db 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -285,6 +285,21 @@ def _load_object(path: str) -> Any: }, grad_input_names=("temb", "weight", "bias"), ), + # The gather's forward is a copy (asserted bitwise in tests/h3), but its VJP + # sums every packed position of a table row, so it is judged as a reduction. + "adaln_row_gather": OperatorSpec( + name="adaln_row_gather", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.h3.adaln_row_gather.NativeH3AdaLNRowGatherOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.kernels.ops.pytorch.h3.adaln_row_gather.NativeH3AdaLNRowGatherOp" + ), + "cuda": "rl_engine.kernels.ops.cuda.h3.adaln_row_gather.H3AdaLNRowGatherCudaOp", + }, + grad_input_names=("rows",), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_modulation.py b/rl_engine/kernels/ops/cuda/h3/adaln_modulation.py new file mode 100644 index 000000000..ea13bd232 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/adaln_modulation.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused H3 AdaLN modulation: ``adaln_projection_3mod`` followed by ``adaln_row_gather``. + +The forward is exactly the two RFC #420 ops back to back (same kernels, same +bits). The difference is the backward. With two separate ops, autograd hands +the gather's table gradient to the projection in the table's dtype (BF16), +which rounds a long FP32 segment sum once more before it reaches the +projection weights and ``temb``. Here the segment sum stays FP32 into the +projection backward, so the parameter gradients carry no extra BF16 rounding. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import ( + _segment_tiles, + adaln_row_gather_available, +) +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.kernels.ops.pytorch.h3.adaln_projection import validate_h3_adaln_projection +from rl_engine.kernels.ops.pytorch.h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_adaln_row_gather_forward" +) +BACKWARD_IMPL = "fp32_segment_sum_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_adaln_row_gather_backward[fp32]" + "+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.kernels.ops.cuda.h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNModulationCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias, timestep_indices, token_tags, hidden): + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias) + rows = table.view(-1, H3_ADALN_CHUNKS * hidden) + packed = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(temb, act, weight, timestep_indices, token_tags) + ctx.table_shape = table.shape + return packed + + @staticmethod + def backward(ctx, grad_packed): + temb, act, weight, timestep_indices, token_tags = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad[:3] + num_timesteps, width = ctx.table_shape + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_timesteps * H3_MODALITY_NUM) + d_rows = _C.h3_adaln_row_gather_backward(grad_packed.contiguous(), *tiles, torch.float32) + d_table = d_rows.view(num_timesteps, width) + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + if need_temb: + d_temb = silu_backward_fp32( + det_linear_backward_input(d_table, weight, torch.float32), temb + ) + record_backward( + "adaln_modulation_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return ( + d_temb, + dw if need_weight else None, + db if need_bias else None, + None, + None, + None, + ) + + +class H3AdaLNModulationCudaOp: + """``adaln_projection_3mod`` + ``adaln_row_gather`` with an FP32 table gradient. + + Returns the six ``(S, H)`` modulation tensors for every packed position, + bitwise equal to running the two ops separately. + """ + + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and adaln_row_gather_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_adaln_row_gather_* ops") + + def __call__(self, temb, weight, bias, timestep_indices, token_tags, *, check_range=True): + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNModulationCudaOp needs CUDA tensors") + _validate_indices(timestep_indices, token_tags, temb.shape[0], check_range) + packed = _H3AdaLNModulationCuda.apply( + temb, + weight, + bias, + timestep_indices.contiguous(), + token_tags.contiguous(), + hidden, + ) + return tuple(packed.unbind(0)) + + +def _validate_indices(timestep_indices, token_tags, num_timesteps, check_range): + # Reuse the gather's checks against a stand-in of the right row count. + stand_in = torch.empty( + (num_timesteps * H3_MODALITY_NUM, H3_ADALN_CHUNKS), device=timestep_indices.device + ) + validate_h3_adaln_row_gather(stand_in, timestep_indices, token_tags, check_range=check_range) diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py b/rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py new file mode 100644 index 000000000..3569d5010 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 AdaLN row gather (RFC #420 ``adaln_row_gather``).""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.kernels.ops.pytorch.h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = "rl_engine._C.h3_adaln_row_gather_forward" +BACKWARD_IMPL = "stable_sorted_segment_tiles_fp32_ordered_fold" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_adaln_row_gather_backward" +# Positions per backward tile; part of the fixed summation order. +BACKWARD_TILE = 256 + + +def adaln_row_gather_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_adaln_row_gather_forward") + and hasattr(_C, "h3_adaln_row_gather_backward") + ) + + +def _segment_tiles(index: torch.Tensor, num_rows: int): + """Sorted positions and fixed-size tiles per row segment (all integer ops).""" + + sorted_pos = torch.argsort(index, stable=True) + counts = torch.bincount(index, minlength=num_rows) + seg_start = torch.cumsum(counts, 0) - counts + tiles_per_seg = (counts + BACKWARD_TILE - 1) // BACKWARD_TILE + seg_first_tile = torch.zeros(num_rows + 1, dtype=torch.int64, device=index.device) + seg_first_tile[1:] = torch.cumsum(tiles_per_seg, 0) + tile_seg = torch.repeat_interleave(torch.arange(num_rows, device=index.device), tiles_per_seg) + tile_local = torch.arange(tile_seg.numel(), device=index.device) - seg_first_tile[tile_seg] + tile_begin = seg_start[tile_seg] + tile_local * BACKWARD_TILE + tile_end = torch.minimum(tile_begin + BACKWARD_TILE, seg_start[tile_seg] + counts[tile_seg]) + return sorted_pos, tile_begin.contiguous(), tile_end.contiguous(), seg_first_tile + + +class _H3AdaLNRowGatherCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, rows, timestep_indices, token_tags): + out = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(timestep_indices, token_tags) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return out + + @staticmethod + def backward(ctx, grad_out): + timestep_indices, token_tags = ctx.saved_tensors + num_rows, rows_dtype = ctx.rows_meta + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_rows) + grad_rows = _C.h3_adaln_row_gather_backward(grad_out.contiguous(), *tiles, rows_dtype) + record_backward( + "adaln_row_gather", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad_rows, None, None + + +class H3AdaLNRowGatherCudaOp: + """CUDA candidate: one launch gathers all six modulation tensors. + + The forward is a byte copy (bitwise equal to six ``index_select`` calls). + The backward replaces ``index_select``'s atomic BF16 scatter-add with a + deterministic FP32 segmented sum in stable sorted order, cast once. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not adaln_row_gather_available(): + raise RuntimeError( + "rl_engine._C lacks h3_adaln_row_gather_*; rebuild with " + "csrc/cuda/h3/adaln_row_gather.cu" + ) + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + if not rows.is_cuda: + raise ValueError("H3AdaLNRowGatherCudaOp needs CUDA tensors") + if rows.stride(1) != 1: + rows = rows.contiguous() + packed = _H3AdaLNRowGatherCuda.apply( + rows, timestep_indices.contiguous(), token_tags.contiguous() + ) + return tuple(packed.unbind(0)) + + def gather_chunks(self, chunks, timestep_indices, token_tags, **kwargs): + """Drop-in for diffusers' six ``(3T, H)`` tensors: concatenate, then gather.""" + + if len(chunks) != H3_ADALN_CHUNKS: + raise ValueError(f"expected {H3_ADALN_CHUNKS} modulation tensors, got {len(chunks)}") + return self.forward(torch.cat(list(chunks), dim=1), timestep_indices, token_tags, **kwargs) + + def forward_fp32(self, rows, timestep_indices, token_tags): + return tuple(out.float() for out in self.forward(rows, timestep_indices, token_tags)) diff --git a/rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py b/rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py new file mode 100644 index 000000000..93341d602 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 AdaLN row gather (RFC #420 row ``adaln_row_gather``). + +H3 computes one row index per packed-sequence position and uses it to pick +every position's six modulation vectors:: + + adaln_indices = timestep_indices * 3 + token_tags (S,) + shift_msa[adaln_indices], scale_msa[adaln_indices], ... six (S, H) + +``rows`` is the ``(3T, 6H)`` view of one block's AdaLN projection +(``table.view(-1, 6H)``); its six column blocks are the six modulation +tensors. ``timestep_indices`` and ``token_tags`` are semantic inputs +(RFC #420 §4): out-of-range values are an error, never clamped. +""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_INDEX_DTYPES = (torch.int64, torch.int32) + + +def validate_h3_adaln_row_gather( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + *, + check_range: bool = True, +) -> int: + """Return the hidden size; raise on anything H3 does not define.""" + + if not isinstance(rows, torch.Tensor) or rows.dim() != 2: + raise ValueError("rows must be a 2-D (3T, 6H) tensor") + if not rows.is_floating_point(): + raise TypeError(f"rows must be floating point, got {rows.dtype}") + num_rows, width = rows.shape + if num_rows == 0 or num_rows % H3_MODALITY_NUM != 0: + raise ValueError(f"rows must hold 3 rows per timestep, got {num_rows}") + if width == 0 or width % H3_ADALN_CHUNKS != 0: + raise ValueError(f"rows width must be 6 * H, got {width}") + for name, index in (("timestep_indices", timestep_indices), ("token_tags", token_tags)): + if not isinstance(index, torch.Tensor) or index.dim() != 1: + raise ValueError(f"{name} must be a 1-D (S,) tensor") + if index.dtype not in _INDEX_DTYPES: + raise TypeError(f"{name} must be int64 or int32, got {index.dtype}") + if index.device != rows.device: + raise ValueError(f"{name} is on {index.device}, rows are on {rows.device}") + if timestep_indices.shape != token_tags.shape: + raise ValueError( + f"timestep_indices {tuple(timestep_indices.shape)} and token_tags " + f"{tuple(token_tags.shape)} must have the same length" + ) + if timestep_indices.numel() == 0: + raise ValueError("the packed sequence must not be empty") + if timestep_indices.dtype != token_tags.dtype: + raise TypeError("timestep_indices and token_tags must share an integer dtype") + if check_range: + num_timesteps = num_rows // H3_MODALITY_NUM + bad = ( + (timestep_indices < 0) + | (timestep_indices >= num_timesteps) + | (token_tags < 0) + | (token_tags >= H3_MODALITY_NUM) + ) + if bool(bad.any()): # one host sync + raise IndexError( + f"timestep_indices must lie in [0, {num_timesteps}) and token_tags in " + f"[0, {H3_MODALITY_NUM}) (0 video, 1 text, 2 audio)" + ) + return width // H3_ADALN_CHUNKS + + +def h3_adaln_indices(timestep_indices: torch.Tensor, token_tags: torch.Tensor) -> torch.Tensor: + return timestep_indices.long() * H3_MODALITY_NUM + token_tags.long() + + +class _DeterministicRowGather(torch.autograd.Function): + """``index_select`` forward; FP32 per-row segment sums for the backward. + + ``index_select``'s own backward is an atomic scatter-add in the grad dtype + (BF16 for H3), which is neither deterministic nor accurate. There are + only 3T table rows, so the reference sums each row's positions with one + FP32 ``torch.sum`` per row (a fixed reduction for a fixed packing) and + rounds once. + """ + + @staticmethod + def forward(ctx, rows, index): + ctx.save_for_backward(index) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return torch.stack( + [chunk.index_select(0, index) for chunk in rows.chunk(H3_ADALN_CHUNKS, dim=-1)] + ) + + @staticmethod + def backward(ctx, grad): + (index,) = ctx.saved_tensors + num_rows, dtype = ctx.rows_meta + grad32 = grad.float() # (6, S, H) + out = grad32.new_zeros((num_rows, grad32.shape[0], grad32.shape[2])) + for row in range(num_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + out[row] = grad32.index_select(1, positions).sum(dim=1) + return out.reshape(num_rows, -1).to(dtype), None + + +class NativeH3AdaLNRowGatherOp: + """PyTorch reference: ``index_select`` on each of the six column blocks. + + The forward is byte-identical to diffusers; the backward is deterministic + (see ``_DeterministicRowGather``). The raw diffusers path, including its + atomic backward, is ``rl_engine.testing.h3_provider.provider_adaln_row_gather``. + ``forward_fp32`` returns the same rows in FP32 with an FP64 gradient. + + The forward is a copy, but the VJP is a segmented reduction over the + packed sequence, so the op is judged with the ``reduction`` tolerance + class; forward bitwise equality is asserted separately. + """ + + op_class = "reduction" + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple(_DeterministicRowGather.apply(rows, index).unbind(0)) + + def forward_fp32(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple( + chunk.index_select(0, index).float() + for chunk in rows.double().chunk(H3_ADALN_CHUNKS, dim=-1) + ) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 22bbab259..ba6c57f05 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -215,6 +215,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_ADALN_PROJECTION = ( "rl_engine.kernels.ops.cuda.h3.adaln_projection.H3AdaLNProjectionCudaOp" ) + PYTORCH_H3_ADALN_ROW_GATHER = ( + "rl_engine.kernels.ops.pytorch.h3.adaln_row_gather.NativeH3AdaLNRowGatherOp" + ) + CUDA_H3_ADALN_ROW_GATHER = ( + "rl_engine.kernels.ops.cuda.h3.adaln_row_gather.H3AdaLNRowGatherCudaOp" + ) def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -637,6 +643,10 @@ def __init__(self): OpBackend.CUDA_H3_ADALN_PROJECTION, OpBackend.PYTORCH_H3_ADALN_PROJECTION, ], + "adaln_row_gather": [ + OpBackend.CUDA_H3_ADALN_ROW_GATHER, + OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, + ], }, "rocm": { "logp": [ @@ -680,6 +690,7 @@ def __init__(self): "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -727,6 +738,7 @@ def __init__(self): "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -755,6 +767,7 @@ def __init__(self): "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 8671ef141..06bfe252b 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -32,6 +32,7 @@ from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps from rl_engine.testing.h3_provider import ( provider_adaln_modulation, + provider_adaln_row_gather, provider_time_embedder, provider_time_proj, ) @@ -94,8 +95,35 @@ def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: golden_atol=5e-2, # reduction / bfloat16 golden_rtol=2e-2, ), + Stage( + name="adaln_row_gather", + op_type="adaln_row_gather", + candidate=lambda op, ctx, up: op.gather_chunks( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + provider=lambda ctx, up: provider_adaln_row_gather( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + golden=lambda op, ctx, up: op.forward_fp32( + torch.cat(list(up), dim=1), ctx["timestep_indices"], ctx["token_tags"] + ), + provider_bitwise_isolated=True, # a copy + golden_atol=5e-2, # carries the projection's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ), ] +# Parameters whose gradients the backward replay reports, in chain order. +GRAD_LEAVES = ( + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +) +BACKWARD_MODES = ("candidate", "candidate_fused", "provider") + def _flat(value: Any) -> list[torch.Tensor]: if isinstance(value, torch.Tensor): @@ -181,6 +209,77 @@ def run_case( return report +def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torch.Tensor]: + """Parameter gradients of the full chain for one execution mode. + + ``candidate`` runs the four RL-Kernel ops separately, ``candidate_fused`` + runs projection + gather as ``H3AdaLNModulationCudaOp``, ``provider`` the + diffusers path, ``golden`` the FP64 references. + """ + + leaves = { + name: ctx["weights"][name].detach().clone().requires_grad_(True) for name in GRAD_LEAVES + } + run_ctx = {**ctx, "weights": {**ctx["weights"], **leaves}} + value = None + if mode == "candidate_fused": + from rl_engine.kernels.ops.cuda.h3.adaln_modulation import H3AdaLNModulationCudaOp + + for stage in STAGES[:2]: + value = stage.candidate(registry.get_op(stage.op_type, device="cuda"), run_ctx, value) + value = H3AdaLNModulationCudaOp()( + value, *adaln_params(run_ctx), ctx["timestep_indices"], ctx["token_tags"] + ) + else: + for stage in STAGES: + if mode == "provider": + value = stage.provider(run_ctx, value) + elif mode == "candidate": + op = registry.get_op(stage.op_type, device="cuda") + value = stage.candidate(op, run_ctx, value) + else: + value = stage.golden(golden_op(registry, stage.op_type), run_ctx, value) + outputs = list(value) + torch.autograd.backward(outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs)]) + return [leaves[name].grad for name in GRAD_LEAVES] + + +def run_backward_case( + registry: KernelRegistry, weights, *, num_timesteps: int, seq_len: int, seed: int = 0 +) -> dict[str, Any]: + """Determinism and accuracy of the chain's parameter gradients, per mode.""" + + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + generator = torch.Generator(device="cuda").manual_seed(seed + 1) + hidden = weights[GRAD_LEAVES[4]].shape[0] // 18 + upstream = [ + torch.randn(seq_len, hidden, device="cuda", generator=generator).bfloat16() + for _ in range(6) + ] + golden = chain_grads("golden", registry, ctx, upstream) + runs = { + mode: [chain_grads(mode, registry, ctx, upstream) for _ in range(2)] + for mode in BACKWARD_MODES + } + report: dict[str, Any] = {"num_timesteps": num_timesteps, "seq_len": seq_len, "leaves": {}} + for index, name in enumerate(GRAD_LEAVES): + gold = golden[index].float() + entry: dict[str, Any] = {"golden_absmax": float(gold.abs().max())} + for mode, (first, second) in runs.items(): + grad = first[index] + err = (grad.float() - gold).abs().max() + entry[mode] = { + "repeat_bitwise_equal": bool(torch.equal(first[index], second[index])), + "max_abs_vs_golden": float(err), + "max_abs_vs_golden_over_absmax": float(err / gold.abs().max().clamp_min(1e-30)), + "correctly_rounded_fraction": float( + (grad == golden[index].to(grad.dtype)).float().mean() + ), + } + report["leaves"][name] = entry + return report + + def git_state() -> dict[str, Any]: def run(*args: str) -> str: return subprocess.check_output(["git", *args], cwd=REPO_ROOT, text=True).strip() diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index f64f7174f..6cedcdf60 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -56,3 +56,12 @@ def provider_adaln_modulation( temb = F.linear(F.silu(temb).to(weight.dtype), weight, bias) temb = temb.view(-1, 6 * hidden_size) return temb.chunk(6, dim=-1) + + +def provider_adaln_row_gather( + modulation: tuple[torch.Tensor, ...], timestep_indices: torch.Tensor, token_tags: torch.Tensor +) -> tuple[torch.Tensor, ...]: + """``adaln_indices = timestep_indices * 3 + token_tags`` then six ``index_select``.""" + + adaln_indices = timestep_indices * 3 + token_tags + return tuple(tensor.index_select(0, adaln_indices) for tensor in modulation) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 3f84eb520..e0a5939b5 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -16,9 +16,10 @@ import torch from rl_engine.kernels.registry import KernelRegistry -from rl_engine.testing.h3_cases import h3_timesteps +from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps from rl_engine.testing.h3_provider import ( provider_adaln_modulation, + provider_adaln_row_gather, provider_time_embedder, provider_time_proj, ) @@ -258,13 +259,105 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, } +# --------------------------------------------------------------------------- # +# adaln_row_gather +# --------------------------------------------------------------------------- # + +GATHER_SEQ_LENS = (4097, 32768, 131072) + + +def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_row_gather", device="cuda") + num_timesteps, hidden = 3, 5376 + rows = torch.randn(3 * num_timesteps, 6 * hidden, device="cuda").bfloat16() + chunks = rows.chunk(6, dim=-1) + cases = [] + for seq in GATHER_SEQ_LENS: + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + grads = [torch.randn(seq, hidden, device="cuda").bfloat16() for _ in range(6)] + + def backward(fn, ti=ti, tags=tags, grads=grads): + leaf = rows.detach().requires_grad_(True) + torch.autograd.backward(list(fn(leaf, ti, tags)), grads) + + cases.append( + { + "op": "adaln_row_gather", + "case": f"S={seq}", + "backend": type(op).__name__, + # bytes written (six (S, H) BF16 outputs); the 3T-row table stays in L2 + "bytes": 6 * seq * hidden * 2, + "candidate": lambda ti=ti, tags=tags: op.forward(rows, ti, tags, check_range=False), + "provider": lambda ti=ti, tags=tags: provider_adaln_row_gather(chunks, ti, tags), + "candidate_backward": lambda b=backward: b( + lambda r, ti, tags: op.forward(r, ti, tags, check_range=False) + ), + "provider_backward": lambda b=backward: b( + lambda r, ti, tags: provider_adaln_row_gather(r.chunk(6, dim=-1), ti, tags) + ), + } + ) + return cases + + +def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: + from rl_engine.testing.h3_chain import run_backward_case + from rl_engine.testing.h3_weights import load_h3_conditioning_weights + + op = registry.get_op("adaln_row_gather", device="cuda") + golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) + rows = torch.randn(9, 6 * 5376, device="cuda").bfloat16() + forward_bitwise = {} + for seq in (1, 257, 4097, 32768): + ti, tags = h3_packed_layout(seq, 3, seed=seq) + ours = op(rows, ti, tags) + theirs = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + forward_bitwise[str(seq)] = all(torch.equal(a, b) for a, b in zip(ours, theirs)) + + # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = [torch.randn(4097, 5376, device="cuda").bfloat16() for _ in range(6)] + + def grad_of(fn): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(fn(leaf)), grads) + return leaf.grad + + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward(list(golden.forward_fp32(ref, ti, tags)), [g.float() for g in grads]) + gold = ref.grad.bfloat16() + backward = {} + for name, fn in ( + ("cuda", lambda leaf: op(leaf, ti, tags)), + ("provider", lambda leaf: provider_adaln_row_gather(leaf.chunk(6, dim=-1), ti, tags)), + ): + first, second = grad_of(fn), grad_of(fn) + backward[name] = { + "repeat_bitwise_equal": bool(torch.equal(first, second)), + "correctly_rounded_fraction": float((first == gold).float().mean()), + } + + weights = load_h3_conditioning_weights("cuda") + chain = [ + run_backward_case(registry, weights, num_timesteps=t, seq_len=s) + for t, s in ((1, 257), (3, 257), (1, 4097), (3, 4097), (4, 32768)) + ] + return { + "forward_bitwise_vs_index_select": forward_bitwise, + "op_backward": backward, + "chain_backward": chain, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, "adaln_projection_3mod": _projection_perf, + "adaln_row_gather": _gather_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, "timestep_mlp_fp32": _mlp_accuracy, "adaln_projection_3mod": _projection_accuracy, + "adaln_row_gather": _gather_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index c20e157e4..c36ade2f2 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -226,7 +226,12 @@ def fake_load_backend(backend): ] -H3_OPS = ("timestep_sinusoid_h3", "timestep_mlp_fp32", "adaln_projection_3mod") +H3_OPS = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", +) @pytest.mark.parametrize("op_name", H3_OPS) diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py index e262ca053..52a1cddde 100644 --- a/scripts/h3_chain_replay.py +++ b/scripts/h3_chain_replay.py @@ -40,6 +40,12 @@ def main() -> None: parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") parser.add_argument("--seed", type=int, default=0) parser.add_argument("--out", type=Path, default=None) + parser.add_argument( + "--backward", + action="store_true", + help="also replay the chain backward (needs every stage): parameter-gradient " + "determinism and accuracy for the separate-op, fused and diffusers chains", + ) args = parser.parse_args() if not torch.cuda.is_available(): @@ -76,6 +82,23 @@ def main() -> None: ) print(f"T={num_timesteps} S={seq_len}: {summary}; first_drift={case['first_drift']}") + backward_cases = [] + if args.backward: + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_backward_case( + registry, weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=args.seed + ) + backward_cases.append(case) + parts = [] + for mode in h3_chain.BACKWARD_MODES: + entries = case["leaves"].values() + det = all(e[mode]["repeat_bitwise_equal"] for e in entries) + worst = max(e[mode]["max_abs_vs_golden_over_absmax"] for e in entries) + parts.append(f"{mode}: {'det' if det else 'NONDET'} {worst:.1e}") + summary = " | ".join(parts) + print(f"backward T={num_timesteps} S={seq_len}: {summary}") + evidence = { "kind": "h3_conditioning_chain_replay", "rfc": manifest["rfc"], @@ -86,6 +109,7 @@ def main() -> None: "environment": h3_chain.environment(), "stages": [s.name for s in stages], "cases": cases, + "backward_cases": backward_cases, } if args.out is not None: args.out.parent.mkdir(parents=True, exist_ok=True) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 6385b4aa4..f2039f0aa 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -46,14 +46,24 @@ def _style(ax, title: str, ylabel: str) -> None: ax.spines["bottom"].set_color(GRID) -def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt: str) -> None: - """Grouped bars with a 2px-style gap and value labels on top.""" +def _bars(ax, labels: list[str], series: list[tuple], fmt: str, log: bool = False) -> None: + """Grouped bars with a gap and value labels on top; a 4th tuple item is a hatch.""" n = len(series) width = 0.8 / n - for i, (name, color, values) in enumerate(series): + for i, (name, color, values, *hatch) in enumerate(series): xs = [x + (i - (n - 1) / 2) * width for x in range(len(labels))] - bars = ax.bar(xs, values, width * 0.92, color=color, label=name, zorder=2) + bars = ax.bar( + xs, + values, + width * 0.92, + color=color if not hatch else SURFACE, + edgecolor=color, + hatch=hatch[0] if hatch else None, + linewidth=1.2, + label=name, + zorder=2, + ) for bar, value in zip(bars, values): ax.annotate( fmt.format(value), @@ -65,9 +75,13 @@ def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt color=INK, ) ax.set_xticks(range(len(labels)), labels) - top = max(v for _, _, values in series for v in values) - ax.set_ylim(0, top * 1.35) - ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=n) + top = max(v for _, _, values, *_ in series for v in values) + if log: + ax.set_yscale("log") + ax.set_ylim(min(v for _, _, values, *_ in series for v in values) / 2, top * 8) + else: + ax.set_ylim(0, top * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=min(n, 4)) # Exact zeros cannot sit on a log axis; draw them on this floor and label them. @@ -261,10 +275,81 @@ def plot_projection(report: dict[str, Any]): return fig +def plot_gather(report: dict[str, Any]): + fig, (left, right) = _figure(report, "adaln_row_gather (T = 3, H = 5376, BF16)") + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("index_select fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "index_select bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + chain = report["accuracy"]["chain_backward"] + _style( + right, + "Whole-chain grads of the FP32 time embedder vs FP64", + "max abs error / golden max", + ) + names = { + "candidate": ("RL-Kernel, separate ops", THIRD, "o"), + "candidate_fused": ("RL-Kernel, fused modulation", RL_KERNEL, "D"), + "provider": ("diffusers", PROVIDER, "s"), + } + labels = [f"T={c['num_timesteps']}\nS={c['seq_len']}" for c in chain] + for j, (mode, (name, color, marker)) in enumerate(names.items()): + leaves = [e for c in chain for e in c["leaves"].values()] + det = sum(e[mode]["repeat_bitwise_equal"] for e in leaves) + adaln = min( + e[mode]["correctly_rounded_fraction"] + for c in chain + for n, e in c["leaves"].items() + if n.startswith("transformer_blocks") + ) + label = ( + f"{name}\n repeat-bitwise {det}/{len(leaves)} · " + f"AdaLN BF16 grads ≥{100 * adaln:.1f}% correctly rounded" + ) + for i, case in enumerate(chain): + values = [ + e[mode]["max_abs_vs_golden_over_absmax"] + for n, e in case["leaves"].items() + if n.startswith("time_embedder") + ] + right.scatter( + [i + (j - 1) * 0.22] * len(values), + values, + s=30, + color=color, + marker=marker, + alpha=0.85, + linewidths=0, + zorder=3, + label=label if i == 0 else None, + ) + right.set_yscale("log") + right.set_ylim(1e-8, 1e4) + right.set_xticks(range(len(labels)), labels, fontsize=7) + right.legend(frameon=False, fontsize=7, labelcolor=INK, loc="upper left") + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, "adaln_projection_3mod": plot_projection, + "adaln_row_gather": plot_gather, } diff --git a/setup.py b/setup.py index b388bdd29..fcefa9016 100644 --- a/setup.py +++ b/setup.py @@ -159,6 +159,7 @@ def get_extensions(): # MiniMax-H3 (RFC #420) conditioning path; plain-arch, runs on SM90 and SM100. cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") cuda_sources.append("csrc/cuda/h3/det_linear.cu") + cuda_sources.append("csrc/cuda/h3/adaln_row_gather.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/test_h3_adaln_modulation.py b/tests/h3/test_h3_adaln_modulation.py new file mode 100644 index 000000000..97b4927d0 --- /dev/null +++ b/tests/h3/test_h3_adaln_modulation.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused AdaLN modulation (projection + row gather) with an FP32 table gradient. + +The forward must be bitwise equal to the two RFC #420 ops run separately; the +backward must be deterministic and must not re-round the table gradient to +BF16 (so ``d_temb`` reaches FP64-golden accuracy). +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.testing.h3_cases import h3_packed_layout + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _ops(): + from rl_engine.kernels.ops.cuda.h3.adaln_modulation import H3AdaLNModulationCudaOp + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks the H3 kernels") + return H3AdaLNModulationCudaOp(), H3AdaLNProjectionCudaOp(), H3AdaLNRowGatherCudaOp() + + +def _params(hidden, dim, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = 18 * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).bfloat16().cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).bfloat16().cuda() + temb = (torch.randn(3, dim, generator=g) * 2).cuda() + return temb, weight, bias + + +def _grads(seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + +def _leaf_grads(fn, temb, weight, bias, grads): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(fn(*leaves)), grads) + return [leaf.grad for leaf in leaves] + + +@requires_cuda +class TestFusedModulation: + def test_forward_bitwise_equal_to_separate_ops(self): + fused, proj, gather = _ops() + temb, weight, bias = _params(64, 2688) + ti, tags = h3_packed_layout(1000, 3) + ours = fused(temb, weight, bias, ti, tags) + separate = gather.gather_chunks(proj(temb, weight, bias), ti, tags) + assert all(torch.equal(a, b) for a, b in zip(ours, separate)) + + def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): + fused, proj, gather = _ops() + weight = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.weight"].cuda() + bias = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.bias"].cuda() + temb = torch.randn(3, 2688, device="cuda") * 2 + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = _grads(4097, 5376, seed=1) + ours = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + separate = _leaf_grads( + lambda t, w, b: gather.gather_chunks(proj(t, w, b), ti, tags), + temb, + weight, + bias, + grads, + ) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + act = act + (act.to(torch.bfloat16).double() - act).detach() # identity-VJP cast + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376) + index = ti * 3 + tags + torch.autograd.backward( + [chunk.index_select(0, index) for chunk in table.chunk(6, dim=-1)], + [g.double() for g in grads], + ) + + def rel(grad, golden): + return ((grad.double() - golden).abs().max() / golden.abs().max()).item() + + # d_temb: no BF16 rounding anywhere on its path, so it reaches FP32 accuracy. + assert rel(ours[0], ref[0].grad) < 1e-5 + assert rel(ours[0], ref[0].grad) < rel(separate[0], ref[0].grad) / 100 + # dW, db: only their own final BF16 rounding remains. + for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad)): + assert (grad == golden.to(torch.bfloat16)).float().mean() > 0.9999 + + def test_backward_repeat_bitwise(self): + fused, _, _ = _ops() + temb, weight, bias = _params(64, 2688, seed=3) + ti, tags = h3_packed_layout(3000, 3, seed=3) + grads = _grads(3000, 64, seed=3) + first = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + again = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + assert all(torch.equal(a, b) for a, b in zip(first, again)) + + def test_rejects_bad_indices(self): + fused, _, _ = _ops() + temb, weight, bias = _params(8, 16) + ti, tags = h3_packed_layout(10, 3) + with pytest.raises(IndexError): + fused(temb, weight, bias, ti + 3, tags) + with pytest.raises(TypeError): + fused(temb.bfloat16(), weight, bias, ti, tags) diff --git a/tests/h3/test_h3_adaln_row_gather.py b/tests/h3/test_h3_adaln_row_gather.py new file mode 100644 index 000000000..0c58c7e49 --- /dev/null +++ b/tests/h3/test_h3_adaln_row_gather.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_row_gather``: ``timestep_index * 3 + token_tag`` lookup of six tensors. + +* forward is a copy: bitwise equal to diffusers' six ``index_select`` calls, + for every packing, length and index dtype; +* semantic indices fail closed (out-of-range tags or timesteps, RFC probes + H2/H3), and flipping a tag or offsetting a timestep picks a different row; +* backward is a deterministic FP32 segmented sum: repeat-bitwise and + correctly rounded, unlike the atomic BF16 ``index_select`` backward. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.h3.adaln_row_gather import NativeH3AdaLNRowGatherOp +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_adaln_row_gather + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks h3_adaln_row_gather_*") + return H3AdaLNRowGatherCudaOp() + + +def _rows(num_timesteps, hidden=5376, dtype=torch.bfloat16, seed=0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(3 * num_timesteps, 6 * hidden, generator=g).to(dtype).to(device) + + +class TestReference: + def test_matches_diffusers_index_select(self): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = h3_packed_layout(11, 2, device="cpu") + index = ti * 3 + tags + chunks = rows.chunk(6, dim=-1) + for ours, chunk in zip(NativeH3AdaLNRowGatherOp().forward(rows, ti, tags), chunks): + assert torch.equal(ours, chunk.index_select(0, index)) + + @pytest.mark.parametrize( + "mutate, error", + [ + (lambda ti, tags: (ti, tags.clone().fill_(3)), IndexError), # H2: no 4th modality + (lambda ti, tags: (ti, tags - 1), IndexError), + (lambda ti, tags: (ti + 2, tags), IndexError), # H3: timestep offset past T + (lambda ti, tags: (ti, tags[:-1]), ValueError), + (lambda ti, tags: (ti[:0], tags[:0]), ValueError), + (lambda ti, tags: (ti.float(), tags), TypeError), + (lambda ti, tags: (ti.int(), tags), TypeError), # mixed index dtypes + ], + ) + def test_rejects_invalid_indices(self, mutate, error): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = mutate(*h3_packed_layout(9, 2, device="cpu")) + with pytest.raises(error): + NativeH3AdaLNRowGatherOp().forward(rows, ti, tags) + + def test_rejects_bad_rows(self): + ti, tags = h3_packed_layout(4, 1, device="cpu") + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(4, 24), ti, tags) # not 3 per t + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(3, 25), ti, tags) # not 6 * H + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("seq", [1, 2, 3, 257, 4097, 32768]) + @pytest.mark.parametrize("num_timesteps", [1, 3]) + def test_bitwise_equal_to_index_select(self, seq, num_timesteps): + rows = _rows(num_timesteps, seed=seq) + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + ours = _cuda_op()(rows, ti, tags) + ref = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + assert len(ours) == 6 + for a, b in zip(ours, ref): + assert a.shape == (seq, 5376) and a.is_contiguous() and torch.equal(a, b) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [24, 5, 5376]) + def test_dtypes_and_unaligned_hidden(self, dtype, hidden): + rows = _rows(2, hidden=hidden, dtype=dtype) + ti, tags = h3_packed_layout(33, 2) + for a, b in zip( + _cuda_op()(rows, ti.int(), tags.int()), NativeH3AdaLNRowGatherOp()(rows, ti, tags) + ): + assert torch.equal(a, b) + + def test_strided_rows_view(self): + table = _rows(4, hidden=16) + rows = table[::2] # row stride 2 * 6H, still 3 rows per timestep after slicing 12 -> 6 + ti, tags = h3_packed_layout(20, 2) + for a, b in zip(_cuda_op()(rows, ti, tags), NativeH3AdaLNRowGatherOp()(rows, ti, tags)): + assert torch.equal(a, b) + + def test_tag_flip_and_timestep_offset_select_other_rows(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(30, 2) + base = _cuda_op()(rows, ti, tags) + flipped = _cuda_op()(rows, ti, (tags + 1) % 3) # H2 + shifted = _cuda_op()(rows, 1 - ti, tags) # H3 within range + assert not torch.equal(base[0], flipped[0]) + assert not torch.equal(base[0], shifted[0]) + # Every output row is exactly the addressed table row, nothing else. + index = ti * 3 + tags + for c, out in enumerate(base): + assert torch.equal(out, rows[index, c * 16 : (c + 1) * 16]) + + def test_rejects_out_of_range_on_device(self): + rows = _rows(1, hidden=16) + ti, tags = h3_packed_layout(8, 1) + with pytest.raises(IndexError): + _cuda_op()(rows, ti, tags + 3) + + def test_gather_chunks_drop_in(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(10, 2) + chunks = rows.chunk(6, dim=-1) + for a, b in zip(_cuda_op().gather_chunks(chunks, ti, tags), _cuda_op()(rows, ti, tags)): + assert torch.equal(a, b) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + def _cuda_grad(self, rows, ti, tags, grads): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(_cuda_op()(leaf, ti, tags)), grads) + return leaf.grad + + def test_correctly_rounded_fp32_segment_sum(self): + rows = _rows(3, seed=1) + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = self._grads(4097, 5376) + ours = self._cuda_grad(rows, ti, tags, grads) + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward( + list(NativeH3AdaLNRowGatherOp().forward_fp32(ref, ti, tags)), [g.float() for g in grads] + ) + assert ours.dtype == torch.bfloat16 + # tolerance_contract.json gradient_accuracy / elementwise / bfloat16. + torch.testing.assert_close(ours.double(), ref.grad, atol=2e-2, rtol=1.6e-2) + assert (ours == ref.grad.bfloat16()).float().mean() > 0.9999 + + def test_repeat_bitwise(self): + rows = _rows(2, hidden=512, seed=2) + ti, tags = h3_packed_layout(5000, 2, seed=2) + grads = self._grads(5000, 512, seed=2) + first = self._cuda_grad(rows, ti, tags, grads) + for _ in range(3): + assert torch.equal(self._cuda_grad(rows, ti, tags, grads), first) + + def test_unreferenced_rows_get_zero_grad(self): + rows = _rows(3, hidden=16) + ti = torch.zeros(7, dtype=torch.long, device="cuda") + tags = torch.zeros(7, dtype=torch.long, device="cuda") # only row 0 is used + grad = self._cuda_grad(rows, ti, tags, self._grads(7, 16)) + assert torch.count_nonzero(grad[1:]) == 0 + assert torch.count_nonzero(grad[0]) > 0 + + def test_segments_longer_than_one_tile(self): + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import BACKWARD_TILE + + seq = 3 * BACKWARD_TILE + 17 + rows = _rows(1, hidden=8, dtype=torch.float32) + ti = torch.zeros(seq, dtype=torch.long, device="cuda") + tags = torch.zeros(seq, dtype=torch.long, device="cuda") + grads = [torch.ones(seq, 8, device="cuda") for _ in range(6)] + grad = self._cuda_grad(rows, ti, tags, grads) + assert torch.equal(grad[0], torch.full((48,), float(seq), device="cuda")) + + def test_reference_backward_is_deterministic_and_accurate(self): + rows = _rows(2, hidden=256, seed=4) + ti, tags = h3_packed_layout(3000, 2, seed=4) + grads = self._grads(3000, 256, seed=4) + + def ref_grad(): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(NativeH3AdaLNRowGatherOp()(leaf, ti, tags)), grads) + return leaf.grad + + first = ref_grad() + assert torch.equal(ref_grad(), first) + torch.testing.assert_close( + first.float(), self._cuda_grad(rows, ti, tags, grads).float(), atol=1e-1, rtol=2e-2 + ) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_row_gather", device="cuda") + assert type(op).__name__ == "H3AdaLNRowGatherCudaOp" diff --git a/tests/h3/test_h3_conditioning_e2e.py b/tests/h3/test_h3_conditioning_e2e.py index edd75a3b2..cdf830f5f 100644 --- a/tests/h3/test_h3_conditioning_e2e.py +++ b/tests/h3/test_h3_conditioning_e2e.py @@ -53,3 +53,22 @@ def test_chain_forward(registry, chain_weights, num_timesteps, seq_len): assert entry["isolated_vs_provider"]["bitwise_equal"], entry["stage"] drift = report["first_isolated_drift"] assert drift is None or not stages[drift].provider_bitwise_isolated + + +@pytest.mark.parametrize("num_timesteps, seq_len", [(1, 257), (3, 4097)]) +def test_chain_backward(registry, chain_weights, num_timesteps, seq_len): + """Parameter gradients of the whole chain: deterministic, and FP32-accurate when fused.""" + + report = h3_chain.run_backward_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + for name, entry in report["leaves"].items(): + for mode in ("candidate", "candidate_fused"): + assert entry[mode]["repeat_bitwise_equal"], (mode, name) + fused = entry["candidate_fused"] + if name.startswith("time_embedder"): + # FP32 parameters: no BF16 rounding anywhere on the fused path. + assert fused["max_abs_vs_golden_over_absmax"] < 1e-5, (name, fused) + else: + # BF16 parameters: only their own final rounding remains. + assert fused["correctly_rounded_fraction"] > 0.99, (name, fused) diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index fdef637f7..455a88638 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -52,6 +52,7 @@ def test_all_ws1_single_ops_are_registered(): "timestep_sinusoid_h3", "timestep_mlp_fp32", "adaln_projection_3mod", + "adaln_row_gather", } <= names From 1f7253c439301cf6105ac4983647a34f1fcc97c6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:45:10 +0000 Subject: [PATCH 08/89] [WS1][MiniMax-H3] adaln_row_gather: B200 evidence report, chain replay and figure Produced by scripts/h3_evidence.py and scripts/h3_chain_replay.py --backward at fa551c2 on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 54636dd09..1fb114792 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -95,7 +95,7 @@ weight and bias) against an FP64 golden in which every declared cast is straight | S >= 257 (72 gradients) | RL-Kernel, separate ops | RL-Kernel, fused | diffusers | | --- | --- | --- | --- | | repeat-bitwise | 72 / 72 | 72 / 72 | 0 / 72 | -| time embedder: max abs error / golden max | 5.3e-4 to 2.4e-3 | 6.5e-7 to 4.9e-6 | 5.6e-3 to 1.2e-1 | +| time embedder: max abs error / golden max | 5.3e-4 to 2.4e-3 | 6.5e-7 to 4.9e-6 | 3.8e-3 to 1.3e-1, varying run to run | | AdaLN weight/bias: correctly rounded BF16, worst case | 59.8% | 99.57% | 1.6% | At S = 3 all three chains are deterministic. Both RL-Kernel variants reach 1.5e-6 on the @@ -111,13 +111,26 @@ B200, T = 3, H = 5376, BF16. "write BW" counts the 6·S·H outputs; the table st | S | CUDA forward | provider forward | CUDA backward | provider backward | | --- | --- | --- | --- | --- | -| 4097 | 62 µs | 92 µs | 1.10 ms | 1.44 ms | -| 32768 | 315 µs | 587 µs | 2.10 ms | 8.66 ms | -| 131072 | 1.24 ms (6.8 TB/s write) | 2.37 ms | 6.84 ms | 34.9 ms | +| 4097 | 0.06 ms | 0.09 ms | 1.15 ms | 1.47 ms | +| 32768 | 0.32 ms | 0.59 ms | 2.18 ms | 8.72 ms | +| 131072 | 1.23 ms (6.8 TB/s write) | 2.37 ms | 6.86 ms | 34.9 ms | The backward's peak memory is about twice the provider's, because autograd stacks the six output gradients and the per-tile FP32 partials are materialised. +## Evidence + +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../usage/evidence/h3-adaln-row-gather-b200/figure.png) + +There are two data files: + +- [`report.json`](../usage/evidence/h3-adaln-row-gather-b200/report.json): op timings, forward + bitwise checks, op-level backward, and the chain-backward cases plotted above. +- [`chain_replay.json`](../usage/evidence/h3-adaln-row-gather-b200/chain_replay.json): the full + stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. + +Both were written from a clean tree at commit `fa551c2`. + ## Tests ```bash From 024ee0813511ab2e430e0b4b28848b42a85c82f2 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 02:37:29 +0000 Subject: [PATCH 09/89] [WS1][CUDA][MiniMax-H3] h3_rmsnorm: RMSNorm with fused AdaLN modulation RFC #420 step 4, row `h3_rmsnorm` (stacked on adaln_row_gather). - CUDA kernel replaying PyTorch's own RMSNorm reduction order (vectorized_layer_norm_kernel: (32, 4) block per row, 4-element vectors, shuffle-down and cross-warp trees, rsqrtf, w * (rstd * x)): bitwise equal to nn.RMSNorm on all four pinned H3 norm weights and batch/position invariant. - Fused AdaLN modulation n * (1 + scale[i]) + shift[i] with the rows gathered in-kernel from strided table views and each step rounded where the eager expression rounds: bitwise equal to diffusers for the block (adaln_indices) and norm_out (timestep_indices), without materialising the (S, H) gathered tensors; 2.5x faster forward at S = 131072. - Deterministic backward: row-local dx, dweight over fixed 256-row tiles, dshift/dscale as sorted segment sums (diffusers: non-deterministic BF16 atomics, ~15x larger error). - Manifest pins the block/refiner-final/norm_out norm weights and norm_out.linear (all in shard 1); chain gains a per-mode history and the block norm1 + modulation stage (covered by the end-to-end test). - Registry, gtest, report/figure hooks, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/rmsnorm_modulate.cu | 403 ++++++++++++++++++++ csrc/ops.cpp | 21 + docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-rmsnorm.md | 107 ++++++ rl_engine/_C.pyi | 21 + rl_engine/kernels/gtest/operator_inputs.py | 15 + rl_engine/kernels/gtest/operator_specs.py | 11 + rl_engine/kernels/ops/cuda/h3/rmsnorm.py | 103 +++++ rl_engine/kernels/ops/pytorch/h3/rmsnorm.py | 143 +++++++ rl_engine/kernels/registry.py | 6 + rl_engine/testing/h3_chain.py | 70 +++- rl_engine/testing/h3_manifest.json | 12 +- rl_engine/testing/h3_provider.py | 16 + rl_engine/testing/h3_report.py | 120 ++++++ rl_engine/tests/test_dispatch.py | 1 + scripts/plot_h3_evidence.py | 65 ++++ setup.py | 1 + tests/h3/test_h3_rmsnorm.py | 203 ++++++++++ tests/test_ws1_gtest_gpu.py | 1 + 20 files changed, 1312 insertions(+), 9 deletions(-) create mode 100644 csrc/cuda/h3/rmsnorm_modulate.cu create mode 100644 docs/operators/h3-rmsnorm.md create mode 100644 rl_engine/kernels/ops/cuda/h3/rmsnorm.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/rmsnorm.py create mode 100644 tests/h3/test_h3_rmsnorm.py diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu new file mode 100644 index 000000000..5212324b1 --- /dev/null +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -0,0 +1,403 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 RMSNorm and fused AdaLN modulation (RFC #420 `h3_rmsnorm`). +// +// n[r] = cast(w * (rstd[r] * x[r])) rstd = rsqrtf(sum(x^2) / N + eps) +// out[r] = cast(cast(n[r] * cast(1 + scale[i])) + shift[i]) i = row_index[r % S] +// +// Forward contract h3-rmsnorm-v1: the statistics replay PyTorch's +// vectorized_layer_norm_kernel (torch cf30153): one +// (32, 4) block per row, 4-element vectors, thread t sums vectors t, t+128, ... +// in order, a shuffle-down tree (16..1), a cross-warp tree, sum / N, rsqrtf, +// then w * (rstd * x) and one cast. nn.RMSNorm therefore matches bitwise, and +// each modulation step rounds to the tensor dtype exactly where the eager +// `n * (1.0 + scale) + shift` expression does. Rows are independent: the result +// does not depend on batch size, position or the other rows. +// +// Backward (FP32, one cast per output, no atomics): +// d_n = g * cast(1 + scale) (g without modulation) +// dx = rstd * w * d_n - x * rstd^3 * sum(w * d_n * x) / N (row-local) +// dweight = sum_r d_n * x * rstd over fixed 256-row tiles, tiles folded in order +// dshift = sum_{r -> i} g, dscale = sum_{r -> i} g * n (sorted tiles) + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +constexpr int kVec = 4; // PyTorch's layer-norm vec_size +constexpr int kWarps = 4; // num_threads() / warp_size = 128 / 32 +constexpr int kRowTile = 256; // dweight row tile + +__device__ __forceinline__ float to_f(float v) { return v; } +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ float from_f(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } + +// Round to T and back: the value an eager T-dtype op would produce. +template +__device__ __forceinline__ float round_t(float v) { return to_f(from_f(v)); } + +// PyTorch's rms statistics: sequential per thread, shuffle-down, cross-warp tree. +// Returns sum(v^2) / N on every thread. +template +__device__ __forceinline__ float row_mean_square(const T* __restrict__ row, int n, float* buf) { + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + float s = 0.0f; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const float v = to_f(row[i * kVec + k]); + s = s + v * v; + } + } + for (int off = 16; off > 0; off >>= 1) s = s + __shfl_down_sync(0xffffffffu, s, off); + for (int off = blockDim.y / 2; off > 0; off /= 2) { + if (threadIdx.x == 0 && threadIdx.y >= off && threadIdx.y < 2 * off) buf[threadIdx.y - off] = s; + __syncthreads(); + if (threadIdx.x == 0 && threadIdx.y < off) s = s + buf[threadIdx.y]; + __syncthreads(); + } + if (thrx == 0) buf[0] = s / float(n); + __syncthreads(); + const float out = buf[0]; + __syncthreads(); + return out; +} + +// Same tree for an arbitrary per-element FP32 term; returns the row sum. +__device__ __forceinline__ float row_sum_tree(float s, float* buf) { + for (int off = 16; off > 0; off >>= 1) s = s + __shfl_down_sync(0xffffffffu, s, off); + for (int off = blockDim.y / 2; off > 0; off /= 2) { + if (threadIdx.x == 0 && threadIdx.y >= off && threadIdx.y < 2 * off) buf[threadIdx.y - off] = s; + __syncthreads(); + if (threadIdx.x == 0 && threadIdx.y < off) s = s + buf[threadIdx.y]; + __syncthreads(); + } + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + if (thrx == 0) buf[0] = s; + __syncthreads(); + const float out = buf[0]; + __syncthreads(); + return out; +} + +struct Modulation { + const void* shift; // (R, N) rows with stride `row_stride` (elements), or nullptr + const void* scale; + int64_t row_stride; + const int64_t* index; // (S,) table row per sequence position + int64_t seq; // S: row r of x uses index[r % S] +}; + +template +__global__ void __launch_bounds__(32 * kWarps) + rmsnorm_modulate_fwd_kernel(const T* __restrict__ x, const T* __restrict__ w, T* __restrict__ y, + float* __restrict__ rstd_out, int n, float eps, Modulation mod) { + __shared__ float buf[kWarps]; + const int64_t r = blockIdx.x; + const T* xr = x + r * n; + const float rstd = rsqrtf(row_mean_square(xr, n, buf) + eps); + const T* shift = nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + shift = static_cast(mod.shift) + i * mod.row_stride; + scale = static_cast(mod.scale) + i * mod.row_stride; + } + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + float v = round_t(__fmul_rn(to_f(w[j]), __fmul_rn(rstd, to_f(xr[j])))); + if (shift != nullptr) { + // Separately rounded ops (no FMA contraction), as the eager expression. + const float t1 = round_t(__fadd_rn(1.0f, to_f(scale[j]))); + v = round_t(__fadd_rn(round_t(__fmul_rn(v, t1)), to_f(shift[j]))); + } + y[r * n + j] = from_f(v); + } + } + if (thrx == 0) rstd_out[r] = rstd; +} + +template +__device__ __forceinline__ float d_norm_out(const T* g_row, const T* scale, int j) { + const float g = to_f(g_row[j]); + return scale == nullptr ? g : g * round_t(1.0f + to_f(scale[j])); +} + +template +__global__ void __launch_bounds__(32 * kWarps) + rmsnorm_modulate_dx_kernel(const T* __restrict__ g, const T* __restrict__ x, + const T* __restrict__ w, const float* __restrict__ rstd, + T* __restrict__ dx, int n, Modulation mod) { + __shared__ float buf[kWarps]; + const int64_t r = blockIdx.x; + const T* xr = x + r * n; + const T* gr = g + r * n; + const T* scale = mod.index != nullptr + ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride + : nullptr; + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + float s = 0.0f; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + s = fmaf(to_f(w[j]) * d_norm_out(gr, scale, j), to_f(xr[j]), s); + } + } + const float dot = row_sum_tree(s, buf); + const float rs = rstd[r]; + const float coef = rs * rs * rs * dot / float(n); + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + dx[r * n + j] = from_f(rs * to_f(w[j]) * d_norm_out(gr, scale, j) - to_f(xr[j]) * coef); + } + } +} + +// partial[tile, j] = sum over rows of the tile (ascending) of d_n * x * rstd +template +__global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* __restrict__ x, + const float* __restrict__ rstd, + float* __restrict__ partial, int64_t rows, int n, + Modulation mod) { + const int j = blockIdx.y * blockDim.x + threadIdx.x; + if (j >= n) return; + const int64_t r0 = static_cast(blockIdx.x) * kRowTile; + const int64_t r1 = min(r0 + kRowTile, rows); + float acc = 0.0f; + for (int64_t r = r0; r < r1; ++r) { + const T* scale = mod.index != nullptr ? static_cast(mod.scale) + + mod.index[r % mod.seq] * mod.row_stride + : nullptr; + acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); + } + partial[static_cast(blockIdx.x) * n + j] = acc; +} + +template +__global__ void fold_tiles_kernel(const float* __restrict__ partial, T* __restrict__ out, + int64_t tiles, int n) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t t = 0; t < tiles; ++t) acc += partial[t * n + j]; + out[j] = from_f(acc); +} + +// Table gradient tiles: positions sorted by (index, position); columns j < n +// accumulate g (shift), columns j >= n accumulate g * n_out (scale). +template +__global__ void modulation_grad_partial_kernel( + const T* __restrict__ g, const T* __restrict__ x, const T* __restrict__ w, + const float* __restrict__ rstd, const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int n) { + const int64_t tile = blockIdx.x; + const int j2 = blockIdx.y * blockDim.x + threadIdx.x; + if (j2 >= 2 * n) return; + const bool is_scale = j2 >= n; + const int j = is_scale ? j2 - n : j2; + const float wj = to_f(w[j]); + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t r = sorted_pos[p]; + const float gv = to_f(g[r * n + j]); + if (is_scale) { + const float nv = round_t(wj * (rstd[r] * to_f(x[r * n + j]))); + acc = fmaf(gv, nv, acc); + } else { + acc += gv; + } + } + partial[tile * 2 * n + j2] = acc; +} + +__global__ void fold_segments_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + float* __restrict__ out, int width) { + const int64_t seg = blockIdx.x; + const int j = blockIdx.y * blockDim.x + threadIdx.x; + if (j >= width) return; + float acc = 0.0f; + for (int64_t t = seg_first_tile[seg]; t < seg_first_tile[seg + 1]; ++t) acc += partial[t * width + j]; + out[seg * width + j] = acc; +} + +void check_rows(const torch::Tensor& t, const char* name, const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device(), name, " must be on ", like.device()); + TORCH_CHECK(t.dim() == 2 && t.is_contiguous(), name, " must be a contiguous 2-D tensor"); + TORCH_CHECK(t.scalar_type() == like.scalar_type(), name, " must have dtype ", like.scalar_type()); +} + +Modulation make_modulation(const c10::optional& shift, + const c10::optional& scale, + const c10::optional& index, const torch::Tensor& x, + int64_t n) { + Modulation mod{nullptr, nullptr, 0, nullptr, 1}; + if (!index.has_value()) { + TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); + return mod; + } + TORCH_CHECK(shift.has_value() && scale.has_value(), "modulation needs both shift and scale"); + const auto& sh = *shift; + const auto& sc = *scale; + const auto& ix = *index; + for (const auto* t : {&sh, &sc}) { + TORCH_CHECK(t->is_cuda() && t->device() == x.device() && t->scalar_type() == x.scalar_type(), + "shift/scale must be CUDA tensors with x's dtype"); + TORCH_CHECK(t->dim() == 2 && t->size(1) == n && t->stride(1) == 1, + "shift/scale must be (R, N) with unit column stride"); + } + TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), + "shift and scale must be views of the same table layout"); + TORCH_CHECK(ix.is_cuda() && ix.scalar_type() == at::kLong && ix.dim() == 1 && ix.is_contiguous(), + "row index must be a contiguous int64 CUDA tensor"); + TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", + ix.size(0), ")"); + mod.shift = sh.data_ptr(); + mod.scale = sc.data_ptr(); + mod.row_stride = sh.stride(0); + mod.index = ix.data_ptr(); + mod.seq = ix.size(0); + return mod; +} + +void check_xw(const torch::Tensor& x, const torch::Tensor& w) { + TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.is_contiguous(), "x must be a contiguous (M, N) CUDA tensor"); + TORCH_CHECK(x.size(0) > 0, "x must have at least one row"); + TORCH_CHECK(w.is_cuda() && w.dim() == 1 && w.size(0) == x.size(1) && w.is_contiguous() && + w.scalar_type() == x.scalar_type(), + "weight must be a contiguous (N,) CUDA tensor with x's dtype"); + TORCH_CHECK(x.size(1) % kVec == 0, "N=", x.size(1), " must be a multiple of ", kVec, + " (PyTorch's vectorized RMSNorm path)"); +} + +// float32 / float16 / bfloat16 only (float64 has no PyTorch-matching vector path here). +#define H3_DISPATCH(TYPE, NAME, ...) \ + [&] { \ + switch (TYPE) { \ + case at::kFloat: { using T = float; __VA_ARGS__(); break; } \ + case at::kHalf: { using T = __half; __VA_ARGS__(); break; } \ + case at::kBFloat16: { using T = __nv_bfloat16; __VA_ARGS__(); break; } \ + default: TORCH_CHECK(false, NAME, ": unsupported dtype ", TYPE); \ + } \ + }() + +} // namespace + +// x (M, N); optional shift/scale (R, N) strided views and index (S,), M % S == 0. +std::vector h3_rmsnorm_forward(torch::Tensor x, torch::Tensor weight, double eps, + c10::optional shift, + c10::optional scale, + c10::optional index) { + check_xw(x, weight); + TORCH_CHECK(x.scalar_type() != at::kDouble, "float64 is not supported"); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + const c10::cuda::CUDAGuard guard(x.device()); + auto y = torch::empty_like(x); + auto rstd = torch::empty({x.size(0)}, x.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_forward", [&] { + rmsnorm_modulate_fwd_kernel<<(x.size(0)), dim3(32, kWarps), 0, stream>>>( + reinterpret_cast(x.data_ptr()), reinterpret_cast(weight.data_ptr()), + reinterpret_cast(y.data_ptr()), rstd.data_ptr(), static_cast(n), + static_cast(eps), mod); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {y, rstd}; +} + +// Returns {dx, dweight} and, with modulation, {d_shift, d_scale} as FP32 (R, N). +std::vector h3_rmsnorm_backward( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, c10::optional sorted_pos, + c10::optional tile_begin, c10::optional tile_end, + c10::optional seg_first_tile) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.scalar_type() == at::kFloat && rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics"); + const int64_t rows = x.size(0); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + const c10::cuda::CUDAGuard guard(x.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto dx = torch::empty_like(x); + const int64_t tiles = (rows + kRowTile - 1) / kRowTile; + auto partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + auto dweight = torch::empty_like(weight); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward", [&] { + const T* g = reinterpret_cast(grad.data_ptr()); + const T* xp = reinterpret_cast(x.data_ptr()); + const T* wp = reinterpret_cast(weight.data_ptr()); + rmsnorm_modulate_dx_kernel<<(rows), dim3(32, kWarps), 0, stream>>>( + g, xp, wp, rstd.data_ptr(), reinterpret_cast(dx.data_ptr()), + static_cast(n), mod); + rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, 0, + stream>>>(g, xp, rstd.data_ptr(), + partial.data_ptr(), rows, + static_cast(n), mod); + fold_tiles_kernel<<>>( + partial.data_ptr(), reinterpret_cast(dweight.data_ptr()), tiles, + static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (mod.index == nullptr) return {dx, dweight}; + + TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && + seg_first_tile.has_value(), + "modulated backward needs the sorted segment tiles"); + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t seg_tiles = tile_begin->numel(); + const int64_t width = 2 * n; + auto seg_partial = torch::empty({std::max(seg_tiles, 1), width}, x.options().dtype(at::kFloat)); + auto table_grad = torch::empty({segments, width}, x.options().dtype(at::kFloat)); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + if (seg_tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_modulation_grad", [&] { + modulation_grad_partial_kernel + <<(seg_tiles), wide_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + sorted_pos->data_ptr(), tile_begin->data_ptr(), + tile_end->data_ptr(), seg_partial.data_ptr(), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + fold_segments_kernel<<(segments), wide_blocks), threads, 0, stream>>>( + seg_partial.data_ptr(), seg_first_tile->data_ptr(), + table_grad.data_ptr(), static_cast(width)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dx, dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 27159d5e9..f2288efe4 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -517,6 +517,16 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor torch::Tensor tile_begin, torch::Tensor tile_end, torch::Tensor seg_first_tile, c10::ScalarType out_dtype); +std::vector h3_rmsnorm_forward(torch::Tensor x, torch::Tensor weight, double eps, + c10::optional shift, + c10::optional scale, + c10::optional index); +std::vector h3_rmsnorm_backward( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, c10::optional sorted_pos, + c10::optional tile_begin, c10::optional tile_end, + c10::optional seg_first_tile); #endif // PyBind11 Module Registration @@ -798,5 +808,16 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Deterministic segmented sum (sorted tiles folded in order) for the row gather", py::arg("grad"), py::arg("sorted_pos"), py::arg("tile_begin"), py::arg("tile_end"), py::arg("seg_first_tile"), py::arg("out_dtype")); + m.def("h3_rmsnorm_forward", &h3_rmsnorm_forward, + "RMSNorm replaying PyTorch's reduction order, with optional fused AdaLN modulation", + py::arg("x"), py::arg("weight"), py::arg("eps"), py::arg("shift") = py::none(), + py::arg("scale") = py::none(), py::arg("index") = py::none()); + m.def("h3_rmsnorm_backward", &h3_rmsnorm_backward, + "Deterministic RMSNorm(+modulation) backward: row-local dx, tiled dweight, sorted table grads", + py::arg("grad"), py::arg("x"), py::arg("weight"), py::arg("rstd"), + py::arg("shift") = py::none(), py::arg("scale") = py::none(), + py::arg("index") = py::none(), py::arg("sorted_pos") = py::none(), + py::arg("tile_begin") = py::none(), py::arg("tile_end") = py::none(), + py::arg("seg_first_tile") = py::none()); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index ff55306a7..af8c1863f 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -34,6 +34,7 @@ nav: - operators/h3-timestep-mlp.md - operators/h3-adaln-projection.md - operators/h3-adaln-row-gather.md + - operators/h3-rmsnorm.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 33e405acd..b687a65b5 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -36,4 +36,5 @@ Every operator page should include: - [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) - [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) - [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) +- [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md new file mode 100644 index 000000000..3bd0b2f6e --- /dev/null +++ b/docs/operators/h3-rmsnorm.md @@ -0,0 +1,107 @@ +# MiniMax-H3 RMSNorm and AdaLN Modulation + +## Summary + +`h3_rmsnorm` covers every RMSNorm in MiniMax-H3 (RFC #420, WS1 step 4). That is the block +`norm1`/`norm2`, the token-refiner norms, the refiner `final_norm` and `norm_out.norm`. All of +them are `nn.RMSNorm(5376, eps=1e-5)` with a BF16 affine weight. In a transformer block and in +`norm_out`, the normalised rows are immediately modulated by per-row AdaLN parameters, and the +op fuses that step: + +```text +n = rms_norm(x, weight, eps) +out = n * (1.0 + scale[index]) + shift[index] +``` + +`index` is `adaln_indices` in a block (rows of the [projection](h3-adaln-projection.md) table) +and `timestep_indices` in `norm_out`. `shift`/`scale` are `(R, H)` row views of the AdaLN table, +gathered inside the kernel, so the `(S, H)` tensors that +[`adaln_row_gather`](h3-adaln-row-gather.md) materialises are never needed. + +## Entry Point + +```python +op = kernel_registry.get_op("h3_rmsnorm", device="cuda") +n = op(x, weight) # plain RMSNorm +out = op.forward_modulated(x, weight, shift_msa, scale_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+, validated on SM100) | `rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp` | `rl_engine._C.h3_rmsnorm_{forward,backward}` | Bitwise equal to `nn.RMSNorm` and to the diffusers modulation | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp` | n/a | `forward*`: provider path; `forward*_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `x` | `(..., S, N)` | bf16 / fp16 / fp32 | `N % 4 == 0`; H3: `N = 5376` | +| `weight` | `(N,)` | `x.dtype` | | +| `shift`, `scale` | `(R, N)` row views, same stride | `x.dtype` | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; one entry per position, shared across the batch | + +Out-of-range indices, mismatched dtypes or shapes, a non-positive eps, and `N % 4 != 0` all fail +closed. + +## Numerics + +- **Statistics (contract `h3-rmsnorm-v1`).** The kernel replays PyTorch's own + `vectorized_layer_norm_kernel` (torch `cf30153`): + - one `(32, 4)` block per row, reading 4-element vectors; + - thread `t` sums vectors `t, t+128, …` in order; + - a shuffle-down tree (16 → 1), then a cross-warp tree; + - `rsqrtf(Σx²/N + eps)`, then `w * (rstd * x)` with one cast. + + The plain norm is therefore **bitwise equal to `nn.RMSNorm`**. Rows are independent of batch + size and position. +- **Modulation.** `1 + scale`, `n * (…)` and `+ shift` are each rounded to the tensor dtype, + exactly where the eager expression rounds. The fused output is **bitwise equal to diffusers' + `norm(x) * (1.0 + scale.index_select(0, i)) + shift.index_select(0, i)`**. +- **Backward.** Everything is in FP32 with one cast per output and no atomics: + - `dx` is row-local, using the same reduction tree; + - `dweight` is folded over fixed 256-row tiles in ascending order; + - `dshift`/`dscale` are segmented sums over positions sorted stably by table row, in the + same scheme as [`adaln_row_gather`](h3-adaln-row-gather.md). + + Diffusers' `index_select` backward accumulates `dshift`/`dscale` with BF16 atomics, so it is + non-deterministic and about 15× less accurate. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op h3_rmsnorm +``` + +B200, block `norm1` + MSA modulation, B = 1, H = 5376, BF16: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.06 ms | 0.11 ms | 1.00 ms | 0.95 ms | +| 32768 | 0.31 ms | 0.75 ms | 2.50 ms | 5.20 ms | +| 131072 | 1.16 ms | 2.90 ms | 7.48 ms | 20.6 ms | + +The forward is a single pass that never materialises the gathered rows. At small S the backward +is dominated by the fixed cost of the stable sort and tile setup. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_rmsnorm.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end, incl. block norm1 + modulation +python scripts/check_operator.py --op h3_rmsnorm --candidate cuda --device cuda \ + --dtype bf16 --batch 2 --seq 257 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op h3_rmsnorm --out docs/usage/evidence/h3-rmsnorm-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-rmsnorm-b200/report.json +``` + +## Known Limitations + +- Bitwise parity with `nn.RMSNorm` is tied to PyTorch's vectorized layer-norm path (`N % 4 == 0`, + aligned tensors). Other shapes are rejected rather than approximated. +- The token-refiner blocks' own norm weights live in shard 14, which is not pinned. They are the + same module (`nn.RMSNorm(5376, eps=1e-5)`) as the four pinned norms tested here. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 144b06e27..01e6e5348 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -356,3 +356,24 @@ def h3_adaln_row_gather_backward( seg_first_tile: torch.Tensor, out_dtype: torch.dtype, ) -> torch.Tensor: ... +def h3_rmsnorm_forward( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_backward( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, + sorted_pos: torch.Tensor | None = None, + tile_begin: torch.Tensor | None = None, + tile_end: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index ef77fe7bc..02f7f2e37 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -50,6 +50,7 @@ def make_operator_inputs( "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, "adaln_row_gather": _make_adaln_row_gather_inputs, + "h3_rmsnorm": _make_h3_rmsnorm_inputs, } try: return builders[op_name](args, dtype, device) @@ -87,6 +88,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: f"x{6 * 3 * _h3_hidden(args)}", "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" f"->{batch * seq}", + "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -431,6 +433,19 @@ def _make_adaln_row_gather_inputs( return {"rows": rows, "timestep_indices": timestep_indices, "token_tags": token_tags} +def _make_h3_rmsnorm_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + weight = _floating_tensor((hidden,), args, torch.float32, device, 1).abs() + 0.5 + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "weight": weight.to(dtype), + "eps": 1e-5, + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 11ce047db..197db8874 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -300,6 +300,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("rows",), ), + "h3_rmsnorm": OperatorSpec( + name="h3_rmsnorm", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp", + "cuda": "rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp", + }, + grad_input_names=("x", "weight"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/rmsnorm.py b/rl_engine/kernels/ops/cuda/h3/rmsnorm.py new file mode 100644 index 000000000..53cc8204a --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/rmsnorm.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 RMSNorm with fused AdaLN modulation (RFC #420 ``h3_rmsnorm``).""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +KERNEL_ID = "rl_engine._C.h3_rmsnorm_forward" +BACKWARD_IMPL = "row_local_dx_tiled_dweight_sorted_table_grads" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_rmsnorm_backward" + + +def h3_rmsnorm_available() -> bool: + return bool( + _EXT_AVAILABLE and hasattr(_C, "h3_rmsnorm_forward") and hasattr(_C, "h3_rmsnorm_backward") + ) + + +class _H3RMSNormCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, index, eps): + hidden = x.shape[-1] + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward(x2, weight.contiguous(), float(eps), shift, scale, index) + ctx.save_for_backward(x2, weight, rstd, shift, scale, index) + ctx.x_shape = x.shape + ctx.modulated = index is not None + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + if not ctx.modulated: + dx, dweight = _C.h3_rmsnorm_backward(g2, x2, weight.contiguous(), rstd) + d_shift = d_scale = None + else: + positions = index.repeat(x2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, shift.shape[0]) + dx, dweight, d_shift, d_scale = _C.h3_rmsnorm_backward( + g2, x2, weight.contiguous(), rstd, shift, scale, index, *tiles + ) + d_shift, d_scale = d_shift.to(shift.dtype), d_scale.to(scale.dtype) + record_backward( + "h3_rmsnorm", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None + + +class H3RMSNormCudaOp: + """CUDA candidate. + + The statistics replay PyTorch's own RMSNorm reduction order, so the plain + norm is bitwise equal to ``nn.RMSNorm``; the modulated form gathers its + ``shift``/``scale`` rows in the kernel (no ``(S, H)`` intermediates) and + rounds each step where the eager expression does, so it is bitwise equal to + diffusers' ``n * (1.0 + scale[i]) + shift[i]`` as well. Rows are + independent of batch size and position. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_rmsnorm_available(): + raise RuntimeError( + "rl_engine._C lacks h3_rmsnorm_*; rebuild with csrc/cuda/h3/rmsnorm_modulate.cu" + ) + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + return _H3RMSNormCuda.apply(x, weight, None, None, None, eps) + + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index, check_range=check_range) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + return _H3RMSNormCuda.apply(x, weight, shift, scale, index.contiguous(), eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + return self.forward(x, weight, eps).float() diff --git a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py new file mode 100644 index 000000000..6ab850a07 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 RMSNorm and its fused AdaLN modulation (RFC #420 row ``h3_rmsnorm``). + +Every H3 RMSNorm is ``nn.RMSNorm(5376, eps=1e-5)`` with an affine weight in the +block dtype: the block ``norm1``/``norm2``, the token-refiner norms and the +final ``norm_out.norm``. Inside a transformer block (and in ``norm_out``) the +normalised rows are immediately modulated by per-row AdaLN parameters:: + + n = rms_norm(x, weight, eps) + out = n * (1.0 + scale[index]) + shift[index] # eager, rounded after every op + +``index`` is ``adaln_indices`` in a block and ``timestep_indices`` in +``norm_out``; ``shift``/``scale`` are ``(R, H)`` rows (views of the AdaLN +table). ``x`` is ``(..., S, H)`` and ``index`` has one entry per position ``S``. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +H3_NORM_EPS = 1e-5 +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> int: + for name, tensor in (("x", x), ("weight", weight)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype not in _DTYPES: + raise TypeError(f"{name} must be fp32, bf16 or fp16, got {tensor.dtype}") + if weight.dtype != x.dtype: + raise TypeError(f"weight dtype {weight.dtype} must match x dtype {x.dtype}") + if weight.device != x.device: + raise ValueError(f"weight is on {weight.device}, x is on {x.device}") + if x.dim() < 2 or x.numel() == 0: + raise ValueError(f"x must be a non-empty (..., N) tensor, got {tuple(x.shape)}") + if weight.shape != (x.shape[-1],): + raise ValueError(f"weight must be ({x.shape[-1]},), got {tuple(weight.shape)}") + if not eps > 0: + raise ValueError(f"eps must be positive, got {eps}") + return x.shape[-1] + + +def validate_h3_modulation( + x: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + hidden = x.shape[-1] + for name, tensor in (("shift", shift), ("scale", scale)): + if tensor.device != x.device or (same_dtype and tensor.dtype != x.dtype): + raise TypeError(f"{name} must match x's dtype and device") + if tensor.dim() != 2 or tensor.shape[1] != hidden: + raise ValueError(f"{name} must be (R, {hidden}), got {tuple(tensor.shape)}") + if shift.shape != scale.shape: + raise ValueError("shift and scale must have the same shape") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != x.device: + raise TypeError("index must be a 1-D int64 tensor on x's device") + if x.dim() < 2 or x.shape[-2] != index.shape[0]: + raise ValueError( + f"x must be (..., S, {hidden}) with S = len(index) = {index.shape[0]}, " + f"got {tuple(x.shape)}" + ) + if check_range and bool(((index < 0) | (index >= shift.shape[0])).any()): + raise IndexError(f"index must lie in [0, {shift.shape[0]})") + + +class _DeterministicModulation(torch.autograd.Function): + """Eager ``n * (1.0 + scale[i]) + shift[i]``; the ``shift``/``scale`` + gradients as FP32 per-row sums instead of ``index_select``'s atomic + scatter-add (``d_n`` is the eager VJP).""" + + @staticmethod + def forward(ctx, n, shift, scale, index): + t1 = 1.0 + scale.index_select(0, index) + ctx.save_for_backward(n, t1, index) + ctx.rows = shift.shape[0] + return n * t1 + shift.index_select(0, index) + + @staticmethod + def backward(ctx, grad): + n, t1, index = ctx.saved_tensors + hidden = n.shape[-1] + g = grad.float().reshape(-1, index.shape[0], hidden) + per_shift = g.sum(dim=0) + per_scale = (g * n.float().reshape(-1, index.shape[0], hidden)).sum(dim=0) + d_shift = per_shift.new_zeros((ctx.rows, hidden)) + d_scale = per_scale.new_zeros((ctx.rows, hidden)) + for row in range(ctx.rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_shift[row] = per_shift.index_select(0, positions).sum(dim=0) + d_scale[row] = per_scale.index_select(0, positions).sum(dim=0) + return grad * t1, d_shift.to(t1.dtype), d_scale.to(t1.dtype), None + + +class NativeH3RMSNormOp: + """PyTorch reference. + + ``forward`` / ``forward_modulated`` are the provider path (``F.rms_norm`` and + the eager modulation expression, with a deterministic shift/scale gradient; + the raw diffusers path is ``h3_provider.provider_norm_modulate``). ``forward_fp32`` / + ``forward_modulated_fp32`` are the golden: the same math in FP64 with no + intermediate rounding, returned in FP32. + """ + + op_class = "reduction" + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + hidden = validate_h3_rmsnorm(x, weight, eps) + return F.rms_norm(x, (hidden,), weight, eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + rstd = torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) + return (x64 * rstd * weight.double()).float() + + def forward_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + validate_h3_modulation(x, shift, scale, index) + n = self.forward(x, weight, eps) + return _DeterministicModulation.apply(n, shift, scale, index) + + def forward_modulated_fp32(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + # The golden may be fed FP32 golden modulation rows for a BF16 x. + validate_h3_modulation(x, shift, scale, index, same_dtype=False) + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) * weight.double() + out = n * (1.0 + scale.double().index_select(0, index)) + shift.double().index_select( + 0, index + ) + return out.float() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index ba6c57f05..8fa1d1453 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -221,6 +221,8 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_ADALN_ROW_GATHER = ( "rl_engine.kernels.ops.cuda.h3.adaln_row_gather.H3AdaLNRowGatherCudaOp" ) + PYTORCH_H3_RMSNORM = "rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp" + CUDA_H3_RMSNORM = "rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -647,6 +649,7 @@ def __init__(self): OpBackend.CUDA_H3_ADALN_ROW_GATHER, OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, ], + "h3_rmsnorm": [OpBackend.CUDA_H3_RMSNORM, OpBackend.PYTORCH_H3_RMSNORM], }, "rocm": { "logp": [ @@ -691,6 +694,7 @@ def __init__(self): "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -739,6 +743,7 @@ def __init__(self): "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -768,6 +773,7 @@ def __init__(self): "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 06bfe252b..fe010c7e3 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -33,11 +33,13 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_norm_modulate, provider_time_embedder, provider_time_proj, ) REPO_ROOT = Path(__file__).resolve().parents[2] +HIDDEN = 5376 @dataclass(frozen=True) @@ -113,6 +115,45 @@ def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: ), ] + +def adaln_indices(ctx: dict[str, Any]) -> torch.Tensor: + return ctx["timestep_indices"] * 3 + ctx["token_tags"] + + +def _norm1_modulated(ctx: dict[str, Any], modulation, call): + shift_msa, scale_msa = modulation[0], modulation[1] + weight = ctx["weights"]["transformer_blocks.0.norm1.weight"] + return call(ctx["hidden"], weight, shift_msa, scale_msa, adaln_indices(ctx)) + + +STAGES.append( + Stage( + name="h3_rmsnorm", + op_type="h3_rmsnorm", + # Block norm1 followed by the MSA shift/scale of the projection stage. + candidate=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated + ), + provider=lambda ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], provider_norm_modulate + ), + golden=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated_fp32 + ), + provider_bitwise_isolated=True, # replays nn.RMSNorm's reduction order + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + +# The backward replay covers the conditioning chain up to the row gather. +CONDITIONING_STAGES = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", +) + # Parameters whose gradients the backward replay reports, in chain order. GRAD_LEAVES = ( "time_embedder.linear_1.weight", @@ -151,6 +192,13 @@ def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dic "timestep_indices": timestep_indices, "token_tags": token_tags, "weights": weights, + # Stand-in for the packed hidden states (the patch/text projections are + # other RFC rows): BF16 (1, S, H), seeded. + "hidden": torch.randn( + (1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 2) + ) + .to(torch.bfloat16) + .cuda(), } @@ -179,16 +227,24 @@ def run_case( } chained = provider = golden = None first_drift = first_isolated = None + # Each mode sees its own earlier outputs (a stage may read any earlier stage). + history = {"candidate": {}, "provider": {}, "golden": {}} + ctx_c = {**ctx, "history": history["candidate"]} + ctx_p = {**ctx, "history": history["provider"]} + ctx_g = {**ctx, "history": history["golden"]} for stage in stages: op = registry.get_op(stage.op_type, device="cuda") gold = golden_op(registry, stage.op_type) with torch.no_grad(): - isolated = stage.candidate(op, ctx, provider) + isolated = stage.candidate(op, ctx_p, provider) chained_input = chained - chained = stage.candidate(op, ctx, chained_input) - repeat = stage.candidate(op, ctx, chained_input) - golden = stage.golden(gold, ctx, golden) - provider = stage.provider(ctx, provider) + chained = stage.candidate(op, ctx_c, chained_input) + repeat = stage.candidate(op, ctx_c, chained_input) + golden = stage.golden(gold, ctx_g, golden) + provider = stage.provider(ctx_p, provider) + history["candidate"][stage.name] = chained + history["provider"][stage.name] = provider + history["golden"][stage.name] = golden entry = { "stage": stage.name, "backend": type(op).__name__, @@ -220,7 +276,7 @@ def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torc leaves = { name: ctx["weights"][name].detach().clone().requires_grad_(True) for name in GRAD_LEAVES } - run_ctx = {**ctx, "weights": {**ctx["weights"], **leaves}} + run_ctx = {**ctx, "weights": {**ctx["weights"], **leaves}, "history": {}} value = None if mode == "candidate_fused": from rl_engine.kernels.ops.cuda.h3.adaln_modulation import H3AdaLNModulationCudaOp @@ -231,7 +287,7 @@ def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torc value, *adaln_params(run_ctx), ctx["timestep_indices"], ctx["token_tags"] ) else: - for stage in STAGES: + for stage in (s for s in STAGES if s.name in CONDITIONING_STAGES): if mode == "provider": value = stage.provider(run_ctx, value) elif mode == "candidate": diff --git a/rl_engine/testing/h3_manifest.json b/rl_engine/testing/h3_manifest.json index 52b1989b1..176ddcf04 100644 --- a/rl_engine/testing/h3_manifest.json +++ b/rl_engine/testing/h3_manifest.json @@ -35,7 +35,13 @@ "time_embedder.linear_2.weight": {"dtype": "float32", "shape": [2688, 5376]}, "time_embedder.linear_2.bias": {"dtype": "float32", "shape": [2688]}, "transformer_blocks.0.adaln_proj.linear.weight": {"dtype": "bfloat16", "shape": [96768, 2688]}, - "transformer_blocks.0.adaln_proj.linear.bias": {"dtype": "bfloat16", "shape": [96768]} + "transformer_blocks.0.adaln_proj.linear.bias": {"dtype": "bfloat16", "shape": [96768]}, + "transformer_blocks.0.norm1.weight": {"dtype": "bfloat16", "shape": [5376]}, + "transformer_blocks.0.norm2.weight": {"dtype": "bfloat16", "shape": [5376]}, + "token_refiner.final_norm.weight": {"dtype": "bfloat16", "shape": [5376]}, + "norm_out.norm.weight": {"dtype": "bfloat16", "shape": [5376]}, + "norm_out.linear.weight": {"dtype": "bfloat16", "shape": [10752, 2688]}, + "norm_out.linear.bias": {"dtype": "bfloat16", "shape": [10752]} }, "reference_implementation": { "repo": "huggingface/diffusers", @@ -53,7 +59,9 @@ "precision": { "timestep_sinusoid": "float32", "timestep_mlp": "float32 weights, float32 activations", - "adaln_projection": "SiLU in float32, cast to bfloat16, bfloat16 weight/bias/output" + "adaln_projection": "SiLU in float32, cast to bfloat16, bfloat16 weight/bias/output", + "rmsnorm": "nn.RMSNorm(5376, eps=1e-5), bfloat16 affine weight, statistics in float32", + "final_modulation": "norm_out.linear(silu(temb).bfloat16) -> (shift, scale), indexed by timestep_indices" } } } diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index 6cedcdf60..05c7ebec9 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -65,3 +65,19 @@ def provider_adaln_row_gather( adaln_indices = timestep_indices * 3 + token_tags return tuple(tensor.index_select(0, adaln_indices) for tensor in modulation) + + +def provider_norm_modulate( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``norm(x) * (1.0 + scale[i]) + shift[i]`` with ``index_select``; block and norm_out.""" + + norm_hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return norm_hidden_states * (1.0 + scale.index_select(0, indices)) + shift.index_select( + 0, indices + ) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index e0a5939b5..dfecec419 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -20,6 +20,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_norm_modulate, provider_time_embedder, provider_time_proj, ) @@ -349,15 +350,134 @@ def grad_of(fn): } +# --------------------------------------------------------------------------- # +# h3_rmsnorm (block norm1 + MSA modulation) +# --------------------------------------------------------------------------- # + +NORM_SEQ_LENS = (4097, 32768, 131072) + + +def _norm_inputs(seq: int, seed: int = 0): + weight = h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0].bfloat16() + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).bfloat16() + shift, scale = table.view(9, 6, 5376)[:, 0], table.view(9, 6, 5376)[:, 1] + ti, tags = h3_packed_layout(seq, 3, seed=seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + return x, weight, shift, scale, ti * 3 + tags + + +def _norm_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("h3_rmsnorm", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, weight, shift, scale, index = _norm_inputs(seq) + grad = torch.randn_like(x) + + def backward(fn, x=x, weight=weight, shift=shift, scale=scale, grad=grad): + leaves = [t.detach().requires_grad_(True) for t in (x, weight, shift, scale)] + fn(*leaves).backward(grad) + + cases.append( + { + "op": "h3_rmsnorm", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 2 * seq * 5376 * 2, # read x, write y + "candidate": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + op.forward_modulated(x, w, sh, sc, i, check_range=False) + ), + "provider": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + provider_norm_modulate(x, w, sh, sc, i) + ), + "candidate_backward": lambda b=backward, i=index: b( + lambda x_, w_, sh_, sc_: op.forward_modulated( + x_, w_, sh_, sc_, i, check_range=False + ) + ), + "provider_backward": lambda b=backward, i=index: b( + lambda x_, w_, sh_, sc_: provider_norm_modulate(x_, w_, sh_, sc_, i) + ), + } + ) + return cases + + +def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("h3_rmsnorm", device="cuda") + weights = load_h3_conditioning_weights("cuda") + plain = {} + for name in ( + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", + ): + x = torch.randn(2, 777, 5376, device="cuda").bfloat16() + ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) + plain[name] = bool(torch.equal(op(x, weights[name]), ref)) + x, weight, shift, scale, index = _norm_inputs(4097, seed=1) + modulated = bool( + torch.equal( + op.forward_modulated(x, weight, shift, scale, index), + provider_norm_modulate(x, weight, shift, scale, index), + ) + ) + invariant = bool( + torch.equal( + op.forward_modulated(x[:, 100:140], weight, shift, scale, index[100:140])[0], + op.forward_modulated(x, weight, shift, scale, index)[0, 100:140], + ) + ) + + # Backward against FP64, for the CUDA op and the diffusers expression. + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, weight, shift, scale)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden(x_, w_, sh_, sc_): + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * w_ + return n * (1 + sc_.index_select(0, index)) + sh_.index_select(0, index) + + ref = grads(golden, torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op.forward_modulated(*t, index)), + ("provider", lambda *t: provider_norm_modulate(*t, index)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref) + }, + } + return { + "plain_bitwise_vs_nn_rmsnorm": plain, + "modulated_bitwise_vs_diffusers": modulated, + "rows_batch_invariant": invariant, + "backward": backward, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, "adaln_projection_3mod": _projection_perf, "adaln_row_gather": _gather_perf, + "h3_rmsnorm": _norm_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, "timestep_mlp_fp32": _mlp_accuracy, "adaln_projection_3mod": _projection_accuracy, "adaln_row_gather": _gather_accuracy, + "h3_rmsnorm": _norm_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index c36ade2f2..49d502f12 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -231,6 +231,7 @@ def fake_load_backend(backend): "timestep_mlp_fp32", "adaln_projection_3mod", "adaln_row_gather", + "h3_rmsnorm", ) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index f2039f0aa..da8bc6774 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -345,11 +345,76 @@ def plot_gather(report: dict[str, Any]): return fig +def plot_rmsnorm(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "h3_rmsnorm (block norm1 + MSA modulation, H = 5376, BF16)" + ) + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("diffusers fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "diffusers bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + acc = report["accuracy"] + bwd = acc["backward"] + keys = ["dx", "dweight", "dshift", "dscale"] + fwd_ok = acc["modulated_bitwise_vs_diffusers"] and all( + acc["plain_bitwise_vs_nn_rmsnorm"].values() + ) + _style(right, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") + + def repeat(mode: str) -> str: + return "yes" if bwd[mode]["repeat_bitwise_equal"] else "no" + + _bars( + right, + keys, + [ + ( + f"RL-Kernel CUDA (repeat-bitwise: {repeat('cuda')})", + RL_KERNEL, + [bwd["cuda"]["rel_error"][k] for k in keys], + ), + ( + f"diffusers (repeat-bitwise: {repeat('provider')})", + PROVIDER, + [bwd["provider"]["rel_error"][k] for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + right.text( + 0.98, + 0.80, + f"forward bitwise equal to diffusers: {'yes' if fwd_ok else 'NO'}", + transform=right.transAxes, + ha="right", + fontsize=8, + color=INK, + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, "adaln_projection_3mod": plot_projection, "adaln_row_gather": plot_gather, + "h3_rmsnorm": plot_rmsnorm, } diff --git a/setup.py b/setup.py index fcefa9016..fb671d0f0 100644 --- a/setup.py +++ b/setup.py @@ -160,6 +160,7 @@ def get_extensions(): cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") cuda_sources.append("csrc/cuda/h3/det_linear.cu") cuda_sources.append("csrc/cuda/h3/adaln_row_gather.cu") + cuda_sources.append("csrc/cuda/h3/rmsnorm_modulate.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/test_h3_rmsnorm.py b/tests/h3/test_h3_rmsnorm.py new file mode 100644 index 000000000..f0d013bfb --- /dev/null +++ b/tests/h3/test_h3_rmsnorm.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``h3_rmsnorm``: RMSNorm (eps 1e-5) and its fused AdaLN modulation. + +* forward is bitwise equal to ``nn.RMSNorm`` and to diffusers' + ``norm(x) * (1.0 + scale[i]) + shift[i]`` (block: ``adaln_indices``; + ``norm_out``: ``timestep_indices``) on the pinned norm weights; +* rows are independent of batch size and position; +* backward is deterministic and checked against an FP64 golden; +* malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import NativeH3RMSNormOp +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_norm_modulate + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NORMS = ( + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", +) + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.rmsnorm import H3RMSNormCudaOp, h3_rmsnorm_available + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3RMSNormCudaOp() + + +def _x(shape, dtype=torch.bfloat16, seed=0, scale=2.0): + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(shape, generator=g) * scale).to(dtype).cuda() + + +def _table(rows, hidden, chunks, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(rows, chunks * hidden, generator=g) * 0.5).to(dtype).cuda() + return table.chunk(chunks, dim=-1) # strided (R, H) views, like the AdaLN outputs + + +def _fp64_modulated(x, weight, shift, scale, index, eps=1e-5): + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * weight.double() + return n * (1 + scale.double().index_select(0, index)) + shift.double().index_select(0, index) + + +class TestReference: + def test_golden_close_to_provider(self): + x, w = _x((3, 64), torch.float32).cpu(), torch.rand(64) + 0.5 + op = NativeH3RMSNormOp() + torch.testing.assert_close(op.forward(x, w), op.forward_fp32(x, w), atol=1e-6, rtol=1e-6) + + def test_rejects_bad_inputs(self): + op = NativeH3RMSNormOp() + x, w = torch.randn(2, 8), torch.ones(8) + with pytest.raises(TypeError): + op.forward(x, w.double()) + with pytest.raises(TypeError): + op.forward(x, w.bfloat16()) + with pytest.raises(ValueError): + op.forward(x, torch.ones(7)) + with pytest.raises(ValueError): + op.forward(x, w, eps=0.0) + shift, scale = torch.zeros(3, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op.forward_modulated(x, w, shift, scale, torch.tensor([0, 3])) + with pytest.raises(ValueError): # one index per position S + op.forward_modulated(x, w, shift, scale, torch.tensor([0])) + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("name", NORMS) + def test_bitwise_equal_to_nn_rmsnorm_on_pinned_weights(self, h3_weights_cpu, name): + weight = h3_weights_cpu[name].cuda() + module = torch.nn.RMSNorm(5376, eps=1e-5).cuda().bfloat16() + module.weight.data.copy_(weight) + x = _x((2, 777, 5376), seed=1) + assert torch.equal(_cuda_op()(x, weight), module(x)) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [12, 256, 4096, 5376]) + def test_bitwise_equal_to_f_rms_norm(self, dtype, hidden): + x = _x((5, 33, hidden), dtype, seed=hidden, scale=30.0) + w = (torch.rand(hidden, device="cuda") * 2).to(dtype) + ref = torch.nn.functional.rms_norm(x, (hidden,), w, 1e-5) + assert torch.equal(_cuda_op()(x, w), ref) + + def test_block_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["transformer_blocks.0.norm1.weight"].cuda() + shift, scale, *_ = _table(9, 5376, 6, seed=2) # (3T, H) rows, T = 3 + ti, tags = h3_packed_layout(4097, 3, seed=2) + index = ti * 3 + tags + x = _x((2, 4097, 5376), seed=2) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, index)) + + def test_norm_out_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["norm_out.norm.weight"].cuda() + shift, scale = _table(3, 5376, 2, seed=3) # (T, H) rows indexed by timestep + ti, _ = h3_packed_layout(1000, 3, seed=3) + x = _x((1, 1000, 5376), seed=3) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, ti) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, ti)) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_modulation_bitwise_in_every_dtype(self, dtype): + """FP32 must not contract the separately rounded eager ops into an FMA.""" + + w = (torch.rand(5376, device="cuda") + 0.5).to(dtype) + shift, scale, *_ = _table(9, 5376, 6, dtype=dtype, seed=9) + ti, tags = h3_packed_layout(333, 3, seed=9) + index = ti * 3 + tags + x = _x((2, 333, 5376), dtype, seed=9) + ours = _cuda_op().forward_modulated(x, w, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, w, shift, scale, index)) + + def test_rows_are_batch_and_position_invariant(self): + w = (torch.rand(5376, device="cuda") + 0.5).bfloat16() + shift, scale, *_ = _table(9, 5376, 6, seed=4) + ti, tags = h3_packed_layout(300, 3, seed=4) + index = ti * 3 + tags + x = _x((3, 300, 5376), seed=4) + full = _cuda_op().forward_modulated(x, w, shift, scale, index) + part = _cuda_op().forward_modulated(x[2:3, 50:90], w, shift, scale, index[50:90]) + assert torch.equal(part[0], full[2, 50:90]) + + def test_rejects_unsupported(self): + w = torch.ones(6, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="multiple of 4"): + _cuda_op()(_x((2, 6)), w) + with pytest.raises(ValueError): + _cuda_op()(_x((2, 8)).cpu(), torch.ones(8).bfloat16()) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, fn, tensors, grad): + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad) + return [leaf.grad for leaf in leaves] + + def test_modulated_backward_against_fp64(self): + g = torch.Generator(device="cpu").manual_seed(5) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.5).bfloat16().cuda() + ti, tags = h3_packed_layout(2049, 3, seed=5) + index = ti * 3 + tags + x = _x((2, 2049, 5376), seed=5) + grad = _x((2, 2049, 5376), seed=6, scale=1.0) + op = _cuda_op() + + def ours(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return op.forward_modulated(x_, w_, shift, scale, index) + + def golden(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return _fp64_modulated(x_, w_, shift, scale, index) + + got = self._grads(ours, (x, w, table), grad) + ref = self._grads(golden, (x.double(), w.double(), table.double()), grad.double()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + rel = ((g_.double() - r).abs().max() / r.abs().max()).item() + # Two BF16 roundings remain: the eager VJP uses round(1 + scale), and + # every gradient is rounded once to BF16 (2^-9 relative each). + assert rel < 1e-2, (name, rel) + again = self._grads(ours, (x, w, table), grad) + assert all(torch.equal(a, b) for a, b in zip(got, again)) + + def test_plain_backward_against_fp64_and_row_local_dx(self): + w = (torch.rand(5376, generator=torch.Generator().manual_seed(7)) + 0.5).bfloat16().cuda() + x = _x((4, 64, 5376), seed=7) + grad = _x((4, 64, 5376), seed=8, scale=1.0) + op = _cuda_op() + got = self._grads(lambda a, b: op(a, b), (x, w), grad) + + def golden(a, b): + return a * torch.rsqrt(a.square().mean(-1, keepdim=True) + 1e-5) * b + + ref = self._grads(golden, (x.double(), w.double()), grad.double()) + for g, r in zip(got, ref): + assert ((g.double() - r).abs().max() / r.abs().max()).item() < 5e-3 + part = self._grads(lambda a, b: op(a, b), (x[1:2, 10:20], w), grad[1:2, 10:20]) + assert torch.equal(part[0][0], got[0][1, 10:20]) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + assert ( + type(KernelRegistry().get_op("h3_rmsnorm", device="cuda")).__name__ == "H3RMSNormCudaOp" + ) diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 455a88638..2b165c088 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -53,6 +53,7 @@ def test_all_ws1_single_ops_are_registered(): "timestep_mlp_fp32", "adaln_projection_3mod", "adaln_row_gather", + "h3_rmsnorm", } <= names From 729978430a95ee8a82d7bdf07147a729ebca88c6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:13:55 +0000 Subject: [PATCH 10/89] [WS1][MiniMax-H3] h3_rmsnorm: B200 evidence report and figure Produced by scripts/h3_evidence.py at f36aa60 on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-rmsnorm.md | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 3bd0b2f6e..957670552 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -67,7 +67,7 @@ closed. same scheme as [`adaln_row_gather`](h3-adaln-row-gather.md). Diffusers' `index_select` backward accumulates `dshift`/`dscale` with BF16 atomics, so it is - non-deterministic and about 15× less accurate. + non-deterministic and about 15–20× less accurate (its error varies from run to run). ## Performance Notes @@ -79,13 +79,24 @@ B200, block `norm1` + MSA modulation, B = 1, H = 5376, BF16: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.06 ms | 0.11 ms | 1.00 ms | 0.95 ms | +| 4097 | 0.06 ms | 0.11 ms | 1.08 ms | 0.95 ms | | 32768 | 0.31 ms | 0.75 ms | 2.50 ms | 5.20 ms | -| 131072 | 1.16 ms | 2.90 ms | 7.48 ms | 20.6 ms | +| 131072 | 1.17 ms | 2.90 ms | 7.40 ms | 20.6 ms | The forward is a single pass that never materialises the gathered rows. At small S the backward is dominated by the fixed cost of the stable sort and tile setup. +## Evidence + +![h3_rmsnorm on B200: latency and backward accuracy](../usage/evidence/h3-rmsnorm-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-rmsnorm-b200/report.json), written by +`scripts/h3_evidence.py` from a clean tree at commit `f36aa60`. It also records: + +- bitwise equality with `nn.RMSNorm` for all four pinned norm weights; +- bitwise equality with diffusers for the modulation; +- row invariance. + ## Tests ```bash From e7bc53aa89f3c4b2c3d6045c0b2ccfec8c091391 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 02:51:41 +0000 Subject: [PATCH 11/89] [WS1][CUDA][MiniMax-H3] adaln_gate_residual: gated residual with in-kernel gate gather RFC #420 step 6, row `adaln_gate_residual` (stacked on h3_rmsnorm). - residual + gate[adaln_index] * sublayer_output in one pass, the gate row read from the AdaLN table view in the kernel; the two roundings sit where the eager expression rounds (__fmul_rn/__fadd_rn, no FMA contraction), so the forward is bitwise equal to diffusers in bf16/fp16/fp32. 1.7x faster forward at S = 131072. - Backward: d_residual = grad; d_sublayer bitwise equal to the eager VJP (exact product, one rounding); d_gate a deterministic FP32 segment sum (diffusers: non-deterministic BF16 atomics, 13-49x larger error, varying run to run). The PyTorch reference gets an equally deterministic gate gradient. - Chain stage residual + gate_msa * (stand-in attention output), covered by the end-to-end test; registry, gtest, report/figure hooks, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 242 ++++++++++++++++++ csrc/ops.cpp | 15 ++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-adaln-gate-residual.md | 91 +++++++ rl_engine/_C.pyi | 13 + rl_engine/kernels/gtest/operator_inputs.py | 18 ++ rl_engine/kernels/gtest/operator_specs.py | 11 + .../kernels/ops/cuda/h3/gate_residual.py | 83 ++++++ .../kernels/ops/pytorch/h3/gate_residual.py | 97 +++++++ rl_engine/kernels/registry.py | 11 + rl_engine/testing/h3_chain.py | 38 +++ rl_engine/testing/h3_provider.py | 8 + rl_engine/testing/h3_report.py | 107 ++++++++ rl_engine/tests/test_dispatch.py | 1 + scripts/plot_h3_evidence.py | 72 ++++++ setup.py | 1 + tests/h3/test_h3_adaln_gate_residual.py | 115 +++++++++ tests/test_ws1_gtest_gpu.py | 1 + 19 files changed, 926 insertions(+) create mode 100644 csrc/cuda/h3/gate_residual.cu create mode 100644 docs/operators/h3-adaln-gate-residual.md create mode 100644 rl_engine/kernels/ops/cuda/h3/gate_residual.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/gate_residual.py create mode 100644 tests/h3/test_h3_adaln_gate_residual.py diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu new file mode 100644 index 000000000..18ede9e72 --- /dev/null +++ b/csrc/cuda/h3/gate_residual.cu @@ -0,0 +1,242 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 gated residual (RFC #420 `adaln_gate_residual`). +// +// out[r, j] = cast(residual[r, j] + cast(gate[i, j] * y[r, j])) i = index[r % S] +// +// The gate row is read from a strided view of the AdaLN table (gate_msa or +// gate_mlp) inside the kernel, and the two roundings sit exactly where the +// eager `residual + gate.index_select(0, i) * y` rounds (__fmul_rn/__fadd_rn +// keep FP32 from contracting them into one FMA), so the result is +// bitwise equal to diffusers. Elementwise: every output depends only on its +// own inputs. +// +// Backward: d_residual = grad; d_y = cast(grad * gate) (the product of two +// 16-bit values is exact in FP32, so this equals the eager VJP bitwise); +// d_gate[i] = sum over positions mapped to i of grad * y, an FP32 segmented +// sum over positions sorted stably by row, in fixed tiles folded in order. + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +__device__ __forceinline__ float to_f(float v) { return v; } +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ float from_f(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } + +template +__device__ __forceinline__ float round_t(float v) { return to_f(from_f(v)); } + +struct Gate { + const void* rows; // (R, N) view, row stride `stride` + int64_t stride; + const int64_t* index; // (S,) + int64_t seq; +}; + +// One block per row (grid-strided); 16-byte vectors when the row allows it. +// `kVecElems` elements are loaded together; each is computed independently. +template +__global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __restrict__ b, + T* __restrict__ out, int64_t rows, int64_t n, + Gate gate, bool vectorized) { + constexpr int kVecElems = 16 / sizeof(T); + for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { + const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const T* ar = a + r * n; + const T* br = kDy ? nullptr : b + r * n; + T* orow = out + r * n; + if (vectorized) { + for (int64_t v = threadIdx.x; v < n / kVecElems; v += blockDim.x) { + const uint4 av = reinterpret_cast(ar)[v]; + const uint4 gv = reinterpret_cast(g)[v]; + uint4 bv = av; + if (!kDy) bv = reinterpret_cast(br)[v]; + const T* ae = reinterpret_cast(&av); + const T* ge = reinterpret_cast(&gv); + const T* be = reinterpret_cast(&bv); + uint4 ov; + T* oe = reinterpret_cast(&ov); +#pragma unroll + for (int k = 0; k < kVecElems; ++k) { + if (kDy) { + oe[k] = from_f(to_f(ae[k]) * to_f(ge[k])); // grad * gate + } else { + const float p = round_t(__fmul_rn(to_f(ge[k]), to_f(be[k]))); // gate * y + oe[k] = from_f(__fadd_rn(to_f(ae[k]), p)); // residual + p + } + } + reinterpret_cast(orow)[v] = ov; + } + } else { + for (int64_t j = threadIdx.x; j < n; j += blockDim.x) { + if (kDy) { + orow[j] = from_f(to_f(ar[j]) * to_f(g[j])); + } else { + const float p = round_t(__fmul_rn(to_f(g[j]), to_f(br[j]))); + orow[j] = from_f(__fadd_rn(to_f(ar[j]), p)); + } + } + } + } +} + +// partial[tile, j] = sum over sorted positions of the tile (ascending) of grad * y +template +__global__ void gate_grad_partial_kernel(const T* __restrict__ grad, const T* __restrict__ y, + const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int64_t n) { + const int64_t tile = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t e = sorted_pos[p] * n + j; + acc = fmaf(to_f(grad[e]), to_f(y[e]), acc); + } + partial[tile * n + j] = acc; +} + +template +__global__ void gate_grad_fold_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + T* __restrict__ out, int64_t n) { + const int64_t seg = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t t = seg_first_tile[seg]; t < seg_first_tile[seg + 1]; ++t) acc += partial[t * n + j]; + out[seg * n + j] = from_f(acc); +} + +#define H3_DISPATCH(TYPE, NAME, ...) \ + [&] { \ + switch (TYPE) { \ + case at::kFloat: { using T = float; __VA_ARGS__(); break; } \ + case at::kHalf: { using T = __half; __VA_ARGS__(); break; } \ + case at::kBFloat16: { using T = __nv_bfloat16; __VA_ARGS__(); break; } \ + default: TORCH_CHECK(false, NAME, ": unsupported dtype ", TYPE); \ + } \ + }() + +Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torch::Tensor& x) { + TORCH_CHECK(gate.is_cuda() && gate.device() == x.device() && gate.scalar_type() == x.scalar_type(), + "gate must be a CUDA tensor with the activations' dtype"); + TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, + "gate must be (R, N) with unit column stride"); + TORCH_CHECK(index.is_cuda() && index.scalar_type() == at::kLong && index.dim() == 1 && + index.is_contiguous() && index.numel() > 0, + "index must be a non-empty contiguous int64 CUDA tensor"); + TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; +} + +void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device() && t.scalar_type() == like.scalar_type(), + name, " must match the residual's device and dtype"); + TORCH_CHECK(t.sizes() == like.sizes() && t.is_contiguous(), name, + " must be contiguous with the residual's shape"); +} + +unsigned row_blocks(int64_t rows) { + return static_cast(std::min(rows, 1 << 30)); +} + +bool can_vectorize(const torch::Tensor& gate, int64_t n, std::initializer_list acts) { + const int64_t vec = 16 / gate.element_size(); + if (n % vec != 0 || (gate.stride(0) * gate.element_size()) % 16 != 0) return false; + if (reinterpret_cast(gate.data_ptr()) % 16 != 0) return false; + for (const auto* t : acts) { + if (reinterpret_cast(t->data_ptr()) % 16 != 0) return false; + } + return true; +} + +} // namespace + +// residual, y: (M, N); gate: (R, N) view; index: (S,), M % S == 0. +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index) { + TORCH_CHECK(residual.is_cuda() && residual.dim() == 2 && residual.is_contiguous() && + residual.numel() > 0, + "residual must be a non-empty contiguous (M, N) CUDA tensor"); + check_act(y, "y", residual); + const Gate g = make_gate(gate, index, residual); + const c10::cuda::CUDAGuard guard(residual.device()); + auto out = torch::empty_like(residual); + auto stream = at::cuda::getCurrentCUDAStream(); + const bool vec = can_vectorize(gate, residual.size(1), {&residual, &y, &out}); + H3_DISPATCH(residual.scalar_type(), "h3_gate_residual_forward", [&] { + gate_residual_rows_kernel<<>>( + reinterpret_cast(residual.data_ptr()), reinterpret_cast(y.data_ptr()), + reinterpret_cast(out.data_ptr()), residual.size(0), residual.size(1), g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// Returns {d_y, d_gate (R, N) in the gate dtype}. d_residual is the incoming grad. +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + torch::Tensor tile_begin, + torch::Tensor tile_end, + torch::Tensor seg_first_tile) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + check_act(y, "y", grad); + const Gate g = make_gate(gate, index, grad); + for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kLong && t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors"); + } + const int64_t rows = grad.size(0); + const int64_t n = grad.size(1); + const int64_t segments = seg_first_tile.numel() - 1; + TORCH_CHECK(segments == gate.size(0), "segments must equal the gate rows"); + const int64_t tiles = tile_begin.numel(); + const c10::cuda::CUDAGuard guard(grad.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto dy = torch::empty_like(y); + auto partial = torch::empty({std::max(tiles, 1), n}, grad.options().dtype(at::kFloat)); + auto dgate = torch::empty({segments, n}, grad.options()); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward", [&] { + const T* gp = reinterpret_cast(grad.data_ptr()); + const T* yp = reinterpret_cast(y.data_ptr()); + const bool vec = can_vectorize(gate, n, {&grad, &dy}); + gate_residual_rows_kernel<<>>( + gp, nullptr, reinterpret_cast(dy.data_ptr()), rows, n, g, vec); + if (tiles > 0) { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, stream>>>( + gp, yp, sorted_pos.data_ptr(), tile_begin.data_ptr(), + tile_end.data_ptr(), partial.data_ptr(), n); + } + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(dgate.data_ptr()), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dy, dgate}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index f2288efe4..08085e721 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -527,6 +527,14 @@ std::vector h3_rmsnorm_backward( c10::optional index, c10::optional sorted_pos, c10::optional tile_begin, c10::optional tile_end, c10::optional seg_first_tile); +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index); +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + torch::Tensor tile_begin, + torch::Tensor tile_end, + torch::Tensor seg_first_tile); #endif // PyBind11 Module Registration @@ -819,5 +827,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("index") = py::none(), py::arg("sorted_pos") = py::none(), py::arg("tile_begin") = py::none(), py::arg("tile_end") = py::none(), py::arg("seg_first_tile") = py::none()); + m.def("h3_gate_residual_forward", &h3_gate_residual_forward, + "residual + gate[index] * y with the gate row gathered in-kernel (eager rounding order)", + py::arg("residual"), py::arg("y"), py::arg("gate"), py::arg("index")); + m.def("h3_gate_residual_backward", &h3_gate_residual_backward, + "Gated-residual backward: exact dy, deterministic sorted segment sum for dgate", + py::arg("grad"), py::arg("y"), py::arg("gate"), py::arg("index"), py::arg("sorted_pos"), + py::arg("tile_begin"), py::arg("tile_end"), py::arg("seg_first_tile")); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index af8c1863f..f6db1bc36 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -35,6 +35,7 @@ nav: - operators/h3-adaln-projection.md - operators/h3-adaln-row-gather.md - operators/h3-rmsnorm.md + - operators/h3-adaln-gate-residual.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index b687a65b5..7042d4561 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -37,4 +37,5 @@ Every operator page should include: - [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) - [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) - [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) +- [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md new file mode 100644 index 000000000..7a3010185 --- /dev/null +++ b/docs/operators/h3-adaln-gate-residual.md @@ -0,0 +1,91 @@ +# MiniMax-H3 AdaLN Gated Residual + +## Summary + +`adaln_gate_residual` adds each sublayer's output back to the residual stream through the +per-row AdaLN gate (RFC #420, WS1 step 6). In every H3 block this happens twice, after attention +with `gate_msa` and after the FFN with `gate_mlp`: + +```text +hidden = residual + gate[adaln_indices] * sublayer_output # RFC #420 §4 order +``` + +`gate` is an `(R, H)` row view of the [AdaLN projection](h3-adaln-projection.md) table. The op +gathers the gate row inside the kernel, so the gathered `(S, H)` gate is never materialised. + +## Entry Point + +```python +op = kernel_registry.get_op("adaln_gate_residual", device="cuda") +hidden = op(residual, attn_output, gate_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA | `rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp` | `rl_engine._C.h3_gate_residual_{forward,backward}` | Forward and `d_sublayer` bitwise equal to diffusers | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp` | n/a | Eager forward with a deterministic gate gradient; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `residual`, `sublayer_output` | `(..., S, H)` | bf16 / fp16 / fp32 | Same shape and dtype | +| `gate` | `(R, H)` row view | same | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; shared across the batch | + +## Numerics + +- **Forward.** Each element computes `p = round(gate * y)` and then `out = round(residual + p)`. + `__fmul_rn`/`__fadd_rn` keep FP32 from contracting the two operations into one FMA. The + roundings fall exactly where the eager expression rounds, so the forward is **bitwise equal to + diffusers** in BF16, FP16 and FP32. Rows are independent. +- **Backward.** + - `d_residual` is the incoming gradient. + - `d_sublayer = round(grad * gate)`. The product of two 16-bit values is exact in FP32, so this + is bitwise equal to the eager VJP. + - `d_gate` is an FP32 segmented sum of `grad * sublayer_output` over positions sorted stably + by table row, using the same tiles as [`adaln_row_gather`](h3-adaln-row-gather.md), rounded + once. + + Diffusers' `d_gate` uses `index_select`'s BF16 atomic scatter-add instead. It is + non-deterministic and about 35× further from the FP64 golden. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op adaln_gate_residual +``` + +B200, B = 1, H = 5376, BF16: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.05 ms | 0.06 ms | 1.21 ms | 0.68 ms | +| 32768 | 0.24 ms | 0.41 ms | 1.40 ms | 2.71 ms | +| 131072 | 0.92 ms | 1.59 ms | 3.31 ms | 10.6 ms | + +The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one +block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and +tile setup. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_adaln_gate_residual.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end, incl. the gated residual +python scripts/check_operator.py --op adaln_gate_residual --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op adaln_gate_residual \ + --out docs/usage/evidence/h3-adaln-gate-residual-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-adaln-gate-residual-b200/report.json +``` + +## Known Limitations + +- Until the attention and FFN rows land, the end-to-end chain feeds a seeded stand-in for the + sublayer output. The op itself is exercised on random and pinned-table inputs. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 01e6e5348..43fb55ff5 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -377,3 +377,16 @@ def h3_rmsnorm_backward( tile_end: torch.Tensor | None = None, seg_first_tile: torch.Tensor | None = None, ) -> list[torch.Tensor]: ... +def h3_gate_residual_forward( + residual: torch.Tensor, y: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... +def h3_gate_residual_backward( + grad: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 02f7f2e37..6afafc3ae 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -51,6 +51,7 @@ def make_operator_inputs( "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, "adaln_row_gather": _make_adaln_row_gather_inputs, "h3_rmsnorm": _make_h3_rmsnorm_inputs, + "adaln_gate_residual": _make_adaln_gate_residual_inputs, } try: return builders[op_name](args, dtype, device) @@ -89,6 +90,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" f"->{batch * seq}", "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", + "adaln_gate_residual": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -446,6 +448,22 @@ def _make_h3_rmsnorm_inputs( } +def _make_adaln_gate_residual_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # gate rows: 3 modality rows for each of --num-timesteps (default --batch) timesteps. + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + rows = 3 * _h3_num_timesteps(args) + generator = _generator(args, device, offset=31) + return { + "residual": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "y": _floating_tensor((batch, seq, hidden), args, dtype, device, 1), + "gate": (_floating_tensor((rows, hidden), args, torch.float32, device, 2) * 0.5).to(dtype), + "index": torch.randint(0, rows, (seq,), generator=generator, device=device), + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 197db8874..5254aba2e 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -311,6 +311,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("x", "weight"), ), + "adaln_gate_residual": OperatorSpec( + name="adaln_gate_residual", + op_class="reduction", # forward is elementwise; the gate VJP sums positions + gold_path="rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp", + "cuda": "rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp", + }, + grad_input_names=("residual", "y", "gate"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/gate_residual.py b/rl_engine/kernels/ops/cuda/h3/gate_residual.py new file mode 100644 index 000000000..2a0fc9e1d --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/gate_residual.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 gated residual (RFC #420 ``adaln_gate_residual``).""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles +from rl_engine.kernels.ops.pytorch.h3.gate_residual import validate_h3_gate_residual + +KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" +BACKWARD_IMPL = "exact_dy_sorted_segment_dgate" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_gate_residual_backward" + + +def h3_gate_residual_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_gate_residual_forward") + and hasattr(_C, "h3_gate_residual_backward") + ) + + +class _H3GateResidualCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, index): + hidden = residual.shape[-1] + res2 = residual.contiguous().view(-1, hidden) + y2 = y.contiguous().view(-1, hidden) + out = _C.h3_gate_residual_forward(res2, y2, gate, index) + ctx.save_for_backward(y2, gate, index) + ctx.shape = residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + positions = index.repeat(y2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, gate.shape[0]) + dy, dgate = _C.h3_gate_residual_backward(g2, y2, gate, index, *tiles) + record_backward( + "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad, dy.view(ctx.shape), dgate, None + + +class H3GateResidualCudaOp: + """CUDA candidate: ``residual + gate[index] * y`` in one elementwise pass. + + The gate row is gathered in the kernel from the AdaLN table view, and the + two roundings match the eager expression, so the output is bitwise equal + to diffusers. ``dy`` is bitwise equal to the eager VJP; ``dgate`` is a + deterministic FP32 segment sum instead of ``index_select``'s BF16 atomics. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_gate_residual_available(): + raise RuntimeError( + "rl_engine._C lacks h3_gate_residual_*; rebuild with csrc/cuda/h3/gate_residual.cu" + ) + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + if not residual.is_cuda: + raise ValueError("H3GateResidualCudaOp needs CUDA tensors") + if gate.stride(1) != 1: + gate = gate.contiguous() + return _H3GateResidualCuda.apply(residual, y, gate, index.contiguous()) + + def forward_fp32(self, residual, y, gate, index): + return self.forward(residual, y, gate, index).float() diff --git a/rl_engine/kernels/ops/pytorch/h3/gate_residual.py b/rl_engine/kernels/ops/pytorch/h3/gate_residual.py new file mode 100644 index 000000000..d2be659d5 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/gate_residual.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 gated residual (RFC #420 row ``adaln_gate_residual``). + +After each sublayer of a block, diffusers adds the sublayer output back to +the residual stream through the per-row AdaLN gate:: + + hidden = residual + gate.index_select(0, adaln_indices) * sublayer_output + +``gate`` is ``gate_msa`` (after attention) or ``gate_mlp`` (after the FFN), +an ``(R, H)`` row view of the AdaLN table; ``residual`` and the sublayer +output are ``(..., S, H)``. The order is fixed by RFC #420 §4: +``residual + gate[row] * sublayer_output``. +""" + +from __future__ import annotations + +import torch + +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_gate_residual( + residual: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + if residual.dtype not in _DTYPES: + raise TypeError(f"residual must be fp32, bf16 or fp16, got {residual.dtype}") + if y.shape != residual.shape or y.dtype != residual.dtype or y.device != residual.device: + raise ValueError("y must match the residual's shape, dtype and device") + if residual.dim() < 2 or residual.numel() == 0: + raise ValueError( + f"residual must be a non-empty (..., S, H) tensor, got {tuple(residual.shape)}" + ) + hidden = residual.shape[-1] + if gate.dim() != 2 or gate.shape[1] != hidden or gate.device != residual.device: + raise ValueError(f"gate must be (R, {hidden}) on the residual's device") + if same_dtype and gate.dtype != residual.dtype: + raise TypeError(f"gate dtype {gate.dtype} must match the residual dtype {residual.dtype}") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != residual.device: + raise TypeError("index must be a 1-D int64 tensor on the residual's device") + if residual.shape[-2] != index.shape[0]: + raise ValueError(f"residual must be (..., S, H) with S = len(index) = {index.shape[0]}") + if check_range and bool(((index < 0) | (index >= gate.shape[0])).any()): + raise IndexError(f"index must lie in [0, {gate.shape[0]})") + + +class _DeterministicGateResidual(torch.autograd.Function): + """Eager forward; the gate gradient as FP32 per-row sums instead of + ``index_select``'s BF16 atomic scatter-add (``d_residual``/``d_y`` are the + eager VJPs).""" + + @staticmethod + def forward(ctx, residual, y, gate, index): + gathered = gate.index_select(0, index) + ctx.save_for_backward(y, gathered, index) + ctx.gate_rows = gate.shape[0] + return residual + gathered * y + + @staticmethod + def backward(ctx, grad): + y, gathered, index = ctx.saved_tensors + contrib = (grad.float() * y.float()).reshape(-1, index.shape[0], y.shape[-1]) + per_position = contrib.sum(dim=0) # (S, H): fixed reduction over the batch + d_gate = per_position.new_zeros((ctx.gate_rows, y.shape[-1])) + for row in range(ctx.gate_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_gate[row] = per_position.index_select(0, positions).sum(dim=0) + return grad, grad * gathered, d_gate.to(gathered.dtype), None + + +class NativeH3GateResidualOp: + """PyTorch reference: ``forward`` is the eager provider expression (with a + deterministic gate gradient), ``forward_fp32`` the FP64 golden returned in + FP32. The raw diffusers path, atomic backward included, is + ``rl_engine.testing.h3_provider.provider_gate_residual``.""" + + op_class = "reduction" # forward is elementwise; the gate VJP sums positions + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + return _DeterministicGateResidual.apply(residual, y, gate, index) + + def forward_fp32(self, residual, y, gate, index): + validate_h3_gate_residual(residual, y, gate, index, same_dtype=False) + out = residual.double() + gate.double().index_select(0, index) * y.double() + return out.float() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 8fa1d1453..bcfc897b8 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -223,6 +223,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ) PYTORCH_H3_RMSNORM = "rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp" CUDA_H3_RMSNORM = "rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp" + PYTORCH_H3_GATE_RESIDUAL = ( + "rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp" + ) + CUDA_H3_GATE_RESIDUAL = "rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -650,6 +654,10 @@ def __init__(self): OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, ], "h3_rmsnorm": [OpBackend.CUDA_H3_RMSNORM, OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [ + OpBackend.CUDA_H3_GATE_RESIDUAL, + OpBackend.PYTORCH_H3_GATE_RESIDUAL, + ], }, "rocm": { "logp": [ @@ -695,6 +703,7 @@ def __init__(self): "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -744,6 +753,7 @@ def __init__(self): "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -774,6 +784,7 @@ def __init__(self): "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index fe010c7e3..7b13fc60f 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -33,6 +33,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_gate_residual, provider_norm_modulate, provider_time_embedder, provider_time_proj, @@ -146,6 +147,35 @@ def _norm1_modulated(ctx: dict[str, Any], modulation, call): ) ) +STAGES.append( + Stage( + name="adaln_gate_residual", + op_type="adaln_gate_residual", + # residual + gate_msa[row] * attention_output (the stand-in sublayer output). + candidate=lambda op, ctx, _up: op( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider=lambda ctx, _up: provider_gate_residual( + ctx["hidden"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ctx["sublayer"], + ), + golden=lambda op, ctx, _up: op.forward_fp32( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider_bitwise_isolated=True, # eager rounding order, elementwise + golden_atol=5e-2, # carries the gate's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ) +) + # The backward replay covers the conditioning chain up to the row gather. CONDITIONING_STAGES = ( "timestep_sinusoid_h3", @@ -199,6 +229,14 @@ def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dic ) .to(torch.bfloat16) .cuda(), + # Stand-in for the attention output the gate multiplies (the attention + # rows belong to other contributors): BF16 (1, S, H), seeded. + "sublayer": ( + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 3)) + * 3.0 + ) + .to(torch.bfloat16) + .cuda(), } diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index 05c7ebec9..8c084afc7 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -81,3 +81,11 @@ def provider_norm_modulate( return norm_hidden_states * (1.0 + scale.index_select(0, indices)) + shift.index_select( 0, indices ) + + +def provider_gate_residual( + residual: torch.Tensor, gate: torch.Tensor, indices: torch.Tensor, sublayer_output: torch.Tensor +) -> torch.Tensor: + """``residual + gate.index_select(0, adaln_indices) * sublayer_output`` (block).""" + + return residual + gate.index_select(0, indices) * sublayer_output diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index dfecec419..021dab4b0 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -20,6 +20,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_gate_residual, provider_norm_modulate, provider_time_embedder, provider_time_proj, @@ -467,12 +468,117 @@ def golden(x_, w_, sh_, sc_): } +# --------------------------------------------------------------------------- # +# adaln_gate_residual (gate_msa after attention) +# --------------------------------------------------------------------------- # + + +def _gate_inputs(seq: int, seed: int = 0, dtype=torch.bfloat16): + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).to(dtype) + residual = torch.randn(1, seq, 5376, device="cuda", generator=g).to(dtype) + y = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 3).to(dtype) + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate_view(table): + return table.view(table.shape[0], 6, -1)[:, 2] + + +def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_gate_residual", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + table, residual, y, index = _gate_inputs(seq) + gate = _gate_view(table) + grad = torch.randn_like(residual) + + def backward(fn, residual=residual, y=y, table=table, grad=grad): + leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] + fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + + cases.append( + { + "op": "adaln_gate_residual", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 3 * seq * 5376 * 2, # read residual and y, write out + "candidate": lambda r=residual, y=y, g=gate, i=index: op.forward( + r, y, g, i, check_range=False + ), + "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( + r, g, i, y + ), + "candidate_backward": lambda b=backward, i=index: b( + lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) + ), + "provider_backward": lambda b=backward, i=index: b( + lambda r, y_, g: provider_gate_residual(r, g, i, y_) + ), + } + ) + return cases + + +def _gate_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("adaln_gate_residual", device="cuda") + forward_bitwise = {} + for dtype in (torch.bfloat16, torch.float16, torch.float32): + table, residual, y, index = _gate_inputs(777, seed=1, dtype=dtype) + gate = _gate_view(table) + forward_bitwise[str(dtype).removeprefix("torch.")] = bool( + torch.equal( + op(residual, y, gate, index), provider_gate_residual(residual, gate, index, y) + ) + ) + table, residual, y, index = _gate_inputs(4097, seed=2) + invariant = bool( + torch.equal( + op(residual[:, 10:60], y[:, 10:60], _gate_view(table), index[10:60])[0], + op(residual, y, _gate_view(table), index)[0, 10:60], + ) + ) + grad = torch.randn( + residual.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(residual.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ref = grads(lambda r, y_, g: provider_gate_residual(r, g, index, y_), torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda r, y_, g: op(r, y_, g, index)), + ("provider", lambda r, y_, g: provider_gate_residual(r, g, index, y_)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("d_residual", "d_sublayer", "d_gate"), first, ref) + }, + } + return { + "forward_bitwise_vs_diffusers": forward_bitwise, + "rows_batch_invariant": invariant, + "backward": backward, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, "adaln_projection_3mod": _projection_perf, "adaln_row_gather": _gather_perf, "h3_rmsnorm": _norm_perf, + "adaln_gate_residual": _gate_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, @@ -480,4 +586,5 @@ def golden(x_, w_, sh_, sc_): "adaln_projection_3mod": _projection_accuracy, "adaln_row_gather": _gather_accuracy, "h3_rmsnorm": _norm_accuracy, + "adaln_gate_residual": _gate_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index 49d502f12..ee8e2c0ff 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -232,6 +232,7 @@ def fake_load_backend(backend): "adaln_projection_3mod", "adaln_row_gather", "h3_rmsnorm", + "adaln_gate_residual", ) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index da8bc6774..2c62208b7 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -409,12 +409,84 @@ def repeat(mode: str) -> str: return fig +def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: + _style(ax, "Forward and backward time (log, lower is better)", "ms") + _bars( + ax, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + (f"{provider} fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + f"{provider} bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + + +def _backward_errors(ax, backward: dict[str, Any], keys: list[str], note: str) -> None: + _style(ax, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") + + def label(name: str, mode: str) -> str: + return ( + f"{name} (repeat-bitwise: {'yes' if backward[mode]['repeat_bitwise_equal'] else 'no'})" + ) + + floor = 1e-9 + _bars( + ax, + keys, + [ + ( + label("RL-Kernel CUDA", "cuda"), + RL_KERNEL, + [max(backward["cuda"]["rel_error"][k], floor) for k in keys], + ), + ( + label("diffusers", "provider"), + PROVIDER, + [max(backward["provider"]["rel_error"][k], floor) for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + for text in ax.texts: + if text.get_text() == f"{floor:.1e}": + text.set_text("exact") + ax.set_ylim(floor / 3, 1e4) + ax.text(0.98, 0.80, note, transform=ax.transAxes, ha="right", fontsize=8, color=INK) + + +def plot_gate_residual(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "adaln_gate_residual (residual + gate_msa[row] * y, H = 5376)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + ok = all(acc["forward_bitwise_vs_diffusers"].values()) + _backward_errors( + right, + acc["backward"], + ["d_residual", "d_sublayer", "d_gate"], + f"forward bitwise equal to diffusers (bf16/fp16/fp32): {'yes' if ok else 'NO'}", + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, "adaln_projection_3mod": plot_projection, "adaln_row_gather": plot_gather, "h3_rmsnorm": plot_rmsnorm, + "adaln_gate_residual": plot_gate_residual, } diff --git a/setup.py b/setup.py index fb671d0f0..3f38e1fbe 100644 --- a/setup.py +++ b/setup.py @@ -161,6 +161,7 @@ def get_extensions(): cuda_sources.append("csrc/cuda/h3/det_linear.cu") cuda_sources.append("csrc/cuda/h3/adaln_row_gather.cu") cuda_sources.append("csrc/cuda/h3/rmsnorm_modulate.cu") + cuda_sources.append("csrc/cuda/h3/gate_residual.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/test_h3_adaln_gate_residual.py b/tests/h3/test_h3_adaln_gate_residual.py new file mode 100644 index 000000000..d67210257 --- /dev/null +++ b/tests/h3/test_h3_adaln_gate_residual.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_gate_residual``: ``residual + gate[row] * sublayer_output``. + +* forward and ``d_sublayer`` are bitwise equal to diffusers' eager + expression in every dtype (the gate row is gathered in the kernel); +* ``d_gate`` is a deterministic FP32 segment sum, checked against FP64; +* rows are elementwise-independent; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.h3.gate_residual import NativeH3GateResidualOp +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_gate_residual + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.gate_residual import ( + H3GateResidualCudaOp, + h3_gate_residual_available, + ) + + if not h3_gate_residual_available(): + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + return H3GateResidualCudaOp() + + +def _case(seq, hidden=5376, dtype=torch.bfloat16, seed=0, batch=2): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(9, 6 * hidden, generator=g) * 0.5).to(dtype).cuda() + residual = torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda() + y = (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate(table, chunk=2): + return table.view(table.shape[0], 6, -1)[:, chunk] # gate_msa (2) / gate_mlp (5) view + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3GateResidualOp() + res, y, gate = torch.zeros(4, 8), torch.zeros(4, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op(res, y, gate, torch.tensor([0, 1, 2, 3])) + with pytest.raises(ValueError): + op(res, y[:, :4], gate, torch.tensor([0, 1, 2, 0])) + with pytest.raises(TypeError): + op(res, y, gate.bfloat16(), torch.tensor([0, 1, 2, 0])) + with pytest.raises(ValueError): # S mismatch + op(res, y, gate, torch.tensor([0, 1])) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("seq, hidden", [(4097, 5376), (33, 12), (17, 13)]) + def test_forward_bitwise_equal_to_diffusers(self, dtype, seq, hidden): + table, residual, y, index = _case(seq, hidden, dtype, seed=seq) + for chunk in (2, 5): + ours = _cuda_op()(residual, y, _gate(table, chunk), index) + assert torch.equal( + ours, provider_gate_residual(residual, _gate(table, chunk), index, y) + ) + + def test_rows_are_batch_and_position_invariant(self): + table, residual, y, index = _case(500, seed=1) + full = _cuda_op()(residual, y, _gate(table), index) + part = _cuda_op()(residual[1:2, 100:150], y[1:2, 100:150], _gate(table), index[100:150]) + assert torch.equal(part[0], full[1, 100:150]) + + def test_backward(self): + table, residual, y, index = _case(4097, seed=2) + grad = torch.randn_like(residual) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ours = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + again = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + theirs = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy)) + golden = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy), torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + assert torch.equal(ours[0], theirs[0]) # d_residual + assert torch.equal(ours[1], theirs[1]) # d_sublayer: exact product, one rounding + rel = ((ours[2].double() - golden[2]).abs().max() / golden[2].abs().max()).item() + assert rel < 5e-3 # only the BF16 rounding of the table gradient + assert torch.count_nonzero(ours[2].view(9, 6, -1)[:, [0, 1, 3, 4, 5]]) == 0 + + def test_rejects_cpu_and_bad_index(self): + table, residual, y, index = _case(10, hidden=16) + with pytest.raises(IndexError): + _cuda_op()(residual, y, _gate(table), index + 9) + with pytest.raises(ValueError): + _cuda_op()(residual.cpu(), y.cpu(), _gate(table).cpu(), index.cpu()) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_gate_residual", device="cuda") + assert type(op).__name__ == "H3GateResidualCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 2b165c088..32c6322c3 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -54,6 +54,7 @@ def test_all_ws1_single_ops_are_registered(): "adaln_projection_3mod", "adaln_row_gather", "h3_rmsnorm", + "adaln_gate_residual", } <= names From cd688b354823109b19b1abd24d8fdfeb5b7a7fc5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:17:29 +0000 Subject: [PATCH 12/89] [WS1][MiniMax-H3] adaln_gate_residual: B200 evidence report and figure Produced by scripts/h3_evidence.py at a26b41a on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 7a3010185..f55412449 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -51,7 +51,7 @@ hidden = op(residual, attn_output, gate_msa, adaln_indices) once. Diffusers' `d_gate` uses `index_select`'s BF16 atomic scatter-add instead. It is - non-deterministic and about 35× further from the FP64 golden. + non-deterministic and 13–49× further from the FP64 golden (its error varies from run to run). ## Performance Notes @@ -63,14 +63,22 @@ B200, B = 1, H = 5376, BF16: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.05 ms | 0.06 ms | 1.21 ms | 0.68 ms | -| 32768 | 0.24 ms | 0.41 ms | 1.40 ms | 2.71 ms | -| 131072 | 0.92 ms | 1.59 ms | 3.31 ms | 10.6 ms | +| 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | +| 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | +| 131072 | 0.92 ms | 1.59 ms | 3.29 ms | 10.6 ms | The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and tile setup. +## Evidence + +![adaln_gate_residual on B200: latency and backward accuracy](../usage/evidence/h3-adaln-gate-residual-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-adaln-gate-residual-b200/report.json), +written by `scripts/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. + ## Tests ```bash From e7428b1e58ea71f0d036f9733faf31b76e0d4d44 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:02:12 +0000 Subject: [PATCH 13/89] [WS1][CUDA][MiniMax-H3] final_adaln_out: norm_out projection + final norm/modulation RFC #420 step 7, row `final_adaln_out` (stacked on adaln_gate_residual). - norm_out as one autograd node: FP32 SiLU -> one BF16 cast -> the deterministic tensor-core GEMV on norm_out.linear (shift first), then the h3_rmsnorm kernel indexed by timestep_indices. Forward equal to diffusers on 99.96% of elements (1-ULP projection ties); rows bitwise batch/position invariant; 2.4x faster forward at S = 131072. - Backward keeps the table gradient in FP32 into the projection backward: d_temb/dW/db 13-57x closer to FP64 than diffusers (whose error varies run to run), all gradients repeat-bitwise. - Golden rounds only at the declared boundaries (SiLU cast, BF16 table), straight-through for the gradient; gtest inputs at realistic scale. - Chain stage: the gated residual stream -> norm_out with the MLP stage's temb (end-to-end test now runs timestep -> ... -> norm_out); registry, report/figure hooks, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-final-adaln-out.md | 89 ++++++++++++ rl_engine/kernels/gtest/operator_inputs.py | 27 ++++ rl_engine/kernels/gtest/operator_specs.py | 11 ++ .../kernels/ops/cuda/h3/final_adaln_out.py | 110 +++++++++++++++ .../kernels/ops/pytorch/h3/final_adaln_out.py | 95 +++++++++++++ rl_engine/kernels/registry.py | 11 ++ rl_engine/testing/h3_chain.py | 28 ++++ rl_engine/testing/h3_provider.py | 18 +++ rl_engine/testing/h3_report.py | 92 +++++++++++++ rl_engine/tests/test_dispatch.py | 1 + scripts/plot_h3_evidence.py | 17 +++ tests/h3/test_h3_final_adaln_out.py | 129 ++++++++++++++++++ tests/test_ws1_gtest_gpu.py | 1 + 15 files changed, 631 insertions(+) create mode 100644 docs/operators/h3-final-adaln-out.md create mode 100644 rl_engine/kernels/ops/cuda/h3/final_adaln_out.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py create mode 100644 tests/h3/test_h3_final_adaln_out.py diff --git a/docs/.nav.yml b/docs/.nav.yml index f6db1bc36..1ef4fa017 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -36,6 +36,7 @@ nav: - operators/h3-adaln-row-gather.md - operators/h3-rmsnorm.md - operators/h3-adaln-gate-residual.md + - operators/h3-final-adaln-out.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 7042d4561..27ffce7d0 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -38,4 +38,5 @@ Every operator page should include: - [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) - [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) - [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) +- [MiniMax-H3 Final AdaLN Output](h3-final-adaln-out.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md new file mode 100644 index 000000000..e05f1d15c --- /dev/null +++ b/docs/operators/h3-final-adaln-out.md @@ -0,0 +1,89 @@ +# MiniMax-H3 Final AdaLN Output + +## Summary + +`final_adaln_out` is H3's `norm_out` (`MiniMaxH3AdaLayerNormOut`), which runs once after the 50 +blocks (RFC #420, WS1 step 7): + +```text +shift, scale = norm_out.linear(silu(temb).to(bf16)).chunk(2) # (T, 5376) each, shift first +out = norm_out.norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] +``` + +The table has one row per distinct timestep and is indexed by `timestep_indices`, not by +`adaln_indices`. Diffusers then upcasts `out` to the FP32 output heads' dtype. That cast is exact +and is left to the caller. + +## Entry Point + +```python +op = kernel_registry.get_op("final_adaln_out", device="cuda") +out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timestep_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+) | `rl_engine.kernels.ops.cuda.h3.final_adaln_out.H3FinalAdaLNOutCudaOp` | `rl_engine._C.h3_det_linear_*`, `rl_engine._C.h3_rmsnorm_*` | One autograd node | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.final_adaln_out.NativeH3FinalAdaLNOutOp` | n/a | `forward`: diffusers replay; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Numerics + +- **Projection.** The [AdaLN projection](h3-adaln-projection.md)'s deterministic tensor-core + GEMV (`h3-det-linear-bf16-mma-v1`), applied to `norm_out.linear`. The FP32 SiLU is rounded + once at the declared cast. BF16 `temb` is rejected (probe H7). +- **Norm and modulation.** The [`h3_rmsnorm`](h3-rmsnorm.md) kernel indexed by + `timestep_indices`. Given the same table, it is bitwise equal to diffusers. Overall, about + 0.04% of output elements differ from diffusers by 1 ULP, all traced to the projection's + summation tree versus cuBLAS. Rows are batch- and position-invariant. +- **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into + the projection backward without an extra BF16 rounding: + - `d_temb`, `dW` and `db` are 13–57× closer to FP64 than diffusers (whose error varies from + run to run), which also rounds the + table gradient to BF16 and accumulates it with `index_select` atomics; + - `dx` and `d_norm_w` are at the BF16 rounding level for both; + - every gradient is repeat-bitwise. +- **Golden.** FP64, rounding only where the model declares a dtype boundary: the SiLU cast, + and `norm_out.linear`'s BF16 output table, which every position of a timestep shares. Both are + rounded straight-through for the gradient. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op final_adaln_out +``` + +B200, B = 1, T = 3, pinned `norm_out` weights: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.13 ms | 0.15 ms | 1.62 ms | 1.09 ms | +| 32768 | 0.39 ms | 0.78 ms | 2.86 ms | 5.32 ms | +| 131072 | 1.24 ms | 2.93 ms | 7.78 ms | 20.7 ms | + +At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the +projection backward. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_final_adaln_out.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end: timestep -> ... -> norm_out +python scripts/check_operator.py --op final_adaln_out --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 64 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op final_adaln_out \ + --out docs/usage/evidence/h3-final-adaln-out-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-final-adaln-out-b200/report.json +``` + +## Known Limitations + +- **gtest at larger S.** With `--seq 257`, a few gradient elements fail the per-element + contract tolerance (BF16 `d_norm_w`, FP32 `d_temb`). Those elements are what is left after + summing hundreds of rows into values around 1e2, which an absolute per-element atol can't + cover. Relative to the gradient's maximum the error is 4e-7 in FP32 and 5e-3 in BF16, and the + PyTorch reference fails the same elements, plus more. `--seq 64` passes in both dtypes. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 6afafc3ae..4abc17a80 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -52,6 +52,7 @@ def make_operator_inputs( "adaln_row_gather": _make_adaln_row_gather_inputs, "h3_rmsnorm": _make_h3_rmsnorm_inputs, "adaln_gate_residual": _make_adaln_gate_residual_inputs, + "final_adaln_out": _make_final_adaln_out_inputs, } try: return builders[op_name](args, dtype, device) @@ -91,6 +92,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: f"->{batch * seq}", "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", "adaln_gate_residual": f"{batch}x{seq}x{_h3_hidden(args)}", + "final_adaln_out": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -464,6 +466,31 @@ def _make_adaln_gate_residual_inputs( } +def _make_final_adaln_out_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # temb is FP32 by contract; norm_out.linear/norm/x use ``dtype`` (BF16 in H3). + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + num_timesteps = _h3_num_timesteps(args) + # Realistic scales (norm weight in [0.5, 1], unit activations): with unscaled + # inputs the parameter gradients sum hundreds of rows into values ~1e3, and + # their cancelling elements cannot meet a per-element FP32 atol. + linear = _floating_tensor((2 * hidden, H3_TIME_EMBED), args, torch.float32, device, 3) + norm = 0.5 + 0.5 * torch.sigmoid(_floating_tensor((hidden,), args, torch.float32, device, 4)) + generator = _generator(args, device, offset=37) + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "norm_weight": norm.to(dtype), + "temb": _floating_tensor((num_timesteps, H3_TIME_EMBED), args, torch.float32, device, 1), + "weight": (linear * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (_floating_tensor((2 * hidden,), args, torch.float32, device, 2) * 0.1).to(dtype), + "timestep_indices": torch.randint( + 0, num_timesteps, (seq,), generator=generator, device=device + ), + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 5254aba2e..d7f53b754 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -322,6 +322,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("residual", "y", "gate"), ), + "final_adaln_out": OperatorSpec( + name="final_adaln_out", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.h3.final_adaln_out.NativeH3FinalAdaLNOutOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.kernels.ops.pytorch.h3.final_adaln_out.NativeH3FinalAdaLNOutOp"), + "cuda": "rl_engine.kernels.ops.cuda.h3.final_adaln_out.H3FinalAdaLNOutCudaOp", + }, + grad_input_names=("x", "norm_weight", "temb", "weight", "bias"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/final_adaln_out.py b/rl_engine/kernels/ops/cuda/h3/final_adaln_out.py new file mode 100644 index 000000000..f2700ef06 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/final_adaln_out.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 final AdaLN output (RFC #420 ``final_adaln_out``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.cuda.h3.rmsnorm import h3_rmsnorm_available +from rl_engine.kernels.ops.pytorch.h3.final_adaln_out import validate_h3_final_adaln_out +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import H3_NORM_EPS + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_rmsnorm_forward[modulated]" +) +BACKWARD_IMPL = "fp32_table_grad_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_rmsnorm_backward+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.kernels.ops.cuda.h3.det_linear.silu_backward_fp32" +) + + +class _H3FinalAdaLNOutCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, norm_weight, temb, weight, bias, timestep_indices, eps): + hidden = x.shape[-1] + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) # declared boundary: one cast after the FP32 SiLU + (table,) = det_linear_forward(act, weight, bias) + shift, scale = table.chunk(2, dim=-1) + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward( + x2, norm_weight.contiguous(), float(eps), shift, scale, timestep_indices + ) + ctx.save_for_backward(x2, norm_weight, rstd, temb, act, weight, table, timestep_indices) + ctx.x_shape = x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, norm_weight, rstd, temb, act, weight, table, timestep_indices = ctx.saved_tensors + shift, scale = table.chunk(2, dim=-1) + positions = timestep_indices.repeat(x2.shape[0] // timestep_indices.shape[0]) + tiles = _segment_tiles(positions, table.shape[0]) + dx, d_norm_weight, d_shift, d_scale = _C.h3_rmsnorm_backward( + grad.contiguous().view_as(x2), + x2, + norm_weight.contiguous(), + rstd, + shift, + scale, + timestep_indices, + *tiles, + ) + # The table gradient stays FP32 into the projection backward. + d_table = torch.cat([d_shift, d_scale], dim=1) + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + d_temb = silu_backward_fp32(det_linear_backward_input(d_table, weight, torch.float32), temb) + record_backward( + "final_adaln_out", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), d_norm_weight, d_temb, dw, db, None, None + + +class H3FinalAdaLNOutCudaOp: + """CUDA candidate: tensor-core shift/scale projection + fused norm/modulation. + + The projection is ``adaln_projection_3mod``'s deterministic GEMV + (``h3-det-linear-bf16-mma-v1``) on ``norm_out.linear``; the norm and + modulation are ``h3_rmsnorm``'s kernel indexed by ``timestep_indices``. + One autograd node, so the table gradient reaches the projection in FP32. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and h3_rmsnorm_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_rmsnorm_* ops") + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + if weight.dtype != x.dtype: + raise TypeError("norm_out.linear must share the activations' dtype") + if not x.is_cuda: + raise ValueError("H3FinalAdaLNOutCudaOp needs CUDA tensors") + return _H3FinalAdaLNOutCuda.apply( + x, norm_weight, temb, weight, bias, timestep_indices.contiguous(), eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps).float() diff --git a/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py b/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py new file mode 100644 index 000000000..94aaa28c1 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 final AdaLN output (RFC #420 row ``final_adaln_out``). + +``MiniMaxH3AdaLayerNormOut`` runs after the 50 blocks:: + + shift, scale = linear(silu(temb).to(bf16)).chunk(2) # (T, H) each, shift first + out = norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] + +``linear`` is ``norm_out.linear`` (``2688 -> 2 * 5376``, BF16), ``norm`` is +``norm_out.norm`` (``nn.RMSNorm(5376, eps=1e-5)``), and the table is indexed +by ``timestep_indices`` (one row per distinct timestep), not by +``adaln_indices``. Diffusers then casts the result to the FP32 output heads' +dtype, an exact upcast left to the caller. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import ( + H3_NORM_EPS, + NativeH3RMSNormOp, + validate_h3_rmsnorm, +) + + +def validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, *, same_dtype: bool = True +) -> int: + hidden = x.shape[-1] + if same_dtype: + validate_h3_rmsnorm(x, norm_weight, eps) + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if temb.dim() != 2 or temb.shape[0] == 0: + raise ValueError(f"temb must be a non-empty (T, D) matrix, got {tuple(temb.shape)}") + if weight.shape != (2 * hidden, temb.shape[1]) or bias.shape != (2 * hidden,): + raise ValueError( + f"norm_out.linear must be ({2 * hidden}, {temb.shape[1]}) with a ({2 * hidden},) " + f"bias, got {tuple(weight.shape)} / {tuple(bias.shape)}" + ) + if bias.dtype != weight.dtype: + raise TypeError("norm_out.linear weight and bias must share a dtype") + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if timestep_indices.dtype != torch.int64 or timestep_indices.dim() != 1: + raise TypeError("timestep_indices must be a 1-D int64 tensor") + if x.dim() < 2 or x.shape[-2] != timestep_indices.shape[0]: + raise ValueError("x must be (..., S, H) with S = len(timestep_indices)") + if bool(((timestep_indices < 0) | (timestep_indices >= temb.shape[0])).any()): + raise IndexError(f"timestep_indices must lie in [0, {temb.shape[0]})") + return hidden + + +class NativeH3FinalAdaLNOutOp: + """PyTorch reference: ``forward`` replays diffusers' ``norm_out``; + ``forward_fp32`` is the FP64 golden, returned in FP32. It rounds only at the + module boundaries the model declares (the SiLU cast and the BF16 table), + straight-through for the gradient.""" + + op_class = "reduction" + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + return NativeH3RMSNormOp().forward_modulated( + x, norm_weight, shift, scale, timestep_indices, eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, same_dtype=False + ) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + act = act + (act.to(weight.dtype).double() - act).detach() # declared cast, identity VJP + table = F.linear(act, weight.double(), bias.double()) + # norm_out.linear is a BF16 module: its (T, 2H) output is rounded to the + # weight dtype, and every position of a timestep shares that row, so the + # rounding is model semantics (straight-through for the gradient). + table = table + (table.to(weight.dtype).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * norm_weight.double() + out = n * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + return out.float() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index bcfc897b8..39255a9d4 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -227,6 +227,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): "rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp" ) CUDA_H3_GATE_RESIDUAL = "rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp" + PYTORCH_H3_FINAL_ADALN_OUT = ( + "rl_engine.kernels.ops.pytorch.h3.final_adaln_out.NativeH3FinalAdaLNOutOp" + ) + CUDA_H3_FINAL_ADALN_OUT = "rl_engine.kernels.ops.cuda.h3.final_adaln_out.H3FinalAdaLNOutCudaOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -658,6 +662,10 @@ def __init__(self): OpBackend.CUDA_H3_GATE_RESIDUAL, OpBackend.PYTORCH_H3_GATE_RESIDUAL, ], + "final_adaln_out": [ + OpBackend.CUDA_H3_FINAL_ADALN_OUT, + OpBackend.PYTORCH_H3_FINAL_ADALN_OUT, + ], }, "rocm": { "logp": [ @@ -704,6 +712,7 @@ def __init__(self): "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -754,6 +763,7 @@ def __init__(self): "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -785,6 +795,7 @@ def __init__(self): "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 7b13fc60f..f4a9599f9 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -33,6 +33,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_final_adaln_out, provider_gate_residual, provider_norm_modulate, provider_time_embedder, @@ -176,6 +177,33 @@ def _norm1_modulated(ctx: dict[str, Any], modulation, call): ) ) + +def _final(ctx: dict[str, Any], call): + weights = ctx["weights"] + history = ctx["history"] + return call( + history["adaln_gate_residual"], # the residual stream after the gated sublayer + weights["norm_out.norm.weight"], + history["timestep_mlp_fp32"], # temb + weights["norm_out.linear.weight"], + weights["norm_out.linear.bias"], + ctx["timestep_indices"], + ) + + +STAGES.append( + Stage( + name="final_adaln_out", + op_type="final_adaln_out", + candidate=lambda op, ctx, _up: _final(ctx, op), + provider=lambda ctx, _up: _final(ctx, provider_final_adaln_out), + golden=lambda op, ctx, _up: _final(ctx, op.forward_fp32), + provider_bitwise_isolated=False, # shift/scale projection: tensor-core tree vs cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + # The backward replay covers the conditioning chain up to the row gather. CONDITIONING_STAGES = ( "timestep_sinusoid_h3", diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index 8c084afc7..c654611ec 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -89,3 +89,21 @@ def provider_gate_residual( """``residual + gate.index_select(0, adaln_indices) * sublayer_output`` (block).""" return residual + gate.index_select(0, indices) * sublayer_output + + +def provider_final_adaln_out( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + temb: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + timestep_indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``MiniMaxH3AdaLayerNormOut.forward``: shift first, indexed by timestep.""" + + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 021dab4b0..41f8d647d 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -20,6 +20,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_final_adaln_out, provider_gate_residual, provider_norm_modulate, provider_time_embedder, @@ -572,6 +573,95 @@ def grads(fn, dtype=None): } +# --------------------------------------------------------------------------- # +# final_adaln_out (norm_out) +# --------------------------------------------------------------------------- # + +FINAL_NAMES = ["norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias"] + + +def _final_inputs(seq: int, seed: int = 0): + norm_weight, weight, bias = h3_params(FINAL_NAMES, [(5376,), (10752, 2688), (10752,)]) + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 + ti, _ = h3_packed_layout(seq, 3, seed=seed) + return x, norm_weight.bfloat16(), temb, weight.bfloat16(), bias.bfloat16(), ti + + +def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("final_adaln_out", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, nw, temb, w, b, ti = _final_inputs(seq) + grad = torch.randn_like(x) + + def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + leaves = [t.detach().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad) + + cases.append( + { + "op": "final_adaln_out", + "case": f"S={seq}", + "backend": type(op).__name__, + # read x + norm_out.linear, write out + "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, + "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), + "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), + "candidate_backward": lambda bw=backward, i=ti: bw(lambda *t: op(*t, i)), + "provider_backward": lambda bw=backward, i=ti: bw( + lambda *t: provider_final_adaln_out(*t, i) + ), + } + ) + return cases + + +def _final_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("final_adaln_out", device="cuda") + golden_op = registry._get_or_create_backend( + registry._priority_map["cpu"]["final_adaln_out"][-1] + ) + x, nw, temb, w, b, ti = _final_inputs(4097, seed=1) + ours = op(x, nw, temb, w, b, ti) + theirs = provider_final_adaln_out(x, nw, temb, w, b, ti) + golden = golden_op.forward_fp32(x, nw, temb, w, b, ti) + invariant = bool( + torch.equal(op(x[:, 100:160], nw, temb, w, b, ti[100:160])[0], ours[0, 100:160]) + ) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn): + leaves = [t.detach().clone().requires_grad_(True) for t in (x, nw, temb, w, b)] + fn(*leaves).backward(grad) + return [leaf.grad for leaf in leaves] + + ref = grads(lambda *t: golden_op.forward_fp32(*t, ti)) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op(*t, ti)), + ("provider", lambda *t: provider_final_adaln_out(*t, ti)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, c) for a, c in zip(first, second)), + "rel_error": { + key: float((g.double() - r.double()).abs().max() / r.double().abs().max()) + for key, g, r in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), first, ref) + }, + } + return { + "equal_to_diffusers_fraction": float((ours == theirs).float().mean()), + "max_abs_vs_golden": float((ours.float() - golden).abs().max()), + "provider_max_abs_vs_golden": float((theirs.float() - golden).abs().max()), + "rows_batch_invariant": invariant, + "backward": backward, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, @@ -579,6 +669,7 @@ def grads(fn, dtype=None): "adaln_row_gather": _gather_perf, "h3_rmsnorm": _norm_perf, "adaln_gate_residual": _gate_perf, + "final_adaln_out": _final_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, @@ -587,4 +678,5 @@ def grads(fn, dtype=None): "adaln_row_gather": _gather_accuracy, "h3_rmsnorm": _norm_accuracy, "adaln_gate_residual": _gate_accuracy, + "final_adaln_out": _final_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index ee8e2c0ff..1e12e3dc5 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -233,6 +233,7 @@ def fake_load_backend(backend): "adaln_row_gather", "h3_rmsnorm", "adaln_gate_residual", + "final_adaln_out", ) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 2c62208b7..af4496904 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -480,6 +480,22 @@ def plot_gate_residual(report: dict[str, Any]): return fig +def plot_final(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "final_adaln_out (norm_out: projection + final norm, T = 3)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + _backward_errors( + right, + acc["backward"], + ["dx", "d_norm_w", "d_temb", "dW", "db"], + f"forward equal to diffusers on {100 * acc['equal_to_diffusers_fraction']:.2f}% " + "of elements (projection 1-ULP ties)", + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, @@ -487,6 +503,7 @@ def plot_gate_residual(report: dict[str, Any]): "adaln_row_gather": plot_gather, "h3_rmsnorm": plot_rmsnorm, "adaln_gate_residual": plot_gate_residual, + "final_adaln_out": plot_final, } diff --git a/tests/h3/test_h3_final_adaln_out.py b/tests/h3/test_h3_final_adaln_out.py new file mode 100644 index 000000000..f322cd93d --- /dev/null +++ b/tests/h3/test_h3_final_adaln_out.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``final_adaln_out``: norm_out shift/scale projection + final norm/modulation. + +* layout: ``shift`` is the first half of ``norm_out.linear``'s output, rows + are indexed by ``timestep_indices`` (not ``adaln_indices``); +* forward matches diffusers' ``norm_out`` to the projection's 1-ULP ties + (the norm and modulation are bitwise; the GEMV tree differs from cuBLAS); +* backward is deterministic, with the table gradient kept in FP32; +* rows are batch/position invariant; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3.final_adaln_out import NativeH3FinalAdaLNOutOp +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_final_adaln_out + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.final_adaln_out import H3FinalAdaLNOutCudaOp + from rl_engine.kernels.ops.cuda.h3.rmsnorm import h3_rmsnorm_available + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3FinalAdaLNOutCudaOp() + + +@pytest.fixture(scope="module") +def norm_out(h3_weights_cpu): + return [ + h3_weights_cpu[name].cuda() + for name in ("norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias") + ] + + +def _inputs(seq, num_timesteps=3, seed=0, batch=1): + g = torch.Generator(device="cpu").manual_seed(seed) + x = (torch.randn(batch, seq, 5376, generator=g) * 2).bfloat16().cuda() + temb = (torch.randn(num_timesteps, 2688, generator=g) * 2).cuda() + ti, _ = h3_packed_layout(seq, num_timesteps, seed=seed) + return x, temb, ti + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3FinalAdaLNOutOp() + x, nw = torch.zeros(1, 4, 8), torch.ones(8) + temb, w, b = torch.zeros(2, 16), torch.zeros(16, 16), torch.zeros(16) + ti = torch.tensor([0, 1, 1, 0]) + with pytest.raises(TypeError): # BF16 temb (probe H7) + op(x, nw, temb.bfloat16(), w, b, ti) + with pytest.raises(ValueError): # not 2H rows + op(x, nw, temb, w[:10], b[:10], ti) + with pytest.raises(IndexError): + op(x, nw, temb, w, b, ti + 2) + + +@requires_cuda +class TestCuda: + def test_shift_first_layout_and_timestep_indexing(self): + """Zero weights: the bias halves are (shift, scale); rows follow timestep_indices.""" + + hidden, seq = 16, 6 + x = torch.randn(1, seq, hidden, device="cuda").bfloat16() + nw = torch.ones(hidden, device="cuda", dtype=torch.bfloat16) + w = torch.zeros(2 * hidden, 32, device="cuda", dtype=torch.bfloat16) + b = torch.cat([torch.full((hidden,), 1.0), torch.full((hidden,), -1.0)]).bfloat16().cuda() + ti = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda") + out = _cuda_op()(x, nw, torch.zeros(2, 32, device="cuda"), w, b, ti) + # scale = -1 -> norm(x) * 0 + shift = 1 everywhere + assert torch.equal(out, torch.ones_like(out)) + + @pytest.mark.parametrize("seq, num_timesteps", [(1, 1), (257, 2), (4097, 3)]) + def test_forward_against_diffusers_and_golden(self, norm_out, seq, num_timesteps): + x, temb, ti = _inputs(seq, num_timesteps, seed=seq) + ours = _cuda_op()(x, *norm_out[:1], temb, *norm_out[1:], ti) + theirs = provider_final_adaln_out(x, norm_out[0], temb, *norm_out[1:], ti) + golden = NativeH3FinalAdaLNOutOp().forward_fp32(x, norm_out[0], temb, *norm_out[1:], ti) + assert (ours == theirs).float().mean() > 0.999 + torch.testing.assert_close(ours.float(), golden, atol=5e-2, rtol=2e-2) + + def test_rows_are_batch_and_position_invariant(self, norm_out): + x, temb, ti = _inputs(600, seed=1, batch=2) + full = _cuda_op()(x, norm_out[0], temb, *norm_out[1:], ti) + part = _cuda_op()(x[1:2, 200:260], norm_out[0], temb, *norm_out[1:], ti[200:260]) + assert torch.equal(part[0], full[1, 200:260]) + + def test_backward_deterministic_and_fp32_table_gradient(self, norm_out): + x, temb, ti = _inputs(4097, seed=2) + grad = torch.randn_like(x) + + def grads(fn, dtype=None): + tensors = [x, norm_out[0], temb, *norm_out[1:]] + leaves = [ + (t if dtype is None else t.to(dtype)).detach().clone().requires_grad_(True) + for t in tensors + ] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden(x_, nw, t, w, b): + act = t * torch.sigmoid(t) + act = act + (act.to(torch.bfloat16).double() - act).detach() + shift, scale = F.linear(act, w, b).chunk(2, dim=-1) + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * nw + return n * (1 + scale.index_select(0, ti)) + shift.index_select(0, ti) + + ours = grads(lambda *t: _cuda_op()(*t, ti)) + again = grads(lambda *t: _cuda_op()(*t, ti)) + ref = grads(golden, torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + for name, g, r in zip(("dx", "d_norm_weight", "d_temb", "dW", "db"), ours, ref): + rel = ((g.double() - r).abs().max() / r.abs().max()).item() + assert rel < 1e-2, (name, rel) # BF16 output rounding and round(1 + scale) only + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("final_adaln_out", device="cuda") + assert type(op).__name__ == "H3FinalAdaLNOutCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 32c6322c3..3679b409c 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -55,6 +55,7 @@ def test_all_ws1_single_ops_are_registered(): "adaln_row_gather", "h3_rmsnorm", "adaln_gate_residual", + "final_adaln_out", } <= names From 623bec886f679f168e17cea916ebbca47e868160 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:21:05 +0000 Subject: [PATCH 14/89] [WS1][MiniMax-H3] final_adaln_out: B200 evidence report, chain replay and figure Produced by scripts/h3_evidence.py and scripts/h3_chain_replay.py --backward at 38d575f on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index e05f1d15c..3e64efbef 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -40,9 +40,9 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste summation tree versus cuBLAS. Rows are batch- and position-invariant. - **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into the projection backward without an extra BF16 rounding: - - `d_temb`, `dW` and `db` are 13–57× closer to FP64 than diffusers (whose error varies from - run to run), which also rounds the - table gradient to BF16 and accumulates it with `index_select` atomics; + - `d_temb`, `dW` and `db` are 13–57× closer to FP64 than diffusers, whose error varies from + run to run because it rounds the table gradient to BF16 and accumulates it with + `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; - every gradient is repeat-bitwise. - **Golden.** FP64, rounding only where the model declares a dtype boundary: the SiLU cast, @@ -59,13 +59,25 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.13 ms | 0.15 ms | 1.62 ms | 1.09 ms | -| 32768 | 0.39 ms | 0.78 ms | 2.86 ms | 5.32 ms | -| 131072 | 1.24 ms | 2.93 ms | 7.78 ms | 20.7 ms | +| 4097 | 0.13 ms | 0.15 ms | 1.64 ms | 1.10 ms | +| 32768 | 0.39 ms | 0.78 ms | 2.88 ms | 5.36 ms | +| 131072 | 1.23 ms | 2.92 ms | 7.80 ms | 20.7 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. +## Evidence + +![final_adaln_out on B200: latency and backward accuracy](../usage/evidence/h3-final-adaln-out-b200/figure.png) + +There are two data files, both written from a clean tree at commit `38d575f`: + +- [`report.json`](../usage/evidence/h3-final-adaln-out-b200/report.json): op timings, the + forward-equality fraction, row invariance and backward accuracy. +- [`chain_replay.json`](../usage/evidence/h3-final-adaln-out-b200/chain_replay.json): the + whole conditioning chain (timestep → … → norm_out) replayed stage by stage over + T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. + ## Tests ```bash From 74a7772d7dbddfd269263fb3fa8a5b86caded83a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:12:05 +0000 Subject: [PATCH 15/89] fix(h3): reject mixed CUDA devices in deterministic linear helpers Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 4 ++++ docs/operators/h3-timestep-mlp.md | 6 +++--- tests/h3/test_h3_timestep_mlp.py | 33 +++++++++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 3 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index cc82acafa..ddaf199ee 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -358,6 +358,8 @@ std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor TORCH_CHECK(bias->is_cuda() && bias->dim() == 1 && bias->size(0) == n_out && bias->is_contiguous() && bias->scalar_type() == weight.scalar_type(), "bias must be a contiguous [N] CUDA tensor with the weight dtype"); + TORCH_CHECK(bias->device() == weight.device(), + "bias and weight must be on the same device"); } const c10::cuda::CUDAGuard device_guard(x.device()); @@ -390,6 +392,7 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei grad.scalar_type() == at::kFloat, "grad must be a contiguous 2-D float32 CUDA tensor"); check_matrix(weight, "weight"); + TORCH_CHECK(grad.device() == weight.device(), "grad and weight must be on the same device"); TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N"); TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, "out_dtype must be float32 or bfloat16"); @@ -434,6 +437,7 @@ std::vector h3_det_linear_backward_weight(torch::Tensor grad, tor grad.scalar_type() == at::kFloat, "grad must be a contiguous 2-D float32 CUDA tensor"); check_matrix(x, "x"); + TORCH_CHECK(grad.device() == x.device(), "grad and x must be on the same device"); TORCH_CHECK(grad.size(0) == x.size(0), "grad rows != x rows"); TORCH_CHECK(w_dtype == at::kFloat || w_dtype == at::kBFloat16, "w_dtype must be float32 or bfloat16"); diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 1da2434ac..a57f98033 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -99,9 +99,9 @@ B200, pinned weights (63.3 MB FP32). End-to-end op times include the Python wrap | T | CUDA op | provider | kernels only (CUDA) | | --- | --- | --- | --- | -| 1 | 32.8 µs | 36.5 µs | 3.8 + 11.6 µs | -| 2 | 38.8 µs | 79.9 µs | 3.9 + 13.7 µs | -| 4 | 42.7 µs | 83.1 µs | 5.0 + 15.8 µs | +| 1 | 35.5 µs | 40.8 µs | 3.8 + 11.6 µs | +| 2 | 40.4 µs | 84.9 µs | 3.9 + 13.7 µs | +| 4 | 44.9 µs | 88.6 µs | 5.0 + 15.8 µs | The 5376→2688 layer reads its weights at about 5.0 TB/s at T = 1. For T >= 2 the provider switches from cuBLAS GEMV to an SGEMM path, which takes about 52 µs of diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py index 7e09da4ba..71f1e694f 100644 --- a/tests/h3/test_h3_timestep_mlp.py +++ b/tests/h3/test_h3_timestep_mlp.py @@ -191,3 +191,36 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("timestep_mlp_fp32", device="cuda") assert type(op).__name__ == "H3TimestepMLPCudaOp" + + +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +class TestCudaDetLinearDeviceContract: + def test_forward_rejects_bias_on_another_device(self, dtype): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_forward + + _cuda_op() + x = torch.randn(2, 8, device="cuda:0", dtype=dtype) + weight = torch.randn(16, 8, device="cuda:0", dtype=dtype) + bias = torch.randn(16, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="bias and weight must be on the same device"): + det_linear_forward(x, weight, bias) + + def test_backward_input_rejects_weight_on_another_device(self, dtype): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_backward_input + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + weight = torch.randn(16, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and weight must be on the same device"): + det_linear_backward_input(grad, weight, dtype) + + def test_backward_weight_rejects_input_on_another_device(self, dtype): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_backward_weight + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + x = torch.randn(2, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and x must be on the same device"): + det_linear_backward_weight(grad, x, dtype) From e0f9dcb174f8ff04ea8bb69cb0116124736925ca Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:14:16 +0000 Subject: [PATCH 16/89] fix(h3): validate native timesteps and extracted weight contracts Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/timestep_sinusoid.cu | 9 ++- csrc/ops.cpp | 7 +- docs/operators/h3-timestep-sinusoid.md | 11 +-- rl_engine/_C.pyi | 5 +- .../kernels/ops/cuda/h3/timestep_sinusoid.py | 13 ++-- scripts/plot_h3_evidence.py | 4 +- scripts/prepare_h3_weights.py | 6 +- tests/h3/test_h3_timestep_sinusoid.py | 34 ++++++++++ tests/h3/test_prepare_h3_weights.py | 67 +++++++++++++++++++ 9 files changed, 138 insertions(+), 18 deletions(-) create mode 100644 tests/h3/test_prepare_h3_weights.py diff --git a/csrc/cuda/h3/timestep_sinusoid.cu b/csrc/cuda/h3/timestep_sinusoid.cu index 778b78965..9322920a4 100644 --- a/csrc/cuda/h3/timestep_sinusoid.cu +++ b/csrc/cuda/h3/timestep_sinusoid.cu @@ -50,7 +50,7 @@ __global__ void h3_timestep_sinusoid_kernel( } // namespace torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, - double max_period) { + double max_period, bool check_range) { TORCH_CHECK(timestep.is_cuda(), "timestep must be a CUDA tensor"); TORCH_CHECK(timestep.scalar_type() == at::kFloat, "timestep must be float32, got ", timestep.scalar_type()); @@ -62,6 +62,13 @@ torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_c const c10::cuda::CUDAGuard device_guard(timestep.device()); auto t = timestep.contiguous(); + // Check at the native boundary, including direct extension calls. The explicit + // opt-out is for already-validated inputs and kernel-only profiling. + if (check_range) { + TORCH_CHECK_VALUE(((t >= 0) & (t <= 1)).all().item(), + "timestep must be finite and lie in [0, 1]: H3 consumes t = 1 - sigma " + "unscaled"); + } const int64_t num_timesteps = t.size(0); const int64_t half = num_channels / 2; auto out = torch::empty({num_timesteps, num_channels}, t.options()); diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 71b5b5502..96a4455e2 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -501,7 +501,7 @@ at::Tensor prefix_shared_attention( #if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \ (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, - double max_period); + double max_period, bool check_range); #endif // PyBind11 Module Registration @@ -762,8 +762,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { #if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \ (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) - m.def("h3_timestep_sinusoid_forward", &h3_timestep_sinusoid_forward, + m.def("h3_timestep_sinusoid_forward", torch::wrap_pybind_function(h3_timestep_sinusoid_forward), "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", - py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0); + py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0, + py::arg("check_range") = true); #endif } diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 514ee651f..316b02132 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -40,12 +40,15 @@ features = op(timestep) # timestep: (T,) in [0, 1] -> (T, 256) floa | --- | --- | --- | --- | | `timestep` | `(T,)`, `T >= 1` | fp32 / bf16 / fp16 | Finite, in `[0, 1]`; any stride; CUDA for the CUDA backend | | `num_channels` | scalar | int | Positive and even; H3 uses 256 | -| output | `(T, num_channels)` | float32 | `[cos | sin]` order | +| output | `(T, num_channels)` | float32 | `[cos \| sin]` order | Inputs that break the contract fail closed: non-1-D, empty, integer or out-of-range timesteps raise. `t * 1000` callers (RFC probe H10) therefore fail instead of -silently producing a different embedding. The range check costs one host read-back; -pass `check_range=False` to `forward` when the caller has already validated `t`. +silently producing a different embedding. The range check costs one host read-back. +The native CUDA entrypoint performs the value check, so the wrapper does not repeat +the host synchronization. Pass `check_range=False` to `forward` or +`rl_engine._C.h3_timestep_sinusoid_forward` only when the caller has already validated +`t`; this trusted path skips value validation and is used for kernel-only profiling. ## Numerics @@ -79,7 +82,7 @@ python benchmarks/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 ``` On a B200 this is a single launch of about 15 µs, independent of `T` for `T <= 64`. With -the range check it takes about 51 µs. The provider path takes about 64 µs, because it is +the range check it takes about 49 µs. The provider path takes about 62–70 µs, because it is six eager kernels plus two concatenations. Either way the op is launch-bound: it moves about 1 KB per timestep. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 970df8515..6c435b850 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -326,5 +326,8 @@ def deterministic_collective_rocm_ipc_all_gather_input( output: torch.Tensor, ) -> None: ... def h3_timestep_sinusoid_forward( - timestep: torch.Tensor, num_channels: int = 256, max_period: float = 10000.0 + timestep: torch.Tensor, + num_channels: int = 256, + max_period: float = 10000.0, + check_range: bool = True, ) -> torch.Tensor: ... diff --git a/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py index a779941b2..72669d7c0 100644 --- a/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py +++ b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py @@ -26,9 +26,11 @@ def h3_sinusoid_cuda_available() -> bool: class _H3TimestepSinusoidCuda(torch.autograd.Function): @staticmethod - def forward(ctx, timestep: torch.Tensor, num_channels: int): + def forward(ctx, timestep: torch.Tensor, num_channels: int, check_range: bool): t32 = timestep.detach().float().contiguous() - out = _C.h3_timestep_sinusoid_forward(t32, int(num_channels), float(H3_MAX_PERIOD)) + out = _C.h3_timestep_sinusoid_forward( + t32, int(num_channels), float(H3_MAX_PERIOD), check_range + ) ctx.save_for_backward(out) ctx.num_channels = int(num_channels) ctx.timestep_dtype = timestep.dtype @@ -53,7 +55,7 @@ def backward(ctx, grad_out: torch.Tensor): impl=BACKWARD_IMPL, family="cuda", ) - return grad_t, None + return grad_t, None, None class H3TimestepSinusoidCudaOp: @@ -84,10 +86,11 @@ def forward( num_channels: int = H3_FREQ_DIM, check_range: bool = True, ) -> torch.Tensor: - validate_h3_timesteps(timestep, num_channels, check_range=check_range) + # Native code owns value validation, avoiding a second host sync here. + validate_h3_timesteps(timestep, num_channels, check_range=False) if not timestep.is_cuda: raise ValueError("H3TimestepSinusoidCudaOp needs a CUDA timestep tensor") - return _H3TimestepSinusoidCuda.apply(timestep, num_channels) + return _H3TimestepSinusoidCuda.apply(timestep, num_channels, check_range) def forward_fp32(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): # The op is FP32 end to end; the FP32 path is the op itself. diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index b0bcc8eb1..1880724a3 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -54,7 +54,7 @@ def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt for i, (name, color, values) in enumerate(series): xs = [x + (i - (n - 1) / 2) * width for x in range(len(labels))] bars = ax.bar(xs, values, width * 0.92, color=color, label=name, zorder=2) - for bar, value in zip(bars, values): + for bar, value in zip(bars, values, strict=True): ax.annotate( fmt.format(value), (bar.get_x() + bar.get_width() / 2, bar.get_height()), @@ -125,7 +125,7 @@ def plot_sinusoid(report: dict[str, Any]): color=MUTED, ha="right", ) - for x, value in zip(xs, cuda): + for x, value in zip(xs, cuda, strict=True): if value == 0: right.annotate( "exact", diff --git a/scripts/prepare_h3_weights.py b/scripts/prepare_h3_weights.py index dc2ed209e..910f34584 100644 --- a/scripts/prepare_h3_weights.py +++ b/scripts/prepare_h3_weights.py @@ -82,8 +82,10 @@ def fetch(filename: str) -> Path: for name, spec in manifest["tensors"].items(): tensor = tensors[name] - assert str(tensor.dtype).removeprefix("torch.") == spec["dtype"], name - assert list(tensor.shape) == spec["shape"], name + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") target = out_dir / EXTRACTED_FILE save_file(tensors, str(target), metadata={"revision": identity["revision"]}) diff --git a/tests/h3/test_h3_timestep_sinusoid.py b/tests/h3/test_h3_timestep_sinusoid.py index 8575041a3..bcd567dc3 100644 --- a/tests/h3/test_h3_timestep_sinusoid.py +++ b/tests/h3/test_h3_timestep_sinusoid.py @@ -139,6 +139,40 @@ def test_rejects_out_of_range(self): with pytest.raises(ValueError): _cuda_op()(torch.tensor([0.0, 999.0], device="cuda")) + @pytest.mark.parametrize("bad", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]) + def test_native_entrypoint_rejects_invalid_values(self, bad): + from rl_engine.kernels.ops.base import _C + + _cuda_op() + # Valid neighbors must not hide an invalid timestep in the native call. + t = torch.tensor([0.0, bad, 1.0], device="cuda") + with pytest.raises(ValueError, match=r"finite and lie in \[0, 1\]"): + _C.h3_timestep_sinusoid_forward(t) + with pytest.raises(ValueError, match=r"finite and lie in \[0, 1\]"): + _cuda_op()(t) + + def test_native_entrypoint_accepts_boundaries(self): + from rl_engine.kernels.ops.base import _C + + _cuda_op() + t = torch.tensor([0.0, 0.5, 1.0], device="cuda") + assert torch.equal(_C.h3_timestep_sinusoid_forward(t), provider_time_proj(t, 256)) + + def test_trusted_path_supports_cuda_graph(self): + op = _cuda_op() + t = h3_timesteps(7) + expected = op(t) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op.forward(t, check_range=False) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op.forward(t, check_range=False) + graph.replay() + assert torch.equal(out, expected) + def test_backward_matches_autograd_reference(self): t = h3_timesteps(17).requires_grad_(True) t_ref = t.detach().clone().requires_grad_(True) diff --git a/tests/h3/test_prepare_h3_weights.py b/tests/h3/test_prepare_h3_weights.py new file mode 100644 index 000000000..826fbd95e --- /dev/null +++ b/tests/h3/test_prepare_h3_weights.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Weight extraction rejects off-manifest tensors even under Python optimization.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest +import torch + +from rl_engine.testing.h3_weights import EXTRACTED_FILE, sha256_file + + +@pytest.mark.parametrize("optimize", [0, 1]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape"]) +def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, mismatch): + import huggingface_hub + from safetensors.torch import load_file, save_file + + tensor = torch.tensor([0.0, 1.0]) + config = tmp_path / "config.json" + config.write_text("{}") + index = tmp_path / "index.json" + index.write_text(json.dumps({"weight_map": {"weight": "shard.safetensors"}})) + shard = tmp_path / "shard.safetensors" + save_file({"weight": tensor}, str(shard)) + manifest = { + "model_identity": { + "hf_repo": "test/h3", + "subfolder": "dit", + "revision": "pinned", + "config_sha256": sha256_file(config), + "index_file": index.name, + "index_sha256": sha256_file(index), + }, + "weight_shards": {shard.name: {"sha256": sha256_file(shard)}}, + "tensors": { + "weight": { + "dtype": "bfloat16" if mismatch == "dtype" else "float32", + "shape": [3] if mismatch == "shape" else [2], + } + }, + } + monkeypatch.setattr( + huggingface_hub, + "hf_hub_download", + lambda repo, filename, **kwargs: str(tmp_path / Path(filename).name), + ) + out_dir = tmp_path / "extracted" + monkeypatch.setattr(sys, "argv", ["prepare_h3_weights.py", "--out", str(out_dir)]) + script = Path(__file__).resolve().parents[2] / "scripts" / "prepare_h3_weights.py" + namespace = {"__name__": "prepare_h3_weights_test", "__file__": str(script)} + # Compiling the actual script at optimize=1 reproduces python -O's assert removal. + exec(compile(script.read_text(), str(script), "exec", optimize=optimize), namespace) + namespace["load_h3_manifest"] = lambda: manifest + target = out_dir / EXTRACTED_FILE + if mismatch: + with pytest.raises(SystemExit, match=f"weight: {mismatch} .* != manifest"): + namespace["main"]() + assert not target.exists() + else: + namespace["main"]() + assert torch.equal(load_file(str(target))["weight"], tensor) From 41ecf74b6300e83179bb020bf8f72331f4cd08ed Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:53:53 +0000 Subject: [PATCH 17/89] fix(h3): address AdaLN projection review findings Validate pinned tensor contents at preparation and load time. Execute chain prerequisites for selected stages and require checkpoint weights for weighted evidence. Reject invalid CUDA tile and backward inputs, align the CUDA wrapper guards, skip unsupported devices, and rotate comparative timing with raw samples. Validation: 143 H3 tests passed on B200 after a fresh CUDA build; illegal FP32 tiles fail compilation; isolated AdaLN replay and pinned-weight evidence generation passed; Black, isort, Ruff, and Flake8 passed. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 3 + .../kernels/ops/cuda/h3/adaln_projection.py | 2 + rl_engine/testing/h3_manifest.json | 30 +++- rl_engine/testing/h3_report.py | 37 +++- rl_engine/testing/h3_weights.py | 22 ++- scripts/h3_chain_replay.py | 12 +- scripts/h3_evidence.py | 10 +- scripts/prepare_h3_weights.py | 10 +- tests/h3/test_h3_cli.py | 160 ++++++++++++++++++ tests/h3/test_h3_conditioning_e2e.py | 5 +- tests/h3/test_h3_det_linear.py | 34 ++++ tests/h3/test_h3_report.py | 68 ++++++++ tests/h3/test_h3_weights.py | 70 ++++++++ tests/h3/test_prepare_h3_weights.py | 5 +- 14 files changed, 443 insertions(+), 25 deletions(-) create mode 100644 tests/h3/test_h3_cli.py create mode 100644 tests/h3/test_h3_det_linear.py create mode 100644 tests/h3/test_h3_report.py create mode 100644 tests/h3/test_h3_weights.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 95e16b234..246d27d60 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -124,6 +124,7 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) using WV = Vec16; using XV = Vec16; constexpr int kVec = WV::kN; + static_assert(kCols * kRows <= kWarp, "one warp cannot store more than kWarp outputs per tile"); static_assert(XV::kN == kVec, "x and weight share a dtype"); constexpr int64_t kStride = static_cast(kWarp) * kVec; @@ -499,6 +500,7 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, "out_dtype must be float32 or bfloat16"); const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); const int64_t n_out = weight.size(0); const int64_t k_in = weight.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); @@ -544,6 +546,7 @@ std::vector h3_det_linear_backward_weight(torch::Tensor grad, tor TORCH_CHECK(w_dtype == at::kFloat || w_dtype == at::kBFloat16, "w_dtype must be float32 or bfloat16"); const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); const int64_t n_out = grad.size(1); const int64_t k_in = x.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py index c8a9b42cc..aee0ea503 100644 --- a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py +++ b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py @@ -100,6 +100,8 @@ def forward_table(self, temb, weight, bias) -> torch.Tensor: """The raw ``(T, 6H*3)`` projection, for consumers that gather it directly.""" validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") return _H3AdaLNProjectionCuda.apply(temb, weight, bias) def forward_fp32(self, temb, weight, bias): diff --git a/rl_engine/testing/h3_manifest.json b/rl_engine/testing/h3_manifest.json index 52b1989b1..5d08b87f3 100644 --- a/rl_engine/testing/h3_manifest.json +++ b/rl_engine/testing/h3_manifest.json @@ -30,12 +30,30 @@ } }, "tensors": { - "time_embedder.linear_1.weight": {"dtype": "float32", "shape": [5376, 256]}, - "time_embedder.linear_1.bias": {"dtype": "float32", "shape": [5376]}, - "time_embedder.linear_2.weight": {"dtype": "float32", "shape": [2688, 5376]}, - "time_embedder.linear_2.bias": {"dtype": "float32", "shape": [2688]}, - "transformer_blocks.0.adaln_proj.linear.weight": {"dtype": "bfloat16", "shape": [96768, 2688]}, - "transformer_blocks.0.adaln_proj.linear.bias": {"dtype": "bfloat16", "shape": [96768]} + "time_embedder.linear_1.weight": { + "dtype": "float32", "shape": [5376, 256], + "sha256": "4c50eabe3c434dc05422e6501cdf7baa5cecbb1a7706ae9599c7a5b0e34e2913" + }, + "time_embedder.linear_1.bias": { + "dtype": "float32", "shape": [5376], + "sha256": "64763a5e39cb078eb0d8c5be5d56ad36cac3dec37be93a87e106d82ea6c071b7" + }, + "time_embedder.linear_2.weight": { + "dtype": "float32", "shape": [2688, 5376], + "sha256": "697819dd9413e0fec3a50987a964f1fd5db568ca545267414238f452294a8154" + }, + "time_embedder.linear_2.bias": { + "dtype": "float32", "shape": [2688], + "sha256": "096a3c7539f0248698cc65376c8da260e7365bf4ccef847c09bb9a36782a743c" + }, + "transformer_blocks.0.adaln_proj.linear.weight": { + "dtype": "bfloat16", "shape": [96768, 2688], + "sha256": "a0000193b30ea36a1996f53fc0d6c4c37ce5fd4585284f0f9bae570d25c6b8cf" + }, + "transformer_blocks.0.adaln_proj.linear.bias": { + "dtype": "bfloat16", "shape": [96768], + "sha256": "d59f9de9e07e6b4e3b57009627e276a26971ce379555fc99dbb06d3a237c8fc0" + } }, "reference_implementation": { "repo": "huggingface/diffusers", diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 3f84eb520..d4b45e09b 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -63,12 +63,35 @@ def peak_mib(fn: Callable[[], Any]) -> float: def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} - for key in TIMED_KEYS: - if key in case: - us = time_us(case[key], warmup, iters) - row[f"{key}_us"] = us - row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 - row[f"{key}_peak_mib"] = peak_mib(case[key]) + keys = [key for key in TIMED_KEYS if key in case] + if not keys: + return row + for _ in range(warmup): + for key in keys: + case[key]() + torch.cuda.synchronize() + samples = {key: [] for key in keys} + orders = [] + for index in range(iters): + # Rotate the first implementation so drift does not favor one backend. + offset = index % len(keys) + order = keys[offset:] + keys[:offset] + orders.append(order) + for key in order: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + case[key]() + end.record() + end.synchronize() + samples[key].append(start.elapsed_time(end) * 1e3) + row["timing_samples_us"] = samples + row["timing_order"] = orders + for key in keys: + us = statistics.median(samples[key]) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key]) return row @@ -244,7 +267,7 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, invariant = all( all( torch.equal(s, f[3 * i : 3 * i + 3]) - for s, f in zip(op(temb[i : i + 1], weight, bias), full) + for s, f in zip(op(temb[i : i + 1], weight, bias), full, strict=True) ) for i in range(9) ) diff --git a/rl_engine/testing/h3_weights.py b/rl_engine/testing/h3_weights.py index 8090c6e83..7d7abdf03 100644 --- a/rl_engine/testing/h3_weights.py +++ b/rl_engine/testing/h3_weights.py @@ -41,6 +41,13 @@ def sha256_file(path: Path, chunk: int = 1 << 24) -> str: return digest.hexdigest() +def sha256_tensor(tensor: torch.Tensor) -> str: + """Hash a CPU tensor's contiguous bytes without safetensors metadata.""" + + tensor_bytes = tensor.contiguous().view(torch.uint8).numpy() + return hashlib.sha256(memoryview(tensor_bytes)).hexdigest() + + def load_h3_conditioning_weights( device: torch.device | str = "cpu", names: list[str] | None = None ) -> dict[str, torch.Tensor]: @@ -55,15 +62,20 @@ def load_h3_conditioning_weights( if not path.is_file(): raise FileNotFoundError(f"{path} missing; run scripts/prepare_h3_weights.py") manifest = load_h3_manifest() - tensors = load_file(str(path), device=str(device)) + tensors = load_file(str(path), device="cpu") wanted = names if names is not None else list(manifest["tensors"]) - out: dict[str, torch.Tensor] = {} - for name in wanted: - spec = manifest["tensors"][name] + for name, spec in manifest["tensors"].items(): tensor = tensors[name] if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: raise ValueError(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") if list(tensor.shape) != spec["shape"]: raise ValueError(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") - out[name] = tensor + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise ValueError(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + out: dict[str, torch.Tensor] = {} + for name in wanted: + if name not in manifest["tensors"]: + raise KeyError(name) + out[name] = tensors[name].to(device=device) return out diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py index e262ca053..512a15f06 100644 --- a/scripts/h3_chain_replay.py +++ b/scripts/h3_chain_replay.py @@ -35,7 +35,11 @@ def _ints(text: str) -> list[int]: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--stages", default=",".join(s.name for s in h3_chain.STAGES)) + parser.add_argument( + "--stages", + default=",".join(s.name for s in h3_chain.STAGES), + help="comma list of stages to report; preceding stages run automatically", + ) parser.add_argument("--timesteps", type=_ints, default=[1, 2, 4], help="comma list of T") parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") parser.add_argument("--seed", type=int, default=0) @@ -49,6 +53,8 @@ def main() -> None: if unknown: raise SystemExit(f"unknown stages: {unknown}") stages = [s for s in h3_chain.STAGES if s.name in wanted] + last_stage = max(i for i, stage in enumerate(h3_chain.STAGES) if stage.name in wanted) + execution_stages = h3_chain.STAGES[: last_stage + 1] torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False @@ -65,8 +71,9 @@ def main() -> None: num_timesteps=num_timesteps, seq_len=seq_len, seed=args.seed, - stages=stages, + stages=execution_stages, ) + case["stages"] = [entry for entry in case["stages"] if entry["stage"] in wanted] cases.append(case) summary = ", ".join( f"{e['stage']}=" @@ -85,6 +92,7 @@ def main() -> None: **h3_chain.git_state(), "environment": h3_chain.environment(), "stages": [s.name for s in stages], + "executed_stages": [s.name for s in execution_stages], "cases": cases, } if args.out is not None: diff --git a/scripts/h3_evidence.py b/scripts/h3_evidence.py index c20778234..8351da245 100644 --- a/scripts/h3_evidence.py +++ b/scripts/h3_evidence.py @@ -26,7 +26,11 @@ from rl_engine.kernels.registry import KernelRegistry # noqa: E402 from rl_engine.testing.h3_chain import environment, git_state # noqa: E402 from rl_engine.testing.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 -from rl_engine.testing.h3_weights import load_h3_manifest # noqa: E402 +from rl_engine.testing.h3_weights import ( # noqa: E402 + WEIGHTS_ENV, + h3_weights_dir, + load_h3_manifest, +) def main() -> None: @@ -34,6 +38,9 @@ def main() -> None: parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) parser.add_argument("--out", type=Path, required=True) args = parser.parse_args() + uses_weights = args.op != "timestep_sinusoid_h3" + if uses_weights and h3_weights_dir() is None: + raise SystemExit(f"set {WEIGHTS_ENV} to the scripts/prepare_h3_weights.py output dir") if not torch.cuda.is_available(): raise SystemExit("needs a CUDA device") torch.backends.cuda.matmul.allow_tf32 = False @@ -46,6 +53,7 @@ def main() -> None: "op": args.op, "rfc": manifest["rfc"], "model_revision": manifest["model_identity"]["revision"], + "weight_source": "pinned_checkpoint" if uses_weights else "not_applicable", "reference_commit": manifest["reference_implementation"]["commit"], **git_state(), "environment": environment(), diff --git a/scripts/prepare_h3_weights.py b/scripts/prepare_h3_weights.py index 910f34584..262408736 100644 --- a/scripts/prepare_h3_weights.py +++ b/scripts/prepare_h3_weights.py @@ -23,7 +23,12 @@ REPO_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(REPO_ROOT)) -from rl_engine.testing.h3_weights import EXTRACTED_FILE, load_h3_manifest, sha256_file # noqa: E402 +from rl_engine.testing.h3_weights import ( # noqa: E402 + EXTRACTED_FILE, + load_h3_manifest, + sha256_file, + sha256_tensor, +) def _check(path: Path, expected: str) -> None: @@ -86,6 +91,9 @@ def fetch(filename: str) -> Path: raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") if list(tensor.shape) != spec["shape"]: raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise SystemExit(f"{name}: sha256 {actual} != manifest {spec['sha256']}") target = out_dir / EXTRACTED_FILE save_file(tensors, str(target), metadata={"revision": identity["revision"]}) diff --git a/tests/h3/test_h3_cli.py b/tests/h3/test_h3_cli.py new file mode 100644 index 000000000..7712fc41a --- /dev/null +++ b/tests/h3/test_h3_cli.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The H3 CLIs preserve chain dependencies and require pinned report weights.""" + +from __future__ import annotations + +import json +import sys +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.testing import h3_chain +from rl_engine.testing.h3_weights import WEIGHTS_ENV +from scripts import h3_chain_replay, h3_evidence + + +@pytest.fixture +def cpu_chain_cli(monkeypatch): + executed = [] + state = {"provider_shift": 0} + + class CpuOp: + def __call__(self, value): + return value + 1 + + forward_fp32 = __call__ + + stages = [] + for index, stage in enumerate(h3_chain.STAGES): + + def stage_input(ctx, upstream, index=index): + if index == 0: + return ctx["timestep"] + assert isinstance(upstream, torch.Tensor), "a preceding stage must run first" + return upstream + + stages.append( + replace( + stage, + candidate=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + golden=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + provider=lambda ctx, up, stage_input=stage_input, index=index: ( + stage_input(ctx, up) + 1 + (state["provider_shift"] if index == 0 else 0) + ), + ) + ) + + def get_op(op_type, *, device): + executed.append(op_type) + return CpuOp() + + monkeypatch.setattr(h3_chain, "STAGES", stages) + monkeypatch.setattr(h3_chain, "make_context", lambda *a, **kw: {"timestep": torch.zeros(1)}) + monkeypatch.setattr(h3_chain, "golden_op", lambda *a: CpuOp()) + monkeypatch.setattr(h3_chain_replay, "KernelRegistry", lambda: SimpleNamespace(get_op=get_op)) + monkeypatch.setattr(h3_chain_replay, "load_h3_conditioning_weights", lambda *a: {}) + monkeypatch.setattr(h3_chain, "environment", lambda: {}) + monkeypatch.setattr(h3_chain, "git_state", lambda: {}) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + return executed, state + + +@pytest.mark.parametrize("provider_shift", [0, 1]) +@pytest.mark.parametrize( + "requested, execution_count", + [ + ("timestep_mlp_fp32", 2), + ("adaln_projection_3mod", 3), + ("timestep_sinusoid_h3,adaln_projection_3mod", 3), + ("timestep_mlp_fp32,timestep_sinusoid_h3", 2), + ], +) +def test_chain_selection_runs_dependencies( + cpu_chain_cli, monkeypatch, tmp_path, capsys, requested, execution_count, provider_shift +): + executed, state = cpu_chain_cli + state["provider_shift"] = provider_shift + out = tmp_path / "chain.json" + monkeypatch.setattr( + sys, + "argv", + [ + "h3_chain_replay.py", + "--stages", + requested, + "--timesteps", + "1", + "--seq-lens", + "3", + "--out", + str(out), + ], + ) + h3_chain_replay.main() + + report = json.loads(out.read_text()) + stage_names = [stage.name for stage in h3_chain.STAGES] + expected_report = [name for name in stage_names if name in requested.split(",")] + assert executed == stage_names[:execution_count] + assert report["executed_stages"] == executed + assert report["stages"] == expected_report + case = report["cases"][0] + assert [entry["stage"] for entry in case["stages"]] == expected_report + assert all(entry["chained_vs_golden"]["bitwise_equal"] for entry in case["stages"]) + assert case["first_drift"] == (stage_names[0] if provider_shift else None) + assert case["first_isolated_drift"] == (stage_names[0] if provider_shift else None) + summary = capsys.readouterr().out.split("; first_drift=", 1)[0] + for name in stage_names: + assert (f"{name}=" in summary) == (name in expected_report) + + +@pytest.mark.parametrize("op", ["timestep_mlp_fp32", "adaln_projection_3mod"]) +@pytest.mark.parametrize("weights_env", [None, " "]) +def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, weights_env): + if weights_env is None: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, weights_env) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + with pytest.raises(SystemExit, match=WEIGHTS_ENV): + h3_evidence.main() + assert not out.exists() + + +@pytest.mark.parametrize( + "op, weight_source", + [ + ("timestep_sinusoid_h3", "not_applicable"), + ("timestep_mlp_fp32", "pinned_checkpoint"), + ("adaln_projection_3mod", "pinned_checkpoint"), + ], +) +def test_evidence_records_weight_source(monkeypatch, tmp_path, op, weight_source): + if weight_source == "not_applicable": + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + monkeypatch.setattr(h3_evidence, "KernelRegistry", lambda: object()) + monkeypatch.setattr( + h3_evidence, + "git_state", + lambda: {"rl_kernel_commit": "test_commit", "tracked_tree_dirty": False}, + ) + monkeypatch.setattr(h3_evidence, "environment", lambda: {}) + monkeypatch.setitem(h3_evidence.ACCURACY, op, lambda registry: {}) + monkeypatch.setitem(h3_evidence.PERF_CASES, op, lambda registry: []) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + h3_evidence.main() + assert json.loads(out.read_text())["weight_source"] == weight_source diff --git a/tests/h3/test_h3_conditioning_e2e.py b/tests/h3/test_h3_conditioning_e2e.py index edd75a3b2..2bd1c47bd 100644 --- a/tests/h3/test_h3_conditioning_e2e.py +++ b/tests/h3/test_h3_conditioning_e2e.py @@ -22,7 +22,10 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing import h3_chain -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) # (T distinct timesteps, S packed rows): single row, short packed, realistic-size packed. CASES = [(1, 3), (3, 257), (4, 4097)] diff --git a/tests/h3/test_h3_det_linear.py b/tests/h3/test_h3_det_linear.py new file mode 100644 index 000000000..f8f90e83c --- /dev/null +++ b/tests/h3/test_h3_det_linear.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Native input validation for the H3 deterministic linear kernels.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.base import _C +from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), + pytest.mark.skipif(not det_linear_available(), reason="rl_engine._C lacks h3_det_linear_*"), +] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_backward_input_rejects_zero_rows(dtype): + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + weight = torch.empty(16, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_input(grad, weight, dtype) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_rejects_zero_rows(dtype, with_bias): + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + x = torch.empty(0, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_weight(grad, x, dtype, with_bias) diff --git a/tests/h3/test_h3_report.py b/tests/h3/test_h3_report.py new file mode 100644 index 000000000..eab1c8807 --- /dev/null +++ b/tests/h3/test_h3_report.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Comparative timing balances backend order and preserves its raw evidence.""" + +from __future__ import annotations + +import pytest + +from rl_engine.testing import h3_report + + +def test_measure_interleaves_backends_and_records_samples(monkeypatch): + elapsed_us = 0.0 + calls = [] + + class Event: + def __init__(self, *, enable_timing): + assert enable_timing + + def record(self): + self.timestamp = elapsed_us + + def synchronize(self): + pass + + def elapsed_time(self, end): + return (end.timestamp - self.timestamp) / 1e3 + + def candidate(): + nonlocal elapsed_us + calls.append("candidate") + elapsed_us += 2.0 + + def provider(): + nonlocal elapsed_us + calls.append("provider") + elapsed_us += 4.0 + + monkeypatch.setattr(h3_report.torch.cuda, "Event", Event) + monkeypatch.setattr(h3_report.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn: 0.0) + report = h3_report.measure( + {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 + ) + assert calls == [ + "candidate", + "provider", # warmup + "candidate", + "provider", + "provider", + "candidate", + "candidate", + "provider", + "provider", + "candidate", + ] + assert report["timing_order"] == [ + ["candidate", "provider"], + ["provider", "candidate"], + ["candidate", "provider"], + ["provider", "candidate"], + ] + assert report["timing_samples_us"] == {"candidate": [2.0] * 4, "provider": [4.0] * 4} + assert report["candidate_us"] == pytest.approx(2.0) + assert report["provider_us"] == pytest.approx(4.0) + assert report["candidate_gbps"] == pytest.approx(0.5) + assert report["provider_gbps"] == pytest.approx(0.25) diff --git a/tests/h3/test_h3_weights.py b/tests/h3/test_h3_weights.py new file mode 100644 index 000000000..015f78954 --- /dev/null +++ b/tests/h3/test_h3_weights.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Extracted H3 weights must retain the pinned tensor contents.""" + +from __future__ import annotations + +import hashlib + +import pytest +import torch + +from rl_engine.testing import h3_weights + + +@pytest.mark.parametrize( + "dtype,raw_bytes", + [ + (torch.float32, bytes.fromhex("000000000000803f")), + (torch.bfloat16, bytes.fromhex("0000803f")), + ], +) +def test_tensor_checksum_hashes_raw_bytes(dtype, raw_bytes): + tensor = torch.tensor([0.0, 1.0], dtype=dtype) + assert h3_weights.sha256_tensor(tensor) == hashlib.sha256(raw_bytes).hexdigest() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) +@pytest.mark.parametrize("names", [None, ["bias"]]) +def test_loading_rejects_modified_weight_contents(tmp_path, monkeypatch, dtype, mismatch, names): + from safetensors.torch import save_file + + tensors = { + "weight": torch.tensor([[0.0, 1.0], [2.0, 3.0]], dtype=dtype), + "bias": torch.tensor([0.0, 1.0], dtype=dtype), + } + manifest = { + "tensors": { + name: { + "dtype": str(tensor.dtype).removeprefix("torch."), + "shape": list(tensor.shape), + "sha256": h3_weights.sha256_tensor(tensor), + } + for name, tensor in tensors.items() + } + } + if mismatch == "dtype": + tensors["weight"] = tensors["weight"].to(torch.float64) + elif mismatch == "shape": + tensors["weight"] = tensors["weight"].reshape(4) + elif mismatch == "sha256": + tensors["weight"][0, 0] = 1.0 + # Tensor checksums do not depend on serializer metadata or tensor key order. + save_file( + dict(reversed(list(tensors.items()))), + str(tmp_path / h3_weights.EXTRACTED_FILE), + metadata={"revision": "a different serializer metadata value"}, + ) + monkeypatch.setenv(h3_weights.WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(h3_weights, "load_h3_manifest", lambda: manifest) + + if mismatch: + with pytest.raises(ValueError, match=f"weight: {mismatch} .* != manifest"): + h3_weights.load_h3_conditioning_weights("cpu", names=names) + else: + actual = h3_weights.load_h3_conditioning_weights("cpu", names=names) + assert list(actual) == (names if names is not None else list(tensors)) + for name, tensor in actual.items(): + assert torch.equal(tensor, tensors[name]) diff --git a/tests/h3/test_prepare_h3_weights.py b/tests/h3/test_prepare_h3_weights.py index 826fbd95e..8a1c94298 100644 --- a/tests/h3/test_prepare_h3_weights.py +++ b/tests/h3/test_prepare_h3_weights.py @@ -12,11 +12,11 @@ import pytest import torch -from rl_engine.testing.h3_weights import EXTRACTED_FILE, sha256_file +from rl_engine.testing.h3_weights import EXTRACTED_FILE, sha256_file, sha256_tensor @pytest.mark.parametrize("optimize", [0, 1]) -@pytest.mark.parametrize("mismatch", [None, "dtype", "shape"]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, mismatch): import huggingface_hub from safetensors.torch import load_file, save_file @@ -42,6 +42,7 @@ def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, misma "weight": { "dtype": "bfloat16" if mismatch == "dtype" else "float32", "shape": [3] if mismatch == "shape" else [2], + "sha256": "0" * 64 if mismatch == "sha256" else sha256_tensor(tensor), } }, } From 2bbef2263b4300fef343726a12e85b7571179d5c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:54:31 +0000 Subject: [PATCH 18/89] fix(h3): validate native gather indices and replay stage dependencies Check timestep and modality bounds on the device before computing gather addresses, including unchecked calls and overflowing int64 indices. Reject non-prefix replay stage selections and partial backward chains before loading weights. Validation: rebuilt the CUDA extension on B200; tests/h3: 175 passed, including 12 isolated native bounds probes and 16 CLI regressions. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 20 +++++-- scripts/h3_chain_replay.py | 26 +++++++-- tests/h3/test_h3_adaln_row_gather.py | 55 ++++++++++++++++++ tests/h3/test_h3_chain_replay_cli.py | 85 ++++++++++++++++++++++++++++ 4 files changed, 174 insertions(+), 12 deletions(-) create mode 100644 tests/h3/test_h3_chain_replay_cli.py diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 192faa187..4b3a168a0 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -63,8 +63,13 @@ struct CudaType { template __device__ __forceinline__ int64_t adaln_row(const idx_t* ti, const idx_t* tags, int64_t s, - int64_t modality_num) { - return static_cast(ti[s]) * modality_num + static_cast(tags[s]); + int64_t modality_num, int64_t num_timesteps) { + const int64_t timestep = static_cast(ti[s]); + const int64_t tag = static_cast(tags[s]); + // Check each semantic index before multiplication, including values that overflow int64. + CUDA_KERNEL_ASSERT(timestep >= 0 && timestep < num_timesteps); + CUDA_KERNEL_ASSERT(tag >= 0 && tag < modality_num); + return timestep * modality_num + tag; } // One block per (s, c); 16-byte copies when every row start is 16-byte aligned. @@ -73,10 +78,10 @@ __global__ void adaln_row_gather_kernel(const T* __restrict__ rows, int64_t row_ const idx_t* __restrict__ ti, const idx_t* __restrict__ tags, T* __restrict__ out, int64_t seq, int64_t hidden, int64_t chunks, - int64_t modality_num) { + int64_t modality_num, int64_t num_timesteps) { for (int64_t s = blockIdx.x; s < seq; s += gridDim.x) { const int64_t c = blockIdx.y; - const int64_t r = adaln_row(ti, tags, s, modality_num); + const int64_t r = adaln_row(ti, tags, s, modality_num, num_timesteps); const T* src = rows + r * row_stride + c * hidden; T* dst = out + (c * seq + s) * hidden; if constexpr (kVector) { @@ -146,6 +151,7 @@ torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor time " is not a multiple of ", chunks); TORCH_CHECK(modality_num > 0 && rows.size(0) % modality_num == 0, "rows count ", rows.size(0), " is not a multiple of ", modality_num); + const int64_t num_timesteps = rows.size(0) / modality_num; const int64_t seq = timestep_indices.numel(); TORCH_CHECK(seq > 0, "the packed sequence must not be empty"); check_index(timestep_indices, "timestep_indices", seq, rows); @@ -174,10 +180,12 @@ torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor time const idx_t* tags = token_tags.data_ptr(); if (vector) { adaln_row_gather_kernel<<>>( - src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num); + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num, + num_timesteps); } else { adaln_row_gather_kernel<<>>( - src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num); + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num, + num_timesteps); } }; if (timestep_indices.scalar_type() == at::kLong) { diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py index 52a1cddde..d5fc6a94d 100644 --- a/scripts/h3_chain_replay.py +++ b/scripts/h3_chain_replay.py @@ -30,12 +30,20 @@ def _ints(text: str) -> list[int]: + """Parse a comma-separated list of integer case sizes.""" + return [int(v) for v in text.split(",")] def main() -> None: + """Validate the chain prefix, replay cases, and optionally save evidence.""" + parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--stages", default=",".join(s.name for s in h3_chain.STAGES)) + parser.add_argument( + "--stages", + default=",".join(s.name for s in h3_chain.STAGES), + help="comma list of stages forming a nonempty ordered prefix of the chain", + ) parser.add_argument("--timesteps", type=_ints, default=[1, 2, 4], help="comma list of T") parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") parser.add_argument("--seed", type=int, default=0) @@ -48,13 +56,19 @@ def main() -> None: ) args = parser.parse_args() - if not torch.cuda.is_available(): - raise SystemExit("the chain replay needs a CUDA device") wanted = args.stages.split(",") - unknown = sorted(set(wanted) - {s.name for s in h3_chain.STAGES}) + stage_names = [s.name for s in h3_chain.STAGES] + unknown = sorted(set(wanted) - set(stage_names)) if unknown: - raise SystemExit(f"unknown stages: {unknown}") - stages = [s for s in h3_chain.STAGES if s.name in wanted] + parser.error(f"unknown stages: {unknown}") + if wanted != stage_names[: len(wanted)]: + parser.error("--stages must form a nonempty ordered prefix of the chain") + stages = h3_chain.STAGES[: len(wanted)] + if args.backward and len(stages) != len(h3_chain.STAGES): + parser.error("--backward requires every stage of the chain") + + if not torch.cuda.is_available(): + raise SystemExit("the chain replay needs a CUDA device") torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False diff --git a/tests/h3/test_h3_adaln_row_gather.py b/tests/h3/test_h3_adaln_row_gather.py index 0c58c7e49..8f6558034 100644 --- a/tests/h3/test_h3_adaln_row_gather.py +++ b/tests/h3/test_h3_adaln_row_gather.py @@ -13,6 +13,12 @@ from __future__ import annotations +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest import torch @@ -123,6 +129,55 @@ def test_rejects_out_of_range_on_device(self): with pytest.raises(IndexError): _cuda_op()(rows, ti, tags + 3) + @pytest.mark.parametrize("entrypoint", ["native", "unchecked_wrapper"]) + @pytest.mark.parametrize( + "timestep, tag, index_dtype", + [ + pytest.param(-1, 0, "int32", id="negative-timestep"), + pytest.param(2, 0, "int64", id="timestep-past-table"), + pytest.param(0, 3, "int32", id="tag-past-modality-valid-row"), + pytest.param(1, -1, "int64", id="negative-tag-valid-row"), + # Multiplication by 3 wraps this value to row 2 in signed int64. + pytest.param((2**64 + 2) // 3, 0, "int64", id="timestep-overflow-valid-row"), + pytest.param(0, 2**63 - 1, "int64", id="int64-max-tag"), + ], + ) + def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): + """Reject semantic index violations without poisoning pytest's CUDA context.""" + + _cuda_op() + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.kernels.ops.base import _C + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import H3AdaLNRowGatherCudaOp + + entrypoint, timestep, tag, index_dtype = sys.argv[1:] + rows = torch.zeros((6, 6), dtype=torch.float32, device="cuda") + dtype = getattr(torch, index_dtype) + ti = torch.tensor([int(timestep)], dtype=dtype, device="cuda") + tags = torch.tensor([int(tag)], dtype=dtype, device="cuda") + if entrypoint == "native": + _C.h3_adaln_row_gather_forward(rows, ti, tags, 6, 3) + else: + H3AdaLNRowGatherCudaOp().forward(rows, ti, tags, check_range=False) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(timestep), str(tag), index_dtype], + cwd=Path(__file__).resolve().parents[2], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + assert completed.returncode != 0, "native gather accepted invalid semantic indices" + assert "device-side assert triggered" in completed.stderr, ( + f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + ) + def test_gather_chunks_drop_in(self): rows = _rows(2, hidden=16) ti, tags = h3_packed_layout(10, 2) diff --git a/tests/h3/test_h3_chain_replay_cli.py b/tests/h3/test_h3_chain_replay_cli.py new file mode 100644 index 000000000..32f695af4 --- /dev/null +++ b/tests/h3/test_h3_chain_replay_cli.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Reject incomplete chain dependencies before checking CUDA or loading weights.""" + +from __future__ import annotations + +import sys + +import pytest + +from scripts import h3_chain_replay + +STAGE_NAMES = [stage.name for stage in h3_chain_replay.h3_chain.STAGES] + + +@pytest.fixture +def no_cuda_check(monkeypatch): + """Fail if invalid CLI arguments reach a CUDA availability check.""" + + def unexpected_check(): + """Report an environment check performed before argument rejection.""" + + pytest.fail("invalid arguments must be rejected before checking CUDA") + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", unexpected_check) + + +@pytest.mark.parametrize( + "stages, message", + [ + ("", "unknown stages"), + ("unknown", "unknown stages"), + (STAGE_NAMES[1], "ordered prefix"), + (STAGE_NAMES[-1], "ordered prefix"), + (",".join(STAGE_NAMES[:1] + STAGE_NAMES[2:]), "ordered prefix"), + (",".join(reversed(STAGE_NAMES)), "ordered prefix"), + (",".join([STAGE_NAMES[0], STAGE_NAMES[0]]), "ordered prefix"), + ], +) +def test_reject_invalid_stages(monkeypatch, capsys, no_cuda_check, stages, message): + """Reject unknown stages and selections that break chain dependencies.""" + + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", "--stages", stages]) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert message in capsys.readouterr().err + + +@pytest.mark.parametrize("count", range(1, len(STAGE_NAMES))) +def test_backward_requires_all_stages(monkeypatch, capsys, no_cuda_check, count): + """Require the full chain for parameter-gradient replay.""" + + monkeypatch.setattr( + sys, + "argv", + ["h3_chain_replay.py", "--stages", ",".join(STAGE_NAMES[:count]), "--backward"], + ) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert "--backward requires every stage" in capsys.readouterr().err + + +@pytest.mark.parametrize("count", range(1, len(STAGE_NAMES) + 1)) +def test_accept_ordered_prefix(monkeypatch, count): + """Allow every nonempty prefix to advance to the device requirement.""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr( + sys, "argv", ["h3_chain_replay.py", "--stages", ",".join(STAGE_NAMES[:count])] + ) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() + + +@pytest.mark.parametrize("options", [[], ["--backward"]]) +def test_accept_all_stages_by_default(monkeypatch, options): + """Keep the complete chain as the default for forward and backward replay.""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", *options]) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() From 50c97dbd5e6b95b9e0222ee7b7d5cdbfa66fac42 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:56:13 +0000 Subject: [PATCH 19/89] ci: enable CodeRabbit auto review for test-h3 CodeRabbit skips automatic reviews targeting non-default branches by default. Allow the test-h3 base branch so pushes to the H3 stack receive incremental reviews. Validation: configuration validated against the official CodeRabbit JSON schema. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .coderabbit.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .coderabbit.yaml diff --git a/.coderabbit.yaml b/.coderabbit.yaml new file mode 100644 index 000000000..cc811057e --- /dev/null +++ b/.coderabbit.yaml @@ -0,0 +1,5 @@ +# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json +reviews: + auto_review: + base_branches: + - "^test-h3$" From d113891c6b67a7262f08adc1837550a47e5c7e3d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:56:45 +0000 Subject: [PATCH 20/89] ci: enable CodeRabbit automatic reviews for test-h3 Include the RFC 420 target branch in CodeRabbit automatic reviews so pushes to the H3 stack are reviewed instead of silently skipped. Validation: configuration parses and validates against the published CodeRabbit v2 schema; the additional branch pattern matches only test-h3. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .coderabbit.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .coderabbit.yaml diff --git a/.coderabbit.yaml b/.coderabbit.yaml new file mode 100644 index 000000000..cc811057e --- /dev/null +++ b/.coderabbit.yaml @@ -0,0 +1,5 @@ +# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json +reviews: + auto_review: + base_branches: + - "^test-h3$" From 5c79216e34368257f14014bcac133be11fb4c216 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:57:41 +0000 Subject: [PATCH 21/89] ci: inherit existing CodeRabbit review settings Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .coderabbit.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.coderabbit.yaml b/.coderabbit.yaml index cc811057e..ebea58ae2 100644 --- a/.coderabbit.yaml +++ b/.coderabbit.yaml @@ -1,4 +1,5 @@ # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json +inheritance: true reviews: auto_review: base_branches: From d1b71b6b8ac6ba0139f74fe97055127258e9edf1 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:30:47 +0000 Subject: [PATCH 22/89] docs(h3): document conditioning operators and review regression tests Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- benchmarks/benchmark_h3_conditioning.py | 2 + rl_engine/_C.pyi | 42 +++++++++++++++++-- rl_engine/kernels/gtest/operator_inputs.py | 8 ++++ .../kernels/ops/cuda/h3/adaln_projection.py | 24 +++++++++++ rl_engine/kernels/ops/cuda/h3/det_linear.py | 25 +++++++++++ rl_engine/kernels/ops/cuda/h3/timestep_mlp.py | 17 ++++++++ .../kernels/ops/cuda/h3/timestep_sinusoid.py | 19 +++++++++ .../ops/pytorch/h3/adaln_projection.py | 22 ++++++++++ .../kernels/ops/pytorch/h3/timestep_mlp.py | 16 +++++++ .../ops/pytorch/h3/timestep_sinusoid.py | 15 +++++++ rl_engine/kernels/registry.py | 1 + rl_engine/testing/h3_cases.py | 13 ++++-- rl_engine/testing/h3_chain.py | 24 ++++++++++- rl_engine/testing/h3_provider.py | 17 ++++++-- rl_engine/testing/h3_report.py | 26 +++++++++++- rl_engine/testing/h3_weights.py | 6 +++ rl_engine/tests/test_dispatch.py | 2 + scripts/h3_chain_replay.py | 4 ++ scripts/h3_evidence.py | 2 + scripts/plot_h3_evidence.py | 18 +++++++- scripts/prepare_h3_weights.py | 6 +++ setup.py | 1 + tests/h3/conftest.py | 2 + tests/h3/test_h3_adaln_projection.py | 36 ++++++++++++++++ tests/h3/test_h3_cli.py | 14 +++++++ tests/h3/test_h3_conditioning_e2e.py | 6 +++ tests/h3/test_h3_det_linear.py | 4 ++ tests/h3/test_h3_report.py | 14 +++++++ tests/h3/test_h3_timestep_mlp.py | 42 +++++++++++++++++++ tests/h3/test_h3_timestep_sinusoid.py | 42 +++++++++++++++++++ tests/h3/test_h3_weights.py | 4 ++ tests/h3/test_prepare_h3_weights.py | 2 + tests/test_ws1_gtest_gpu.py | 2 + 33 files changed, 463 insertions(+), 15 deletions(-) diff --git a/benchmarks/benchmark_h3_conditioning.py b/benchmarks/benchmark_h3_conditioning.py index de1aaceb7..dc5aed949 100644 --- a/benchmarks/benchmark_h3_conditioning.py +++ b/benchmarks/benchmark_h3_conditioning.py @@ -27,6 +27,8 @@ def main() -> None: + """Benchmark selected CUDA operators, print summaries, and optionally save JSON.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--op", default="all", choices=["all", *PERF_CASES]) parser.add_argument("--warmup", type=int, default=20) diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fbe7558a8..29a889f27 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -330,17 +330,51 @@ def h3_timestep_sinusoid_forward( num_channels: int = 256, max_period: float = 10000.0, check_range: bool = True, -) -> torch.Tensor: ... +) -> torch.Tensor: + """Return FP32 CUDA cosine-then-sine features of shape ``(T, num_channels)``. + + Require a nonempty 1-D FP32 CUDA timestep tensor, positive even channels + and positive ``max_period``. Check finite values in ``[0, 1]`` when enabled; + the Python operator supplies the analytic backward for this native forward. + """ + ... + def h3_det_linear_forward( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None, activation: int = 0, save_pre_activation: bool = False, -) -> list[torch.Tensor]: ... +) -> list[torch.Tensor]: + """Return a deterministic CUDA projection and optional FP32 pre-activation. + + Require same-device/dtype FP32 or BF16 contiguous, 16-byte-aligned matrices + ``x (T, K)`` and ``weight (N, K)``, ``T > 0`` and ``K`` divisible by the + 16-byte vector width. Optional contiguous bias is ``(N,)`` in that dtype. + Activation 0 is identity and 1 is SiLU; BF16 forward requires SM80 or newer. + Return ``[out]`` in the input dtype or ``[out, pre]`` when saving activation. + """ + ... + def h3_det_linear_backward_input( grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype -) -> torch.Tensor: ... +) -> torch.Tensor: + """Return a deterministic ``(T, K)`` CUDA input gradient in ``out_dtype``. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``weight (N, K)``. Output dtype + must be FP32 or BF16; output-column chunks accumulate in fixed order. + """ + ... + def h3_det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True -) -> list[torch.Tensor]: ... +) -> list[torch.Tensor]: + """Return CUDA weight and optional bias gradients using an ascending row fold. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``x (T, K)``. Return ``[dweight]`` + with shape ``(N, K)``, or ``[dweight, dbias]`` with ``dbias (N,)`` when + requested, in FP32 or BF16 as specified by ``w_dtype``. + """ + ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 255a48c0c..ced084383 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -28,6 +28,7 @@ def make_operator_inputs( dtype: torch.dtype, device: torch.device, ) -> dict[str, Any]: + """Build operator keyword inputs on the requested device from CLI shape and seed options.""" builders = { "rms_norm": _make_rms_norm_inputs, "qk_norm": _make_qk_norm_inputs, @@ -57,6 +58,7 @@ def make_operator_inputs( def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: + """Return the operator's dimension label for benchmark and evidence reports.""" batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) names = { @@ -344,6 +346,7 @@ def _make_kv_cache_attention_inputs( def _h3_num_timesteps(args: argparse.Namespace) -> int: + """Read the packed H3 timestep count, falling back to the CLI batch size.""" # H3 packs a handful of distinct timesteps; reuse --batch as their count. return _arg_int(args, "num_timesteps", _arg_int(args, "batch", 2)) @@ -351,6 +354,7 @@ def _h3_num_timesteps(args: argparse.Namespace) -> int: def _h3_timesteps( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> torch.Tensor: + """Build seeded timesteps in [0, 1] with endpoint cases in the requested storage dtype.""" num = _h3_num_timesteps(args) mode = _arg_str(args, "input_mode", "random") if mode == "constant": @@ -366,12 +370,14 @@ def _h3_timesteps( def _make_timestep_sinusoid_h3_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: + """Supply a one-dimensional packed timestep tensor to the H3 sinusoid operator.""" return {"timestep": _h3_timesteps(args, dtype, device)} def _make_timestep_mlp_fp32_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: + """Build [T, 256] sinusoid features and seeded FP32 256→5376→2688 MLP parameters.""" # The H3 time_embedder is declared FP32; ``dtype`` does not apply to it. del dtype from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import NativeH3TimestepSinusoidOp @@ -390,12 +396,14 @@ def _make_timestep_mlp_fp32_inputs( def _h3_hidden(args: argparse.Namespace) -> int: + """Read the AdaLN channel width, defaulting to the checkpoint's 5376 channels.""" return _arg_int(args, "normalized_dim", H3_HIDDEN) def _make_adaln_projection_3mod_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: + """Build FP32 [T, 2688] embeddings and 18H projection parameters in the chosen dtype.""" # temb is FP32 by contract (the SiLU runs before the cast); ``dtype`` is the # projection's weight dtype, BF16 in the checkpoint. n_out = 6 * 3 * _h3_hidden(args) diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py index aee0ea503..fde0480e6 100644 --- a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py +++ b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py @@ -35,6 +35,11 @@ class _H3AdaLNProjectionCuda(torch.autograd.Function): @staticmethod def forward(ctx, temb, weight, bias): + """Project FP32 ``(T, D)`` embeddings to a weight-dtype ``(T, 18H)`` table. + + Save the embedding, rounded SiLU activation and weight for the VJP. + """ + temb = temb.contiguous() # Declared mixed-precision boundary: SiLU at temb's FP32 precision, # then exactly one rounding to the projection dtype. These are the @@ -46,6 +51,12 @@ def forward(ctx, temb, weight, bias): @staticmethod def backward(ctx, grad_table): + """Return requested embedding, weight and bias gradients in their input dtypes. + + Accumulate projection gradients deterministically and evaluate the + straight-through cast and SiLU VJP in FP32 for the embedding gradient. + """ + temb, act, weight = ctx.saved_tensors need_temb, need_weight, need_bias = ctx.needs_input_grad grad_table = grad_table.float().contiguous() @@ -81,15 +92,26 @@ class H3AdaLNProjectionCudaOp: contract = CONTRACT def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + if not det_linear_available(): raise RuntimeError( "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" ) def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` modulation views for CUDA inputs.""" + return self.forward(temb, weight, bias) def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six ``(3T, H)`` views after FP32 SiLU and a deterministic projection. + + Require CUDA FP32 ``temb`` of shape ``(T, D)`` with ``T > 0`` and + same-device FP32 or BF16 weight/bias of shapes ``(18H, D)``/``(18H,)``. + Outputs share one table in the weight dtype; autograd covers all inputs. + """ + hidden = validate_h3_adaln_projection(temb, weight, bias) if not temb.is_cuda: raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") @@ -105,4 +127,6 @@ def forward_table(self, temb, weight, bias) -> torch.Tensor: return _H3AdaLNProjectionCuda.apply(temb, weight, bias) def forward_fp32(self, temb, weight, bias): + """Run the CUDA projection with its declared weight-dtype output boundary.""" + return self.forward(temb, weight, bias) diff --git a/rl_engine/kernels/ops/cuda/h3/det_linear.py b/rl_engine/kernels/ops/cuda/h3/det_linear.py index f5dab854e..f12acabba 100644 --- a/rl_engine/kernels/ops/cuda/h3/det_linear.py +++ b/rl_engine/kernels/ops/cuda/h3/det_linear.py @@ -24,10 +24,14 @@ def det_linear_available() -> bool: + """Return whether the extension exposes all three H3 deterministic linear APIs.""" + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) def _require() -> None: + """Raise when a deterministic linear entry point is unavailable.""" + if not det_linear_available(): raise RuntimeError( "rl_engine._C lacks the h3_det_linear_* symbols; rebuild the CUDA extension " @@ -43,6 +47,14 @@ def det_linear_forward( activation: int = ACT_NONE, save_pre_activation: bool = False, ) -> list[torch.Tensor]: + """Project CUDA ``x`` of shape ``(T, K)`` with ``(N, K)`` weights and optional bias. + + Inputs share an FP32 or BF16 dtype/device, ``T > 0``, and ``K`` supports + 16-byte vector loads; BF16 forward needs SM80 or newer. Return ``[out]`` + in the input dtype, plus an FP32 pre-activation when requested. Activation + is identity or SiLU; contiguous copies preserve the reduction contract. + """ + _require() return _C.h3_det_linear_forward( x.contiguous(), @@ -56,6 +68,12 @@ def det_linear_forward( def det_linear_backward_input( grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype ) -> torch.Tensor: + """Compute a deterministic ``(T, K)`` input VJP on the inputs' CUDA device. + + Convert nonempty ``(T, N)`` gradients to FP32 and combine them with FP32 + or BF16 ``(N, K)`` weights; return FP32 or BF16 as specified by ``out_dtype``. + """ + _require() return _C.h3_det_linear_backward_input( grad.float().contiguous(), weight.contiguous(), out_dtype @@ -65,6 +83,13 @@ def det_linear_backward_input( def det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True ) -> list[torch.Tensor]: + """Fold nonempty CUDA ``(T, N)`` gradients and ``(T, K)`` inputs in row order. + + Convert gradients to FP32; ``x`` is FP32 or BF16 on the same device. + Return ``[dweight]`` of shape ``(N, K)``, optionally followed by ``dbias`` + of shape ``(N,)``, both in the requested FP32 or BF16 ``w_dtype``. + """ + _require() return _C.h3_det_linear_backward_weight( grad.float().contiguous(), x.contiguous(), w_dtype, bool(with_bias) diff --git a/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py b/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py index da85e87d8..041116c70 100644 --- a/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py +++ b/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py @@ -32,6 +32,8 @@ class _H3TimestepMLPCuda(torch.autograd.Function): @staticmethod def forward(ctx, x, w1, b1, w2, b2): + """Evaluate the FP32 CUDA linear-SiLU-linear graph and save VJP inputs.""" + x = x.contiguous() hidden, pre = det_linear_forward(x, w1, b1, activation=ACT_SILU, save_pre_activation=True) (out,) = det_linear_forward(hidden, w2, b2, activation=ACT_NONE) @@ -40,6 +42,8 @@ def forward(ctx, x, w1, b1, w2, b2): @staticmethod def backward(ctx, grad_out): + """Return requested FP32 input and parameter VJPs with fixed reduction order.""" + x, w1, w2, pre, hidden = ctx.saved_tensors need_x, need_w1, need_b1, need_w2, need_b2 = ctx.needs_input_grad grad_out = grad_out.float().contiguous() @@ -81,19 +85,32 @@ class H3TimestepMLPCudaOp: contract = CONTRACT def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + if not det_linear_available(): raise RuntimeError( "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" ) def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` timestep embeddings for same-device CUDA inputs.""" + return self.forward(x, w1, b1, w2, b2) def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Apply two deterministic FP32 CUDA projections with an intervening SiLU. + + Require same-device FP32 tensors: ``x`` is nonempty ``(T, K)``, + weights are ``(H, K)`` and ``(D, H)``, and biases are ``(H,)`` and + ``(D,)``. Return ``(T, D)`` with autograd support for all five inputs. + """ + validate_h3_timestep_mlp(x, w1, b1, w2, b2) if not x.is_cuda: raise ValueError("H3TimestepMLPCudaOp needs CUDA tensors") return _H3TimestepMLPCuda.apply(x, w1, b1, w2, b2) def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Use the standard CUDA MLP, whose inputs and output are already FP32.""" + return self.forward(x, w1, b1, w2, b2) diff --git a/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py index 72669d7c0..8601ef8f8 100644 --- a/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py +++ b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py @@ -21,12 +21,16 @@ def h3_sinusoid_cuda_available() -> bool: + """Return whether the extension provides the CUDA H3 sinusoid entry point.""" + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) class _H3TimestepSinusoidCuda(torch.autograd.Function): @staticmethod def forward(ctx, timestep: torch.Tensor, num_channels: int, check_range: bool): + """Compute FP32 CUDA features and save them for the analytic timestep VJP.""" + t32 = timestep.detach().float().contiguous() out = _C.h3_timestep_sinusoid_forward( t32, int(num_channels), float(H3_MAX_PERIOD), check_range @@ -38,6 +42,8 @@ def forward(ctx, timestep: torch.Tensor, num_channels: int, check_range: bool): @staticmethod def backward(ctx, grad_out: torch.Tensor): + """Reduce channel derivatives with a fixed FP32 tree into the timestep dtype.""" + (out,) = ctx.saved_tensors half = ctx.num_channels // 2 freq = NativeH3TimestepSinusoidOp.frequencies_fp32(ctx.num_channels, out.device) @@ -70,6 +76,8 @@ class H3TimestepSinusoidCudaOp: backward_impl = BACKWARD_IMPL def __init__(self) -> None: + """Require the native H3 CUDA sinusoid symbol before creating this operator.""" + if not h3_sinusoid_cuda_available(): raise RuntimeError( "rl_engine._C.h3_timestep_sinusoid_forward is unavailable; rebuild the " @@ -77,6 +85,8 @@ def __init__(self) -> None: ) def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 CUDA features for timesteps in ``[0, 1]`` with gradients.""" + return self.forward(timestep, num_channels=num_channels) def forward( @@ -86,6 +96,13 @@ def forward( num_channels: int = H3_FREQ_DIM, check_range: bool = True, ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on CUDA. + + Require nonempty 1-D FP32, FP16 or BF16 timesteps and a positive even + channel count. Native validation rejects nonfinite/out-of-range values + when ``check_range`` is true; backward returns the original input dtype. + """ + # Native code owns value validation, avoiding a second host sync here. validate_h3_timesteps(timestep, num_channels, check_range=False) if not timestep.is_cuda: @@ -93,5 +110,7 @@ def forward( return _H3TimestepSinusoidCuda.apply(timestep, num_channels, check_range) def forward_fp32(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Use the standard CUDA sinusoid path, which already returns FP32 features.""" + # The op is FP32 end to end; the FP32 path is the op itself. return self.forward(timestep, num_channels=num_channels) diff --git a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py index 75ec2b8bd..b6d013961 100644 --- a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py +++ b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py @@ -26,6 +26,8 @@ def adaln_hidden_size(weight: torch.Tensor) -> int: + """Infer ``H`` from a 2-D projection weight whose row count is divisible by 18.""" + rows_per_hidden = H3_ADALN_CHUNKS * H3_MODALITY_NUM if weight.dim() != 2 or weight.shape[0] % rows_per_hidden != 0: raise ValueError( @@ -36,6 +38,12 @@ def adaln_hidden_size(weight: torch.Tensor) -> int: def validate_h3_adaln_projection(temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): + """Check the AdaLN input shapes, shared device and precision boundary; return ``H``. + + Require nonempty FP32 ``temb`` of shape ``(T, D)`` and FP32 or BF16 weight + of shape ``(18H, D)`` with matching-dtype bias of shape ``(18H,)``. + """ + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): if not isinstance(tensor, torch.Tensor): raise TypeError(f"{name} must be a torch.Tensor") @@ -79,14 +87,28 @@ class NativeH3AdaLNProjectionOp: op_class = "reduction" def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` views using the PyTorch provider path.""" + return self.forward(temb, weight, bias) def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Apply FP32 SiLU, cast to the weight dtype and project on the input device. + + Validate ``(T, D)`` embeddings and ``(18H, D)``/``(18H,)`` parameters; + return six weight-dtype ``(3T, H)`` views with PyTorch autograd support. + """ + hidden = validate_h3_adaln_projection(temb, weight, bias) table = F.linear(F.silu(temb).to(weight.dtype), weight, bias) return split_adaln_table(table, hidden) def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six FP32 golden views from FP64 SiLU and projection arithmetic. + + Preserve the declared activation rounding to the weight dtype while + using an identity VJP at that cast boundary for the FP64 golden graph. + """ + hidden = validate_h3_adaln_projection(temb, weight, bias) t64 = temb.double() act = t64 * torch.sigmoid(t64) diff --git a/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py b/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py index ac54d295f..51ef48db7 100644 --- a/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py +++ b/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py @@ -27,6 +27,12 @@ def validate_h3_timestep_mlp( w2: torch.Tensor, b2: torch.Tensor, ) -> None: + """Check a same-device FP32 ``(T, K) -> (T, H) -> (T, D)`` MLP contract. + + Require ``T > 0``, weights of shapes ``(H, K)``/``(D, H)`` and biases + of shapes ``(H,)``/``(D,)``; reject lower-precision parameters or inputs. + """ + named = {"x": x, "w1": w1, "b1": b1, "w2": w2, "b2": b2} for name, tensor in named.items(): if not isinstance(tensor, torch.Tensor): @@ -64,13 +70,23 @@ class NativeH3TimestepMLPOp: op_class = "reduction" def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` embeddings using the PyTorch provider graph.""" + return self.forward(x, w1, b1, w2, b2) def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Validate and evaluate linear-SiLU-linear in FP32 on the input device. + + Inputs follow the ``(T, K) -> (T, H) -> (T, D)`` contract and retain + PyTorch autograd support; CUDA callers must disable TF32 for this path. + """ + validate_h3_timestep_mlp(x, w1, b1, w2, b2) return F.linear(F.silu(F.linear(x, w1, b1)), w2, b2) def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Evaluate the validated MLP in FP64 and round its ``(T, D)`` output to FP32.""" + validate_h3_timestep_mlp(x, w1, b1, w2, b2) z = F.linear(x.double(), w1.double(), b1.double()) h = z * torch.sigmoid(z) diff --git a/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py b/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py index 8ca0c0ccb..54da13527 100644 --- a/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py +++ b/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py @@ -68,6 +68,8 @@ class NativeH3TimestepSinusoidOp: op_class = "elementwise" def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 sinusoidal features for nonempty timesteps in ``[0, 1]``.""" + return self.forward(timestep, num_channels=num_channels) def forward( @@ -77,6 +79,13 @@ def forward( num_channels: int = H3_FREQ_DIM, check_range: bool = True, ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on the input device. + + Require nonempty 1-D FP32, BF16 or FP16 timesteps and positive even + channels. Validate finite values in ``[0, 1]`` when ``check_range`` is + true and preserve PyTorch autograd through the provider formula. + """ + validate_h3_timesteps(timestep, num_channels, check_range=check_range) half = num_channels // 2 # Same op sequence and dtypes as diffusers get_timestep_embedding with @@ -96,6 +105,12 @@ def forward_fp32( num_channels: int = H3_FREQ_DIM, check_range: bool = True, ) -> torch.Tensor: + """Evaluate the validated sinusoid formula in FP64 and round once to FP32. + + Return ``(T, num_channels)`` on the timestep device with cosine channels + followed by sine channels and gradients through the FP64 golden graph. + """ + validate_h3_timesteps(timestep, num_channels, check_range=check_range) half = num_channels // 2 k = torch.arange(half, dtype=torch.float64, device=timestep.device) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 22bbab259..880363b6c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -456,6 +456,7 @@ class KernelRegistry: """ def __init__(self): + """Initialize backend caches, capability contracts, and device-specific priorities.""" self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() self.semantic = SemanticOperatorCatalog(_default_semantic_descriptors()) diff --git a/rl_engine/testing/h3_cases.py b/rl_engine/testing/h3_cases.py index 527a820ec..8eb509d33 100644 --- a/rl_engine/testing/h3_cases.py +++ b/rl_engine/testing/h3_cases.py @@ -9,7 +9,11 @@ def h3_timesteps(num: int, *, seed: int = 0, device: str = "cuda") -> torch.Tensor: - """Distinct-looking timesteps in [0, 1] with both endpoints present.""" + """Return seeded FP32 timesteps of shape ``(num,)`` on ``device``. + + Requires positive ``num``; the first value is 0 and the last is 1 when + ``num > 1``, with all remaining values sampled uniformly in [0, 1]. + """ generator = torch.Generator(device="cpu").manual_seed(seed) t = torch.rand(num, generator=generator) @@ -22,10 +26,11 @@ def h3_timesteps(num: int, *, seed: int = 0, device: str = "cuda") -> torch.Tens def h3_packed_layout( seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" ) -> tuple[torch.Tensor, torch.Tensor]: - """(timestep_indices, token_tags) for a packed sequence with every modality. + """Return seeded int64 timestep indices and modality tags, each shape ``(seq_len,)``. - Rows are video, text and audio blocks in random order, each block mapped - to a random timestep, like H3's conditioning/target packing. + Each row independently samples a timestep in ``[0, num_timesteps)`` and + a video/text/audio tag in ``[0, 3)``. The first three rows cover all tags + when available; both tensors are returned on ``device``. """ generator = torch.Generator(device="cpu").manual_seed(seed) diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 8671ef141..11ca1c0fb 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -54,11 +54,15 @@ class Stage: def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: + """Return the context's MLP weights and biases in linear_1, linear_2 call order.""" + weights = ctx["weights"] return [weights[f"time_embedder.linear_{i}.{p}"] for i in (1, 2) for p in ("weight", "bias")] def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + """Return the context's block-0 AdaLN projection weight and bias.""" + prefix = "transformer_blocks.0.adaln_proj.linear" return ctx["weights"][f"{prefix}.weight"], ctx["weights"][f"{prefix}.bias"] @@ -98,12 +102,16 @@ def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: def _flat(value: Any) -> list[torch.Tensor]: + """Collect tensors from nested sequences in their original order.""" + if isinstance(value, torch.Tensor): return [value] return [tensor for item in value for tensor in _flat(item)] def compare(lhs: Any, rhs: Any, atol: float = 0.0, rtol: float = 0.0) -> dict[str, Any]: + """Report tensor equality, maximum FP32 absolute error, and tolerance checks.""" + lhs_t, rhs_t = _flat(lhs), _flat(rhs) bitwise = all(torch.equal(a, b) for a, b in zip(lhs_t, rhs_t, strict=True)) max_abs = 0.0 @@ -117,6 +125,8 @@ def compare(lhs: Any, rhs: Any, atol: float = 0.0, rtol: float = 0.0) -> dict[st def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dict[str, Any]: + """Build seeded CUDA timesteps and packed row labels with the supplied weights.""" + timestep_indices, token_tags = h3_packed_layout(seq_len, num_timesteps, seed=seed) return { "timestep": h3_timesteps(num_timesteps, seed=seed), @@ -127,7 +137,7 @@ def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dic def golden_op(registry: KernelRegistry, op_type: str): - """The PyTorch reference: the last entry of the CPU priority list.""" + """Return the PyTorch reference at the end of the operator's CPU priority list.""" return registry._get_or_create_backend(registry._priority_map["cpu"][op_type][-1]) @@ -141,6 +151,12 @@ def run_case( seed: int = 0, stages: list[Stage] | None = None, ) -> dict[str, Any]: + """Replay CUDA candidate, provider, and golden paths and report each stage's drift. + + Supplied stages must include their prerequisites in execution order. + Omitting ``stages`` runs the full registered chain without gradients. + """ + stages = STAGES if stages is None else stages ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) report: dict[str, Any] = { @@ -182,7 +198,11 @@ def run_case( def git_state() -> dict[str, Any]: + """Return HEAD and tracked-tree dirtiness, defaulting to unknown/dirty if Git fails.""" + def run(*args: str) -> str: + """Run Git in the repository root and return stripped stdout.""" + return subprocess.check_output(["git", *args], cwd=REPO_ROOT, text=True).strip() try: @@ -194,6 +214,8 @@ def run(*args: str) -> str: def environment() -> dict[str, Any]: + """Record the active CUDA device, software versions, and matmul TF32 setting.""" + return { "gpu": torch.cuda.get_device_name(), "capability": list(torch.cuda.get_device_capability()), diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index f64f7174f..eaf31ba92 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -19,7 +19,11 @@ def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch.Tensor: - """``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0)``.""" + """Replay provider FP32 cosine/sine features for ``(T,)`` timesteps. + + Matches ``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0)`` + with the default channel count, returning shape ``(T, 256)``. + """ half_dim = num_channels // 2 exponent = -math.log(10000) * torch.arange( @@ -41,7 +45,10 @@ def provider_time_embedder( w2: torch.Tensor, b2: torch.Tensor, ) -> torch.Tensor: - """``TimestepEmbedding(256, 5376, out_dim=2688)``: linear_1 -> SiLU -> linear_2, FP32.""" + """Replay linear_1, SiLU, and linear_2 using the checkpoint weight dtype. + + H3's FP32 weights map features of shape ``(T, 256)`` to ``(T, 2688)``. + """ sample = F.linear(features.to(w1.dtype), w1, b1) sample = F.silu(sample) @@ -51,7 +58,11 @@ def provider_time_embedder( def provider_adaln_modulation( temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, hidden_size: int = 5376 ) -> tuple[torch.Tensor, ...]: - """``MiniMaxH3AdaLayerNormModulation.forward``: six (3T, H) views.""" + """Replay provider SiLU, the weight-dtype cast, and the AdaLN projection. + + Return six ``(3T, hidden_size)`` views in timestep-major modality order + for checkpoint weights shaped ``(18 * hidden_size, 2688)``. + """ temb = F.linear(F.silu(temb).to(weight.dtype), weight, bias) temb = temb.view(-1, 6 * hidden_size) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index d4b45e09b..ad9eb87fc 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -53,6 +53,8 @@ def time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 200) -> float: def peak_mib(fn: Callable[[], Any]) -> float: + """Synchronize one CUDA call and report peak allocated memory above baseline in MiB.""" + torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() base = torch.cuda.memory_allocated() @@ -62,6 +64,12 @@ def peak_mib(fn: Callable[[], Any]) -> float: def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + """Time case callables in rotating order after warmup and return their evidence. + + Record raw CUDA-event samples, executed orders, median microseconds, + decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. + """ + row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} keys = [key for key in TIMED_KEYS if key in case] if not keys: @@ -101,6 +109,8 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st def _sinusoid_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA sinusoid timing cases with checked, unchecked, and provider calls.""" + op = registry.get_op("timestep_sinusoid_h3", device="cuda") cases = [] for num in (1, 2, 4, 64): @@ -120,6 +130,8 @@ def _sinusoid_perf(registry: KernelRegistry) -> list[dict[str, Any]]: def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: + """Compare sinusoid outputs to provider and FP64 golden across timestep counts.""" + op = registry.get_op("timestep_sinusoid_h3", device="cuda") golden = registry._get_or_create_backend( registry._priority_map["cpu"]["timestep_sinusoid_h3"][-1] @@ -140,7 +152,7 @@ def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Tensor]: - """Pinned checkpoint tensors when available, otherwise same-shape random ones.""" + """Return CUDA checkpoint parameters, or seeded random tensors when weights are unset.""" if h3_weights_dir() is not None: weights = load_h3_conditioning_weights("cuda", names) @@ -161,6 +173,8 @@ def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Ten def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA MLP candidate/provider timing cases for one to four timesteps.""" + op = registry.get_op("timestep_mlp_fp32", device="cuda") params = h3_params(MLP_NAMES, MLP_SHAPES) weight_bytes = sum(p.numel() * p.element_size() for p in params) @@ -181,6 +195,8 @@ def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: + """Measure per-draw MLP error against FP64 and test single-row/batched equality.""" + op = registry.get_op("timestep_mlp_fp32", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["timestep_mlp_fp32"][-1]) sinusoid = registry.get_op("timestep_sinusoid_h3", device="cuda") @@ -216,11 +232,15 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: def _adaln_params() -> tuple[torch.Tensor, torch.Tensor]: + """Return block-0 AdaLN projection weights and bias as BF16 CUDA tensors.""" + weight, bias = h3_params(ADALN_NAMES, [(96768, 2688), (96768,)]) return weight.bfloat16(), bias.bfloat16() def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA AdaLN candidate/provider timing cases for one to four timesteps.""" + op = registry.get_op("adaln_projection_3mod", device="cuda") weight, bias = _adaln_params() cases = [] @@ -240,11 +260,13 @@ def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: def _flat_cat(outputs) -> torch.Tensor: + """Concatenate flattened modulation tensors in their output order.""" + return torch.cat([out.reshape(-1) for out in outputs]) def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, Any]: - """Fraction of BF16 outputs equal to the correctly rounded FP64 golden, per draw.""" + """Report per-draw BF16 equality to FP64 goldens and single-row/batch invariance.""" op = registry.get_op("adaln_projection_3mod", device="cuda") golden = registry._get_or_create_backend( diff --git a/rl_engine/testing/h3_weights.py b/rl_engine/testing/h3_weights.py index 7d7abdf03..6b2dc995c 100644 --- a/rl_engine/testing/h3_weights.py +++ b/rl_engine/testing/h3_weights.py @@ -25,15 +25,21 @@ def load_h3_manifest() -> dict[str, Any]: + """Read the bundled checkpoint identity and pinned tensor contracts.""" + return json.loads(MANIFEST_PATH.read_text()) def h3_weights_dir() -> Path | None: + """Expand the configured extraction directory, or return None when unset or blank.""" + value = os.environ.get(WEIGHTS_ENV, "").strip() return Path(value).expanduser() if value else None def sha256_file(path: Path, chunk: int = 1 << 24) -> str: + """Return a file SHA256 while reading bounded chunks.""" + digest = hashlib.sha256() with path.open("rb") as handle: while block := handle.read(chunk): diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index c20e157e4..7f810d815 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -231,6 +231,8 @@ def fake_load_backend(backend): @pytest.mark.parametrize("op_name", H3_OPS) def test_h3_ops_prefer_cuda_and_fall_back_to_pytorch_reference(op_name): + """Prefer H3 CUDA backends on CUDA and use their PyTorch fallbacks elsewhere.""" + registry = KernelRegistry() cuda_chain = registry._priority_map["cuda"][op_name] assert cuda_chain[0].name.startswith("CUDA_H3_") diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py index 512a15f06..395d3e17c 100644 --- a/scripts/h3_chain_replay.py +++ b/scripts/h3_chain_replay.py @@ -30,10 +30,14 @@ def _ints(text: str) -> list[int]: + """Parse comma-separated integers for the replay CLI's case dimensions.""" + return [int(v) for v in text.split(",")] def main() -> None: + """Replay requested CUDA stages with prerequisites and optionally save JSON evidence.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--stages", diff --git a/scripts/h3_evidence.py b/scripts/h3_evidence.py index 8351da245..426b0848d 100644 --- a/scripts/h3_evidence.py +++ b/scripts/h3_evidence.py @@ -34,6 +34,8 @@ def main() -> None: + """Collect CUDA accuracy and performance evidence, requiring weights for weighted ops.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) parser.add_argument("--out", type=Path, required=True) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 64edd56fe..7ebd228d0 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -35,6 +35,8 @@ def _style(ax, title: str, ylabel: str) -> None: + """Apply the shared report axis styling, title, and vertical-axis label.""" + ax.set_facecolor(SURFACE) ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) ax.set_ylabel(ylabel, color=MUTED, fontsize=9) @@ -47,7 +49,7 @@ def _style(ax, title: str, ylabel: str) -> None: def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt: str) -> None: - """Grouped bars with a 2px-style gap and value labels on top.""" + """Draw grouped series on the axis with formatted value labels and a legend.""" n = len(series) width = 0.8 / n @@ -75,10 +77,14 @@ def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt def _log_values(values: list[float]) -> list[float]: + """Replace nonpositive values with the plotting floor for logarithmic axes.""" + return [v if v > 0 else LOG_FLOOR for v in values] def _figure(report: dict[str, Any], suptitle: str): + """Create two report axes with a title and GPU, software, and revision provenance.""" + fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), facecolor=SURFACE) env = report["environment"] fig.suptitle(suptitle, x=0.01, ha="left", fontsize=12, color=INK, fontweight="bold") @@ -94,6 +100,8 @@ def _figure(report: dict[str, Any], suptitle: str): def plot_sinusoid(report: dict[str, Any]): + """Return a sinusoid report figure showing latency and errors against FP64.""" + fig, (left, right) = _figure(report, "timestep_sinusoid_h3") perf = report["perf"] _style(left, "Latency per call (lower is better)", "µs") @@ -155,6 +163,8 @@ def plot_sinusoid(report: dict[str, Any]): def _latency_panel(ax, perf: list[dict[str, Any]], provider_name: str) -> None: + """Plot candidate and provider latency summaries in microseconds for each case.""" + _style(ax, "Latency per call (lower is better)", "µs") _bars( ax, @@ -209,6 +219,8 @@ def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], con def plot_mlp(report: dict[str, Any]): + """Return an MLP report figure showing latency and per-draw FP64 error distributions.""" + fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") _latency_panel(left, report["perf"], "diffusers (cuBLAS)") acc = report["accuracy"] @@ -225,6 +237,8 @@ def plot_mlp(report: dict[str, Any]): def plot_projection(report: dict[str, Any]): + """Return an AdaLN figure with latency, BF16 rounding accuracy, and the early-cast probe.""" + fig, (left, right) = _figure(report, "adaln_projection_3mod (pinned block-0 AdaLN weights)") _latency_panel(left, report["perf"], "diffusers (cuBLAS)") acc = report["accuracy"] @@ -269,6 +283,8 @@ def plot_projection(report: dict[str, Any]): def main() -> None: + """Read an operator report and save its figure as a PNG beside it or at ``--out``.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("report", type=Path) parser.add_argument("--out", type=Path, default=None) diff --git a/scripts/prepare_h3_weights.py b/scripts/prepare_h3_weights.py index 262408736..a9af15c37 100644 --- a/scripts/prepare_h3_weights.py +++ b/scripts/prepare_h3_weights.py @@ -32,6 +32,8 @@ def _check(path: Path, expected: str) -> None: + """Verify a downloaded file's SHA-256, exiting on mismatch or printing success.""" + actual = sha256_file(path) if actual != expected: raise SystemExit(f"sha256 mismatch for {path.name}: {actual} != {expected}") @@ -39,6 +41,8 @@ def _check(path: Path, expected: str) -> None: def main() -> None: + """Download pinned model artifacts, verify tensor identities, and write extracted weights.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--out", type=Path, required=True, help="output directory") parser.add_argument( @@ -60,6 +64,8 @@ def main() -> None: out_dir.mkdir(parents=True, exist_ok=True) def fetch(filename: str) -> Path: + """Download one file from the manifest's model subfolder and pinned revision.""" + return Path( hf_hub_download( identity["hf_repo"], diff --git a/setup.py b/setup.py index b388bdd29..2211604c5 100644 --- a/setup.py +++ b/setup.py @@ -96,6 +96,7 @@ def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: def get_extensions(): + """Configure the optional PyTorch native extension for the active CUDA or ROCm toolchain.""" torch, _, CUDAExtension = _load_torch_extension_tools() if torch is None: message = ( diff --git a/tests/h3/conftest.py b/tests/h3/conftest.py index 44c8fe117..1c3ba8caa 100644 --- a/tests/h3/conftest.py +++ b/tests/h3/conftest.py @@ -17,6 +17,8 @@ @pytest.fixture(scope="session") def h3_weights_cpu() -> dict[str, torch.Tensor]: + """Load pinned CPU weights once per session, skipping unavailable artifacts.""" + if h3_weights_dir() is None: pytest.skip(f"{WEIGHTS_ENV} not set; run scripts/prepare_h3_weights.py") try: diff --git a/tests/h3/test_h3_adaln_projection.py b/tests/h3/test_h3_adaln_projection.py index c9aac92d1..0b031fa4e 100644 --- a/tests/h3/test_h3_adaln_projection.py +++ b/tests/h3/test_h3_adaln_projection.py @@ -29,6 +29,8 @@ def _cuda_op(): + """Construct the CUDA projection operator, skipping missing native linear kernels.""" + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available @@ -38,11 +40,15 @@ def _cuda_op(): def _temb(num, dim=2688, seed=0, device="cuda"): + """Generate seeded FP32 timestep embeddings on CPU and move them to the test device.""" + g = torch.Generator(device="cpu").manual_seed(seed) return (torch.randn(num, dim, generator=g) * 2).to(device) def _synthetic(hidden, dim, dtype=torch.bfloat16, seed=0, device="cuda"): + """Create seeded projection weights for six chunks and three modalities per channel.""" + g = torch.Generator(device="cpu").manual_seed(seed) n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).to(dtype) @@ -51,11 +57,15 @@ def _synthetic(hidden, dim, dtype=torch.bfloat16, seed=0, device="cuda"): def _flat(outputs): + """Flatten and concatenate all modulation chunks for elementwise comparisons.""" + return torch.cat([out.reshape(-1) for out in outputs]) class TestReference: def test_provider_layout_matches_diffusers(self): + """Match the provider values and timestep-major modality layout on CPU.""" + temb = _temb(2, dim=16, device="cpu") weight, bias = _synthetic(8, 16, device="cpu") ours = NativeH3AdaLNProjectionOp().forward(temb, weight, bias) @@ -66,11 +76,15 @@ def test_provider_layout_matches_diffusers(self): def test_rejects_bf16_temb(self): # RFC probe H7: casting temb before the SiLU. + """Reject an early BF16 cast so SiLU receives the declared FP32 embeddings.""" + weight, bias = _synthetic(8, 16, device="cpu") with pytest.raises(TypeError, match="float32"): NativeH3AdaLNProjectionOp().forward(_temb(2, 16, device="cpu").bfloat16(), weight, bias) def test_rejects_bad_shapes(self): + """Reject invalid projection dimensions, empty batches, and mismatched bias dtype.""" + weight, bias = _synthetic(8, 16, device="cpu") op = NativeH3AdaLNProjectionOp() with pytest.raises(ValueError): # not a multiple of 6 * 3 rows @@ -104,6 +118,8 @@ def test_channel_identity_lands_in_its_chunk_and_modality_row(self): assert torch.equal(out[t * 3 + m].float(), expected), (c, t, m) def test_outputs_are_views_of_one_table_like_diffusers(self): + """Require all six modulation chunks to share storage with the expected offsets.""" + weight, bias = _synthetic(16, 32) outputs = _cuda_op()(_temb(2, 32), weight, bias) base = outputs[0].untyped_storage().data_ptr() @@ -113,6 +129,8 @@ def test_outputs_are_views_of_one_table_like_diffusers(self): @pytest.fixture(scope="module") def block0(h3_weights_cpu): + """Move the pinned block-zero projection weight and bias to CUDA for this module.""" + return h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() @@ -120,6 +138,8 @@ def block0(h3_weights_cpu): class TestCudaRealWeights: @pytest.mark.parametrize("num", [1, 2, 3, 4]) def test_forward_is_correctly_rounded_declared_golden(self, block0, num): + """Check checkpoint accuracy and near-exact BF16 rounding of the declared golden.""" + weight, bias = block0 temb = _temb(num, seed=num) outputs = _cuda_op()(temb, weight, bias) @@ -146,6 +166,8 @@ def test_early_cast_is_numerically_detectable(self, block0): assert (ours == early).float().mean() < 0.8 def test_rows_are_batch_and_position_invariant(self, block0): + """Keep each timestep modulation row identical alone and under batch reordering.""" + weight, bias = block0 op = _cuda_op() temb = _temb(5, seed=3) @@ -159,6 +181,8 @@ def test_rows_are_batch_and_position_invariant(self, block0): assert torch.equal(s.view(5, 3, -1).flip(0), f.view(5, 3, -1)) def test_backward_against_golden(self, block0): + """Check embedding and parameter gradients against a high-precision declared-cast VJP.""" + weight, bias = block0 temb = _temb(2, seed=5) g = torch.Generator(device="cuda").manual_seed(0) @@ -184,6 +208,8 @@ class TestCudaSynthetic: @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("hidden, dim", [(24, 16), (40, 72), (8, 2688)]) def test_other_sizes(self, dtype, hidden, dim): + """Match the projection golden at non-checkpoint dimensions in FP32 and BF16.""" + weight, bias = _synthetic(hidden, dim, dtype=dtype, seed=hidden) temb = _temb(3, dim, seed=hidden) tol = dict(atol=5e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-4, rtol=1e-4) @@ -194,6 +220,8 @@ def test_other_sizes(self, dtype, hidden, dim): torch.testing.assert_close(out.float(), gold, **tol) def test_repeat_forward_backward_bitwise(self): + """Require repeated projection outputs and all input gradients to match bitwise.""" + weight, bias = _synthetic(64, 2688, seed=9) temb = _temb(3, seed=9) grads = [torch.randn(9, 64, device="cuda").bfloat16() for _ in range(6)] @@ -208,11 +236,15 @@ def test_repeat_forward_backward_bitwise(self): assert torch.equal(a, b) def test_temb_grad_rows_are_batch_invariant(self): + """Keep embedding gradients bitwise equal in full batches and single-row calls.""" + weight, bias = _synthetic(64, 2688, seed=11) temb = _temb(4, seed=11) grads = [torch.randn(12, 64, device="cuda").bfloat16() for _ in range(6)] def d_temb(rows): + """Compute selected embedding-row gradients with their matching modality gradients.""" + leaf = temb[rows].detach().clone().requires_grad_(True) outs = _cuda_op()(leaf, weight, bias) torch.autograd.backward( @@ -237,11 +269,15 @@ def test_rows_invariant_across_tensor_core_row_tiles(self, num): assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) def test_rejects_cpu(self): + """Reject CPU tensors at the CUDA projection boundary.""" + weight, bias = _synthetic(8, 16, device="cpu") with pytest.raises(ValueError): _cuda_op()(_temb(2, 16, device="cpu"), weight, bias) def test_registry_dispatches_cuda(self): + """Resolve the projection registry entry to its dedicated CUDA implementation.""" + from rl_engine.kernels.registry import KernelRegistry _cuda_op() diff --git a/tests/h3/test_h3_cli.py b/tests/h3/test_h3_cli.py index 7712fc41a..7288d795a 100644 --- a/tests/h3/test_h3_cli.py +++ b/tests/h3/test_h3_cli.py @@ -20,11 +20,15 @@ @pytest.fixture def cpu_chain_cli(monkeypatch): + """Replace CUDA chain stages with CPU stubs that track dependencies and provider drift.""" + executed = [] state = {"provider_shift": 0} class CpuOp: def __call__(self, value): + """Advance the synthetic chain value by one for candidate and golden replay.""" + return value + 1 forward_fp32 = __call__ @@ -33,6 +37,8 @@ def __call__(self, value): for index, stage in enumerate(h3_chain.STAGES): def stage_input(ctx, upstream, index=index): + """Use raw timesteps for stage zero and require predecessor outputs thereafter.""" + if index == 0: return ctx["timestep"] assert isinstance(upstream, torch.Tensor), "a preceding stage must run first" @@ -50,6 +56,8 @@ def stage_input(ctx, upstream, index=index): ) def get_op(op_type, *, device): + """Record each dispatched stage and return its deterministic CPU stub.""" + executed.append(op_type) return CpuOp() @@ -79,6 +87,8 @@ def get_op(op_type, *, device): def test_chain_selection_runs_dependencies( cpu_chain_cli, monkeypatch, tmp_path, capsys, requested, execution_count, provider_shift ): + """Execute prerequisite stages but report requested stages and preserve first drift.""" + executed, state = cpu_chain_cli state["provider_shift"] = provider_shift out = tmp_path / "chain.json" @@ -118,6 +128,8 @@ def test_chain_selection_runs_dependencies( @pytest.mark.parametrize("op", ["timestep_mlp_fp32", "adaln_projection_3mod"]) @pytest.mark.parametrize("weights_env", [None, " "]) def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, weights_env): + """Reject weighted evidence without a configured checkpoint before writing a report.""" + if weights_env is None: monkeypatch.delenv(WEIGHTS_ENV, raising=False) else: @@ -138,6 +150,8 @@ def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, we ], ) def test_evidence_records_weight_source(monkeypatch, tmp_path, op, weight_source): + """Label weighted evidence as pinned and sinusoid evidence as weight-independent.""" + if weight_source == "not_applicable": monkeypatch.delenv(WEIGHTS_ENV, raising=False) else: diff --git a/tests/h3/test_h3_conditioning_e2e.py b/tests/h3/test_h3_conditioning_e2e.py index 2bd1c47bd..a24ceebb9 100644 --- a/tests/h3/test_h3_conditioning_e2e.py +++ b/tests/h3/test_h3_conditioning_e2e.py @@ -33,16 +33,22 @@ @pytest.fixture(scope="module") def chain_weights(h3_weights_cpu): + """Move all pinned conditioning tensors to CUDA once for the end-to-end module.""" + return {name: tensor.cuda() for name, tensor in h3_weights_cpu.items()} @pytest.fixture(scope="module") def registry(): + """Provide the production registry used to dispatch every conditioning stage.""" + return KernelRegistry() @pytest.mark.parametrize("num_timesteps, seq_len", CASES) def test_chain_forward(registry, chain_weights, num_timesteps, seq_len): + """Check CUDA dispatch, golden tolerance, repeatability, and declared provider parity.""" + report = h3_chain.run_case( registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len ) diff --git a/tests/h3/test_h3_det_linear.py b/tests/h3/test_h3_det_linear.py index f8f90e83c..63ae03a8e 100644 --- a/tests/h3/test_h3_det_linear.py +++ b/tests/h3/test_h3_det_linear.py @@ -19,6 +19,8 @@ @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_backward_input_rejects_zero_rows(dtype): + """Reject empty gradient batches in native input-gradient kernels for both dtypes.""" + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) weight = torch.empty(16, 8, device="cuda", dtype=dtype) with pytest.raises(RuntimeError, match="grad must have at least one row"): @@ -28,6 +30,8 @@ def test_backward_input_rejects_zero_rows(dtype): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("with_bias", [False, True]) def test_backward_weight_rejects_zero_rows(dtype, with_bias): + """Reject empty gradient batches before parameter-gradient or bias-gradient reduction.""" + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) x = torch.empty(0, 8, device="cuda", dtype=dtype) with pytest.raises(RuntimeError, match="grad must have at least one row"): diff --git a/tests/h3/test_h3_report.py b/tests/h3/test_h3_report.py index eab1c8807..be62daa38 100644 --- a/tests/h3/test_h3_report.py +++ b/tests/h3/test_h3_report.py @@ -11,28 +11,42 @@ def test_measure_interleaves_backends_and_records_samples(monkeypatch): + """Balance candidate/provider timing order and preserve samples and derived bandwidth.""" + elapsed_us = 0.0 calls = [] class Event: def __init__(self, *, enable_timing): + """Require timing-enabled events in the deterministic timing stub.""" + assert enable_timing def record(self): + """Capture the synthetic elapsed clock at this event.""" + self.timestamp = elapsed_us def synchronize(self): + """Leave synchronization inert because the timing stub has no pending GPU work.""" + pass def elapsed_time(self, end): + """Return elapsed synthetic event time in CUDA milliseconds.""" + return (end.timestamp - self.timestamp) / 1e3 def candidate(): + """Record a candidate call and advance the synthetic clock by two microseconds.""" + nonlocal elapsed_us calls.append("candidate") elapsed_us += 2.0 def provider(): + """Record a provider call and advance the synthetic clock by four microseconds.""" + nonlocal elapsed_us calls.append("provider") elapsed_us += 4.0 diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py index 71f1e694f..f378a3b35 100644 --- a/tests/h3/test_h3_timestep_mlp.py +++ b/tests/h3/test_h3_timestep_mlp.py @@ -26,6 +26,8 @@ def _cuda_op(): + """Construct the CUDA MLP operator, skipping missing native linear kernels.""" + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available from rl_engine.kernels.ops.cuda.h3.timestep_mlp import H3TimestepMLPCudaOp @@ -35,10 +37,14 @@ def _cuda_op(): def _features(num, device="cuda", seed=0): + """Generate seeded H3 sinusoidal features on the requested test device.""" + return NativeH3TimestepSinusoidOp().forward(h3_timesteps(num, seed=seed, device=device)) def _synthetic_params(k_in, hidden, out, device="cuda", seed=0): + """Create seeded FP32 parameters for a two-layer SiLU MLP of the requested sizes.""" + g = torch.Generator(device="cpu").manual_seed(seed) w1 = torch.randn(hidden, k_in, generator=g) / k_in**0.5 b1 = torch.randn(hidden, generator=g) * 0.1 @@ -48,12 +54,16 @@ def _synthetic_params(k_in, hidden, out, device="cuda", seed=0): def _real_params(weights): + """Return pinned timestep MLP weights and biases on CUDA in layer order.""" + return [ weights[f"time_embedder.linear_{i}.{p}"].cuda() for i in (1, 2) for p in ("weight", "bias") ] def _fp64_grads(x, params, grad_out): + """Compute gradients of both linear layers and SiLU using FP64 autograd.""" + leaves = [t.detach().double().requires_grad_(True) for t in (x, *params)] z = F.linear(leaves[0], leaves[1], leaves[2]) F.linear(z * torch.sigmoid(z), leaves[3], leaves[4]).backward(grad_out.double()) @@ -62,6 +72,8 @@ def _fp64_grads(x, params, grad_out): class TestReference: def test_provider_and_golden_agree(self): + """Keep the FP32 provider close to the high-precision MLP replay on CPU.""" + x = _features(5, device="cpu") params = _synthetic_params(256, 64, 32, device="cpu") op = NativeH3TimestepMLPOp() @@ -72,12 +84,16 @@ def test_provider_and_golden_agree(self): @pytest.mark.parametrize("which", range(5)) def test_rejects_bf16_anywhere(self, which): # RFC probe H7: casting a declared FP32 path to BF16 early. + """Reject BF16 in every argument of the declared FP32 MLP path.""" + args = [_features(2, device="cpu"), *_synthetic_params(256, 64, 32, device="cpu")] args[which] = args[which].bfloat16() with pytest.raises(TypeError, match="float32"): NativeH3TimestepMLPOp().forward(*args) def test_rejects_shape_mismatches(self): + """Reject incompatible layer dimensions, malformed biases, and empty feature batches.""" + x = _features(2, device="cpu") w1, b1, w2, b2 = _synthetic_params(256, 64, 32, device="cpu") op = NativeH3TimestepMLPOp() @@ -95,6 +111,8 @@ def test_rejects_shape_mismatches(self): class TestCudaRealWeights: @pytest.mark.parametrize("num", [1, 2, 3, 4, 7]) def test_forward_against_fp64_golden(self, h3_weights_cpu, num): + """Check checkpoint forward accuracy, FP32 dtype, and the timestep output shape.""" + params = _real_params(h3_weights_cpu) x = _features(num, seed=num) out = _cuda_op()(x, *params) @@ -105,6 +123,8 @@ def test_forward_against_fp64_golden(self, h3_weights_cpu, num): assert (out - gold).abs().max().item() < 1e-5 def test_rows_are_batch_and_position_invariant(self, h3_weights_cpu): + """Keep each checkpoint embedding row bitwise unchanged across batch sizes and positions.""" + op = _cuda_op() params = _real_params(h3_weights_cpu) x_all = _features(9, seed=1) @@ -118,6 +138,8 @@ def test_rows_are_batch_and_position_invariant(self, h3_weights_cpu): assert torch.equal(op(x, *params)[pos], full[4]), (num, pos) def test_backward_against_fp64_golden(self, h3_weights_cpu): + """Match all five checkpoint input gradients to the FP64 autograd golden.""" + params = _real_params(h3_weights_cpu) x = _features(3, seed=2) leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] @@ -134,6 +156,8 @@ def test_backward_against_fp64_golden(self, h3_weights_cpu): class TestCudaSynthetic: @pytest.mark.parametrize("k_in, hidden, out", [(8, 24, 16), (12, 40, 20), (256, 5376, 2688)]) def test_other_sizes(self, k_in, hidden, out): + """Match the high-precision forward golden for small and checkpoint-sized MLPs.""" + params = _synthetic_params(k_in, hidden, out, seed=k_in) x = torch.randn(5, k_in, device="cuda") torch.testing.assert_close( @@ -144,11 +168,15 @@ def test_other_sizes(self, k_in, hidden, out): ) def test_rejects_unaligned_k(self): + """Reject inner dimensions that violate the native FP32 four-element alignment.""" + params = _synthetic_params(6, 8, 8) with pytest.raises(RuntimeError, match="multiple of 4"): _cuda_op()(torch.randn(2, 6, device="cuda"), *params) def test_rejects_cpu_and_mixed_devices(self): + """Reject CPU-only and mixed-device inputs at the CUDA MLP boundary.""" + params = _synthetic_params(8, 24, 16) with pytest.raises(ValueError): _cuda_op()(torch.randn(2, 8), *[p.cpu() for p in params]) @@ -156,6 +184,8 @@ def test_rejects_cpu_and_mixed_devices(self): _cuda_op()(torch.randn(2, 8), *params) def test_repeat_and_backward_are_deterministic(self): + """Require repeated MLP outputs and all input gradients to be bitwise equal.""" + op = _cuda_op() params = _synthetic_params(256, 5376, 2688, seed=3) x = _features(4, seed=3) @@ -171,12 +201,16 @@ def test_repeat_and_backward_are_deterministic(self): assert torch.equal(a, b) def test_dx_rows_are_batch_invariant(self): + """Keep feature gradients bitwise unchanged when rows are evaluated independently.""" + op = _cuda_op() params = _synthetic_params(256, 5376, 2688, seed=4) x = _features(6, seed=4) grad_out = torch.randn(6, 2688, device="cuda") def dx(rows): + """Compute selected feature-row gradients using matching upstream rows.""" + leaf = x[rows].detach().clone().requires_grad_(True) op(leaf, *params).backward(grad_out[rows]) return leaf.grad @@ -186,6 +220,8 @@ def dx(rows): assert torch.equal(dx(slice(i, i + 1))[0], full[i]) def test_registry_dispatches_cuda(self): + """Resolve the timestep MLP registry entry to its dedicated CUDA implementation.""" + from rl_engine.kernels.registry import KernelRegistry _cuda_op() @@ -198,6 +234,8 @@ def test_registry_dispatches_cuda(self): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) class TestCudaDetLinearDeviceContract: def test_forward_rejects_bias_on_another_device(self, dtype): + """Reject a projection bias on a different CUDA device from its weight.""" + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_forward _cuda_op() @@ -208,6 +246,8 @@ def test_forward_rejects_bias_on_another_device(self, dtype): det_linear_forward(x, weight, bias) def test_backward_input_rejects_weight_on_another_device(self, dtype): + """Reject native input-gradient operands on different CUDA devices.""" + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_backward_input _cuda_op() @@ -217,6 +257,8 @@ def test_backward_input_rejects_weight_on_another_device(self, dtype): det_linear_backward_input(grad, weight, dtype) def test_backward_weight_rejects_input_on_another_device(self, dtype): + """Reject native parameter-gradient operands on different CUDA devices.""" + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_backward_weight _cuda_op() diff --git a/tests/h3/test_h3_timestep_sinusoid.py b/tests/h3/test_h3_timestep_sinusoid.py index bcd567dc3..2f9747a00 100644 --- a/tests/h3/test_h3_timestep_sinusoid.py +++ b/tests/h3/test_h3_timestep_sinusoid.py @@ -28,6 +28,8 @@ def _cuda_op(): + """Construct the CUDA sinusoid operator, skipping a missing native forward entrypoint.""" + from rl_engine.kernels.ops.cuda.h3.timestep_sinusoid import ( H3TimestepSinusoidCudaOp, h3_sinusoid_cuda_available, @@ -40,6 +42,8 @@ def _cuda_op(): class TestReference: def test_layout_is_cos_then_sin(self): + """Check FP32 cosine-then-sine layout and the unit-frequency boundary values.""" + t = torch.tensor([0.0, 1.0]) out = NativeH3TimestepSinusoidOp().forward(t) assert out.shape == (2, 256) and out.dtype == torch.float32 @@ -50,11 +54,15 @@ def test_layout_is_cos_then_sin(self): assert out[1, 128].item() == pytest.approx(math.sin(1.0), abs=1e-7) def test_provider_path_matches_diffusers_cpu(self): + """Require CPU sinusoid values to match the diffusers provider bitwise.""" + t = h3_timesteps(9, device="cpu") ours = NativeH3TimestepSinusoidOp().forward(t) assert torch.equal(ours, provider_time_proj(t, 256)) def test_golden_close_to_provider_path(self): + """Keep the CPU provider within FP32 error of the high-precision sinusoid golden.""" + t = h3_timesteps(33, device="cpu") op = NativeH3TimestepSinusoidOp() torch.testing.assert_close(op.forward(t), op.forward_fp32(t), atol=1e-6, rtol=0) @@ -71,10 +79,14 @@ def test_golden_close_to_provider_path(self): ], ) def test_rejects_invalid_timesteps(self, bad, error): + """Reject empty, malformed, non-floating, non-finite, or out-of-range timesteps.""" + with pytest.raises(error): NativeH3TimestepSinusoidOp().forward(bad) def test_rejects_odd_channels(self): + """Reject odd feature counts that cannot split into equal cosine and sine halves.""" + with pytest.raises(ValueError): NativeH3TimestepSinusoidOp().forward(torch.tensor([0.5]), num_channels=255) @@ -83,6 +95,8 @@ def test_rejects_odd_channels(self): class TestCuda: @pytest.mark.parametrize("num", [1, 2, 3, 7, 64, 1000, 4097]) def test_bitwise_equal_to_diffusers_cuda_path(self, num): + """Require CUDA sinusoid values to match both provider replays for varied batch sizes.""" + t = h3_timesteps(num, seed=num) out = _cuda_op()(t) assert out.dtype == torch.float32 and out.shape == (num, 256) @@ -91,12 +105,16 @@ def test_bitwise_equal_to_diffusers_cuda_path(self, num): @pytest.mark.parametrize("num_channels", [2, 8, 96, 256, 320]) def test_other_channel_counts_match_reference(self, num_channels): + """Preserve bitwise reference parity across supported even channel counts.""" + t = h3_timesteps(5) out = _cuda_op()(t, num_channels=num_channels) ref = NativeH3TimestepSinusoidOp().forward(t, num_channels=num_channels) assert torch.equal(out, ref) def test_fp32_contract_against_fp64_golden(self): + """Bound CUDA sinusoid error against the high-precision golden by FP32 precision.""" + t = h3_timesteps(257) out = _cuda_op()(t) gold = NativeH3TimestepSinusoidOp().forward_fp32(t) @@ -106,10 +124,14 @@ def test_fp32_contract_against_fp64_golden(self): @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) def test_low_precision_timesteps_are_upcast_first(self, dtype): + """Match provider FP32 upcasting for BF16 and FP16 timesteps.""" + t = h3_timesteps(6).to(dtype) assert torch.equal(_cuda_op()(t), provider_time_proj(t, 256)) def test_batch_position_and_repeat_invariance(self): + """Keep row bits unchanged across batch positions, repeated calls, and permutations.""" + op = _cuda_op() base = torch.tensor(0.3141592, device="cuda") single = op(base.view(1)) @@ -126,21 +148,29 @@ def test_batch_position_and_repeat_invariance(self): assert torch.equal(op(t[perm]), first[perm]) def test_non_contiguous_input(self): + """Treat strided timesteps identically to their contiguous copies.""" + t = h3_timesteps(20) strided = t[::2] assert not strided.is_contiguous() assert torch.equal(_cuda_op()(strided), _cuda_op()(strided.contiguous())) def test_rejects_cpu_input(self): + """Reject CPU timesteps at the CUDA sinusoid boundary.""" + with pytest.raises(ValueError): _cuda_op()(torch.tensor([0.5])) def test_rejects_out_of_range(self): + """Reject scaled timesteps outside the declared unit interval.""" + with pytest.raises(ValueError): _cuda_op()(torch.tensor([0.0, 999.0], device="cuda")) @pytest.mark.parametrize("bad", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]) def test_native_entrypoint_rejects_invalid_values(self, bad): + """Reject invalid values in native and wrapper calls even when neighbors are valid.""" + from rl_engine.kernels.ops.base import _C _cuda_op() @@ -152,6 +182,8 @@ def test_native_entrypoint_rejects_invalid_values(self, bad): _cuda_op()(t) def test_native_entrypoint_accepts_boundaries(self): + """Accept both interval endpoints and match provider values in the native call.""" + from rl_engine.kernels.ops.base import _C _cuda_op() @@ -159,6 +191,8 @@ def test_native_entrypoint_accepts_boundaries(self): assert torch.equal(_C.h3_timestep_sinusoid_forward(t), provider_time_proj(t, 256)) def test_trusted_path_supports_cuda_graph(self): + """Capture and replay the prevalidated sinusoid path without changing its output.""" + op = _cuda_op() t = h3_timesteps(7) expected = op(t) @@ -174,6 +208,8 @@ def test_trusted_path_supports_cuda_graph(self): assert torch.equal(out, expected) def test_backward_matches_autograd_reference(self): + """Match timestep gradients to the high-precision autograd sinusoid replay.""" + t = h3_timesteps(17).requires_grad_(True) t_ref = t.detach().clone().requires_grad_(True) grad = torch.randn(17, 256, device="cuda", generator=torch.Generator("cuda").manual_seed(1)) @@ -182,10 +218,14 @@ def test_backward_matches_autograd_reference(self): torch.testing.assert_close(t.grad, t_ref.grad, atol=1e-4, rtol=1e-4) def test_backward_is_row_invariant(self): + """Keep a timestep gradient bitwise unchanged when unrelated rows join the batch.""" + op = _cuda_op() grad_row = torch.randn(256, device="cuda", generator=torch.Generator("cuda").manual_seed(2)) def grad_of(t, pos): + """Differentiate one selected timestep using a fixed upstream feature-gradient row.""" + t = t.clone().requires_grad_(True) grad = torch.zeros(t.shape[0], 256, device="cuda") grad[pos] = grad_row @@ -198,6 +238,8 @@ def grad_of(t, pos): assert torch.equal(grad_of(t, 13), single) def test_registry_dispatches_cuda(self): + """Resolve the sinusoid registry entry to its dedicated CUDA implementation.""" + from rl_engine.kernels.registry import KernelRegistry _cuda_op() diff --git a/tests/h3/test_h3_weights.py b/tests/h3/test_h3_weights.py index 015f78954..c081234e2 100644 --- a/tests/h3/test_h3_weights.py +++ b/tests/h3/test_h3_weights.py @@ -21,6 +21,8 @@ ], ) def test_tensor_checksum_hashes_raw_bytes(dtype, raw_bytes): + """Hash the exact FP32 and BF16 bit patterns without numerical conversion.""" + tensor = torch.tensor([0.0, 1.0], dtype=dtype) assert h3_weights.sha256_tensor(tensor) == hashlib.sha256(raw_bytes).hexdigest() @@ -29,6 +31,8 @@ def test_tensor_checksum_hashes_raw_bytes(dtype, raw_bytes): @pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) @pytest.mark.parametrize("names", [None, ["bias"]]) def test_loading_rejects_modified_weight_contents(tmp_path, monkeypatch, dtype, mismatch, names): + """Reject altered shape, dtype, or bytes even outside the requested tensor subset.""" + from safetensors.torch import save_file tensors = { diff --git a/tests/h3/test_prepare_h3_weights.py b/tests/h3/test_prepare_h3_weights.py index 8a1c94298..1405e044e 100644 --- a/tests/h3/test_prepare_h3_weights.py +++ b/tests/h3/test_prepare_h3_weights.py @@ -18,6 +18,8 @@ @pytest.mark.parametrize("optimize", [0, 1]) @pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, mismatch): + """Enforce tensor contracts before extraction writes, including under Python optimization.""" + import huggingface_hub from safetensors.torch import load_file, save_file diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index fdef637f7..8e77797ce 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -33,6 +33,8 @@ def _run(script: str, *args: str, timeout: int = 300) -> None: def test_all_ws1_single_ops_are_registered(): + """Require gtest specs for all WS1 single operators and the H3 conditioning stages.""" + names = set(operator_names()) assert { "rms_norm", From 59476c9788984e2a257e639ed69d3d75d95690ee Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:32:23 +0000 Subject: [PATCH 23/89] test(h3): enforce complete tensor pairings in regression checks Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- scripts/plot_h3_evidence.py | 4 ++-- tests/h3/test_h3_adaln_projection.py | 15 ++++++++------- tests/h3/test_h3_timestep_mlp.py | 4 ++-- 3 files changed, 12 insertions(+), 11 deletions(-) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 7ebd228d0..988f68c7d 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -181,7 +181,7 @@ def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], con """Per-draw max |error| as boxes with every draw overlaid, log scale.""" _style(ax, title, "max abs error per draw") - for i, (name, color, values) in enumerate(series): + for i, (_name, color, values) in enumerate(series): ax.boxplot( [values], positions=[i], @@ -251,7 +251,7 @@ def plot_projection(report: dict[str, Any]): ("RL-Kernel CUDA", RL_KERNEL, [100 * v for v in acc["cuda_correctly_rounded"]]), ("diffusers (cuBLAS)", PROVIDER, [100 * v for v in acc["provider_correctly_rounded"]]), ] - for i, (name, color, values) in enumerate(series): + for i, (_name, color, values) in enumerate(series): jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] ax_vals = sorted(values) median = ax_vals[len(ax_vals) // 2] diff --git a/tests/h3/test_h3_adaln_projection.py b/tests/h3/test_h3_adaln_projection.py index 0b031fa4e..d052651a3 100644 --- a/tests/h3/test_h3_adaln_projection.py +++ b/tests/h3/test_h3_adaln_projection.py @@ -71,7 +71,7 @@ def test_provider_layout_matches_diffusers(self): ours = NativeH3AdaLNProjectionOp().forward(temb, weight, bias) theirs = provider_adaln_modulation(temb, weight, bias, hidden_size=8) assert len(ours) == 6 - for a, b in zip(ours, theirs): + for a, b in zip(ours, theirs, strict=True): assert a.shape == (6, 8) and torch.equal(a, b) def test_rejects_bf16_temb(self): @@ -144,7 +144,7 @@ def test_forward_is_correctly_rounded_declared_golden(self, block0, num): temb = _temb(num, seed=num) outputs = _cuda_op()(temb, weight, bias) golden = NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias) - for out, gold in zip(outputs, golden): + for out, gold in zip(outputs, golden, strict=True): assert out.dtype == torch.bfloat16 and out.shape == (3 * num, 5376) # tolerance_contract.json forward_accuracy / reduction / bfloat16. torch.testing.assert_close(out.float(), gold, atol=5e-2, rtol=2e-2) @@ -174,10 +174,10 @@ def test_rows_are_batch_and_position_invariant(self, block0): full = op(temb, weight, bias) for i in range(5): single = op(temb[i : i + 1], weight, bias) - for s, f in zip(single, full): + for s, f in zip(single, full, strict=True): assert torch.equal(s, f[3 * i : 3 * i + 3]) swapped = op(temb.flip(0), weight, bias) - for s, f in zip(swapped, full): + for s, f in zip(swapped, full, strict=True): assert torch.equal(s.view(5, 3, -1).flip(0), f.view(5, 3, -1)) def test_backward_against_golden(self, block0): @@ -198,7 +198,7 @@ def test_backward_against_golden(self, block0): # tolerance_contract.json gradient_accuracy / reduction: FP32 temb grad # and BF16 weight/bias grads (rounded once from an FP32 accumulation). torch.testing.assert_close(leaves[0].grad.double(), ref[0].grad, atol=1e-4, rtol=1e-4) - for leaf, r in zip(leaves[1:], ref[1:]): + for leaf, r in zip(leaves[1:], ref[1:], strict=True): torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) assert (leaf.grad == r.grad.to(torch.bfloat16)).float().mean() > 0.999 @@ -216,6 +216,7 @@ def test_other_sizes(self, dtype, hidden, dim): for out, gold in zip( _cuda_op()(temb, weight, bias), NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias), + strict=True, ): torch.testing.assert_close(out.float(), gold, **tol) @@ -232,7 +233,7 @@ def test_repeat_forward_backward_bitwise(self): torch.autograd.backward(list(outs), grads) runs.append([*[o.detach() for o in outs], *[leaf.grad for leaf in leaves]]) for later in runs[1:]: - for a, b in zip(runs[0], later): + for a, b in zip(runs[0], later, strict=True): assert torch.equal(a, b) def test_temb_grad_rows_are_batch_invariant(self): @@ -265,7 +266,7 @@ def test_rows_invariant_across_tensor_core_row_tiles(self, num): full = _cuda_op()(temb, weight, bias) for i in (0, 7, num - 1): single = _cuda_op()(temb[i : i + 1], weight, bias) - for s_out, f_out in zip(single, full): + for s_out, f_out in zip(single, full, strict=True): assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) def test_rejects_cpu(self): diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py index f378a3b35..c61d4faac 100644 --- a/tests/h3/test_h3_timestep_mlp.py +++ b/tests/h3/test_h3_timestep_mlp.py @@ -147,7 +147,7 @@ def test_backward_against_fp64_golden(self, h3_weights_cpu): 3, 2688, device="cuda", generator=torch.Generator("cuda").manual_seed(0) ) _cuda_op()(*leaves).backward(grad_out) - for name, leaf, ref in zip(NAMES, leaves, _fp64_grads(x, params, grad_out)): + for name, leaf, ref in zip(NAMES, leaves, _fp64_grads(x, params, grad_out), strict=True): # tolerance_contract.json gradient_accuracy / reduction / float32. torch.testing.assert_close(leaf.grad.double(), ref, atol=1e-4, rtol=1e-4, msg=name) @@ -197,7 +197,7 @@ def test_repeat_and_backward_are_deterministic(self): out.backward(grad_out) runs.append([out.detach(), *[leaf.grad for leaf in leaves]]) for later in runs[1:]: - for a, b in zip(runs[0], later): + for a, b in zip(runs[0], later, strict=True): assert torch.equal(a, b) def test_dx_rows_are_batch_invariant(self): From 5f5cf945b42556653474e1070a8389b5b2450b03 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:44:51 +0000 Subject: [PATCH 24/89] fix(h3): preserve gradients for empty linear dimensions Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 8 ++++- .../ops/pytorch/h3/adaln_projection.py | 2 +- tests/h3/test_h3_det_linear.py | 30 +++++++++++++++++++ tests/h3/test_h3_timestep_mlp.py | 24 +++++++++++++++ 4 files changed, 62 insertions(+), 2 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 246d27d60..636a00414 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -504,9 +504,11 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei const int64_t n_out = weight.size(0); const int64_t k_in = weight.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); + auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); + if (k_in == 0) return out; + if (n_out == 0) return out.zero_(); const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; auto partial = torch::empty({chunks, rows, k_in}, grad.options()); - auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); auto stream = at::cuda::getCurrentCUDAStream(); const int threads = 256; dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); @@ -551,6 +553,10 @@ std::vector h3_det_linear_backward_weight(torch::Tensor grad, tor const int64_t k_in = x.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); auto dw = torch::empty({n_out, k_in}, grad.options().dtype(w_dtype)); + if (n_out == 0) { + if (!with_bias) return {dw}; + return {dw, torch::empty({n_out}, grad.options().dtype(w_dtype))}; + } auto stream = at::cuda::getCurrentCUDAStream(); const int threads = 256; const int64_t total = n_out * k_in; diff --git a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py index b6d013961..e9cac5ff0 100644 --- a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py +++ b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py @@ -103,7 +103,7 @@ def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: return split_adaln_table(table, hidden) def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: - """Return six FP32 golden views from FP64 SiLU and projection arithmetic. + """Return six FP32 golden outputs from FP64 SiLU and projection arithmetic. Preserve the declared activation rounding to the weight dtype while using an identity VJP at that cast boundary for the FP64 golden graph. diff --git a/tests/h3/test_h3_det_linear.py b/tests/h3/test_h3_det_linear.py index 63ae03a8e..a84184811 100644 --- a/tests/h3/test_h3_det_linear.py +++ b/tests/h3/test_h3_det_linear.py @@ -36,3 +36,33 @@ def test_backward_weight_rejects_zero_rows(dtype, with_bias): x = torch.empty(0, 8, device="cuda", dtype=dtype) with pytest.raises(RuntimeError, match="grad must have at least one row"): _C.h3_det_linear_backward_weight(grad, x, dtype, with_bias) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +def test_backward_input_handles_empty_columns(dtype, out_dtype, k_in, n_out): + """Match the input VJP of empty linear dimensions in the requested output dtype.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + weight = torch.empty(n_out, k_in, device="cuda", dtype=dtype) + expected = (grad @ weight.float()).to(out_dtype) + actual = _C.h3_det_linear_backward_input(grad, weight, out_dtype) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_handles_empty_columns(dtype, w_dtype, k_in, n_out, with_bias): + """Return empty weight gradients while preserving any nonempty bias reduction.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + x = torch.arange(3 * k_in, device="cuda", dtype=dtype).reshape(3, k_in) + expected_weight = (grad.T @ x.float()).to(w_dtype) + actual = _C.h3_det_linear_backward_weight(grad, x, w_dtype, with_bias) + assert len(actual) == (2 if with_bias else 1) + torch.testing.assert_close(actual[0], expected_weight, atol=0, rtol=0) + if with_bias: + torch.testing.assert_close(actual[1], grad.sum(dim=0).to(w_dtype), atol=0, rtol=0) diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py index c61d4faac..fa74344ab 100644 --- a/tests/h3/test_h3_timestep_mlp.py +++ b/tests/h3/test_h3_timestep_mlp.py @@ -167,6 +167,30 @@ def test_other_sizes(self, k_in, hidden, out): rtol=1e-4, ) + def test_zero_hidden_size_matches_cpu_forward_backward(self): + """Match CPU forward values and all input gradients with a zero-width hidden layer.""" + + generator = torch.Generator(device="cpu").manual_seed(0) + cpu_inputs = [ + torch.randn(3, 8, generator=generator), + torch.empty(0, 8), + torch.empty(0), + torch.empty(16, 0), + torch.linspace(-1.0, 1.0, 16), + ] + cpu_inputs = [tensor.requires_grad_(True) for tensor in cpu_inputs] + cuda_inputs = [tensor.detach().cuda().requires_grad_(True) for tensor in cpu_inputs] + cpu_out = NativeH3TimestepMLPOp()(*cpu_inputs) + cuda_out = _cuda_op()(*cuda_inputs) + torch.testing.assert_close(cuda_out.cpu(), cpu_out, atol=0, rtol=0) + grad = torch.randn(3, 16, generator=generator) + cpu_out.backward(grad) + cuda_out.backward(grad.cuda()) + for name, cpu_input, cuda_input in zip(NAMES, cpu_inputs, cuda_inputs, strict=True): + torch.testing.assert_close( + cuda_input.grad.cpu(), cpu_input.grad, atol=0, rtol=0, msg=name + ) + def test_rejects_unaligned_k(self): """Reject inner dimensions that violate the native FP32 four-element alignment.""" From d2a6546cbb424f8201428ed9183fd513bc33f9f5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:21:02 +0000 Subject: [PATCH 25/89] fix(h3): validate RMSNorm inputs and isolate backward benchmarks Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- benchmarks/benchmark_h3_conditioning.py | 13 +- csrc/cuda/h3/rmsnorm_modulate.cu | 43 +++- rl_engine/kernels/ops/pytorch/h3/rmsnorm.py | 6 +- rl_engine/testing/h3_report.py | 120 +++++++--- tests/h3/test_h3_benchmark_timing.py | 211 ++++++++++++++++++ tests/h3/test_h3_rmsnorm.py | 17 +- tests/h3/test_h3_rmsnorm_native_validation.py | 158 +++++++++++++ 7 files changed, 519 insertions(+), 49 deletions(-) create mode 100644 tests/h3/test_h3_benchmark_timing.py create mode 100644 tests/h3/test_h3_rmsnorm_native_validation.py diff --git a/benchmarks/benchmark_h3_conditioning.py b/benchmarks/benchmark_h3_conditioning.py index de1aaceb7..5a39fc175 100644 --- a/benchmarks/benchmark_h3_conditioning.py +++ b/benchmarks/benchmark_h3_conditioning.py @@ -4,7 +4,8 @@ """Benchmark the MiniMax-H3 conditioning ops (RFC #420) against the provider path. CUDA-event medians, bandwidth and peak memory per case, plus the backend -the registry dispatched. Cases live in ``rl_engine/testing/h3_report.py``. +the registry dispatched. Timings alternate candidate/provider execution order; +backward timings exclude forward setup. Cases live in ``rl_engine/testing/h3_report.py``. python benchmarks/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 python benchmarks/benchmark_h3_conditioning.py --op all --json out.json @@ -40,6 +41,7 @@ def main() -> None: registry = KernelRegistry() env = environment() print(f"device={env['gpu']} torch={env['torch']} cuda={env['cuda']}") + print("timing order alternates each iteration; backward timings exclude forward setup") results = [] for name in list(PERF_CASES) if args.op == "all" else [args.op]: for case in PERF_CASES[name](registry): @@ -51,7 +53,14 @@ def main() -> None: for key in TIMED_KEYS if f"{key}_us" in row ] - print(f"{name} {row['case']} [{row['backend']}]: " + "; ".join(parts)) + order = row["execution_order"] + phases = [" -> ".join(order[f"iteration_{i}"]) for i in (0, 1)] + print( + f"{name} {row['case']} [{row['backend']}]: " + + "; ".join(parts) + + "; alternating order: " + + " / ".join(phases) + ) if args.json is not None: args.json.write_text(json.dumps({"environment": env, "results": results}, indent=2) + "\n") diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 5212324b1..4eef07be1 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -255,6 +255,12 @@ void check_rows(const torch::Tensor& t, const char* name, const torch::Tensor& l TORCH_CHECK(t.scalar_type() == like.scalar_type(), name, " must have dtype ", like.scalar_type()); } +void check_index_tensor(const torch::Tensor& t, const char* name, const torch::Tensor& x) { + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.scalar_type() == at::kLong && + t.dim() == 1 && t.is_contiguous(), + name, " must be a contiguous 1-D int64 CUDA tensor on x's device"); +} + Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, @@ -276,8 +282,8 @@ Modulation make_modulation(const c10::optional& shift, } TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), "shift and scale must be views of the same table layout"); - TORCH_CHECK(ix.is_cuda() && ix.scalar_type() == at::kLong && ix.dim() == 1 && ix.is_contiguous(), - "row index must be a contiguous int64 CUDA tensor"); + check_index_tensor(ix, "row index", x); + TORCH_CHECK(ix.numel() > 0, "row index must be nonempty"); TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", ix.size(0), ")"); mod.shift = sh.data_ptr(); @@ -291,9 +297,11 @@ Modulation make_modulation(const c10::optional& shift, void check_xw(const torch::Tensor& x, const torch::Tensor& w) { TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.is_contiguous(), "x must be a contiguous (M, N) CUDA tensor"); TORCH_CHECK(x.size(0) > 0, "x must have at least one row"); - TORCH_CHECK(w.is_cuda() && w.dim() == 1 && w.size(0) == x.size(1) && w.is_contiguous() && + TORCH_CHECK(w.is_cuda() && w.device() == x.device() && w.dim() == 1 && + w.size(0) == x.size(1) && w.is_contiguous() && w.scalar_type() == x.scalar_type(), - "weight must be a contiguous (N,) CUDA tensor with x's dtype"); + "weight must be a contiguous (N,) CUDA tensor with x's dtype and device"); + TORCH_CHECK(x.size(1) > 0, "x must have at least one column"); TORCH_CHECK(x.size(1) % kVec == 0, "N=", x.size(1), " must be a multiple of ", kVec, " (PyTorch's vectorized RMSNorm path)"); } @@ -344,11 +352,31 @@ std::vector h3_rmsnorm_backward( check_xw(x, weight); check_rows(grad, "grad", x); TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); - TORCH_CHECK(rstd.is_cuda() && rstd.scalar_type() == at::kFloat && rstd.numel() == x.size(0), - "rstd must be the forward's float32 (M,) statistics"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && + rstd.scalar_type() == at::kFloat && rstd.dim() == 1 && + rstd.is_contiguous() && rstd.size(0) == x.size(0), + "rstd must be contiguous float32 (M,) statistics on x's device"); const int64_t rows = x.size(0); const int64_t n = x.size(1); const Modulation mod = make_modulation(shift, scale, index, x, n); + if (mod.index != nullptr) { + TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && + seg_first_tile.has_value(), + "modulated backward needs the sorted segment tiles"); + check_index_tensor(*sorted_pos, "sorted_pos", x); + check_index_tensor(*tile_begin, "tile_begin", x); + check_index_tensor(*tile_end, "tile_end", x); + check_index_tensor(*seg_first_tile, "seg_first_tile", x); + TORCH_CHECK(sorted_pos->size(0) == rows, "sorted_pos must have M entries"); + TORCH_CHECK(seg_first_tile->size(0) == shift->size(0) + 1, + "seg_first_tile must have R + 1 entries"); + TORCH_CHECK(tile_begin->size(0) == tile_end->size(0), + "tile_begin and tile_end must have the same number of entries"); + } else { + TORCH_CHECK(!sorted_pos.has_value() && !tile_begin.has_value() && !tile_end.has_value() && + !seg_first_tile.has_value(), + "sorted segment tiles require modulation"); + } const c10::cuda::CUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); auto dx = torch::empty_like(x); @@ -375,9 +403,6 @@ std::vector h3_rmsnorm_backward( C10_CUDA_KERNEL_LAUNCH_CHECK(); if (mod.index == nullptr) return {dx, dweight}; - TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && - seg_first_tile.has_value(), - "modulated backward needs the sorted segment tiles"); const int64_t segments = seg_first_tile->numel() - 1; const int64_t seg_tiles = tile_begin->numel(); const int64_t width = 2 * n; diff --git a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py index 6ab850a07..aeee51754 100644 --- a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py +++ b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py @@ -126,8 +126,10 @@ def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: rstd = torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) return (x64 * rstd * weight.double()).float() - def forward_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): - validate_h3_modulation(x, shift, scale, index) + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ): + validate_h3_modulation(x, shift, scale, index, check_range=check_range) n = self.forward(x, weight, eps) return _DeterministicModulation.apply(n, shift, scale, index) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index dfecec419..c8dfcba78 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -36,41 +36,83 @@ ) -def time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 200) -> float: - """Median CUDA-event time of ``fn`` in microseconds (includes the Python wrapper).""" +def _prepared_call( + fn: Callable[..., Any], setup: Callable[[], Any] | None = None +) -> Callable[[], Any]: + if setup is None: + return fn + prepared = setup() + return lambda: fn(prepared) + + +def _sample_us(fn: Callable[..., Any], setup: Callable[[], Any] | None = None) -> float: + run = _prepared_call(fn, setup) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1e3 + + +def time_us( + fn: Callable[..., Any], + warmup: int = 20, + iters: int = 200, + *, + setup: Callable[[], Any] | None = None, +) -> float: + """Median CUDA-event microseconds; optional setup runs outside the timed region. + + When supplied, ``setup`` runs before every call and its return value is passed + to ``fn``. Backward measurements use it to build a fresh forward graph. + """ for _ in range(warmup): - fn() + _prepared_call(fn, setup)() torch.cuda.synchronize() - samples = [] - for _ in range(iters): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - fn() - end.record() - end.synchronize() - samples.append(start.elapsed_time(end) * 1e3) - return statistics.median(samples) - - -def peak_mib(fn: Callable[[], Any]) -> float: + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) + + +def peak_mib(fn: Callable[..., Any], *, setup: Callable[[], Any] | None = None) -> float: + """Incremental peak CUDA memory of the call, excluding optional setup.""" + + run = _prepared_call(fn, setup) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() base = torch.cuda.memory_allocated() - fn() + run() torch.cuda.synchronize() return (torch.cuda.max_memory_allocated() - base) / 2**20 def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + """Interleave timed keys, reversing their execution order on alternate iterations.""" + row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} - for key in TIMED_KEYS: - if key in case: - us = time_us(case[key], warmup, iters) - row[f"{key}_us"] = us - row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 - row[f"{key}_peak_mib"] = peak_mib(case[key]) + keys = [key for key in TIMED_KEYS if key in case] + orders = (keys, list(reversed(keys))) + row["execution_order"] = { + "policy": "alternating", + "iteration_0": orders[0], + "iteration_1": orders[1], + } + if any(key.endswith("_backward") for key in keys): + row["backward_timing_scope"] = "backward_only" + for iteration in range(warmup): + for key in orders[iteration % 2]: + _prepared_call(case[key], case.get(f"{key}_setup"))() + torch.cuda.synchronize() + samples = {key: [] for key in keys} + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) + for key in keys: + us = statistics.median(samples[key]) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) return row @@ -277,9 +319,12 @@ def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) grads = [torch.randn(seq, hidden, device="cuda").bfloat16() for _ in range(6)] - def backward(fn, ti=ti, tags=tags, grads=grads): + def prepare_backward(fn, ti=ti, tags=tags): leaf = rows.detach().requires_grad_(True) - torch.autograd.backward(list(fn(leaf, ti, tags)), grads) + return list(fn(leaf, ti, tags)) + + def backward(outputs, grads=grads): + torch.autograd.backward(outputs, grads) cases.append( { @@ -290,11 +335,15 @@ def backward(fn, ti=ti, tags=tags, grads=grads): "bytes": 6 * seq * hidden * 2, "candidate": lambda ti=ti, tags=tags: op.forward(rows, ti, tags, check_range=False), "provider": lambda ti=ti, tags=tags: provider_adaln_row_gather(chunks, ti, tags), - "candidate_backward": lambda b=backward: b( - lambda r, ti, tags: op.forward(r, ti, tags, check_range=False) + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: op.forward(r, indices, row_tags, check_range=False) ), - "provider_backward": lambda b=backward: b( - lambda r, ti, tags: provider_adaln_row_gather(r.chunk(6, dim=-1), ti, tags) + "provider_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: provider_adaln_row_gather( + r.chunk(6, dim=-1), indices, row_tags + ) ), } ) @@ -374,9 +423,12 @@ def _norm_perf(registry: KernelRegistry) -> list[dict[str, Any]]: x, weight, shift, scale, index = _norm_inputs(seq) grad = torch.randn_like(x) - def backward(fn, x=x, weight=weight, shift=shift, scale=scale, grad=grad): + def prepare_backward(fn, x=x, weight=weight, shift=shift, scale=scale): leaves = [t.detach().requires_grad_(True) for t in (x, weight, shift, scale)] - fn(*leaves).backward(grad) + return fn(*leaves) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -390,12 +442,14 @@ def backward(fn, x=x, weight=weight, shift=shift, scale=scale, grad=grad): "provider": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( provider_norm_modulate(x, w, sh, sc, i) ), - "candidate_backward": lambda b=backward, i=index: b( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( lambda x_, w_, sh_, sc_: op.forward_modulated( x_, w_, sh_, sc_, i, check_range=False ) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda x_, w_, sh_, sc_: provider_norm_modulate(x_, w_, sh_, sc_, i) ), } diff --git a/tests/h3/test_h3_benchmark_timing.py b/tests/h3/test_h3_benchmark_timing.py new file mode 100644 index 000000000..e6517bd70 --- /dev/null +++ b/tests/h3/test_h3_benchmark_timing.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU regressions for benchmark ordering and backward timing boundaries.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.testing import h3_report + + +@pytest.fixture +def cuda_clock(monkeypatch): + clock = SimpleNamespace(time=0.0, active=False, log=[], allocated=0, peak=0, events=0) + + class Event: + def __init__(self, *, enable_timing): + assert enable_timing + self.start = clock.events % 2 == 0 + clock.events += 1 + + def record(self): + clock.active = self.start + clock.log.append("start" if self.start else "end") + self.time = clock.time + + def synchronize(self): + assert not clock.active + + def elapsed_time(self, end): + return end.time - self.time + + def reset_peak(): + clock.log.append("reset_peak") + clock.peak = clock.allocated + + def allocated(): + clock.log.append("memory_baseline") + return clock.allocated + + monkeypatch.setattr(torch.cuda, "Event", Event) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: clock.log.append("synchronize")) + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", reset_peak) + monkeypatch.setattr(torch.cuda, "memory_allocated", allocated) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: clock.peak) + return clock + + +def test_setup_is_outside_backward_events_and_memory_baseline(cuda_clock): + states = [] + + def setup(): + assert not cuda_clock.active + cuda_clock.log.append("setup") + cuda_clock.time += 50.0 + cuda_clock.allocated += 64 * 2**20 + return object() + + def backward(state): + states.append(state) + cuda_clock.log.append("backward") + cuda_clock.time += 2.0 + cuda_clock.allocated += 3 * 2**20 + cuda_clock.peak = cuda_clock.allocated + + assert h3_report.time_us(backward, warmup=1, iters=2, setup=setup) == 2000.0 + assert len({id(state) for state in states}) == 3 + assert cuda_clock.log.count("start") == 2 + for index, entry in enumerate(cuda_clock.log): + if entry == "start": + assert cuda_clock.log[index - 1] == "setup" + assert cuda_clock.log[index + 1 : index + 3] == ["backward", "end"] + + cuda_clock.log.clear() + assert h3_report.peak_mib(backward, setup=setup) == 3.0 + assert cuda_clock.log == [ + "setup", + "synchronize", + "reset_peak", + "memory_baseline", + "backward", + "synchronize", + ] + + +@pytest.mark.parametrize("keys", [("candidate", "provider"), h3_report.TIMED_KEYS]) +def test_measure_interleaves_and_records_both_orders(cuda_clock, keys): + calls = [] + case = {"op": "test", "case": "tiny", "backend": "cpu", "bytes": 4096} + for position, key in enumerate(keys, start=1): + + def run(key=key, duration=position): + calls.append(key) + cuda_clock.time += duration + + case[key] = run + + row = h3_report.measure(case, warmup=2, iters=4) + first, second = list(keys), list(reversed(keys)) + assert calls == (first + second) * 3 + first # warmup, samples, then memory calls + assert row["execution_order"] == { + "policy": "alternating", + "iteration_0": first, + "iteration_1": second, + } + for position, key in enumerate(keys, start=1): + assert row[f"{key}_us"] == position * 1000.0 + assert row[f"{key}_gbps"] == pytest.approx(4096 / (position * 1e-3) / 1e9) + assert row[f"{key}_peak_mib"] == 0.0 + assert not any(key.endswith("_setup") for key in row) + + +@pytest.mark.parametrize("operator", ["norm", "gather"]) +def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( + monkeypatch, cuda_clock, operator +): + prepared_leaves = [] + previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + + def track_backward(outputs): + tensors = outputs if isinstance(outputs, tuple) else (outputs,) + + def backward_hook(grad): + cuda_clock.time += 2.0 / len(tensors) + return grad + + for tensor in tensors: + tensor.register_hook(backward_hook) + return outputs + + def track(fn, leaves, *args): + if leaves[0].requires_grad: + assert not cuda_clock.active + prepared_leaves.append(leaves) + cuda_clock.time += 1.0 + outputs = fn(*leaves, *args) + return track_backward(outputs) if leaves[0].requires_grad else outputs + + if operator == "norm": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 4), + torch.randn(3, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_norm_modulate + + def forward(x, weight, shift, scale, index, **kwargs): + return track(provider, [x, weight, shift, scale], index) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_norm_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) + op = SimpleNamespace(forward_modulated=forward) + build = h3_report._norm_perf + else: + randn = torch.randn + + def cpu_randn(*args, **kwargs): + kwargs["device"] = "cpu" + return randn(*args, **kwargs) + + provider = h3_report.provider_adaln_row_gather + + def forward(rows, ti, tags, **kwargs): + def gather(leaf, ti, tags): + return provider(leaf.chunk(6, dim=-1), ti, tags) + + return track(gather, [rows], ti, tags) + + def provider_forward(chunks, ti, tags): + if chunks[0].requires_grad: + assert not cuda_clock.active + prepared_leaves.append([chunks[0]._base]) + cuda_clock.time += 1.0 + outputs = provider(chunks, ti, tags) + return track_backward(outputs) if chunks[0].requires_grad else outputs + + monkeypatch.setattr(torch, "randn", cpu_randn) + monkeypatch.setattr(h3_report, "GATHER_SEQ_LENS", (2, 3)) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: (torch.arange(seq) % num, torch.arange(seq) % 3), + ) + monkeypatch.setattr(h3_report, "provider_adaln_row_gather", provider_forward) + op = SimpleNamespace(forward=forward) + build = h3_report._gather_perf + + try: + registry = SimpleNamespace(get_op=lambda *args, **kwargs: op) + for case in build(registry): + row = h3_report.measure(case, warmup=1, iters=2) + assert row["backward_timing_scope"] == "backward_only" + assert row["candidate_backward_us"] == pytest.approx(2000.0) + assert row["provider_backward_us"] == pytest.approx(2000.0) + finally: + torch.set_num_threads(previous_threads) + + assert prepared_leaves + assert len({id(leaves[0]) for leaves in prepared_leaves}) == len(prepared_leaves) + assert all(leaf.grad is not None for leaves in prepared_leaves for leaf in leaves) diff --git a/tests/h3/test_h3_rmsnorm.py b/tests/h3/test_h3_rmsnorm.py index f0d013bfb..53e46b16b 100644 --- a/tests/h3/test_h3_rmsnorm.py +++ b/tests/h3/test_h3_rmsnorm.py @@ -37,9 +37,9 @@ def _cuda_op(): return H3RMSNormCudaOp() -def _x(shape, dtype=torch.bfloat16, seed=0, scale=2.0): +def _x(shape, dtype=torch.bfloat16, seed=0, scale=2.0, device="cuda"): g = torch.Generator(device="cpu").manual_seed(seed) - return (torch.randn(shape, generator=g) * scale).to(dtype).cuda() + return (torch.randn(shape, generator=g) * scale).to(device=device, dtype=dtype) def _table(rows, hidden, chunks, dtype=torch.bfloat16, seed=0): @@ -56,10 +56,21 @@ def _fp64_modulated(x, weight, shift, scale, index, eps=1e-5): class TestReference: def test_golden_close_to_provider(self): - x, w = _x((3, 64), torch.float32).cpu(), torch.rand(64) + 0.5 + x, w = _x((3, 64), torch.float32, device="cpu"), torch.rand(64) + 0.5 op = NativeH3RMSNormOp() torch.testing.assert_close(op.forward(x, w), op.forward_fp32(x, w), atol=1e-6, rtol=1e-6) + def test_modulated_with_prevalidated_indices(self): + x, w = _x((2, 8), torch.float32, device="cpu"), torch.ones(8) + shift, scale = torch.randn(3, 8), torch.randn(3, 8) + index = torch.tensor([2, 0]) + op = NativeH3RMSNormOp() + expected = op.forward(x, w) * (1 + scale[index]) + shift[index] + torch.testing.assert_close( + op.forward_modulated(x, w, shift, scale, index, check_range=False), + expected, + ) + def test_rejects_bad_inputs(self): op = NativeH3RMSNormOp() x, w = torch.randn(2, 8), torch.ones(8) diff --git a/tests/h3/test_h3_rmsnorm_native_validation.py b/tests/h3/test_h3_rmsnorm_native_validation.py new file mode 100644 index 000000000..a0d6f4df7 --- /dev/null +++ b/tests/h3/test_h3_rmsnorm_native_validation.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Direct native calls reject malformed pointers and layouts before launching CUDA work.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = [ + pytest.mark.cuda_only, + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), +] + +TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") + + +@pytest.fixture +def native(): + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not all( + hasattr(_C, name) for name in ("h3_rmsnorm_forward", "h3_rmsnorm_backward") + ): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return _C + + +@pytest.fixture +def inputs(native): + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device=x.device) + _, rstd = native.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device=x.device).chunk(2, dim=1) + return { + "grad": torch.ones_like(x), + "x": x, + "weight": weight, + "rstd": rstd, + "shift": shift, + "scale": scale, + "index": torch.tensor([0, 1], dtype=torch.int64, device=x.device), + "sorted_pos": torch.tensor([0, 2, 1, 3], dtype=torch.int64, device=x.device), + "tile_begin": torch.tensor([0, 2], dtype=torch.int64, device=x.device), + "tile_end": torch.tensor([2, 4], dtype=torch.int64, device=x.device), + "seg_first_tile": torch.tensor([0, 1, 2], dtype=torch.int64, device=x.device), + } + + +def _call(native, inputs, entrypoint): + if entrypoint == "forward": + return native.h3_rmsnorm_forward( + inputs["x"], + inputs["weight"], + 1e-5, + inputs["shift"], + inputs["scale"], + inputs["index"], + ) + return native.h3_rmsnorm_backward(**inputs) + + +def _malformed(tensor, defect): + if defect == "cpu": + return tensor.cpu() + if defect == "dtype": + dtype = torch.float16 if tensor.dtype == torch.float32 else torch.int32 + return tensor.to(dtype) + if defect == "rank": + return tensor.unsqueeze(0) + if defect == "stride": + return torch.stack((tensor, tensor), dim=1)[:, 0] + if defect == "length": + return tensor[:-1] + raise ValueError(f"unknown defect: {defect}") + + +def test_valid_plain_and_modulated_backward(native, inputs): + plain = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + assert len(native.h3_rmsnorm_backward(**plain)) == 2 + dx, dweight, dshift, dscale = native.h3_rmsnorm_backward(**inputs) + assert dx.shape == inputs["x"].shape + assert dweight.shape == inputs["weight"].shape + assert dscale.shape == inputs["scale"].shape + torch.testing.assert_close(dshift, torch.full_like(inputs["shift"], 2.0)) + torch.cuda.synchronize(inputs["x"].device) + + +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_empty_row_index(native, inputs, entrypoint): + inputs["index"] = inputs["index"][:0] + with pytest.raises(RuntimeError, match="row index must be nonempty"): + _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_zero_hidden_size(native, inputs, entrypoint): + inputs["x"] = inputs["x"][:, :0].contiguous() + inputs["weight"] = inputs["weight"][:0] + with pytest.raises(RuntimeError, match="at least one column"): + _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +@pytest.mark.parametrize("mode", ["plain", "modulated"]) +def test_rejects_invalid_rstd(native, inputs, defect, mode): + inputs["rstd"] = _malformed(inputs["rstd"], defect) + if mode == "plain": + inputs = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.parametrize("name", TILE_NAMES) +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +def test_rejects_invalid_segment_tiles(native, inputs, name, defect): + inputs[name] = _malformed(inputs[name], defect) + message = "tile_begin and tile_end" if defect == "length" and name.startswith("tile_") else name + with pytest.raises(RuntimeError, match=message): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.parametrize("name", TILE_NAMES) +def test_rejects_missing_segment_tiles(native, inputs, name): + inputs[name] = None + with pytest.raises(RuntimeError, match="needs the sorted segment tiles"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.parametrize("name", TILE_NAMES) +def test_plain_backward_rejects_segment_tiles(native, inputs, name): + plain = {key: inputs[key] for key in ("grad", "x", "weight", "rstd", name)} + with pytest.raises(RuntimeError, match="sorted segment tiles require modulation"): + native.h3_rmsnorm_backward(**plain) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize( + ("entrypoint", "name"), + [ + ("forward", "weight"), + ("forward", "index"), + ("backward", "weight"), + ("backward", "index"), + ("backward", "rstd"), + ("backward", "sorted_pos"), + ("backward", "tile_begin"), + ("backward", "tile_end"), + ("backward", "seg_first_tile"), + ], +) +def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name): + other_device = (inputs["x"].device.index + 1) % torch.cuda.device_count() + inputs[name] = inputs[name].to(f"cuda:{other_device}") + message = "row index" if name == "index" else name + with pytest.raises(RuntimeError, match=message): + _call(native, inputs, entrypoint) From df519d8e23786a66c0f7afb50fc2afb453e7a984 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:24:11 +0000 Subject: [PATCH 26/89] docs(h3): refresh B200 evidence with backward-only alternating samples Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 18 +++++++++++------- docs/operators/h3-rmsnorm.md | 12 +++++++----- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 1fb114792..f60b2371f 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -70,7 +70,7 @@ Measured on a B200 (torch 2.13.0+cu130): | --- | --- | --- | | Forward vs `index_select` | bitwise equal | | | `d_rows` repeat-bitwise | yes | no | -| `d_rows` correctly rounded from the FP64 sum | 99.99999% | 7.8% | +| `d_rows` correctly rounded from the FP64 sum | 99.99965% | 7.8% | | `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | ### Fused modulation: `H3AdaLNModulationCudaOp` @@ -111,12 +111,14 @@ B200, T = 3, H = 5376, BF16. "write BW" counts the 6·S·H outputs; the table st | S | CUDA forward | provider forward | CUDA backward | provider backward | | --- | --- | --- | --- | --- | -| 4097 | 0.06 ms | 0.09 ms | 1.15 ms | 1.47 ms | -| 32768 | 0.32 ms | 0.59 ms | 2.18 ms | 8.72 ms | -| 131072 | 1.23 ms (6.8 TB/s write) | 2.37 ms | 6.86 ms | 34.9 ms | +| 4097 | 0.07 ms | 0.11 ms | 0.88 ms | 1.28 ms | +| 32768 | 0.33 ms | 0.62 ms | 1.75 ms | 8.04 ms | +| 131072 | 1.17 ms | 2.42 ms | 5.62 ms | 32.45 ms | -The backward's peak memory is about twice the provider's, because autograd stacks the -six output gradients and the per-tile FP32 partials are materialised. +Backward timings and peak memory exclude leaf creation and the forward pass. The CUDA +backward stacks the six output gradients and materialises per-tile FP32 partials, adding +8130 MiB at S = 131072; the provider adds under 2 MiB. Candidate/provider execution order +alternates each iteration and is recorded in the report. ## Evidence @@ -129,7 +131,9 @@ There are two data files: - [`chain_replay.json`](../usage/evidence/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -Both were written from a clean tree at commit `fa551c2`. +`report.json` was regenerated from a clean tree at commit `80e4609` with backward-only +timings and alternating execution order. `chain_replay.json` was written from a clean tree +at commit `fa551c2`. ## Tests diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 957670552..3c69706e7 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -79,19 +79,21 @@ B200, block `norm1` + MSA modulation, B = 1, H = 5376, BF16: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.06 ms | 0.11 ms | 1.08 ms | 0.95 ms | -| 32768 | 0.31 ms | 0.75 ms | 2.50 ms | 5.20 ms | -| 131072 | 1.17 ms | 2.90 ms | 7.40 ms | 20.6 ms | +| 4097 | 0.08 ms | 0.12 ms | 1.23 ms | 0.85 ms | +| 32768 | 0.35 ms | 0.76 ms | 2.18 ms | 4.40 ms | +| 131072 | 1.21 ms | 2.90 ms | 6.17 ms | 17.55 ms | The forward is a single pass that never materialises the gathered rows. At small S the backward -is dominated by the fixed cost of the stable sort and tile setup. +is dominated by the fixed cost of the stable sort and tile setup. Backward timings and peak +memory exclude leaf creation and the forward pass. Candidate/provider execution order alternates +each iteration and is recorded in the report. ## Evidence ![h3_rmsnorm on B200: latency and backward accuracy](../usage/evidence/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](../usage/evidence/h3-rmsnorm-b200/report.json), written by -`scripts/h3_evidence.py` from a clean tree at commit `f36aa60`. It also records: +`scripts/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; From cc867401e434b8e3d33daf28f3e8a64cd548dcef Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:27:16 +0000 Subject: [PATCH 27/89] fix(h3): reject mixed-device RMSNorm native inputs Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/rmsnorm_modulate.cu | 16 ++++++------ tests/h3/test_h3_rmsnorm.py | 43 ++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 7 deletions(-) diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 5212324b1..d42fe3525 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -276,8 +276,9 @@ Modulation make_modulation(const c10::optional& shift, } TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), "shift and scale must be views of the same table layout"); - TORCH_CHECK(ix.is_cuda() && ix.scalar_type() == at::kLong && ix.dim() == 1 && ix.is_contiguous(), - "row index must be a contiguous int64 CUDA tensor"); + TORCH_CHECK(ix.is_cuda() && ix.device() == x.device() && ix.scalar_type() == at::kLong && + ix.dim() == 1 && ix.is_contiguous(), + "row index must be a contiguous int64 CUDA tensor on x's device"); TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", ix.size(0), ")"); mod.shift = sh.data_ptr(); @@ -291,9 +292,9 @@ Modulation make_modulation(const c10::optional& shift, void check_xw(const torch::Tensor& x, const torch::Tensor& w) { TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.is_contiguous(), "x must be a contiguous (M, N) CUDA tensor"); TORCH_CHECK(x.size(0) > 0, "x must have at least one row"); - TORCH_CHECK(w.is_cuda() && w.dim() == 1 && w.size(0) == x.size(1) && w.is_contiguous() && - w.scalar_type() == x.scalar_type(), - "weight must be a contiguous (N,) CUDA tensor with x's dtype"); + TORCH_CHECK(w.is_cuda() && w.device() == x.device() && w.dim() == 1 && w.size(0) == x.size(1) && + w.is_contiguous() && w.scalar_type() == x.scalar_type(), + "weight must be a contiguous (N,) CUDA tensor with x's dtype on x's device"); TORCH_CHECK(x.size(1) % kVec == 0, "N=", x.size(1), " must be a multiple of ", kVec, " (PyTorch's vectorized RMSNorm path)"); } @@ -344,8 +345,9 @@ std::vector h3_rmsnorm_backward( check_xw(x, weight); check_rows(grad, "grad", x); TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); - TORCH_CHECK(rstd.is_cuda() && rstd.scalar_type() == at::kFloat && rstd.numel() == x.size(0), - "rstd must be the forward's float32 (M,) statistics"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics on x's device"); const int64_t rows = x.size(0); const int64_t n = x.size(1); const Modulation mod = make_modulation(shift, scale, index, x, n); diff --git a/tests/h3/test_h3_rmsnorm.py b/tests/h3/test_h3_rmsnorm.py index f0d013bfb..fffb7e3e0 100644 --- a/tests/h3/test_h3_rmsnorm.py +++ b/tests/h3/test_h3_rmsnorm.py @@ -143,6 +143,49 @@ def test_rejects_unsupported(self): _cuda_op()(_x((2, 8)).cpu(), torch.ones(8).bfloat16()) +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +class TestCudaNativeDeviceValidation: + @pytest.mark.parametrize("argument", ["weight", "index"]) + def test_forward_rejects_different_device(self, argument): + from rl_engine.kernels.ops.base import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + if argument == "weight": + weight = weight.to("cuda:1") + else: + index = index.to("cuda:1") + with pytest.raises(RuntimeError, match="on x's device"): + _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + @pytest.mark.parametrize("argument", ["weight", "index", "rstd"]) + def test_backward_rejects_different_device(self, argument): + from rl_engine.kernels.ops.base import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + if argument == "weight": + weight = weight.to("cuda:1") + elif argument == "index": + index = index.to("cuda:1") + else: + rstd = rstd.to("cuda:1") + with pytest.raises(RuntimeError, match="on x's device"): + _C.h3_rmsnorm_backward(torch.ones_like(x), x, weight, rstd, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + @requires_cuda class TestCudaBackward: def _grads(self, fn, tensors, grad): From 40a9b445bbdc46527bd5f8b1a99c7c1f08f0b77a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:27:16 +0000 Subject: [PATCH 28/89] fix(h3): retain FP64 final AdaLN golden gradients Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/testing/h3_report.py | 21 ++++++++--- tests/h3/test_h3_final_adaln_out.py | 58 +++++++++++++++++++++++++---- 2 files changed, 66 insertions(+), 13 deletions(-) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 41f8d647d..4b331e3fc 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -634,12 +634,23 @@ def _final_accuracy(registry: KernelRegistry) -> dict[str, Any]: x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) ).to(x.dtype) - def grads(fn): - leaves = [t.detach().clone().requires_grad_(True) for t in (x, nw, temb, w, b)] - fn(*leaves).backward(grad) + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, nw, temb, w, b)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) return [leaf.grad for leaf in leaves] - ref = grads(lambda *t: golden_op.forward_fp32(*t, ti)) + def golden_backward(x_, nw_, t_, w_, b_): + # Keep the declared BF16 boundaries with FP64 leaves and identity VJPs. + act = t_ * torch.sigmoid(t_) + act = act + (act.to(torch.bfloat16).double() - act).detach() + table = torch.nn.functional.linear(act, w_, b_) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * nw_ + return n * (1.0 + scale.index_select(0, ti)) + shift.index_select(0, ti) + + ref = grads(golden_backward, torch.float64) backward = {} for name, fn in ( ("cuda", lambda *t: op(*t, ti)), @@ -649,7 +660,7 @@ def grads(fn): backward[name] = { "repeat_bitwise_equal": all(torch.equal(a, c) for a, c in zip(first, second)), "rel_error": { - key: float((g.double() - r.double()).abs().max() / r.double().abs().max()) + key: float((g.double() - r).abs().max() / r.abs().max()) for key, g, r in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), first, ref) }, } diff --git a/tests/h3/test_h3_final_adaln_out.py b/tests/h3/test_h3_final_adaln_out.py index f322cd93d..17a3843b8 100644 --- a/tests/h3/test_h3_final_adaln_out.py +++ b/tests/h3/test_h3_final_adaln_out.py @@ -49,6 +49,18 @@ def _inputs(seq, num_timesteps=3, seed=0, batch=1): return x, temb, ti +def _golden(x, norm_weight, temb, weight, bias, timestep_indices): + act = temb * torch.sigmoid(temb) + act = act + (act.to(torch.bfloat16).double() - act).detach() + table = F.linear(act, weight, bias) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) * norm_weight + return n * (1 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + class TestReference: def test_rejects_bad_inputs(self): op = NativeH3FinalAdaLNOutOp() @@ -106,21 +118,51 @@ def grads(fn, dtype=None): fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) return [leaf.grad for leaf in leaves] - def golden(x_, nw, t, w, b): - act = t * torch.sigmoid(t) - act = act + (act.to(torch.bfloat16).double() - act).detach() - shift, scale = F.linear(act, w, b).chunk(2, dim=-1) - n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * nw - return n * (1 + scale.index_select(0, ti)) + shift.index_select(0, ti) - ours = grads(lambda *t: _cuda_op()(*t, ti)) again = grads(lambda *t: _cuda_op()(*t, ti)) - ref = grads(golden, torch.float64) + ref = grads(lambda *t: _golden(*t, ti), torch.float64) assert all(torch.equal(a, b) for a, b in zip(ours, again)) for name, g, r in zip(("dx", "d_norm_weight", "d_temb", "dW", "db"), ours, ref): rel = ((g.double() - r).abs().max() / r.abs().max()).item() assert rel < 1e-2, (name, rel) # BF16 output rounding and round(1 + scale) only + def test_report_backward_uses_fp64_golden(self, monkeypatch): + from types import SimpleNamespace + + from rl_engine.testing import h3_report + + rng = torch.Generator(device="cuda").manual_seed(3) + x = torch.randn(1, 160, 8, device="cuda", generator=rng).bfloat16() + nw = torch.randn(8, device="cuda", generator=rng).bfloat16() + temb = torch.randn(3, 4, device="cuda", generator=rng) + w = torch.randn(16, 4, device="cuda", generator=rng).bfloat16() + b = torch.randn(16, device="cuda", generator=rng).bfloat16() + ti = torch.arange(160, device="cuda") % 3 + inputs = (x, nw, temb, w, b) + native = NativeH3FinalAdaLNOutOp() + registry = SimpleNamespace( + get_op=lambda *args, **kwargs: native, + _get_or_create_backend=lambda _: native, + _priority_map={"cpu": {"final_adaln_out": [None]}}, + ) + monkeypatch.setattr(h3_report, "_final_inputs", lambda *args, **kwargs: (*inputs, ti)) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", native) + report = h3_report._final_accuracy(registry) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + leaves = [t.detach().clone().requires_grad_(True) for t in inputs] + native(*leaves, ti).backward(grad) + ref = [t.double().detach().requires_grad_(True) for t in inputs] + _golden(*ref, ti).backward(grad.double()) + for name, leaf, reference in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), leaves, ref): + assert reference.grad.dtype == torch.float64 + expected = float( + (leaf.grad.double() - reference.grad).abs().max() / reference.grad.abs().max() + ) + for backend in ("cuda", "provider"): + assert report["backward"][backend]["rel_error"][name] == pytest.approx(expected) + def test_registry_dispatches_cuda(self): from rl_engine.kernels.registry import KernelRegistry From 7090f00fc986130391ebe91541f404586313df91 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:28:04 +0000 Subject: [PATCH 29/89] fix(h3): address native validation and evidence review comments Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 11 +- csrc/cuda/h3/rmsnorm_modulate.cu | 18 ++- docs/operators/h3-adaln-gate-residual.md | 2 + docs/operators/h3-rmsnorm.md | 2 + rl_engine/testing/h3_report.py | 61 +++++--- scripts/plot_h3_evidence.py | 14 ++ tests/h3/test_h3_native_validation.py | 124 ++++++++++++++++ tests/h3/test_h3_report.py | 177 +++++++++++++++++++++++ 8 files changed, 380 insertions(+), 29 deletions(-) create mode 100644 tests/h3/test_h3_native_validation.py create mode 100644 tests/h3/test_h3_report.py diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 18ede9e72..c2bd7921a 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -144,9 +144,9 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc "gate must be a CUDA tensor with the activations' dtype"); TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, "gate must be (R, N) with unit column stride"); - TORCH_CHECK(index.is_cuda() && index.scalar_type() == at::kLong && index.dim() == 1 && - index.is_contiguous() && index.numel() > 0, - "index must be a non-empty contiguous int64 CUDA tensor"); + TORCH_CHECK(index.is_cuda() && index.device() == x.device() && index.scalar_type() == at::kLong && + index.dim() == 1 && index.is_contiguous() && index.numel() > 0, + "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; } @@ -207,8 +207,9 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: check_act(y, "y", grad); const Gate g = make_gate(gate, index, grad); for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { - TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kLong && t->dim() == 1 && t->is_contiguous(), - "tile metadata must be contiguous int64 CUDA tensors"); + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 tensors on ", grad.device()); } const int64_t rows = grad.size(0); const int64_t n = grad.size(1); diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 5212324b1..2775572e6 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -276,8 +276,9 @@ Modulation make_modulation(const c10::optional& shift, } TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), "shift and scale must be views of the same table layout"); - TORCH_CHECK(ix.is_cuda() && ix.scalar_type() == at::kLong && ix.dim() == 1 && ix.is_contiguous(), - "row index must be a contiguous int64 CUDA tensor"); + TORCH_CHECK(ix.is_cuda() && ix.device() == x.device() && ix.scalar_type() == at::kLong && + ix.dim() == 1 && ix.is_contiguous() && ix.numel() > 0, + "row index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", ix.size(0), ")"); mod.shift = sh.data_ptr(); @@ -349,6 +350,16 @@ std::vector h3_rmsnorm_backward( const int64_t rows = x.size(0); const int64_t n = x.size(1); const Modulation mod = make_modulation(shift, scale, index, x, n); + if (mod.index != nullptr) { + TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && + seg_first_tile.has_value(), + "modulated backward needs the sorted segment tiles"); + for (const auto* t : {&*sorted_pos, &*tile_begin, &*tile_end, &*seg_first_tile}) { + TORCH_CHECK(t->is_cuda() && t->device() == x.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 tensors on ", x.device()); + } + } const c10::cuda::CUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); auto dx = torch::empty_like(x); @@ -375,9 +386,6 @@ std::vector h3_rmsnorm_backward( C10_CUDA_KERNEL_LAUNCH_CHECK(); if (mod.index == nullptr) return {dx, dweight}; - TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && - seg_first_tile.has_value(), - "modulated backward needs the sorted segment tiles"); const int64_t segments = seg_first_tile->numel() - 1; const int64_t seg_tiles = tile_begin->numel(); const int64_t width = 2 * n; diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index f55412449..2ef31ba54 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -16,6 +16,8 @@ gathers the gate row inside the kernel, so the gathered `(S, H)` gate is never m ## Entry Point ```python +from rl_engine.kernels.registry import kernel_registry + op = kernel_registry.get_op("adaln_gate_residual", device="cuda") hidden = op(residual, attn_output, gate_msa, adaln_indices) ``` diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 957670552..faaed48e2 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -21,6 +21,8 @@ gathered inside the kernel, so the `(S, H)` tensors that ## Entry Point ```python +from rl_engine.kernels.registry import kernel_registry + op = kernel_registry.get_op("h3_rmsnorm", device="cuda") n = op(x, weight) # plain RMSNorm out = op.forward_modulated(x, weight, shift_msa, scale_msa, adaln_indices) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 021dab4b0..489404bb6 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -25,7 +25,11 @@ provider_time_embedder, provider_time_proj, ) -from rl_engine.testing.h3_weights import h3_weights_dir, load_h3_conditioning_weights +from rl_engine.testing.h3_weights import ( + WEIGHTS_ENV, + h3_weights_dir, + load_h3_conditioning_weights, +) # Keys of a perf case that hold a timed callable, in display order. TIMED_KEYS = ( @@ -247,7 +251,7 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, invariant = all( all( torch.equal(s, f[3 * i : 3 * i + 3]) - for s, f in zip(op(temb[i : i + 1], weight, bias), full) + for s, f in zip(op(temb[i : i + 1], weight, bias), full, strict=True) ) for i in range(9) ) @@ -304,7 +308,6 @@ def backward(fn, ti=ti, tags=tags, grads=grads): def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: from rl_engine.testing.h3_chain import run_backward_case - from rl_engine.testing.h3_weights import load_h3_conditioning_weights op = registry.get_op("adaln_row_gather", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) @@ -314,7 +317,9 @@ def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: ti, tags = h3_packed_layout(seq, 3, seed=seq) ours = op(rows, ti, tags) theirs = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) - forward_bitwise[str(seq)] = all(torch.equal(a, b) for a, b in zip(ours, theirs)) + forward_bitwise[str(seq)] = all( + torch.equal(a, b) for a, b in zip(ours, theirs, strict=True) + ) # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. ti, tags = h3_packed_layout(4097, 3, seed=1) @@ -339,16 +344,22 @@ def grad_of(fn): "correctly_rounded_fraction": float((first == gold).float().mean()), } - weights = load_h3_conditioning_weights("cuda") - chain = [ - run_backward_case(registry, weights, num_timesteps=t, seq_len=s) - for t, s in ((1, 257), (3, 257), (1, 4097), (3, 4097), (4, 32768)) - ] - return { + result: dict[str, Any] = { "forward_bitwise_vs_index_select": forward_bitwise, "op_backward": backward, - "chain_backward": chain, + "chain_backward": [], } + if h3_weights_dir() is None: + result["chain_backward_skipped"] = ( + f"{WEIGHTS_ENV} not set; run scripts/prepare_h3_weights.py for chain measurements" + ) + else: + weights = load_h3_conditioning_weights("cuda") + result["chain_backward"] = [ + run_backward_case(registry, weights, num_timesteps=t, seq_len=s) + for t, s in ((1, 257), (3, 257), (1, 4097), (3, 4097), (4, 32768)) + ] + return result # --------------------------------------------------------------------------- # @@ -406,14 +417,19 @@ def backward(fn, x=x, weight=weight, shift=shift, scale=scale, grad=grad): def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("h3_rmsnorm", device="cuda") - weights = load_h3_conditioning_weights("cuda") - plain = {} - for name in ( + weight_source = "pinned_checkpoint" if h3_weights_dir() is not None else "synthetic" + names = [ "transformer_blocks.0.norm1.weight", "transformer_blocks.0.norm2.weight", "token_refiner.final_norm.weight", "norm_out.norm.weight", - ): + ] + weights = { + name: weight.bfloat16() + for name, weight in zip(names, h3_params(names, [(5376,)] * len(names)), strict=True) + } + plain = {} + for name in names: x = torch.randn(2, 777, 5376, device="cuda").bfloat16() ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) plain[name] = bool(torch.equal(op(x, weights[name]), ref)) @@ -454,13 +470,16 @@ def golden(x_, w_, sh_, sc_): ): first, second = grads(fn), grads(fn) backward[name] = { - "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), "rel_error": { key: float((g.double() - r).abs().max() / r.abs().max()) - for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref) + for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref, strict=True) }, } return { + "weight_source": weight_source, "plain_bitwise_vs_nn_rmsnorm": plain, "modulated_bitwise_vs_diffusers": modulated, "rows_batch_invariant": invariant, @@ -559,10 +578,14 @@ def grads(fn, dtype=None): ): first, second = grads(fn), grads(fn) backward[name] = { - "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), "rel_error": { key: float((g.double() - r).abs().max() / r.abs().max()) - for key, g, r in zip(("d_residual", "d_sublayer", "d_gate"), first, ref) + for key, g, r in zip( + ("d_residual", "d_sublayer", "d_gate"), first, ref, strict=True + ) }, } return { diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 15c800b5f..d8bd5ea43 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -302,6 +302,20 @@ def plot_gather(report: dict[str, Any]): "Whole-chain grads of the FP32 time embedder vs FP64", "max abs error / golden max", ) + if not chain: + right.text( + 0.5, + 0.5, + report["accuracy"].get("chain_backward_skipped", "No chain measurements available"), + transform=right.transAxes, + ha="center", + va="center", + color=INK, + fontsize=9, + wrap=True, + ) + right.set_axis_off() + return fig names = { "candidate": ("RL-Kernel, separate ops", THIRD, "o"), "candidate_fused": ("RL-Kernel, fused modulation", RL_KERNEL, "D"), diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py new file mode 100644 index 000000000..8dd919028 --- /dev/null +++ b/tests/h3/test_h3_native_validation.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Exercise H3 validation at the native entrypoints, bypassing Python guards.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@pytest.fixture +def native_extension(): + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + names = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward", + "h3_gate_residual_forward", + "h3_gate_residual_backward", + ) + if not _EXT_AVAILABLE or not all(hasattr(_C, name) for name in names): + pytest.skip("rl_engine._C lacks native H3 entrypoints") + return _C + + +@pytest.fixture +def native_case(): + device = "cuda:0" + x = torch.arange(16, device=device, dtype=torch.float32).view(2, 8) / 10 + return { + "x": x, + "weight": torch.ones(8, device=device), + "shift": torch.zeros(2, 8, device=device), + "scale": torch.zeros(2, 8, device=device), + "gate": torch.ones(2, 8, device=device), + "grad": torch.ones_like(x), + "rstd": torch.rsqrt(x.square().mean(-1) + 1e-5), + "index": torch.arange(2, device=device, dtype=torch.int64), + "sorted_pos": torch.arange(2, device=device, dtype=torch.int64), + "tile_begin": torch.arange(2, device=device, dtype=torch.int64), + "tile_end": torch.arange(1, 3, device=device, dtype=torch.int64), + "seg_first_tile": torch.arange(3, device=device, dtype=torch.int64), + } + + +def _native_call(extension, operation, boundary, case): + x, index = case["x"], case["index"] + tiles = tuple(case[name] for name in ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile")) + if operation == "rmsnorm": + modulation = (case["shift"], case["scale"], index) + if boundary == "forward": + return extension.h3_rmsnorm_forward(x, case["weight"], 1e-5, *modulation) + return extension.h3_rmsnorm_backward( + case["grad"], x, case["weight"], case["rstd"], *modulation, *tiles + ) + if boundary == "forward": + return extension.h3_gate_residual_forward(x, x, case["gate"], index) + return extension.h3_gate_residual_backward(case["grad"], x, case["gate"], index, *tiles) + + +def _other_cuda(tensor): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + return tensor.to("cuda:1") + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("boundary", ["forward", "backward"]) +@pytest.mark.parametrize("index_case", ["empty", "cpu", "other_cuda"]) +def test_native_rejects_invalid_row_index( + native_extension, native_case, operation, boundary, index_case +): + index = native_case["index"] + if index_case == "empty": + native_case["index"] = index[:0] + elif index_case == "cpu": + native_case["index"] = index.cpu() + else: + native_case["index"] = _other_cuda(index) + with pytest.raises( + RuntimeError, match="index must be a non-empty contiguous int64 tensor on cuda:0" + ): + _native_call(native_extension, operation, boundary, native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +@pytest.mark.parametrize("device", ["cpu", "other_cuda"]) +def test_native_rejects_mixed_device_tiles( + native_extension, native_case, operation, metadata, device +): + tensor = native_case[metadata] + native_case[metadata] = tensor.cpu() if device == "cpu" else _other_cuda(tensor) + with pytest.raises( + RuntimeError, match="tile metadata must be contiguous int64 tensors on cuda:0" + ): + _native_call(native_extension, operation, "backward", native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +def test_native_accepts_same_device_inputs(native_extension, native_case, operation): + x = native_case["x"].detach().clone().requires_grad_(True) + if operation == "rmsnorm": + weight = native_case["weight"].detach().clone().requires_grad_(True) + shift = native_case["shift"].detach().clone().requires_grad_(True) + scale = native_case["scale"].detach().clone().requires_grad_(True) + norm = torch.nn.functional.rms_norm(x, (8,), weight, 1e-5) + ref = norm * (1 + scale) + shift + ref.backward(native_case["grad"]) + expected_grads = (x.grad, weight.grad, shift.grad, scale.grad) + got, _ = _native_call(native_extension, operation, "forward", native_case) + else: + gate = native_case["gate"].detach().clone().requires_grad_(True) + ref = native_case["x"] + gate * x + ref.backward(native_case["grad"]) + expected_grads = (x.grad, gate.grad) + got = _native_call(native_extension, operation, "forward", native_case) + torch.testing.assert_close(got, ref) + actual_grads = _native_call(native_extension, operation, "backward", native_case) + for actual, expected in zip(actual_grads, expected_grads, strict=True): + torch.testing.assert_close(actual, expected) diff --git a/tests/h3/test_h3_report.py b/tests/h3/test_h3_report.py new file mode 100644 index 000000000..b6bb6ec3f --- /dev/null +++ b/tests/h3/test_h3_report.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Optional checkpoint handling preserves standalone operator measurements.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.testing import h3_chain, h3_report +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_adaln_row_gather, provider_norm_modulate +from rl_engine.testing.h3_weights import WEIGHTS_ENV + + +@pytest.fixture +def small_cpu_report(monkeypatch): + """Exercise actual report/autograd logic with small CPU tensors and provider ops.""" + + def dimension(size): + return {5376: 4, 6 * 5376: 24, 257: 160, 777: 160, 4097: 160, 32768: 160}.get(size, size) + + class CpuTorch: + nn = SimpleNamespace( + functional=SimpleNamespace( + rms_norm=lambda x, _shape, weight, eps: F.rms_norm(x, (x.shape[-1],), weight, eps) + ) + ) + + def __getattr__(self, name): + return getattr(torch, name) + + def Generator(self, device): + return torch.Generator(device="cpu") + + def randn(self, *shape, **kwargs): + if len(shape) == 1 and isinstance(shape[0], (tuple, list, torch.Size)): + shape = shape[0] + kwargs["device"] = "cpu" + return torch.randn(*(dimension(size) for size in shape), **kwargs) + + class Gather: + def __call__(self, rows, ti, tags): + return provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + + forward_fp32 = __call__ + + class Norm: + def __call__(self, x, weight): + return F.rms_norm(x, (x.shape[-1],), weight, 1e-5) + + forward_modulated = staticmethod(provider_norm_modulate) + + gather = Gather() + ops = {"adaln_row_gather": gather, "h3_rmsnorm": Norm()} + registry = SimpleNamespace( + get_op=lambda name, device: ops[name], + _priority_map={"cpu": {"adaln_row_gather": ["gather"]}}, + _get_or_create_backend=lambda _backend: gather, + ) + monkeypatch.setattr(h3_report, "torch", CpuTorch()) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: h3_packed_layout(dimension(seq), num, seed=seed, device="cpu"), + ) + + def norm_inputs(seq, seed=0): + generator = torch.Generator().manual_seed(seed) + weight = h3_report.h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0] + shift, scale = [torch.randn(9, 4, generator=generator).bfloat16() for _ in range(2)] + x = torch.randn(1, dimension(seq), 4, generator=generator).bfloat16() + ti, tags = h3_packed_layout(dimension(seq), 3, seed=seed, device="cpu") + return x, weight.bfloat16(), shift, scale, ti * 3 + tags + + monkeypatch.setattr(h3_report, "_norm_inputs", norm_inputs) + return registry + + +def test_gather_without_weights_preserves_op_results(small_cpu_report, monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + + def unexpected_load(*args, **kwargs): + pytest.fail("an unset checkpoint must not be loaded") + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", unexpected_load) + result = h3_report._gather_accuracy(small_cpu_report) + assert all(result["forward_bitwise_vs_index_select"].values()) + assert set(result["op_backward"]) == {"cuda", "provider"} + assert all(row["repeat_bitwise_equal"] for row in result["op_backward"].values()) + assert result["chain_backward"] == [] + assert WEIGHTS_ENV in result["chain_backward_skipped"] + + +@pytest.mark.parametrize("configured", [False, True]) +def test_norm_reports_weight_source(small_cpu_report, monkeypatch, tmp_path, configured): + if configured: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + else: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + calls = [] + + def load_weights(device, names): + assert configured, "an unset checkpoint must not be loaded" + calls.append(names) + return {name: torch.ones(4, dtype=torch.bfloat16) for name in names} + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", load_weights) + result = h3_report._norm_accuracy(small_cpu_report) + assert result["weight_source"] == ("pinned_checkpoint" if configured else "synthetic") + assert len(result["plain_bitwise_vs_nn_rmsnorm"]) == 4 + assert all(result["plain_bitwise_vs_nn_rmsnorm"].values()) + assert result["modulated_bitwise_vs_diffusers"] + assert result["rows_batch_invariant"] + assert set(result["backward"]) == {"cuda", "provider"} + assert bool(calls) is configured + + +def test_gather_with_weights_runs_chain(small_cpu_report, monkeypatch, tmp_path): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + weights = {"weight": torch.ones(1)} + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", lambda device: weights) + + def run_chain(registry, loaded, **case): + assert registry is small_cpu_report + assert loaded is weights + return case + + monkeypatch.setattr(h3_chain, "run_backward_case", run_chain) + result = h3_report._gather_accuracy(small_cpu_report) + assert len(result["chain_backward"]) == 5 + assert result["chain_backward"][-1] == {"num_timesteps": 4, "seq_len": 32768} + assert "chain_backward_skipped" not in result + + +@pytest.mark.parametrize("accuracy", [h3_report._gather_accuracy, h3_report._norm_accuracy]) +def test_configured_missing_weights_raise(small_cpu_report, monkeypatch, tmp_path, accuracy): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + with pytest.raises(FileNotFoundError, match="missing; run scripts/prepare_h3_weights.py"): + accuracy(small_cpu_report) + + +def test_plot_gather_without_chain(small_cpu_report, monkeypatch, tmp_path): + pytest.importorskip("matplotlib") + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + script = Path(__file__).resolve().parents[2] / "scripts" / "plot_h3_evidence.py" + spec = importlib.util.spec_from_file_location("plot_h3_evidence_test", script) + plot = importlib.util.module_from_spec(spec) + spec.loader.exec_module(plot) + accuracy = h3_report._gather_accuracy(small_cpu_report) + report = { + "environment": {"gpu": "CPU fixture", "torch": torch.__version__, "cuda": "none"}, + "rl_kernel_commit": "testcommit", + "model_revision": "pinned_revision", + "accuracy": accuracy, + "perf": [ + { + "case": "S=160", + "candidate_us": 1, + "provider_us": 2, + "candidate_backward_us": 3, + "provider_backward_us": 4, + } + ], + } + figure = plot.plot_gather(report) + output = tmp_path / "figure.png" + figure.savefig(output) + assert output.stat().st_size > 0 + assert figure.axes[1].texts[0].get_text() == accuracy["chain_backward_skipped"] + plot.plt.close(figure) From cff41e857503e6ba3d7147c7d43dfd0d5a3cb93b Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:28:41 +0000 Subject: [PATCH 30/89] docs(h3): regenerate final AdaLN evidence against FP64 gradients Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 3e64efbef..59b4ab0a9 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -40,8 +40,8 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste summation tree versus cuBLAS. Rows are batch- and position-invariant. - **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into the projection backward without an extra BF16 rounding: - - `d_temb`, `dW` and `db` are 13–57× closer to FP64 than diffusers, whose error varies from - run to run because it rounds the table gradient to BF16 and accumulates it with + - `d_temb`, `dW` and `db` are 22–43× closer to FP64 in the recorded run than diffusers, + whose error varies from run to run because it rounds the table gradient to BF16 and accumulates it with `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; - every gradient is repeat-bitwise. @@ -59,9 +59,9 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.13 ms | 0.15 ms | 1.64 ms | 1.10 ms | -| 32768 | 0.39 ms | 0.78 ms | 2.88 ms | 5.36 ms | -| 131072 | 1.23 ms | 2.92 ms | 7.80 ms | 20.7 ms | +| 4097 | 0.14 ms | 0.15 ms | 1.67 ms | 1.12 ms | +| 32768 | 0.40 ms | 0.78 ms | 2.97 ms | 5.39 ms | +| 131072 | 1.25 ms | 2.94 ms | 7.87 ms | 22.2 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -70,13 +70,15 @@ projection backward. ![final_adaln_out on B200: latency and backward accuracy](../usage/evidence/h3-final-adaln-out-b200/figure.png) -There are two data files, both written from a clean tree at commit `38d575f`: +There are two data files: - [`report.json`](../usage/evidence/h3-final-adaln-out-b200/report.json): op timings, the - forward-equality fraction, row invariance and backward accuracy. + forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean + tree at `372d6d7`, with FP64 leaves and upstream gradients and explicit BF16 forward boundaries. - [`chain_replay.json`](../usage/evidence/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over - T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. + T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean + tree at `38d575f`. ## Tests From 507c9ad86540dd0aaab9376996d85d8c1c5f60a8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 10:39:17 +0000 Subject: [PATCH 31/89] [WS2][CUDA][MiniMax-H3] tp_adaln_3mod: column-parallel AdaLN projection, byte-equal to WS1 RFC #420 WS2 row `tp_adaln_3mod` (stacked on final_adaln_out). - Ownership: rank r owns the contiguous output columns [r*N/tp, (r+1)*N/tp) of the 3 x 6 x 5376 table (the same weight/bias rows); N/tp must be a multiple of the 64-row d_input chunk, so a shard never splits a WS1 partial. Unsupported TP sizes fail closed. - Every floating-point op is the WS1 kernel: the local tensor-core GEMV gives the WS1 columns, a rank-ordered all-gather (a copy) rebuilds the table, dW/db are shard-local, and d_temb all-gathers the WS1 chunk partials and runs the WS1 ascending fold on every rank. No collective reduction arithmetic. - det_linear: d_input split into h3_det_linear_backward_input_partials and h3_det_linear_fold_chunks (same kernels; backward_input composes them). - ws2_comm.gather_rows: rank-ordered row gather that pads outputs of at most 256 KiB past DeterministicCollective's single-block byte-copy path (~4 us/KiB); keeps the T = 1 forward at ~0.12 ms. - Readback: tp, rank, owned columns and (modality, chunk) slots, collective backend, kernel ids, fallback. - Tests: ownership, single-GPU shard arithmetic (columns and chunk partials equal the WS1 slice, the fold equals WS1 d_input), fail-closed inputs, and real NCCL TP2/4/8 on the pinned block-0 weights. - scripts/h3_ws2_evidence.py + figure; operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 41 +++- csrc/ops.cpp | 8 + docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-tp-adaln-3mod.md | 94 ++++++++ rl_engine/_C.pyi | 4 + rl_engine/kernels/ops/cuda/h3/det_linear.py | 17 ++ .../ops/cuda/h3/tp_adaln_projection.py | 211 ++++++++++++++++++ rl_engine/kernels/ops/cuda/h3/ws2_comm.py | 33 +++ rl_engine/testing/h3_ws2.py | 187 ++++++++++++++++ scripts/h3_ws2_evidence.py | 90 ++++++++ scripts/plot_h3_evidence.py | 96 +++++++- tests/h3/test_h3_tp_adaln.py | 148 ++++++++++++ 13 files changed, 923 insertions(+), 8 deletions(-) create mode 100644 docs/operators/h3-tp-adaln-3mod.md create mode 100644 rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py create mode 100644 rl_engine/kernels/ops/cuda/h3/ws2_comm.py create mode 100644 rl_engine/testing/h3_ws2.py create mode 100644 scripts/h3_ws2_evidence.py create mode 100644 tests/h3/test_h3_tp_adaln.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 95e16b234..66d62931d 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -26,6 +26,9 @@ // * N is split into fixed chunks of 64 rows; each chunk is an ascending // fmaf chain into FP32 from 0; the chunk partials are then left-folded in // ascending chunk order and cast once. +// The partials and the fold are also bound separately: a contiguous N +// slice starting on a chunk boundary yields exactly the full call's +// partials for those chunks, which is what tensor-parallel shards gather. // d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] // * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one // cross-row reduction; its order is the logical row order. @@ -487,24 +490,24 @@ std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor return {out}; } -// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype. -torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, - c10::ScalarType out_dtype) { +// grad [T, N] (float32), weight [N, K] -> partial [ceil(N / 64), T, K] float32: +// the per-chunk sums of contract h3-det-linear-v1, before the ascending fold. +// A contiguous slice of N that starts on a chunk boundary yields exactly the +// matching rows of the full call (tensor-parallel shards rely on this). +torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight) { TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && grad.scalar_type() == at::kFloat, "grad must be a contiguous 2-D float32 CUDA tensor"); check_matrix(weight, "weight"); TORCH_CHECK(grad.device() == weight.device(), "grad and weight must be on the same device"); TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N"); - TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, - "out_dtype must be float32 or bfloat16"); const int64_t rows = grad.size(0); const int64_t n_out = weight.size(0); const int64_t k_in = weight.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; auto partial = torch::empty({chunks, rows, k_in}, grad.options()); - auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); + if (partial.numel() == 0) return partial; auto stream = at::cuda::getCurrentCUDAStream(); const int threads = 256; dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); @@ -518,7 +521,23 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei n_out, k_in); } C10_CUDA_KERNEL_LAUNCH_CHECK(); - const int64_t elems = rows * k_in; + return partial; +} + +// partial [C, T, K] float32 -> left fold over C in ascending order, cast once. +torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype) { + TORCH_CHECK(partial.is_cuda() && partial.dim() == 3 && partial.is_contiguous() && + partial.scalar_type() == at::kFloat && partial.size(0) > 0, + "partial must be a non-empty contiguous 3-D float32 CUDA tensor"); + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + const c10::cuda::CUDAGuard device_guard(partial.device()); + const int64_t chunks = partial.size(0); + auto out = torch::empty({partial.size(1), partial.size(2)}, partial.options().dtype(out_dtype)); + const int64_t elems = out.numel(); + if (elems == 0) return out; + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); if (out_dtype == at::kFloat) { det_fold_chunks_kernel<<>>( @@ -531,6 +550,14 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei return out; } +// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype. +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype) { + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + return h3_det_linear_fold_chunks(h3_det_linear_backward_input_partials(grad, weight), out_dtype); +} + // grad [T, N] float32, x [T, K] (float32 or bfloat16) -> (dW [N, K], dbias [N]) in w_dtype. std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, c10::ScalarType w_dtype, diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 52ece3f13..f3dc5cca2 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -507,6 +507,8 @@ std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor int64_t activation, bool save_pre_activation); torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, c10::ScalarType out_dtype); +torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight); +torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype); std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, c10::ScalarType w_dtype, bool with_bias); @@ -806,6 +808,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("h3_det_linear_backward_input", &h3_det_linear_backward_input, "Deterministic grad @ weight with fixed 64-row N chunks folded in order", py::arg("grad"), py::arg("weight"), py::arg("out_dtype")); + m.def("h3_det_linear_backward_input_partials", &h3_det_linear_backward_input_partials, + "Per-64-row-chunk FP32 partials of grad @ weight, before the ascending fold", + py::arg("grad"), py::arg("weight")); + m.def("h3_det_linear_fold_chunks", &h3_det_linear_fold_chunks, + "Ascending left fold of h3_det_linear_backward_input_partials, cast once", + py::arg("partial"), py::arg("out_dtype")); m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight, "Deterministic dW/dbias as ascending-row FP32 folds", py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true); diff --git a/docs/.nav.yml b/docs/.nav.yml index 1ef4fa017..48f12fe3b 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -37,6 +37,7 @@ nav: - operators/h3-rmsnorm.md - operators/h3-adaln-gate-residual.md - operators/h3-final-adaln-out.md + - operators/h3-tp-adaln-3mod.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 27ffce7d0..99e02cf0c 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -39,4 +39,5 @@ Every operator page should include: - [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) - [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) - [MiniMax-H3 Final AdaLN Output](h3-final-adaln-out.md) +- [MiniMax-H3 Tensor-Parallel AdaLN Projection](h3-tp-adaln-3mod.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md new file mode 100644 index 000000000..da7b0ae42 --- /dev/null +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -0,0 +1,94 @@ +# MiniMax-H3 Tensor-Parallel AdaLN Projection + +## Summary + +`tp_adaln_3mod` shards the [AdaLN projection](h3-adaln-projection.md) (`2688 -> 96768`, 520 MB of +BF16 weight per block) across tensor-parallel ranks (RFC #420, WS2). Every rank ends with the +**same bytes as the single-GPU (WS1) op**: the six modulation tensors, the `3T` modality rows, +`d_temb`, and its shard of `dW`/`db`. + +## Entry Point + +```python +from rl_engine.distributed.collectives import DeterministicCollective +from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, +) + +collective = DeterministicCollective(group=tp_group, device=local_rank) +op = H3TPAdaLNProjectionCudaOp(collective, n_total=96768) +w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, w_shard, b_shard) +op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, fallback +``` + +## Ownership + +Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the +`N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. +`N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard +never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 qualify. Anything else fails closed when +the op is built. + +| TP | Columns per rank | Rank 0 owns | +| --- | --- | --- | +| 2 | 48384 | video: all six chunks; text: shift/scale/gate_msa | +| 4 | 24192 | video: shift_msa .. shift_mlp, half of scale_mlp | +| 8 | 12096 | video: shift_msa, scale_msa, 1344 channels of gate_msa | + +`AdaLNColumnShard.slots()` lists each rank's `(modality, chunk, h_begin, h_end)` pieces. The +all-gathered table is the WS1 table, so the six tensors and the `t * 3 + modality` rows are views +of it exactly as in WS1. Nothing downstream needs to know the TP size. + +## Numerics + +Every floating-point operation is the WS1 kernel itself. The only collective is a rank-ordered +all-gather, which is a copy. No collective reduction arithmetic is used. + +- **Forward.** An output column depends only on `silu(temb)` and its own weight row, so the local + tensor-core GEMV yields exactly the WS1 columns. The all-gather rebuilds the `(T, N)` table. +- **`dW`, `db`.** These are rows of the shard, computed locally with the WS1 ascending-`t` fold. +- **`d_temb`.** WS1 sums `N` in 64-row chunks and then left-folds the chunk partials in ascending + order. Each rank computes the partials of its own chunks + (`h3_det_linear_backward_input_partials`), the all-gather puts them in global chunk order, and + every rank runs the WS1 fold (`h3_det_linear_fold_chunks`) and the FP32 SiLU VJP. + +The backward reads its own columns of the table gradient. It therefore requires that gradient to +be identical on every TP rank, as it is when every rank applies the modulation to the same rows. + +## Communication + +| Direction | Payload per rank | At T = 3, TP = 8 | +| --- | --- | --- | +| forward | `T x N / tp` BF16 | 72.6 KB | +| backward | `(N / 64 / tp) x T x 2688` FP32 | 6.1 MB | + +`DeterministicCollective` copies all-gathers whose output is at most 256 KiB with a single-block +byte loop, at about 4 µs per KiB on B200. `ws2_comm.gather_rows` pads such shards just past that +size, where the multi-block path takes about 45 µs. This matters for the forward at small `T`. + +## Evidence + +The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 +node, one GPU per rank. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_tp_adaln.py -v # ownership, shard arithmetic, NCCL TP2/4/8 +python scripts/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \ + --out docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json +``` + +The NCCL tests need as many GPUs as ranks and skip otherwise. The shard-arithmetic tests check +the same property on one GPU: each rank's columns and chunk partials equal the matching slice of +the WS1 call, and the fold of the gathered partials equals WS1's `d_input`. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- `norm_out.linear` (`2688 -> 10752`) is not sharded here. It is small and stays replicated. +- There is no ROCm path. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 89843f979..672bf6049 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -341,6 +341,10 @@ def h3_det_linear_forward( def h3_det_linear_backward_input( grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype ) -> torch.Tensor: ... +def h3_det_linear_backward_input_partials( + grad: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: ... +def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: ... def h3_det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True ) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/ops/cuda/h3/det_linear.py b/rl_engine/kernels/ops/cuda/h3/det_linear.py index f5dab854e..17e7bce0f 100644 --- a/rl_engine/kernels/ops/cuda/h3/det_linear.py +++ b/rl_engine/kernels/ops/cuda/h3/det_linear.py @@ -16,10 +16,13 @@ ACT_NONE = 0 ACT_SILU = 1 CONTRACT = "h3-det-linear-v1" +DINPUT_CHUNK = 64 # N rows per d_input partial (kDInputChunk) _SYMBOLS = ( "h3_det_linear_forward", "h3_det_linear_backward_input", "h3_det_linear_backward_weight", + "h3_det_linear_backward_input_partials", + "h3_det_linear_fold_chunks", ) @@ -62,6 +65,20 @@ def det_linear_backward_input( ) +def det_linear_backward_input_partials(grad: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``(ceil(N / 64), T, K)`` FP32 chunk sums of ``grad @ weight``, before the fold.""" + + _require() + return _C.h3_det_linear_backward_input_partials(grad.float().contiguous(), weight.contiguous()) + + +def det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Ascending left fold of the chunk partials, cast once (the second half of d_input).""" + + _require() + return _C.h3_det_linear_fold_chunks(partial.contiguous(), out_dtype) + + def det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True ) -> list[torch.Tensor]: diff --git a/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py new file mode 100644 index 000000000..0eb5a4709 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tensor-parallel H3 AdaLN projection (RFC #420 ``tp_adaln_3mod``). + +Ownership: rank ``r`` of ``tp`` owns the contiguous output columns +``[r * N / tp, (r + 1) * N / tp)`` of the ``N = 3 * 6 * H`` projection, i.e. the +same rows of ``adaln_proj.linear.weight``/``bias``. ``N / tp`` must be a +multiple of the 64-row ``d_input`` chunk of contract ``h3-det-linear-v1``. + +Every floating-point operation is the WS1 kernel itself, so every rank ends +with the WS1 bytes: + +* forward: each column depends only on ``x`` and its own weight row, so the + local GEMV produces the WS1 columns; a rank-ordered all-gather (a copy) + rebuilds the ``(T, N)`` table, and the six tensors and ``3T`` modality rows + are views of it exactly as in WS1; +* ``dW``/``db``: rows of the shard, computed locally (no cross-rank sum); +* ``d_temb``: each rank computes the WS1 chunk partials of its columns, the + all-gather puts them in global chunk order, and every rank runs the WS1 + ascending fold. No collective reduction arithmetic is used. + +The backward expects the table gradient to be identical on every TP rank +(the modulation consumes the full replicated table); it reads its own columns. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + CONTRACT, + DINPUT_CHUNK, + det_linear_available, + det_linear_backward_input_partials, + det_linear_backward_weight, + det_linear_fold_chunks, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.cuda.h3.ws2_comm import gather_rows +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.kernels.ops.pytorch.h3.adaln_projection import split_adaln_table + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[column_shard]+all_gather[rank_order]" +BACKWARD_IMPL = "local_dweight+all_gather_chunk_partials+ascending_fold" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input_partials" + "+all_gather+rl_engine._C.h3_det_linear_fold_chunks" +) + + +@dataclass(frozen=True) +class AdaLNColumnShard: + """The output columns a TP rank owns, and the WS1 objects they map to.""" + + tp: int + rank: int + n_total: int + begin: int + end: int + + @property + def hidden(self) -> int: + return self.n_total // (H3_MODALITY_NUM * H3_ADALN_CHUNKS) + + def slots(self) -> list[tuple[int, int, int, int]]: + """``(modality, chunk, h_begin, h_end)`` pieces covered by this shard.""" + + out, col = [], self.begin + while col < self.end: + slot, h0 = divmod(col, self.hidden) + h1 = min(self.hidden, h0 + self.end - col) + out.append((*divmod(slot, H3_ADALN_CHUNKS), h0, h1)) + col += h1 - h0 + return out + + +def adaln_column_shard(n_total: int, tp: int, rank: int) -> AdaLNColumnShard: + if tp < 1 or not 0 <= rank < tp: + raise ValueError(f"need 0 <= rank < tp, got rank={rank}, tp={tp}") + if n_total % (H3_MODALITY_NUM * H3_ADALN_CHUNKS): + raise ValueError(f"N={n_total} is not 3 modalities x 6 chunks x H") + if n_total % (tp * DINPUT_CHUNK): + # A shard boundary inside a d_input chunk would split one WS1 partial. + raise ValueError( + f"tp={tp} does not split N={n_total} on {DINPUT_CHUNK}-column chunk boundaries" + ) + width = n_total // tp + return AdaLNColumnShard(tp, rank, n_total, rank * width, (rank + 1) * width) + + +def shard_adaln_projection(weight: torch.Tensor, bias: torch.Tensor, tp: int, rank: int): + """This rank's rows of the full projection weight and bias (views).""" + + shard = adaln_column_shard(weight.shape[0], tp, rank) + return weight[shard.begin : shard.end], bias[shard.begin : shard.end] + + +def _gather_columns(collective, local: torch.Tensor) -> torch.Tensor: + """``(T, N / tp)`` column shards -> the ``(T, N)`` table, in rank order (a copy).""" + + num_t, width = local.shape + rows = gather_rows(collective, local) # (tp * T, N / tp), rank-major + return rows.view(-1, num_t, width).transpose(0, 1).reshape(num_t, -1) + + +def _validate_shard_inputs(temb, weight, bias, shard: AdaLNColumnShard) -> None: + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if not tensor.is_cuda or tensor.device != temb.device: + raise ValueError(f"{name} must be on this rank's CUDA device, got {tensor.device}") + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in (torch.bfloat16, torch.float32) or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share bfloat16 or float32, got {weight.dtype}/{bias.dtype}" + ) + width = shard.end - shard.begin + if weight.dim() != 2 or weight.shape[0] != width or bias.shape != (width,): + raise ValueError( + f"rank {shard.rank} of {shard.tp} owns {width} rows, got weight " + f"{tuple(weight.shape)} and bias {tuple(bias.shape)}" + ) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + + +class _TPAdaLNProjection(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight_shard, bias_shard, collective, shard): + temb = temb.contiguous() + act = F.silu(temb).to(weight_shard.dtype) + (local,) = det_linear_forward(act, weight_shard, bias_shard) + ctx.save_for_backward(temb, act, weight_shard) + ctx.collective, ctx.shard = collective, shard + return _gather_columns(collective, local) + + @staticmethod + def backward(ctx, grad_table): + temb, act, weight = ctx.saved_tensors + shard = ctx.shard + grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() + dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) + partial = det_linear_backward_input_partials(grad_local, weight) + d_act = det_linear_fold_chunks(gather_rows(ctx.collective, partial), torch.float32) + record_backward( + "tp_adaln_3mod", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return silu_backward_fp32(d_act, temb), dw, db, None, None + + +class H3TPAdaLNProjectionCudaOp: + """Column-parallel ``adaln_projection_3mod``, byte-equal to WS1 on every rank. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` + and ``backend_id`` (e.g. ``DeterministicCollective``). Inputs are the replicated + FP32 ``temb`` and this rank's weight/bias rows. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self, collective, n_total: int) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + self.collective = collective + self.shard = adaln_column_shard(n_total, collective.world_size, collective.rank) + + def __call__(self, temb, weight_shard, bias_shard): + return self.forward(temb, weight_shard, bias_shard) + + def forward(self, temb, weight_shard, bias_shard) -> tuple[torch.Tensor, ...]: + table = self.forward_table(temb, weight_shard, bias_shard) + return split_adaln_table(table, self.shard.hidden) + + def forward_table(self, temb, weight_shard, bias_shard) -> torch.Tensor: + _validate_shard_inputs(temb, weight_shard, bias_shard, self.shard) + return _TPAdaLNProjection.apply(temb, weight_shard, bias_shard, self.collective, self.shard) + + def readback(self) -> dict: + """Runtime identity of this rank's projection, for evidence and strict traces.""" + + return { + "op": "tp_adaln_3mod", + "kernel_id": KERNEL_ID, + "backward_kernel_id": BACKWARD_KERNEL_ID, + "contract": CONTRACT, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_ascending_chunk_fold_after_rank_ordered_gather", + "tp": self.shard.tp, + "rank": self.shard.rank, + "columns": [self.shard.begin, self.shard.end], + "slots": [list(s) for s in self.shard.slots()], + "fallback": None, + } diff --git a/rl_engine/kernels/ops/cuda/h3/ws2_comm.py b/rl_engine/kernels/ops/cuda/h3/ws2_comm.py new file mode 100644 index 000000000..f03597b7d --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/ws2_comm.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rank-ordered row gather shared by the H3 WS2 ops (a copy: no arithmetic).""" + +from __future__ import annotations + +import torch + +# DeterministicCollective sends all-gathers whose output is at most this many bytes +# through a single-block byte-copy kernel (csrc/cuda/distributed/ +# deterministic_collective.cu, kSingleBlockFastPathMaxBytes) that costs about 4 us per +# KiB on B200; larger outputs take the multi-block path (~45 us). Padding the shard +# past the threshold is cheaper for every non-trivial size. +_SINGLE_BLOCK_GATHER_MAX_BYTES = 256 * 1024 + + +def gather_rows(collective, shard: torch.Tensor) -> torch.Tensor: + """``(world * rows, ...)``: every rank's ``(rows, ...)`` shard, in rank order.""" + + shard = shard.contiguous() + rows, world = shard.shape[0], collective.world_size + row_bytes = shard[0].numel() * shard.element_size() if rows else 0 + if world == 1: + return shard + if row_bytes == 0: + return collective.all_gather(shard) + padded = max(rows, _SINGLE_BLOCK_GATHER_MAX_BYTES // (world * row_bytes) + 1) + if padded == rows: + return collective.all_gather(shard) + staged = torch.cat([shard, shard.new_zeros((padded - rows, *shard.shape[1:]))]) + out = collective.all_gather(staged).view(world, padded, *shard.shape[1:]) + return out[:, :rows].reshape(world * rows, *shard.shape[1:]) diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py new file mode 100644 index 000000000..a0dc116f6 --- /dev/null +++ b/rl_engine/testing/h3_ws2.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS1 vs WS2 comparisons for the H3 conditioning path (tests and evidence). + +The per-rank helpers run inside real multi-process groups (tests and +``scripts/h3_ws2_evidence.py``); the WS1 side runs on one GPU. +""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def ws1_projection(temb, weight, bias, grad) -> dict[str, torch.Tensor]: + """The single-GPU ``adaln_projection_3mod`` table and its three gradients.""" + + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + table = H3AdaLNProjectionCudaOp().forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight": leaves[1].grad, + "d_bias": leaves[2].grad, + } + + +def tp_projection_rank(collective, temb, weight, bias, grad) -> dict[str, Any]: + """One TP rank: shard the full weight, run the TP op, return its view of the results.""" + + from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + op = H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]) + w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, w_shard, b_shard)] + table = op.forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight_shard": leaves[1].grad, + "d_bias_shard": leaves[2].grad, + "readback": op.readback(), + } + + +def tp_matches_ws1(ws1: dict[str, torch.Tensor], ranks: list[dict[str, Any]]) -> dict[str, bool]: + """Byte equality of every rank's replicated outputs and of the concatenated shards.""" + + return { + "table": all(torch.equal(r["table"], ws1["table"]) for r in ranks), + "d_temb": all(torch.equal(r["d_temb"], ws1["d_temb"]) for r in ranks), + "d_weight": torch.equal(torch.cat([r["d_weight_shard"] for r in ranks]), ws1["d_weight"]), + "d_bias": torch.equal(torch.cat([r["d_bias_shard"] for r in ranks]), ws1["d_bias"]), + } + + +def tp_projection_sweep(collective, temb, weight, bias, grad, iters: int = 100) -> dict[str, Any]: + """Evidence for one rank: per-T byte equality against WS1 on this GPU, timing, readback. + + ``temb``/``grad`` hold the largest T; every T in ``1..T`` uses their first + rows. Each rank compares its own replicated outputs and its own shard of + ``dW``/``db`` with the WS1 result computed here, so nothing large is shipped. + """ + + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + adaln_column_shard, + ) + from rl_engine.testing.h3_report import time_us + + shard = adaln_column_shard(weight.shape[0], collective.world_size, collective.rank) + cols = slice(shard.begin, shard.end) + tp_op, ws1_op = ( + H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]), + H3AdaLNProjectionCudaOp(), + ) + equality, perf = [], [] + for num_t in range(1, temb.shape[0] + 1): + t, g = temb[:num_t], grad[:num_t] + ws1 = ws1_projection(t, weight, bias, g) + mine = tp_projection_rank(collective, t, weight, bias, g) + equality.append( + { + "num_timesteps": num_t, + "table": torch.equal(mine["table"], ws1["table"]), + "d_temb": torch.equal(mine["d_temb"], ws1["d_temb"]), + "d_weight_shard": torch.equal(mine["d_weight_shard"], ws1["d_weight"][cols]), + "d_bias_shard": torch.equal(mine["d_bias_shard"], ws1["d_bias"][cols]), + } + ) + leaves = [x.detach().clone().requires_grad_(True) for x in (t, weight[cols], bias[cols])] + full = [x.detach().clone().requires_grad_(True) for x in (t, weight, bias)] + + def step(op, args): + for x in args: + x.grad = None + op.forward_table(*args).backward(g) + + perf.append( + { + "num_timesteps": num_t, + "tp_forward_us": time_us( + lambda: tp_op.forward_table(t, weight[cols], bias[cols]), 30, iters + ), + "tp_forward_backward_us": time_us(lambda: step(tp_op, leaves), 10, iters), + "ws1_forward_us": time_us(lambda: ws1_op.forward_table(t, weight, bias), 30, iters), + "ws1_forward_backward_us": time_us(lambda: step(ws1_op, full), 10, iters), + } + ) + return {"equality": equality, "perf": perf, "readback": tp_op.readback()} + + +# --- real multi-process groups ------------------------------------------------- + + +def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): + import traceback + + import torch.distributed as dist + + from rl_engine.distributed.collectives import DeterministicCollective + + try: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", init_method=init_method, rank=rank, world_size=world, device_id=rank + ) + inputs = {k: v.cuda() for k, v in torch.load(inputs_path).items()} + collective = DeterministicCollective(device=rank, max_size_bytes=256 << 20) + out = target(collective, **inputs, **kwargs) + # Results go through a file: tensors sent over the queue are shared by + # file descriptor and vanish if this process exits before the parent reads. + torch.save( + {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, + inputs_path.with_name(f"rank{rank}.pt"), + ) + queue.put({"rank": rank}) + collective.close() + dist.destroy_process_group() + except Exception: # forwarded to the parent + queue.put({"rank": rank, "error": traceback.format_exc()}) + + +def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: float = 900, **kwargs): + """Run ``target(collective, **inputs, **kwargs)`` on ranks ``0..world-1`` (one GPU each). + + ``target`` must be importable (it is pickled into spawned processes). Each + rank gets a NCCL group and a ``DeterministicCollective``; results come back + on the CPU, ordered by rank. Any rank's exception is re-raised here. + """ + + import tempfile + from pathlib import Path + + import torch.multiprocessing as mp + + if torch.cuda.device_count() < world: + raise RuntimeError(f"needs {world} GPUs, found {torch.cuda.device_count()}") + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "inputs.pt" + torch.save({k: v.cpu() for k, v in inputs.items()}, path) + queue = ctx.Queue() + init = (Path(tmp) / "init").as_uri() + procs = [ + ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) + for r in range(world) + ] + for proc in procs: + proc.start() + status = sorted((queue.get(timeout=timeout) for _ in procs), key=lambda r: r["rank"]) + for proc in procs: + proc.join(timeout=120) + for result in status: + if "error" in result: + raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") + return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] diff --git a/scripts/h3_ws2_evidence.py b/scripts/h3_ws2_evidence.py new file mode 100644 index 000000000..c27332f77 --- /dev/null +++ b/scripts/h3_ws2_evidence.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write an RFC #420 WS2 evidence report (byte equality vs WS1, timing, readback). + +Each TP size runs as real NCCL processes, one GPU per rank, with the +deterministic CUDA collective. Run it on a clean tree on one 8-GPU node; +``scripts/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python scripts/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \\ + --out docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from rl_engine.testing.h3_chain import environment, git_state # noqa: E402 +from rl_engine.testing.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) +from rl_engine.testing.h3_ws2 import run_world, tp_projection_sweep # noqa: E402 + +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _tp_adaln(worlds: list[int], max_t: int) -> dict: + weights = load_h3_conditioning_weights("cpu") + weight, bias = weights[WEIGHT], weights[BIAS] + g = torch.Generator(device="cpu").manual_seed(0) + inputs = { + "temb": torch.randn(max_t, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(max_t, weight.shape[0], generator=g).bfloat16(), + } + runs = [] + for world in worlds: + ranks = run_world(world, tp_projection_sweep, inputs) + runs.append({"tp": world, "ranks": ranks}) + equal = all( + v for r in ranks for e in r["equality"] for k, v in e.items() if k != "num_timesteps" + ) + print(f"tp={world}: byte-equal to WS1 on every rank and T: {equal}") + return {"weights": [WEIGHT, BIAS], "runs": runs} + + +OPS = {"tp_adaln_3mod": _tp_adaln} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--worlds", default="1,2,4,8") + parser.add_argument("--max-timesteps", type=int, default=4) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + worlds = [int(w) for w in args.worlds.split(",")] + if torch.cuda.device_count() < max(worlds): + raise SystemExit(f"needs {max(worlds)} GPUs, found {torch.cuda.device_count()}") + + manifest = load_h3_manifest() + report = { + "kind": "h3_ws2_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + **git_state(), + "environment": {**environment(), "gpus": torch.cuda.device_count()}, + **OPS[args.op](worlds, args.max_timesteps), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + commit, dirty = report["rl_kernel_commit"][:7], report["tracked_tree_dirty"] + print(f"wrote {args.out} (commit {commit}, dirty={dirty})") + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 3159c7379..ba4b3b4ef 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -2,7 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Render an RFC #420 operator report (``scripts/h3_evidence.py``) as a PNG. +"""Render an RFC #420 operator or WS2 report (``scripts/h3_evidence.py``, +``scripts/h3_ws2_evidence.py``) as a PNG. Needs only the report JSON and matplotlib (not torch or a GPU), so it can run anywhere: @@ -22,6 +23,7 @@ import matplotlib matplotlib.use("Agg") +import matplotlib.patches # noqa: E402 import matplotlib.pyplot as plt # noqa: E402 # Fixed series identity across every H3 figure (categorical slots 1-3). @@ -496,6 +498,97 @@ def plot_final(report: dict[str, Any]): return fig +TP_TENSORS = ( + ("table", "table (6 x 3T rows)"), + ("d_temb", "d_temb"), + ("d_weight_shard", "dW shard"), + ("d_bias_shard", "db shard"), +) + + +def plot_tp_adaln(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "tp_adaln_3mod (column-parallel AdaLN projection, 2688 -> 96768, NCCL)" + ) + runs = report["runs"] + # Left: every rank, every T, every tensor byte-equal to WS1? + left.set_facecolor(SURFACE) + left.set_title( + "Byte-equal to WS1 (every rank, every T)", loc="left", fontsize=11, color=INK, pad=10 + ) + for y, run in enumerate(runs): + checks = [e for r in run["ranks"] for e in r["equality"]] + for x, (key, _) in enumerate(TP_TENSORS): + ok = sum(bool(e[key]) for e in checks) + good = ok == len(checks) + left.add_patch( + matplotlib.patches.FancyBboxPatch( + (x + 0.06, y + 0.08), + 0.88, + 0.84, + boxstyle="round,pad=0,rounding_size=0.06", + facecolor=RL_KERNEL if good else PROVIDER, + edgecolor=SURFACE, + linewidth=2, + ) + ) + left.text( + x + 0.5, + y + 0.5, + f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} rank x T", + ha="center", + va="center", + fontsize=8, + color="white", + ) + left.set_xlim(0, len(TP_TENSORS)) + left.set_ylim(len(runs), 0) + left.set_xticks([i + 0.5 for i in range(len(TP_TENSORS))], [label for _, label in TP_TENSORS]) + left.set_yticks([i + 0.5 for i in range(len(runs))], [f"TP{run['tp']}" for run in runs]) + left.tick_params(colors=MUTED, labelsize=8, length=0) + for side in left.spines.values(): + side.set_visible(False) + # Right: slowest rank's time at the largest T; WS1 on one GPU as reference lines. + num_t = max(p["num_timesteps"] for p in runs[0]["ranks"][0]["perf"]) + + def worst(run, key): + perf = [p for r in run["ranks"] for p in r["perf"] if p["num_timesteps"] == num_t] + return max(p[key] for p in perf) / 1e3 + + tp_runs = [run for run in runs if run["tp"] > 1] + _style(right, f"Time per call, T = {num_t} (slowest rank; lower is better)", "ms") + _bars( + right, + [f"TP{run['tp']}" for run in tp_runs], + [ + ("TP fwd", RL_KERNEL, [worst(run, "tp_forward_us") for run in tp_runs]), + ( + "TP fwd+bwd", + RL_KERNEL, + [worst(run, "tp_forward_backward_us") for run in tp_runs], + "////", + ), + ], + "{:.2f}", + log=True, + ) + for key, label, style in ( + ("ws1_forward_us", "WS1 fwd", ":"), + ("ws1_forward_backward_us", "WS1 fwd+bwd", "--"), + ): + value = worst(runs[0], key) + right.axhline( + value, + color=PROVIDER, + linestyle=style, + linewidth=1.4, + zorder=1, + label=f"{label}, 1 GPU: {value:.2f}", + ) + right.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, @@ -504,6 +597,7 @@ def plot_final(report: dict[str, Any]): "h3_rmsnorm": plot_rmsnorm, "adaln_gate_residual": plot_gate_residual, "final_adaln_out": plot_final, + "tp_adaln_3mod": plot_tp_adaln, } diff --git a/tests/h3/test_h3_tp_adaln.py b/tests/h3/test_h3_tp_adaln.py new file mode 100644 index 000000000..8b341c8d3 --- /dev/null +++ b/tests/h3/test_h3_tp_adaln.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``tp_adaln_3mod``: column-parallel AdaLN projection, byte-equal to WS1. + +* ownership: contiguous 64-aligned column shards cover the 3 x 6 x H table + exactly once; unsupported TP sizes and wrong shards fail closed; +* shard arithmetic (one GPU): a rank's forward columns and ``d_input`` chunk + partials are bitwise the matching slice of the full WS1 call, and folding + the rank-ordered partials reproduces WS1's ``d_input``; +* real NCCL (2/4/8 GPUs): every rank's table and ``d_temb`` and the + concatenated ``dW``/``db`` shards are bitwise equal to WS1. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import adaln_column_shard +from rl_engine.testing.h3_ws2 import run_world, tp_matches_ws1, tp_projection_rank, ws1_projection + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" +N_H3 = 3 * 6 * 5376 +ALL_TRUE = dict.fromkeys(("table", "d_temb", "d_weight", "d_bias"), True) + + +def _require_ext(): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + + +def _inputs(num_t, n_out, k, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + temb = (torch.randn(num_t, k, generator=g) * 2).cuda() + weight = (torch.randn(n_out, k, generator=g) / k**0.5).to(dtype).cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + grad = torch.randn(num_t, n_out, generator=g).to(dtype).cuda() + return temb, weight, bias, grad + + +class TestOwnership: + @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8]) + def test_shards_tile_the_table_once(self, tp): + shards = [adaln_column_shard(N_H3, tp, r) for r in range(tp)] + assert [s.begin for s in shards] == [0, *[s.end for s in shards[:-1]]] + assert shards[-1].end == N_H3 + assert sum(h1 - h0 for s in shards for (_, _, h0, h1) in s.slots()) == N_H3 + + def test_tp2_slots(self): + # rank 0: video's six chunks and text's first three; rank 1: the rest. + assert adaln_column_shard(N_H3, 2, 0).slots() == [ + *[(0, c, 0, 5376) for c in range(6)], + *[(1, c, 0, 5376) for c in range(3)], + ] + + @pytest.mark.parametrize("tp, rank", [(5, 0), (16, 0), (0, 0), (2, 2), (2, -1)]) + def test_rejects_unsupported(self, tp, rank): + with pytest.raises(ValueError): + adaln_column_shard(N_H3, tp, rank) + + +@requires_cuda +class TestShardArithmetic: + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_pinned_shards_reproduce_ws1(self, h3_weights_cpu, tp): + _require_ext() + from rl_engine.kernels.ops.cuda.h3 import det_linear as dl + + weight, bias = h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + g = torch.Generator(device="cpu").manual_seed(tp) + act = F.silu(torch.randn(3, weight.shape[1], generator=g).cuda() * 2).bfloat16() + grad = torch.randn(3, weight.shape[0], generator=g).cuda() + (full,) = dl.det_linear_forward(act, weight, bias) + full_partials = dl.det_linear_backward_input_partials(grad, weight) + pieces = [] + for rank in range(tp): + s = adaln_column_shard(weight.shape[0], tp, rank) + (cols,) = dl.det_linear_forward(act, weight[s.begin : s.end], bias[s.begin : s.end]) + assert torch.equal(cols, full[:, s.begin : s.end]) + part = dl.det_linear_backward_input_partials( + grad[:, s.begin : s.end].contiguous(), weight[s.begin : s.end] + ) + c0 = s.begin // dl.DINPUT_CHUNK + assert torch.equal(part, full_partials[c0 : c0 + part.shape[0]]) + pieces.append(part) + folded = dl.det_linear_fold_chunks(torch.cat(pieces), torch.float32) + assert torch.equal(folded, dl.det_linear_backward_input(grad, weight, torch.float32)) + + def test_tp1_op_is_ws1(self): + _require_ext() + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(3, 3 * 6 * 256, 136) + ws1 = ws1_projection(temb, weight, bias, grad) + assert tp_matches_ws1(ws1, [tp_projection_rank(one, temb, weight, bias, grad)]) == ALL_TRUE + + def test_fails_closed(self): + _require_ext() + from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + temb, weight, bias, _ = _inputs(2, 3 * 6 * 256, 136) + rank1 = SimpleNamespace(rank=1, world_size=2, backend_id="unused", all_gather=None) + op = H3TPAdaLNProjectionCudaOp(rank1, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 2, 1) + with pytest.raises(TypeError): # RFC probe H7: cast before the SiLU + op(temb.bfloat16(), w, b) + with pytest.raises(ValueError): # the full weight instead of this rank's rows + op(temb, weight, bias) + with pytest.raises(ValueError): + op(temb.cpu(), w, b) + with pytest.raises(ValueError): # 4608 / 16 is not a multiple of 64 + H3TPAdaLNProjectionCudaOp(SimpleNamespace(rank=0, world_size=16), weight.shape[0]) + assert op.readback()["columns"] == [2304, 4608] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_byte_equal_to_ws1(h3_weights_cpu, world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + weight, bias = h3_weights_cpu[WEIGHT], h3_weights_cpu[BIAS] + g = torch.Generator(device="cpu").manual_seed(world) + inputs = { + "temb": torch.randn(3, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(3, weight.shape[0], generator=g).bfloat16(), + } + ws1 = { + k: v.cpu() for k, v in ws1_projection(**{k: v.cuda() for k, v in inputs.items()}).items() + } + torch.cuda.empty_cache() + ranks = run_world(world, tp_projection_rank, inputs) + assert tp_matches_ws1(ws1, ranks) == ALL_TRUE + assert [r["readback"]["rank"] for r in ranks] == list(range(world)) + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} From ecde5bb76a19bdc47457b67dad19e5423225fe55 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 10:58:54 +0000 Subject: [PATCH 32/89] [WS2][CUDA][MiniMax-H3] sp_norm_adaln: sequence-parallel norm/AdaLN/gated residual, byte-equal to WS1 RFC #420 WS2 row `sp_norm_adaln` (stacked on tp_adaln_3mod). - Ownership: rank r holds packed positions [r*S//sp, (r+1)*S//sp) of every batch item with the matching metadata; any S >= sp. The (S,) row index is replicated and every rank derives the same reduction plan from it. - Forward, dx, d_sublayer and d_residual: the WS1 kernels on local rows. - d_norm_weight, d_shift/d_scale, d_gate keep the WS1 two-level order (256-element tiles, then an ascending fold): each tile is computed by the rank holding its first row, after all-gathering the tile's rows from other ranks (only tiles crossing a shard boundary move rows); the partials are all-gathered into WS1 tile order and every rank runs the WS1 fold. Rank-ordered all-gathers only; no collective arithmetic. - New bindings around the unchanged WS1 kernels: h3_rmsnorm_backward_{dx,partials}, h3_rmsnorm_fold_partials, h3_gate_residual_backward_dy, h3_gate_grad_{partials,fold}; the WS1 dweight tile kernel also takes explicit row lists (same arithmetic). - The plan is built once per index object and shared by every norm and gated residual that uses it. - Tests: layout, per-rank backward vs WS1 for SP 2-8 on one GPU (odd S, B = 2, shards smaller than a tile, block and interleaved packing, bf16/fp32), fail-closed inputs, and the autograd block region over real NCCL SP2/4/8. - h3_block_layout (H3's block packing), evidence sweep with rows exchanged, timing and a naive-SP comparison; figure; operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 80 ++++ csrc/cuda/h3/rmsnorm_modulate.cu | 165 ++++++- csrc/ops.cpp | 42 ++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-sp-norm-adaln.md | 96 +++++ rl_engine/_C.pyi | 43 ++ .../kernels/ops/cuda/h3/sp_norm_adaln.py | 404 ++++++++++++++++++ rl_engine/testing/h3_cases.py | 24 ++ rl_engine/testing/h3_ws2.py | 181 ++++++++ scripts/h3_ws2_evidence.py | 42 +- scripts/plot_h3_evidence.py | 183 ++++++-- tests/h3/test_h3_sp_norm_adaln.py | 210 +++++++++ 13 files changed, 1428 insertions(+), 44 deletions(-) create mode 100644 docs/operators/h3-sp-norm-adaln.md create mode 100644 rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py create mode 100644 tests/h3/test_h3_sp_norm_adaln.py diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 18ede9e72..9a03660f1 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -16,6 +16,8 @@ // 16-bit values is exact in FP32, so this equals the eager VJP bitwise); // d_gate[i] = sum over positions mapped to i of grad * y, an FP32 segmented // sum over positions sorted stably by row, in fixed tiles folded in order. +// The tile partials and the fold are also bound separately for sequence +// parallel callers (each tile computed where its rows live, folded in order). #include #include @@ -240,3 +242,81 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: C10_CUDA_KERNEL_LAUNCH_CHECK(); return {dy, dgate}; } + +// Sequence parallel: the WS1 d_gate tile partials over explicit row lists of +// (grad, y) rows -> (tiles, N) FP32. +torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch::Tensor rows, + torch::Tensor tile_begin, torch::Tensor tile_end) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + check_act(y, "y", grad); + for (const auto* t : {&rows, &tile_begin, &tile_end}) { + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors on grad's device"); + } + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin and tile_end differ in length"); + const int64_t n = grad.size(1); + const int64_t tiles = tile_begin.numel(); + const c10::cuda::CUDAGuard guard(grad.device()); + auto partial = torch::empty({tiles, n}, grad.options().dtype(at::kFloat)); + if (tiles == 0) return partial; + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(grad.scalar_type(), "h3_gate_grad_partials", [&] { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(y.data_ptr()), + rows.data_ptr(), tile_begin.data_ptr(), tile_end.data_ptr(), + partial.data_ptr(), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return partial; +} + +// The WS1 d_gate fold: each segment's tiles in ascending order, cast once to dtype. +torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_tile, + c10::ScalarType dtype) { + TORCH_CHECK(partial.is_cuda() && partial.dim() == 2 && partial.is_contiguous() && + partial.scalar_type() == at::kFloat, + "partial must be a contiguous float32 (tiles, N) CUDA tensor"); + TORCH_CHECK(seg_first_tile.is_cuda() && seg_first_tile.device() == partial.device() && + seg_first_tile.scalar_type() == at::kLong && seg_first_tile.dim() == 1 && + seg_first_tile.is_contiguous(), + "seg_first_tile must be a contiguous int64 CUDA tensor on partial's device"); + const int64_t n = partial.size(1); + const int64_t segments = seg_first_tile.numel() - 1; + const c10::cuda::CUDAGuard guard(partial.device()); + auto out = torch::empty({segments, n}, partial.options().dtype(dtype)); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(dtype, "h3_gate_grad_fold", [&] { + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(out.data_ptr()), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// Row-local part of the WS1 backward only: d_y = cast(grad * gate[index]). +torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate, + torch::Tensor index) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + const Gate g = make_gate(gate, index, grad); + const int64_t rows = grad.size(0); + const int64_t n = grad.size(1); + const c10::cuda::CUDAGuard guard(grad.device()); + auto dy = torch::empty_like(grad); + H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward_dy", [&] { + const bool vec = can_vectorize(gate, n, {&grad, &dy}); + gate_residual_rows_kernel<<>>( + reinterpret_cast(grad.data_ptr()), nullptr, reinterpret_cast(dy.data_ptr()), + rows, n, g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dy; +} diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index d42fe3525..0c7a1bbfe 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -20,6 +20,10 @@ // dx = rstd * w * d_n - x * rstd^3 * sum(w * d_n * x) / N (row-local) // dweight = sum_r d_n * x * rstd over fixed 256-row tiles, tiles folded in order // dshift = sum_{r -> i} g, dscale = sum_{r -> i} g * n (sorted tiles) +// +// The tile partials and the folds are also bound separately (`*_partials`, +// `*_fold_*`) so a sequence-parallel caller can compute each WS1 tile on the +// rank that holds its rows and fold the gathered partials in WS1 order. #include #include @@ -180,18 +184,25 @@ __global__ void __launch_bounds__(32 * kWarps) } } -// partial[tile, j] = sum over rows of the tile (ascending) of d_n * x * rstd +// partial[tile, j] = sum over rows of the tile (ascending) of d_n * x * rstd. +// Tiles are the fixed 256-row blocks, or (sequence parallel) explicit row lists +// rows[tile_begin[tile] .. tile_end[tile]) that hold the same rows in the same order. template __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* __restrict__ x, const float* __restrict__ rstd, float* __restrict__ partial, int64_t rows, int n, - Modulation mod) { + Modulation mod, + const int64_t* __restrict__ row_list = nullptr, + const int64_t* __restrict__ tile_begin = nullptr, + const int64_t* __restrict__ tile_end = nullptr) { const int j = blockIdx.y * blockDim.x + threadIdx.x; if (j >= n) return; - const int64_t r0 = static_cast(blockIdx.x) * kRowTile; - const int64_t r1 = min(r0 + kRowTile, rows); + const int64_t tile = blockIdx.x; + const int64_t p0 = row_list != nullptr ? tile_begin[tile] : tile * kRowTile; + const int64_t p1 = row_list != nullptr ? tile_end[tile] : min(p0 + kRowTile, rows); float acc = 0.0f; - for (int64_t r = r0; r < r1; ++r) { + for (int64_t p = p0; p < p1; ++p) { + const int64_t r = row_list != nullptr ? row_list[p] : p; const T* scale = mod.index != nullptr ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride : nullptr; @@ -403,3 +414,147 @@ std::vector h3_rmsnorm_backward( C10_CUDA_KERNEL_LAUNCH_CHECK(); return {dx, dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; } + +namespace { + +void check_tiles(const torch::Tensor& rows, const torch::Tensor& begin, const torch::Tensor& end, + const torch::Tensor& like) { + for (const auto* t : {&rows, &begin, &end}) { + TORCH_CHECK(t->is_cuda() && t->device() == like.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors on x's device"); + } + TORCH_CHECK(begin.numel() == end.numel(), "tile_begin and tile_end differ in length"); +} + +} // namespace + +// Sequence parallel: the dweight and (with modulation) table-gradient tile +// partials of the WS1 backward, for explicit tiles over the rows of x. `index` +// gives each row's table row (one entry per row of x). Returns +// {dweight_partial (tiles, N)} or {dweight_partial, table_partial (seg_tiles, 2N)}. +std::vector h3_rmsnorm_backward_partials( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, torch::Tensor dw_rows, torch::Tensor dw_begin, + torch::Tensor dw_end, c10::optional seg_rows, + c10::optional seg_begin, c10::optional seg_end) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics on x's device"); + check_tiles(dw_rows, dw_begin, dw_end, x); + const int64_t rows = x.size(0); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + TORCH_CHECK(mod.index == nullptr || mod.seq == rows, "index needs one entry per row of x"); + const c10::cuda::CUDAGuard guard(x.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + const int64_t tiles = dw_begin.numel(); + auto dw_partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + if (tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_partials", [&] { + rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, + 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + rstd.data_ptr(), dw_partial.data_ptr(), rows, static_cast(n), mod, + dw_rows.data_ptr(), dw_begin.data_ptr(), dw_end.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + if (mod.index == nullptr) return {dw_partial}; + + TORCH_CHECK(seg_rows.has_value() && seg_begin.has_value() && seg_end.has_value(), + "modulated partials need the segment tiles"); + check_tiles(*seg_rows, *seg_begin, *seg_end, x); + const int64_t seg_tiles = seg_begin->numel(); + const int64_t width = 2 * n; + auto seg_partial = torch::empty({seg_tiles, width}, x.options().dtype(at::kFloat)); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + if (seg_tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_modulation_partials", [&] { + modulation_grad_partial_kernel + <<(seg_tiles), wide_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + seg_rows->data_ptr(), seg_begin->data_ptr(), + seg_end->data_ptr(), seg_partial.data_ptr(), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + return {dw_partial, seg_partial}; +} + +// The WS1 folds: dweight = ascending fold of all tiles, cast to weight's dtype; +// table rows = ascending fold of each segment's tiles, FP32 (R, 2N). +std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight, + c10::optional seg_partial, + c10::optional seg_first_tile) { + TORCH_CHECK(dw_partial.is_cuda() && dw_partial.dim() == 2 && dw_partial.is_contiguous() && + dw_partial.scalar_type() == at::kFloat, + "dw_partial must be a contiguous float32 (tiles, N) CUDA tensor"); + TORCH_CHECK(weight.is_cuda() && weight.device() == dw_partial.device() && weight.dim() == 1 && + weight.size(0) == dw_partial.size(1), + "weight must be (N,) on dw_partial's device"); + const int64_t n = dw_partial.size(1); + const c10::cuda::CUDAGuard guard(dw_partial.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + auto dweight = torch::empty_like(weight); + H3_DISPATCH(weight.scalar_type(), "h3_rmsnorm_fold_partials", [&] { + fold_tiles_kernel<<>>( + dw_partial.data_ptr(), reinterpret_cast(dweight.data_ptr()), + dw_partial.size(0), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!seg_partial.has_value()) return {dweight}; + + TORCH_CHECK(seg_first_tile.has_value() && seg_first_tile->is_cuda() && + seg_first_tile->device() == dw_partial.device() && + seg_first_tile->scalar_type() == at::kLong && seg_first_tile->is_contiguous(), + "segment fold needs seg_first_tile (int64, on dw_partial's device)"); + TORCH_CHECK(seg_partial->is_cuda() && seg_partial->device() == dw_partial.device() && + seg_partial->dim() == 2 && seg_partial->is_contiguous() && + seg_partial->scalar_type() == at::kFloat && seg_partial->size(1) == 2 * n, + "seg_partial must be a contiguous float32 (tiles, 2N) CUDA tensor"); + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t width = 2 * n; + auto table_grad = torch::empty({segments, width}, dw_partial.options()); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + fold_segments_kernel<<(segments), wide_blocks), threads, 0, stream>>>( + seg_partial->data_ptr(), seg_first_tile->data_ptr(), + table_grad.data_ptr(), static_cast(width)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; +} + +// Row-local part of the WS1 backward only: dx (sequence-parallel callers +// compute the cross-row reductions with h3_rmsnorm_backward_partials). +torch::Tensor h3_rmsnorm_backward_dx(torch::Tensor grad, torch::Tensor x, torch::Tensor weight, + torch::Tensor rstd, c10::optional shift, + c10::optional scale, + c10::optional index) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics on x's device"); + const Modulation mod = make_modulation(shift, scale, index, x, x.size(1)); + const c10::cuda::CUDAGuard guard(x.device()); + auto dx = torch::empty_like(x); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_dx", [&] { + rmsnorm_modulate_dx_kernel<<(x.size(0)), dim3(32, kWarps), 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + reinterpret_cast(dx.data_ptr()), static_cast(x.size(1)), mod); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dx; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index f3dc5cca2..342d09516 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -537,6 +537,25 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: torch::Tensor tile_begin, torch::Tensor tile_end, torch::Tensor seg_first_tile); +std::vector h3_rmsnorm_backward_partials( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, torch::Tensor dw_rows, torch::Tensor dw_begin, + torch::Tensor dw_end, c10::optional seg_rows, + c10::optional seg_begin, c10::optional seg_end); +std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight, + c10::optional seg_partial, + c10::optional seg_first_tile); +torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch::Tensor rows, + torch::Tensor tile_begin, torch::Tensor tile_end); +torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_tile, + c10::ScalarType dtype); +torch::Tensor h3_rmsnorm_backward_dx(torch::Tensor grad, torch::Tensor x, torch::Tensor weight, + torch::Tensor rstd, c10::optional shift, + c10::optional scale, + c10::optional index); +torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate, + torch::Tensor index); #endif // PyBind11 Module Registration @@ -843,5 +862,28 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Gated-residual backward: exact dy, deterministic sorted segment sum for dgate", py::arg("grad"), py::arg("y"), py::arg("gate"), py::arg("index"), py::arg("sorted_pos"), py::arg("tile_begin"), py::arg("tile_end"), py::arg("seg_first_tile")); + m.def("h3_rmsnorm_backward_partials", &h3_rmsnorm_backward_partials, + "WS1 RMSNorm dweight / table-gradient tile partials over explicit row lists (SP)", + py::arg("grad"), py::arg("x"), py::arg("weight"), py::arg("rstd"), py::arg("shift"), + py::arg("scale"), py::arg("index"), py::arg("dw_rows"), py::arg("dw_begin"), + py::arg("dw_end"), py::arg("seg_rows") = py::none(), py::arg("seg_begin") = py::none(), + py::arg("seg_end") = py::none()); + m.def("h3_rmsnorm_fold_partials", &h3_rmsnorm_fold_partials, + "WS1 ascending folds of RMSNorm dweight and per-segment table-gradient partials", + py::arg("dw_partial"), py::arg("weight"), py::arg("seg_partial") = py::none(), + py::arg("seg_first_tile") = py::none()); + m.def("h3_gate_grad_partials", &h3_gate_grad_partials, + "WS1 d_gate tile partials over explicit row lists (SP)", py::arg("grad"), py::arg("y"), + py::arg("rows"), py::arg("tile_begin"), py::arg("tile_end")); + m.def("h3_gate_grad_fold", &h3_gate_grad_fold, + "WS1 per-segment ascending fold of d_gate partials, cast once", py::arg("partial"), + py::arg("seg_first_tile"), py::arg("dtype")); + m.def("h3_rmsnorm_backward_dx", &h3_rmsnorm_backward_dx, + "Row-local dx of the WS1 RMSNorm(+modulation) backward", py::arg("grad"), py::arg("x"), + py::arg("weight"), py::arg("rstd"), py::arg("shift") = py::none(), + py::arg("scale") = py::none(), py::arg("index") = py::none()); + m.def("h3_gate_residual_backward_dy", &h3_gate_residual_backward_dy, + "Row-local d_y = grad * gate[index] of the WS1 gated-residual backward", + py::arg("grad"), py::arg("gate"), py::arg("index")); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index 48f12fe3b..cfc833fe2 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -38,6 +38,7 @@ nav: - operators/h3-adaln-gate-residual.md - operators/h3-final-adaln-out.md - operators/h3-tp-adaln-3mod.md + - operators/h3-sp-norm-adaln.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 99e02cf0c..065d754b6 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -40,4 +40,5 @@ Every operator page should include: - [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) - [MiniMax-H3 Final AdaLN Output](h3-final-adaln-out.md) - [MiniMax-H3 Tensor-Parallel AdaLN Projection](h3-tp-adaln-3mod.md) +- [MiniMax-H3 Sequence-Parallel Norm and AdaLN](h3-sp-norm-adaln.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md new file mode 100644 index 000000000..167b63f7c --- /dev/null +++ b/docs/operators/h3-sp-norm-adaln.md @@ -0,0 +1,96 @@ +# MiniMax-H3 Sequence-Parallel Norm, AdaLN Modulation and Gated Residual + +## Summary + +`sp_norm_adaln` runs the row-wise part of an H3 block with the packed sequence split across +sequence-parallel ranks (RFC #420, WS2). It covers the [RMSNorm and AdaLN modulation](h3-rmsnorm.md) +and the [gated residual](h3-adaln-gate-residual.md). Every rank's output and row gradients are the +**WS1 bytes of its rows**. The gradients of the replicated tensors (`d_norm_weight`, `d_shift`, +`d_scale`, `d_gate`) are **the WS1 bytes, identical on every rank**. + +## Entry Point + +```python +from rl_engine.distributed.collectives import DeterministicCollective +from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + +collective = DeterministicCollective(group=sp_group, device=local_rank) +op = H3SPNormAdaLNCudaOp(collective, seq_len=S, batch=B) +rows = slice(op.layout.lo, op.layout.hi) # this rank's packed positions +hidden = op.gate_residual(residual[:, rows], attn_out_local, gate_msa, adaln_indices) +normed = op.norm_modulated(hidden, norm2_weight, shift_mlp, scale_mlp, adaln_indices) +final = op.norm_modulated(h, norm_out_weight, shift, scale, timestep_indices) # norm_out +op.readback() # sp, rank, positions, collective backend, reduction order, fallback +``` + +The table views, the norm weight and the full `(S,)` row index are replicated. Activations are +the rank's `(B, S_local, H)` rows. + +## Ownership + +Rank `r` of `sp` holds packed positions `[r * S // sp, (r + 1) * S // sp)` of every batch item, +together with the matching rows of the metadata. Any `S >= sp` works, including `S` not divisible +by `sp`. Shapes that do not match the rank's slice, a sharded index, an out-of-range index, CPU +tensors and malformed tables all fail closed. + +## Numerics + +- **Forward and row gradients.** `out`, `dx`, `d_sublayer` and `d_residual` are the WS1 kernels on + the local rows, with the local slice of the index. Rows are independent, so they are byte-equal + to WS1. +- **Cross-row gradients.** In WS1, `d_norm_weight` sums `d_n * x * rstd` over fixed 256-row tiles + of the `B x S` rows. `d_shift`/`d_scale`/`d_gate` sum over 256-element tiles of each table row's + positions, sorted stably. In both cases the tile partials are then folded in ascending order. SP + keeps exactly that two-level order: + 1. Every rank builds the global WS1 tile list from the replicated index (`sp_plan`). + 2. Each tile is computed by the rank that holds its first row. Rows of that tile that live on + other ranks are all-gathered first. Only tiles that cross a shard boundary move rows. + 3. The tile partials (`h3_rmsnorm_backward_partials`, `h3_gate_grad_partials`, the WS1 partial + kernels over explicit row lists) are all-gathered and put back in WS1 tile order. + 4. Every rank runs the WS1 folds (`h3_rmsnorm_fold_partials`, `h3_gate_grad_fold`). + + The only collectives are rank-ordered all-gathers, which are copies. No collective reduction + arithmetic is used. + +A naive SP backward runs the WS1 backward on each rank's rows and sums the per-rank results. It +changes the reduction tree, so it is not WS1. The evidence measures how often it differs. + +The plan depends only on the layout and the index. `H3SPNormAdaLNCudaOp` builds it once per index +object and reuses it for every norm and gated residual that uses the same index, which in a forward +pass is every block. + +## Rows Exchanged + +How many rows move depends on how the table rows interleave along the sequence: + +| Packing | Rows sent per rank (S = 32768, SP8, 4096 rows per rank) | +| --- | --- | +| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 174 | +| interleaved (modality and timestep random per position; stress case) | 0 – 1536 | + +## Evidence + +The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 +node, one GPU per rank. + +## Tests + +```bash +python -m pytest tests/h3/test_h3_sp_norm_adaln.py -v # layout, rank backward, NCCL SP2/4/8 +python scripts/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \ + --out docs/usage/evidence/h3-sp-norm-adaln-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-sp-norm-adaln-b200/report.json +``` + +The rank-backward tests run SP 2–8 on one GPU, with ranks as threads around the plain backward +functions. They cover odd `S`, `B = 2`, shards smaller than a tile, block and interleaved packing, +and bf16 and fp32, and compare every rank against WS1. The NCCL tests run the autograd region end +to end and need as many GPUs as ranks. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- With interleaved packing, up to a full shard of rows can move in the backward. The result is + still exact. +- Attention and FFN are other rows. Until they land, the region uses a seeded stand-in for the + sublayer output. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 672bf6049..d9948d98d 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -397,3 +397,46 @@ def h3_gate_residual_backward( tile_end: torch.Tensor, seg_first_tile: torch.Tensor, ) -> list[torch.Tensor]: ... +def h3_rmsnorm_backward_partials( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + index: torch.Tensor | None, + dw_rows: torch.Tensor, + dw_begin: torch.Tensor, + dw_end: torch.Tensor, + seg_rows: torch.Tensor | None = None, + seg_begin: torch.Tensor | None = None, + seg_end: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_fold_partials( + dw_partial: torch.Tensor, + weight: torch.Tensor, + seg_partial: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_gate_grad_partials( + grad: torch.Tensor, + y: torch.Tensor, + rows: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, +) -> torch.Tensor: ... +def h3_gate_grad_fold( + partial: torch.Tensor, seg_first_tile: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: ... +def h3_rmsnorm_backward_dx( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> torch.Tensor: ... +def h3_gate_residual_backward_dy( + grad: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py new file mode 100644 index 000000000..19a867b3e --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py @@ -0,0 +1,404 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sequence-parallel H3 RMSNorm / AdaLN modulation / gated residual (RFC #420 ``sp_norm_adaln``). + +Ownership: rank ``r`` of ``sp`` holds packed positions ``[r * S // sp, (r + 1) * S // sp)`` +of every batch item. The full ``(S,)`` row index is replicated; it is small, and every +rank derives the same reduction plan from it. + +Forward and the row-local gradients (``dx``, ``d_sublayer``, ``d_residual``) are the WS1 +kernels on the local rows. Rows are independent, so they are byte-equal by construction. + +The cross-row reductions (``d_norm_weight``, ``d_shift``/``d_scale``, ``d_gate``) are WS1 +two-level folds: FP32 sums over fixed 256-element tiles, then an ascending fold of the +tiles. Every rank builds the global WS1 tile list. A tile is computed by the rank that +holds its first row; the tile's rows from other ranks are all-gathered beforehand (only +rows of tiles that straddle a shard boundary move). The tile partials are all-gathered, +put back in WS1 tile order, and every rank runs the WS1 fold. The partial and fold kernels +are the WS1 kernels, so every gradient equals WS1 on every rank. The only collectives are +rank-ordered all-gathers (copies). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import BACKWARD_TILE, _segment_tiles +from rl_engine.kernels.ops.cuda.h3.ws2_comm import gather_rows +from rl_engine.kernels.ops.pytorch.h3.gate_residual import validate_h3_gate_residual +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +_SYMBOLS = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward_dx", + "h3_rmsnorm_backward_partials", + "h3_rmsnorm_fold_partials", + "h3_gate_residual_forward", + "h3_gate_residual_backward_dy", + "h3_gate_grad_partials", + "h3_gate_grad_fold", +) +BACKWARD_IMPL = "row_local_ws1+gather_straddling_rows+ws1_tile_partials+gather+ws1_fold" + + +def sp_norm_adaln_available() -> bool: + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +@dataclass(frozen=True) +class SPRowLayout: + """Which packed positions each sequence-parallel rank holds.""" + + seq_len: int + batch: int + sp: int + rank: int + + def bounds(self, rank: int) -> tuple[int, int]: + return rank * self.seq_len // self.sp, (rank + 1) * self.seq_len // self.sp + + @property + def lo(self) -> int: + return self.bounds(self.rank)[0] + + @property + def hi(self) -> int: + return self.bounds(self.rank)[1] + + @property + def local_len(self) -> int: + return self.hi - self.lo + + +def sp_row_layout(seq_len: int, batch: int, sp: int, rank: int) -> SPRowLayout: + if sp < 1 or not 0 <= rank < sp: + raise ValueError(f"need 0 <= rank < sp, got rank={rank}, sp={sp}") + if batch < 1 or seq_len < sp: + raise ValueError(f"every rank needs a row: seq_len={seq_len}, sp={sp}, batch={batch}") + return SPRowLayout(seq_len, batch, sp, rank) + + +def _ordinal_by_group(group: torch.Tensor, groups: int) -> tuple[torch.Tensor, torch.Tensor]: + """Per-group counts, and each element's position within its group (in input order).""" + + counts = torch.bincount(group, minlength=groups) + order = torch.argsort(group, stable=True) + ordinal = torch.empty_like(group) + first = torch.cumsum(counts, 0) - counts + ordinal[order] = torch.arange(group.numel(), device=group.device) - first[group[order]] + return counts, ordinal + + +class _Plan: + """Global WS1 tiles of one backward, and this rank's part in computing them. + + ``families`` are ``(rows, tile_begin, tile_end)`` over global flattened rows + ``b * S + s``; consecutive tiles cover ``rows`` in order. Every rank builds the same + plan from replicated metadata, so collective calls and shapes agree. + """ + + def __init__(self, layout: SPRowLayout, families, device) -> None: + self.layout, sp = layout, layout.sp + self._his = torch.tensor([layout.bounds(r)[1] for r in range(sp)], device=device) + + # Rows of a tile computed by another rank: the only activation rows that move. + tile_owner, shared = [], [] + for rows, begin, end in families: + owner = self.owner(rows[begin]) + tile_owner.append(owner) + shared.append(rows[torch.repeat_interleave(owner, end - begin) != self.owner(rows)]) + shared = torch.unique(torch.cat(shared)) + shared_owner = self.owner(shared) + counts, slot = _ordinal_by_group(shared_owner, sp) + self.max_send = int(counts.max()) if shared.numel() else 0 + + # Buffer row of every global row this rank reads: its own rows first, then + # every rank's send block (padded to max_send) in rank order. + own = self.global_rows(layout.rank) + buffer_index = torch.full((layout.batch * layout.seq_len,), -1, device=device) + buffer_index[own] = torch.arange(own.numel(), device=device) + foreign = shared_owner != layout.rank + buffer_index[shared[foreign]] = ( + own.numel() + shared_owner[foreign] * self.max_send + slot[foreign] + ) + mine = ~foreign + self.send_local = self.local_index(shared[mine][torch.argsort(slot[mine])]) + self.recv_global = torch.zeros(sp * self.max_send, dtype=torch.long, device=device) + self.recv_global[shared_owner * self.max_send + slot] = shared # padding: never read + + self.families = [ + self._my_tiles(rows, begin, end, owner, buffer_index) + for (rows, begin, end), owner in zip(families, tile_owner) + ] + + def owner(self, rows: torch.Tensor) -> torch.Tensor: + return torch.searchsorted(self._his, rows % self.layout.seq_len, right=True) + + def global_rows(self, rank: int) -> torch.Tensor: + lo, hi = self.layout.bounds(rank) + s = torch.arange(lo, hi, device=self._his.device) + b = torch.arange(self.layout.batch, device=self._his.device) + return (b[:, None] * self.layout.seq_len + s[None, :]).reshape(-1) + + def local_index(self, rows: torch.Tensor) -> torch.Tensor: + b, s = rows // self.layout.seq_len, rows % self.layout.seq_len + return b * self.layout.local_len + s - self.layout.lo + + def _my_tiles(self, rows, begin, end, owner, buffer_index): + mine = torch.nonzero(owner == self.layout.rank).flatten() + lengths = (end - begin)[mine] + my_end = torch.cumsum(lengths, 0) + my_begin = my_end - lengths + elems = torch.repeat_interleave(begin[mine] - my_begin, lengths) + elems = elems + torch.arange(elems.numel(), device=rows.device) + counts, ordinal = _ordinal_by_group(owner, self.layout.sp) + width = int(counts.max()) if owner.numel() else 0 + return { + "tiles": (buffer_index[rows[elems]].contiguous(), my_begin, my_end), + "width": width, # tiles per rank in the padded partial gather + "gather_index": owner * width + ordinal, # WS1 tile order within that gather + } + + def exchange(self, collective, *tensors: torch.Tensor) -> list[torch.Tensor]: + """For each ``(M_local, ...)`` tensor: local rows, then the gathered send blocks.""" + + out = [] + for t in tensors: + if self.max_send == 0: + out.append(t) + continue + send = t.new_zeros((self.max_send, *t.shape[1:])) + send[: self.send_local.numel()] = t[self.send_local] + out.append(torch.cat([t, gather_rows(collective, send)])) + return out + + def gather_partials(self, collective, family: int, partial: torch.Tensor) -> torch.Tensor: + """This rank's tile partials -> every tile's partial, in WS1 tile order.""" + + fam = self.families[family] + padded = partial.new_zeros((fam["width"], *partial.shape[1:])) + padded[: partial.shape[0]] = partial + return gather_rows(collective, padded).index_select(0, fam["gather_index"]).contiguous() + + +def _dweight_tiles(total: int, device): + begin = torch.arange(0, total, BACKWARD_TILE, device=device) + return torch.arange(total, device=device), begin, torch.clamp(begin + BACKWARD_TILE, max=total) + + +@dataclass(frozen=True) +class SPPlan: + """Everything an SP backward needs that depends only on the layout and the row index. + + Built once per ``(row index, table rows)`` and shared by every norm and gated + residual that uses that index (every block of a forward pass). + """ + + layout: SPRowLayout + plan: _Plan + local_index: torch.Tensor | None # (S_local,) table row of each local position + index_buf: torch.Tensor | None # table row of each exchange-buffer row + seg_first_tile: torch.Tensor | None + num_rows: int | None + + +def sp_plan(layout: SPRowLayout, index_full=None, num_rows: int | None = None) -> SPPlan: + """The reduction plan for one row index (``None``: the plain norm's ``dweight`` only).""" + + device = index_full.device if index_full is not None else torch.device("cuda") + families = [_dweight_tiles(layout.batch * layout.seq_len, device)] + if index_full is None: + return SPPlan(layout, _Plan(layout, families, device), None, None, None, None) + *seg_tiles, seg_first_tile = _segment_tiles(index_full.repeat(layout.batch), num_rows) + plan = _Plan(layout, [*families, tuple(seg_tiles)], device) + local_index = index_full[layout.lo : layout.hi].contiguous() + received = index_full[plan.recv_global % layout.seq_len] + index_buf = torch.cat([local_index.repeat(layout.batch), received]).contiguous() + return SPPlan(layout, plan, local_index, index_buf, seg_first_tile, num_rows) + + +def sp_rmsnorm_backward(grad, x, weight, rstd, shift, scale, sp: SPPlan, collective): + """One rank's backward on its ``(B * S_local, H)`` rows: ``(dx, dweight, d_shift, d_scale)``. + + ``dx`` is row-local; the other gradients are the full WS1 reductions, identical on + every rank. ``shift``/``scale`` are ``None`` for the plain norm. + """ + + weight, plan, modulated = weight.contiguous(), sp.plan, sp.local_index is not None + dx = _C.h3_rmsnorm_backward_dx(grad, x, weight, rstd, shift, scale, sp.local_index) + g_buf, x_buf, rstd_buf = plan.exchange(collective, grad, x, rstd) + seg = plan.families[1]["tiles"] if modulated else (None, None, None) + partials = _C.h3_rmsnorm_backward_partials( + g_buf, x_buf, weight, rstd_buf, shift, scale, sp.index_buf, + *plan.families[0]["tiles"], *seg, + ) # fmt: skip + gathered = [plan.gather_partials(collective, i, p) for i, p in enumerate(partials)] + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_rmsnorm_backward_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + if not modulated: + (dweight,) = _C.h3_rmsnorm_fold_partials(gathered[0], weight) + return dx, dweight, None, None + dweight, d_shift, d_scale = _C.h3_rmsnorm_fold_partials( + gathered[0], weight, gathered[1], sp.seg_first_tile + ) + return dx, dweight, d_shift.to(shift.dtype), d_scale.to(scale.dtype) + + +def sp_gate_residual_backward(grad, y, gate, sp: SPPlan, collective): + """One rank's backward on its rows: ``(d_sublayer, d_gate)``; ``d_residual`` is ``grad``.""" + + plan = sp.plan + dy = _C.h3_gate_residual_backward_dy(grad, gate, sp.local_index) + g_buf, y_buf = plan.exchange(collective, grad, y) + partial = _C.h3_gate_grad_partials(g_buf, y_buf, *plan.families[1]["tiles"]) + dgate = _C.h3_gate_grad_fold( + plan.gather_partials(collective, 1, partial), sp.seg_first_tile, gate.dtype + ) + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_gate_grad_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + return dy, dgate + + +class _SPRMSNorm(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, sp, collective, eps): + x2 = x.contiguous().view(-1, x.shape[-1]) + y, rstd = _C.h3_rmsnorm_forward( + x2, weight.contiguous(), float(eps), shift, scale, sp.local_index + ) + ctx.save_for_backward(x2, weight, rstd, shift, scale) + ctx.sp, ctx.collective, ctx.x_shape = sp, collective, x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + dx, dweight, d_shift, d_scale = sp_rmsnorm_backward( + g2, x2, weight, rstd, shift, scale, ctx.sp, ctx.collective + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None, None + + +class _SPGateResidual(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, sp, collective): + hidden = residual.shape[-1] + res2 = residual.contiguous().view(-1, hidden) + y2 = y.contiguous().view(-1, hidden) + out = _C.h3_gate_residual_forward(res2, y2, gate, sp.local_index) + ctx.save_for_backward(y2, gate) + ctx.sp, ctx.collective, ctx.shape = sp, collective, residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + dy, dgate = sp_gate_residual_backward(g2, y2, gate, ctx.sp, ctx.collective) + return grad, dy.view(ctx.shape), dgate, None, None + + +class H3SPNormAdaLNCudaOp: + """Sequence-parallel ``h3_rmsnorm`` (+ modulation) and ``adaln_gate_residual``. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` and + ``backend_id`` (e.g. ``DeterministicCollective``). Activations are this rank's + ``(B, S_local, H)`` rows; the norm weight, the table views and the ``(S,)`` row + index are the full, replicated tensors. Gradients of the replicated tensors come + back identical on every rank, equal to WS1. + """ + + op_class = "reduction" + backward_impl = BACKWARD_IMPL + + def __init__(self, collective, seq_len: int, batch: int) -> None: + if not sp_norm_adaln_available(): + raise RuntimeError("rl_engine._C lacks the H3 SP symbols; rebuild the CUDA extension") + self.collective = collective + self.layout = sp_row_layout(seq_len, batch, collective.world_size, collective.rank) + self._plans: list[tuple[object, int, int | None, SPPlan]] = [] + + def plan(self, index: torch.Tensor | None, num_rows: int | None) -> SPPlan: + """The cached plan for this index object (rebuilt if it was modified in place).""" + + version = -1 if index is None else index._version + for obj, ver, rows, plan in self._plans: + if obj is index and ver == version and rows == num_rows: + return plan + plan = sp_plan(self.layout, index, num_rows) + # Holding the index keeps its storage alive, so a new tensor can never alias it. + self._plans = [(index, version, num_rows, plan), *self._plans[:3]] + return plan + + def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | None) -> None: + expected = (self.layout.batch, self.layout.local_len) + if not isinstance(x, torch.Tensor) or x.dim() != 3 or tuple(x.shape[:2]) != expected: + shape = tuple(x.shape) if isinstance(x, torch.Tensor) else type(x).__name__ + raise ValueError(f"activations must be this rank's {expected} + (H,) rows, got {shape}") + if not x.is_cuda: + raise ValueError("H3SPNormAdaLNCudaOp needs CUDA tensors") + if index is None: + return + if index.dim() != 1 or index.shape[0] != self.layout.seq_len: + raise ValueError( + f"index must be the full ({self.layout.seq_len},) row index, " + f"got {tuple(index.shape)}" + ) + if index.numel() and (int(index.min()) < 0 or int(index.max()) >= num_rows): + raise IndexError(f"index must be in [0, {num_rows})") + + def norm(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + self._check(x, None, None) + validate_h3_rmsnorm(x, weight, eps) + return _SPRMSNorm.apply(x, weight, None, None, self.plan(None, None), self.collective, eps) + + def norm_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + self._check(x, index, shift.shape[0]) + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index[self.layout.lo : self.layout.hi]) + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + sp = self.plan(index, shift.shape[0]) + return _SPRMSNorm.apply(x, weight, shift, scale, sp, self.collective, eps) + + def gate_residual(self, residual, y, gate, index) -> torch.Tensor: + self._check(residual, index, gate.shape[0]) + validate_h3_gate_residual(residual, y, gate, index[self.layout.lo : self.layout.hi]) + if gate.stride(1) != 1: + gate = gate.contiguous() + sp = self.plan(index, gate.shape[0]) + return _SPGateResidual.apply(residual, y, gate, sp, self.collective) + + def readback(self) -> dict: + return { + "op": "sp_norm_adaln", + "backward_impl": BACKWARD_IMPL, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_256_element_tiles_then_ascending_fold", + "sp": self.layout.sp, + "rank": self.layout.rank, + "positions": [self.layout.lo, self.layout.hi], + "batch": self.layout.batch, + "fallback": None, + } diff --git a/rl_engine/testing/h3_cases.py b/rl_engine/testing/h3_cases.py index 527a820ec..512d3ca4a 100644 --- a/rl_engine/testing/h3_cases.py +++ b/rl_engine/testing/h3_cases.py @@ -33,3 +33,27 @@ def h3_packed_layout( timestep_indices = torch.randint(0, num_timesteps, (seq_len,), generator=generator) token_tags[: min(3, seq_len)] = torch.arange(min(3, seq_len)) return timestep_indices.to(device), token_tags.to(device) + + +def h3_block_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """(timestep_indices, token_tags) for block-structured packing, as H3 packs requests. + + Each timestep owns one contiguous run of text, video and audio blocks (about + 5% / 80% / 15% of its tokens), so table rows form contiguous position runs. + ``h3_packed_layout`` is the interleaved stress case. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + cuts = torch.sort(torch.randperm(seq_len - 1, generator=generator)[: num_timesteps - 1] + 1) + edges = [0, *cuts.values.tolist(), seq_len] + timestep_indices = torch.empty(seq_len, dtype=torch.long) + token_tags = torch.empty(seq_len, dtype=torch.long) + for t, (lo, hi) in enumerate(zip(edges, edges[1:])): + text, audio = (hi - lo) // 20, (hi - lo) * 3 // 20 + timestep_indices[lo:hi] = t + token_tags[lo:hi] = 0 # video + token_tags[lo : lo + text] = 1 + token_tags[hi - audio : hi] = 2 + return timestep_indices.to(device), token_tags.to(device) diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py index a0dc116f6..b2e3b448f 100644 --- a/rl_engine/testing/h3_ws2.py +++ b/rl_engine/testing/h3_ws2.py @@ -120,6 +120,187 @@ def step(op, args): return {"equality": equality, "perf": perf, "readback": tp_op.readback()} +def _region(gate_fn, norm_fn, residual, y, weight, table): + """One block's MLP-side SP region: ``norm2(residual + gate_msa * y)`` with modulation.""" + + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = table.chunk(6, dim=-1) + hidden = gate_fn(residual, y, gate_msa) + return norm_fn(hidden, weight, shift_mlp, scale_mlp) + + +def ws1_sp_region(residual, y, weight, table, index, grad) -> dict[str, torch.Tensor]: + """The SP region on one GPU with the WS1 ops (the reference for ``sp_norm_adaln``).""" + + from rl_engine.kernels.ops.cuda.h3.gate_residual import H3GateResidualCudaOp + from rl_engine.kernels.ops.cuda.h3.rmsnorm import H3RMSNormCudaOp + + gate_op, norm_op = H3GateResidualCudaOp(), H3RMSNormCudaOp() + leaves = [t.detach().clone().requires_grad_(True) for t in (residual, y, weight, table)] + out = _region( + lambda r, yy, g: gate_op(r, yy, g, index), + lambda h, w, sh, sc: norm_op.forward_modulated(h, w, sh, sc, index), + *leaves, + ) + out.backward(grad) + names = ("d_residual", "d_y", "d_weight", "d_table") + return {"out": out.detach(), **{n: leaf.grad for n, leaf in zip(names, leaves)}} + + +def sp_region_rank(collective, residual, y, weight, table, index, grad) -> dict[str, Any]: + """One SP rank: its rows of the region and its (replicated) parameter gradients.""" + + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + op = H3SPNormAdaLNCudaOp(collective, seq_len=residual.shape[1], batch=residual.shape[0]) + rows = slice(op.layout.lo, op.layout.hi) + local = [t[:, rows].contiguous() for t in (residual, y)] + leaves = [t.detach().clone().requires_grad_(True) for t in (*local, weight, table)] + out = _region( + lambda r, yy, g: op.gate_residual(r, yy, g, index), + lambda h, w, sh, sc: op.norm_modulated(h, w, sh, sc, index), + *leaves, + ) + out.backward(grad[:, rows].contiguous()) + names = ("d_residual", "d_y", "d_weight", "d_table") + return { + "out": out.detach(), + **{n: leaf.grad for n, leaf in zip(names, leaves)}, + "readback": op.readback(), + } + + +def sp_matches_ws1(ws1: dict[str, torch.Tensor], ranks: list[dict[str, Any]]) -> dict[str, bool]: + """Rows: each rank's slice equals WS1's; parameters: every rank's copy equals WS1's.""" + + def rows(key): + return all( + torch.equal(r[key], ws1[key][:, slice(*r["readback"]["positions"])]) for r in ranks + ) + + return { + "out": rows("out"), + "d_residual": rows("d_residual"), + "d_y": rows("d_y"), + "d_weight": all(torch.equal(r["d_weight"], ws1["d_weight"]) for r in ranks), + "d_table": all(torch.equal(r["d_table"], ws1["d_table"]) for r in ranks), + } + + +def sp_case(batch, seq, hidden=5376, num_t=3, *, layout="block", dtype=torch.bfloat16, seed=0): + """Inputs of the SP region; identical on every rank (CPU generator).""" + + from rl_engine.testing.h3_cases import h3_block_layout, h3_packed_layout + + g = torch.Generator(device="cpu").manual_seed(seed) + make = h3_block_layout if layout == "block" else h3_packed_layout + ti, tags = make(seq, num_t, seed=seed) + return { + "residual": torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda(), + "y": (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda(), + "weight": (torch.rand(hidden, generator=g) + 0.5).to(dtype).cuda(), + "table": (torch.randn(3 * num_t, 6 * hidden, generator=g) * 0.5).to(dtype).cuda(), + "index": (ti * 3 + tags).cuda(), + "grad": torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda(), + } + + +def naive_sp_mismatch(case: dict[str, torch.Tensor], sp: int) -> dict[str, float]: + """Fraction of parameter-gradient elements where a naive SP backward differs from WS1. + + Naive: each rank runs the WS1 backward on its own rows, then the per-rank results + are summed in rank order (what a plain all-reduce of local gradients does). + """ + + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import sp_row_layout + + ws1 = ws1_sp_region(**case) + parts = [] + for rank in range(sp): + lay = sp_row_layout(case["residual"].shape[1], case["residual"].shape[0], sp, rank) + part = slice(lay.lo, lay.hi) + local = {k: (v[:, part] if k in ("residual", "y", "grad") else v) for k, v in case.items()} + local["index"] = case["index"][part] + parts.append(ws1_sp_region(**local)) + out = {} + for key in ("d_weight", "d_table"): + total = parts[0][key].float() + for p in parts[1:]: + total = total + p[key].float() + out[key] = (total.to(ws1[key].dtype) != ws1[key]).float().mean().item() + return out + + +def sp_region_sweep(collective, cases: list[dict[str, Any]], iters: int = 20) -> dict[str, Any]: + """Evidence for one SP rank: byte equality, rows exchanged and time, per case.""" + + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import _dweight_tiles, _Plan, sp_row_layout + + results = [] + for spec in cases: + case = sp_case(**spec) + batch, seq = case["residual"].shape[:2] + ws1 = ws1_sp_region(**case) + mine = sp_region_rank(collective, **case) + equal = sp_matches_ws1(ws1, [mine]) + lay = sp_row_layout(seq, batch, collective.world_size, collective.rank) + tiles = _segment_tiles(case["index"].repeat(batch), case["table"].shape[0]) + plan = _Plan(lay, [_dweight_tiles(batch * seq, "cuda"), tuple(tiles[:3])], "cuda") + del ws1, mine + torch.cuda.empty_cache() + results.append( + { + **spec, + "equal": equal, + "local_rows": batch * lay.local_len, + "rows_sent": int(plan.send_local.numel()), + "rows_gathered_per_rank": plan.max_send, + **_region_times(collective, case, lay, iters), + } + ) + return {"cases": results} + + +def _region_times(collective, case, lay, iters) -> dict[str, float]: + """Forward + backward time of the region (leaves prepared once, grads reset per call).""" + + from rl_engine.kernels.ops.cuda.h3.gate_residual import H3GateResidualCudaOp + from rl_engine.kernels.ops.cuda.h3.rmsnorm import H3RMSNormCudaOp + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + from rl_engine.testing.h3_report import time_us + + index, rows = case["index"], slice(lay.lo, lay.hi) + sp_op = H3SPNormAdaLNCudaOp(collective, seq_len=lay.seq_len, batch=lay.batch) + gate_op, norm_op = H3GateResidualCudaOp(), H3RMSNormCudaOp() + + def leaves(local): + acts = [case[k][:, rows] if local else case[k] for k in ("residual", "y")] + return [ + t.detach().contiguous().requires_grad_(True) + for t in (*acts, case["weight"], case["table"]) + ] + + def step(fns, args, grad): + for leaf in args: + leaf.grad = None + _region(*fns, *args).backward(grad) + + sp_fns = ( + lambda r, yy, g: sp_op.gate_residual(r, yy, g, index), + lambda h, w, sh, sc: sp_op.norm_modulated(h, w, sh, sc, index), + ) + ws1_fns = ( + lambda r, yy, g: gate_op(r, yy, g, index), + lambda h, w, sh, sc: norm_op.forward_modulated(h, w, sh, sc, index), + ) + sp_args, ws1_args = leaves(True), leaves(False) + sp_grad = case["grad"][:, rows].contiguous() + return { + "sp_forward_backward_us": time_us(lambda: step(sp_fns, sp_args, sp_grad), 3, iters), + "ws1_forward_backward_us": time_us(lambda: step(ws1_fns, ws1_args, case["grad"]), 3, iters), + } + + # --- real multi-process groups ------------------------------------------------- diff --git a/scripts/h3_ws2_evidence.py b/scripts/h3_ws2_evidence.py index c27332f77..021cfc1f8 100644 --- a/scripts/h3_ws2_evidence.py +++ b/scripts/h3_ws2_evidence.py @@ -11,6 +11,8 @@ export RL_KERNEL_H3_WEIGHTS= python scripts/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \\ --out docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json + python scripts/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \\ + --out docs/usage/evidence/h3-sp-norm-adaln-b200/report.json """ from __future__ import annotations @@ -29,7 +31,13 @@ load_h3_conditioning_weights, load_h3_manifest, ) -from rl_engine.testing.h3_ws2 import run_world, tp_projection_sweep # noqa: E402 +from rl_engine.testing.h3_ws2 import ( # noqa: E402 + naive_sp_mismatch, + run_world, + sp_case, + sp_region_sweep, + tp_projection_sweep, +) WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" BIAS = "transformer_blocks.0.adaln_proj.linear.bias" @@ -56,7 +64,37 @@ def _tp_adaln(worlds: list[int], max_t: int) -> dict: return {"weights": [WEIGHT, BIAS], "runs": runs} -OPS = {"tp_adaln_3mod": _tp_adaln} +SP_CASES = [ + {"batch": 1, "seq": 4097, "layout": "block", "seed": 1}, + {"batch": 1, "seq": 4097, "layout": "interleaved", "seed": 2}, + {"batch": 2, "seq": 4097, "layout": "block", "seed": 3}, + {"batch": 1, "seq": 32768, "layout": "block", "seed": 4}, + {"batch": 1, "seq": 32768, "layout": "interleaved", "seed": 5}, + {"batch": 1, "seq": 131072, "layout": "block", "seed": 6}, +] + + +def _sp_norm(worlds: list[int], _max_t: int) -> dict: + runs = [] + for world in worlds: + ranks = run_world(world, sp_region_sweep, {}, cases=SP_CASES) + runs.append({"sp": world, "ranks": ranks}) + equal = all(all(c["equal"].values()) for r in ranks for c in r["cases"]) + print(f"sp={world}: byte-equal to WS1 on every rank and case: {equal}") + naive = { + layout: naive_sp_mismatch(sp_case(1, 32768, layout=layout, seed=4), max(worlds)) + for layout in ("block", "interleaved") + } + region = "norm2(residual + gate_msa[row] * y) * (1 + scale_mlp[row]) + shift_mlp[row]" + return { + "region": region, + "hidden": 5376, + "runs": runs, + "naive_sp": {"sp": max(worlds), **naive}, + } + + +OPS = {"tp_adaln_3mod": _tp_adaln, "sp_norm_adaln": _sp_norm} def main() -> None: diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index ba4b3b4ef..6ee3de114 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -511,43 +511,13 @@ def plot_tp_adaln(report: dict[str, Any]): report, "tp_adaln_3mod (column-parallel AdaLN projection, 2688 -> 96768, NCCL)" ) runs = report["runs"] - # Left: every rank, every T, every tensor byte-equal to WS1? - left.set_facecolor(SURFACE) - left.set_title( - "Byte-equal to WS1 (every rank, every T)", loc="left", fontsize=11, color=INK, pad=10 - ) - for y, run in enumerate(runs): - checks = [e for r in run["ranks"] for e in r["equality"]] - for x, (key, _) in enumerate(TP_TENSORS): - ok = sum(bool(e[key]) for e in checks) - good = ok == len(checks) - left.add_patch( - matplotlib.patches.FancyBboxPatch( - (x + 0.06, y + 0.08), - 0.88, - 0.84, - boxstyle="round,pad=0,rounding_size=0.06", - facecolor=RL_KERNEL if good else PROVIDER, - edgecolor=SURFACE, - linewidth=2, - ) - ) - left.text( - x + 0.5, - y + 0.5, - f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} rank x T", - ha="center", - va="center", - fontsize=8, - color="white", - ) - left.set_xlim(0, len(TP_TENSORS)) - left.set_ylim(len(runs), 0) - left.set_xticks([i + 0.5 for i in range(len(TP_TENSORS))], [label for _, label in TP_TENSORS]) - left.set_yticks([i + 0.5 for i in range(len(runs))], [f"TP{run['tp']}" for run in runs]) - left.tick_params(colors=MUTED, labelsize=8, length=0) - for side in left.spines.values(): - side.set_visible(False) + _equality_grid( + left, + "Byte-equal to WS1 (every rank, every T)", + [(f"TP{run['tp']}", [e for r in run["ranks"] for e in r["equality"]]) for run in runs], + TP_TENSORS, + "rank x T", + ) # Right: slowest rank's time at the largest T; WS1 on one GPU as reference lines. num_t = max(p["num_timesteps"] for p in runs[0]["ranks"][0]["perf"]) @@ -589,6 +559,144 @@ def worst(run, key): return fig +def _equality_grid(ax, title, rows, tensors, unit) -> None: + """One cell per (configuration, tensor): how many checks were byte-equal to WS1.""" + + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) + for y, (_, checks) in enumerate(rows): + for x, (key, _) in enumerate(tensors): + ok = sum(bool(e[key]) for e in checks) + good = ok == len(checks) + ax.add_patch( + matplotlib.patches.FancyBboxPatch( + (x + 0.06, y + 0.08), + 0.88, + 0.84, + boxstyle="round,pad=0,rounding_size=0.06", + facecolor=RL_KERNEL if good else PROVIDER, + edgecolor=SURFACE, + linewidth=2, + ) + ) + label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} {unit}" + ax.text(x + 0.5, y + 0.5, label, ha="center", va="center", fontsize=8, color="white") + ax.set_xlim(0, len(tensors)) + ax.set_ylim(len(rows), 0) + ax.set_xticks([i + 0.5 for i in range(len(tensors))], [label for _, label in tensors]) + ax.set_yticks([i + 0.5 for i in range(len(rows))], [name for name, _ in rows]) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + for side in ax.spines.values(): + side.set_visible(False) + + +SP_TENSORS = ( + ("out", "output rows"), + ("d_residual", "d_residual"), + ("d_y", "d_sublayer"), + ("d_weight", "d_norm_w"), + ("d_table", "d_table"), +) + + +def plot_sp_norm(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(15, 3.9), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.35, 1, 0.9]} + ) + env = report["environment"] + fig.suptitle( + "sp_norm_adaln (sequence-parallel norm2(residual + gate * y) with AdaLN modulation, " + "H = 5376, NCCL)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.01, + f"{env['gpu']} x {env['gpus']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · bf16", + fontsize=7, + color=MUTED, + ) + runs = report["runs"] + _equality_grid( + axes[0], + "Byte-equal to WS1 (every rank, every case)", + [ + (f"SP{run['sp']}", [c["equal"] for r in run["ranks"] for c in r["cases"]]) + for run in runs + ], + SP_TENSORS, + "rank x case", + ) + + # Middle: the largest block-layout case, slowest rank. + big = max( + (c for c in runs[0]["ranks"][0]["cases"] if c["layout"] == "block"), + key=lambda c: c["seq"] * c["batch"], + ) + + def worst(run, key): + return ( + max( + c[key] + for r in run["ranks"] + for c in r["cases"] + if (c["seq"], c["batch"], c["layout"]) == (big["seq"], big["batch"], big["layout"]) + ) + / 1e3 + ) + + sp_runs = [run for run in runs if run["sp"] > 1] + ax = axes[1] + _style(ax, f"Region fwd+bwd, S = {big['seq']}, block layout", "ms (slowest rank)") + _bars( + ax, + [f"SP{run['sp']}" for run in sp_runs], + [("SP", RL_KERNEL, [worst(run, "sp_forward_backward_us") for run in sp_runs])], + "{:.2f}", + ) + ws1 = worst(runs[0], "ws1_forward_backward_us") + ax.axhline( + ws1, color=PROVIDER, linestyle="--", linewidth=1.4, zorder=1, label=f"WS1, 1 GPU: {ws1:.2f}" + ) + ax.set_ylim(0, ws1 * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper right") + + # Right: what a naive SP reduction would do instead. + naive = report["naive_sp"] + ax = axes[2] + _style(ax, f"Naive SP{naive['sp']} vs WS1 (S = 32768)", "% of elements that differ") + _bars( + ax, + ["d_norm_w", "d_table"], + [ + ("naive, block", PROVIDER, [100 * naive["block"][k] for k in ("d_weight", "d_table")]), + ( + "naive, interleaved", + PROVIDER, + [100 * naive["interleaved"][k] for k in ("d_weight", "d_table")], + "////", + ), + ("RL-Kernel SP", RL_KERNEL, [0.0, 0.0]), + ], + "{:.1f}", + ) + ax.text( + 0.98, + 0.62, + "naive: each rank's WS1 backward,\nthen a rank-order sum", + transform=ax.transAxes, + ha="right", + fontsize=7, + color=MUTED, + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, @@ -598,6 +706,7 @@ def worst(run, key): "adaln_gate_residual": plot_gate_residual, "final_adaln_out": plot_final, "tp_adaln_3mod": plot_tp_adaln, + "sp_norm_adaln": plot_sp_norm, } diff --git a/tests/h3/test_h3_sp_norm_adaln.py b/tests/h3/test_h3_sp_norm_adaln.py new file mode 100644 index 000000000..04d58395c --- /dev/null +++ b/tests/h3/test_h3_sp_norm_adaln.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``sp_norm_adaln``: sequence-parallel norm / modulation / gated residual. + +* ownership: contiguous position shards per batch item; bad layouts fail closed; +* backward (one GPU, ranks as threads around the plain backward functions): every + rank's row gradients equal WS1's rows, and its ``dweight``/``d_shift``/``d_scale``/ + ``d_gate`` equal WS1's, for SP 2..8, odd S, B > 1, interleaved modality rows and + bf16/fp32; +* end to end (real NCCL, autograd): the MLP-side block region + ``norm2(residual + gate_msa * y)`` with modulation is byte-equal to WS1 on every rank. +""" + +from __future__ import annotations + +import os +import threading +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import SPRowLayout, sp_row_layout +from rl_engine.testing.h3_ws2 import ( + run_world, + sp_case, + sp_matches_ws1, + sp_region_rank, + ws1_sp_region, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +ALL_TRUE = dict.fromkeys(("out", "d_residual", "d_y", "d_weight", "d_table"), True) + + +def _require_ext(): + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import sp_norm_adaln_available + + if not sp_norm_adaln_available(): + pytest.skip("rl_engine._C lacks the H3 SP symbols") + + +class _ThreadGroup: + """Ranks as threads on one GPU; ``all_gather`` is the rank-ordered concatenation. + + Only for the plain backward functions: autograd runs every backward of a device + on one worker thread, so autograd-driven ranks need real processes. + """ + + def __init__(self, world: int) -> None: + self.barrier, self.slots = threading.Barrier(world), [None] * world + + def handle(self, rank: int): + def all_gather(t): + self.slots[rank] = t + self.barrier.wait() + out = torch.cat(self.slots) + self.barrier.wait() + return out + + return SimpleNamespace(rank=rank, world_size=len(self.slots), all_gather=all_gather) + + def run(self, fn): + results, errors = [None] * len(self.slots), [] + + def body(rank): + try: + results[rank] = fn(self.handle(rank)) + except BaseException as exc: # re-raised below + errors.append(exc) + self.barrier.abort() + + threads = [threading.Thread(target=body, args=(r,)) for r in range(len(self.slots))] + for t in threads: + t.start() + for t in threads: + t.join() + if errors: + raise next( + (e for e in errors if not isinstance(e, threading.BrokenBarrierError)), errors[0] + ) + return results + + +class TestLayout: + @pytest.mark.parametrize("seq, sp", [(8, 8), (257, 2), (1000, 3), (4097, 8)]) + def test_shards_cover_every_position_once(self, seq, sp): + bounds = [SPRowLayout(seq, 1, sp, r).bounds(r) for r in range(sp)] + assert bounds[0][0] == 0 and bounds[-1][1] == seq + assert all(a[1] == b[0] and a[0] < a[1] for a, b in zip(bounds, bounds[1:])) + + @pytest.mark.parametrize("seq, batch, sp, rank", [(3, 1, 4, 0), (8, 0, 2, 0), (8, 1, 2, 2)]) + def test_rejects(self, seq, batch, sp, rank): + with pytest.raises(ValueError): + sp_row_layout(seq, batch, sp, rank) + + +@requires_cuda +class TestBackwardRanks: + @pytest.mark.parametrize( + "sp, batch, seq, hidden, layout, dtype", + [ + (2, 1, 4097, 5376, "block", torch.bfloat16), + (2, 1, 4097, 5376, "interleaved", torch.bfloat16), + (3, 2, 1000, 256, "interleaved", torch.bfloat16), + (4, 2, 777, 512, "block", torch.float32), + (5, 1, 1301, 128, "interleaved", torch.bfloat16), + (8, 2, 2049, 256, "block", torch.bfloat16), + (8, 1, 40, 64, "interleaved", torch.float32), # shards far smaller than a tile + ], + ) + def test_ranks_reproduce_ws1(self, sp, batch, seq, hidden, layout, dtype): + _require_ext() + from rl_engine import _C + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import ( + sp_gate_residual_backward, + sp_plan, + sp_rmsnorm_backward, + ) + + c = sp_case(batch, seq, hidden, layout=layout, dtype=dtype, seed=sp * 100 + seq) + chunks = c["table"].chunk(6, dim=-1) + shift, scale, gate = chunks[3], chunks[4], chunks[2] # shift/scale_mlp, gate_msa + x2, g2, y2 = (c[k].view(-1, hidden) for k in ("residual", "grad", "y")) + idx, rows = c["index"], shift.shape[0] + tiles = _segment_tiles(idx.repeat(batch), rows) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5, shift, scale, idx) + ws1_norm = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd, shift, scale, idx, *tiles) + ws1_gate = _C.h3_gate_residual_backward(g2, y2, gate, idx, *tiles) + + def rank(coll): + lay = sp_row_layout(seq, batch, sp, coll.rank) + part = slice(lay.lo, lay.hi) + local = { + k: c[k][:, part].reshape(-1, hidden).contiguous() for k in ("residual", "grad", "y") + } + rs = rstd.view(batch, seq)[:, part].reshape(-1).contiguous() + plan = sp_plan(lay, idx, rows) + norm = sp_rmsnorm_backward( + local["grad"], local["residual"], c["weight"], rs, shift, scale, plan, coll + ) + gate_grads = sp_gate_residual_backward(local["grad"], local["y"], gate, plan, coll) + return lay, norm, gate_grads + + for lay, (dx, dw, dsh, dsc), (dy, dgate) in _ThreadGroup(sp).run(rank): + part = slice(lay.lo, lay.hi) + assert torch.equal( + dx.view(batch, -1, hidden), ws1_norm[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dw, ws1_norm[1]) + assert torch.equal(dsh, ws1_norm[2].to(dtype)) and torch.equal( + dsc, ws1_norm[3].to(dtype) + ) + assert torch.equal( + dy.view(batch, -1, hidden), ws1_gate[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dgate, ws1_gate[1]) + + def test_plain_norm_ranks_reproduce_ws1(self): + _require_ext() + from rl_engine import _C + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import sp_plan, sp_rmsnorm_backward + + c = sp_case(2, 1500, 256, seed=5) + x2, g2 = c["residual"].view(-1, 256), c["grad"].view(-1, 256) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5) + ws1 = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd) + + def rank(coll): + lay = sp_row_layout(1500, 2, 4, coll.rank) + part = slice(lay.lo, lay.hi) + loc = [c[k][:, part].reshape(-1, 256).contiguous() for k in ("grad", "residual")] + rs = rstd.view(2, 1500)[:, part].reshape(-1).contiguous() + return sp_rmsnorm_backward(*loc, c["weight"], rs, None, None, sp_plan(lay), coll) + + assert all(torch.equal(dw, ws1[1]) for _, dw, _, _ in _ThreadGroup(4).run(rank)) + + def test_fails_closed(self): + _require_ext() + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + c = sp_case(1, 100, 64) + shift, scale = c["table"].chunk(6, dim=-1)[3:5] + op = H3SPNormAdaLNCudaOp(SimpleNamespace(rank=1, world_size=2), seq_len=100, batch=1) + local = c["residual"][:, 50:] + with pytest.raises(ValueError): # the full sequence instead of this rank's rows + op.norm_modulated(c["residual"], c["weight"], shift, scale, c["index"]) + with pytest.raises(ValueError): # a shard of the index instead of the full index + op.norm_modulated(local, c["weight"], shift, scale, c["index"][50:]) + with pytest.raises(IndexError): + op.gate_residual(local, c["y"][:, 50:], shift, c["index"] + 9) + with pytest.raises(ValueError): + op.norm(local.cpu(), c["weight"].cpu()) + assert op.readback()["positions"] == [50, 100] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_region_byte_equal_to_ws1(world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + case = sp_case(2, 4097, layout="interleaved", seed=world) + ws1 = {k: v.cpu() for k, v in ws1_sp_region(**case).items()} + torch.cuda.empty_cache() + ranks = run_world(world, sp_region_rank, {k: v.cpu() for k, v in case.items()}) + assert sp_matches_ws1(ws1, ranks) == ALL_TRUE + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} From ef0c9ee24bf2924af0cbdea5e6d64022118a5366 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 19:43:01 +0000 Subject: [PATCH 33/89] [WS2][MiniMax-H3] tp_adaln_3mod: 8 x B200 NCCL evidence report and figure Produced by scripts/h3_ws2_evidence.py --worlds 1,2,4,8 at 4947996 on a clean tree (real NCCL, one GPU per rank); rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-tp-adaln-3mod.md | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index da7b0ae42..6f525efe6 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -70,8 +70,24 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ## Evidence -The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 -node, one GPU per rank. +![tp_adaln_3mod on 8 x B200: byte equality and time](../usage/evidence/h3-tp-adaln-3mod-b200/figure.png) + +The report is written by `scripts/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For +TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are +byte-equal to WS1 computed on the same GPU. + +| T | | WS1 (1 GPU) | TP2 | TP4 | TP8 | +| --- | --- | --- | --- | --- | --- | +| 1 | forward | 0.10 ms | 0.12 ms | 0.11 ms | 0.12 ms | +| 1 | forward + backward | 1.90 ms | 1.49 ms | 1.19 ms | 1.04 ms | +| 3 | forward | 0.10 ms | 0.11 ms | 0.11 ms | 0.12 ms | +| 3 | forward + backward | 2.25 ms | 1.77 ms | 1.48 ms | 1.32 ms | + +Times are for the slowest rank. The WS1 forward already streams the 520 MB weight at about +5 TB/s, so the TP forward is bounded by the all-gather. `DeterministicCollective` synchronises the +host on every call, which costs about 0.1 ms. The backward is dominated by the shard-local `dW` +and gains 1.7-1.8x at TP8. Each rank also holds only `1 / tp` of the weight and its gradient. ## Tests From 63b8caa226e2af3aa49b7e9a2a913f8f235692be Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 19:43:48 +0000 Subject: [PATCH 34/89] [WS2][MiniMax-H3] sp_norm_adaln: 8 x B200 NCCL evidence report and figure Produced by scripts/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 at 6c0d380 on a clean tree (real NCCL, one GPU per rank); rendered with scripts/plot_h3_evidence.py (equality-grid label moved to the panel title) and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-sp-norm-adaln.md | 25 +++++++++++++++++++++++-- scripts/plot_h3_evidence.py | 12 ++++++------ 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 167b63f7c..8506060ea 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -70,8 +70,29 @@ How many rows move depends on how the table rows interleave along the sequence: ## Evidence -The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 -node, one GPU per rank. +![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](../usage/evidence/h3-sp-norm-adaln-b200/figure.png) + +The report is written by `scripts/h3_ws2_evidence.py` from a clean tree at commit `6c0d380`, with +real NCCL processes on one 8 x B200 node, one GPU per rank. The region is +`norm2(residual + gate_msa[row] * y)` with `shift_mlp`/`scale_mlp` modulation, H = 5376, BF16. It +covers six cases: S = 4097 (block, interleaved, B = 2), 32768 (block, interleaved) and 131072 +(block). For SP 1, 2, 4 and 8, every rank's output rows, `d_residual` and `d_sublayer`, and its +`d_norm_w` and `d_table`, are byte-equal to WS1 computed on the same GPU. + +Forward + backward of the region (slowest rank): + +| S, packing | WS1 (1 GPU) | SP2 | SP4 | SP8 | +| --- | --- | --- | --- | --- | +| 4097, block | 2.04 ms | 1.73 ms | 1.88 ms | 2.43 ms | +| 32768, block | 3.72 ms | 2.85 ms | 2.40 ms | 2.72 ms | +| 32768, interleaved | 3.72 ms | 3.45 ms | 3.88 ms | 6.82 ms | +| 131072, block | 10.73 ms | 7.27 ms | 4.66 ms | 3.92 ms | + +At long sequences the rows split and the gain grows: 2.74x at SP8 for S = 131072. At short +sequences, each backward's handful of host-synchronising all-gathers (about 0.1 ms each) +dominates. Interleaved packing also moves more rows. A naive SP8 backward (each rank's WS1 +backward, then a rank-order sum) differs from WS1 on 40% of `d_norm_w` elements and on 6% (block) +or 21% (interleaved) of `d_table` elements. ## Tests diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 6ee3de114..c4f7c05c6 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -513,10 +513,10 @@ def plot_tp_adaln(report: dict[str, Any]): runs = report["runs"] _equality_grid( left, - "Byte-equal to WS1 (every rank, every T)", + "Byte-equal to WS1", [(f"TP{run['tp']}", [e for r in run["ranks"] for e in r["equality"]]) for run in runs], TP_TENSORS, - "rank x T", + "ranks x T", ) # Right: slowest rank's time at the largest T; WS1 on one GPU as reference lines. num_t = max(p["num_timesteps"] for p in runs[0]["ranks"][0]["perf"]) @@ -563,7 +563,7 @@ def _equality_grid(ax, title, rows, tensors, unit) -> None: """One cell per (configuration, tensor): how many checks were byte-equal to WS1.""" ax.set_facecolor(SURFACE) - ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) + ax.set_title(f"{title}; checks = {unit}", loc="left", fontsize=11, color=INK, pad=10) for y, (_, checks) in enumerate(rows): for x, (key, _) in enumerate(tensors): ok = sum(bool(e[key]) for e in checks) @@ -579,7 +579,7 @@ def _equality_grid(ax, title, rows, tensors, unit) -> None: linewidth=2, ) ) - label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} {unit}" + label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)}" ax.text(x + 0.5, y + 0.5, label, ha="center", va="center", fontsize=8, color="white") ax.set_xlim(0, len(tensors)) ax.set_ylim(len(rows), 0) @@ -624,13 +624,13 @@ def plot_sp_norm(report: dict[str, Any]): runs = report["runs"] _equality_grid( axes[0], - "Byte-equal to WS1 (every rank, every case)", + "Byte-equal to WS1", [ (f"SP{run['sp']}", [c["equal"] for r in run["ranks"] for c in r["cases"]]) for run in runs ], SP_TENSORS, - "rank x case", + "ranks x cases", ) # Middle: the largest block-layout case, slowest rank. From e4b8a62afa5980330fa73a0085e0c09b3cd58a92 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:50:47 +0000 Subject: [PATCH 35/89] fix(h3): store norm_out.norm and 1 + scale in BF16 in the final AdaLN golden The gtest gold skipped two roundings diffusers' norm_out performs, so the BF16 check missed d_norm_weight from S=257 and d_temb/dW from S=1024 for CUDA and the PyTorch reference alike. With them declared straight-through, the BF16 gtest passes at S=257/1024/4097; FP32 is unchanged and its long-reduction misses are documented. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 37 ++++++++++++++----- .../kernels/ops/pytorch/h3/final_adaln_out.py | 19 +++++++--- tests/h3/test_h3_final_adaln_out.py | 28 ++++++++++++++ 3 files changed, 70 insertions(+), 14 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 59b4ab0a9..7439e84a1 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -45,9 +45,12 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; - every gradient is repeat-bitwise. -- **Golden.** FP64, rounding only where the model declares a dtype boundary: the SiLU cast, - and `norm_out.linear`'s BF16 output table, which every position of a timestep shares. Both are - rounded straight-through for the gradient. +- **Golden.** FP64, rounding only where the model stores a value in its own dtype: the SiLU + cast, `norm_out.linear`'s BF16 output table, `norm_out.norm`'s BF16 output and the BF16 + `1 + scale`. All four are rounded straight-through for the gradient. The table and `1 + scale` + are shared by every position of a timestep. Without them the golden's `d_norm_w` drifts + systematically with S, and the gtest missed it from S = 257. `norm_out.norm`'s rounding + enters `d_scale`, `dW` and `d_temb` summed over S, and the gtest missed those from S = 1024. ## Performance Notes @@ -87,7 +90,7 @@ export RL_KERNEL_H3_WEIGHTS= python -m pytest tests/h3/test_h3_final_adaln_out.py -v # operator python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end: timestep -> ... -> norm_out python scripts/check_operator.py --op final_adaln_out --candidate cuda --device cuda \ - --dtype bf16 --batch 3 --seq 64 --normalized-dim 5376 --check-grad + --dtype bf16 --batch 3 --seq 4097 --normalized-dim 5376 --check-grad # also 257, 1024 python scripts/h3_evidence.py --op final_adaln_out \ --out docs/usage/evidence/h3-final-adaln-out-b200/report.json python scripts/plot_h3_evidence.py docs/usage/evidence/h3-final-adaln-out-b200/report.json @@ -95,9 +98,25 @@ python scripts/plot_h3_evidence.py docs/usage/evidence/h3-final-adaln-out-b200/r ## Known Limitations -- **gtest at larger S.** With `--seq 257`, a few gradient elements fail the per-element - contract tolerance (BF16 `d_norm_w`, FP32 `d_temb`). Those elements are what is left after - summing hundreds of rows into values around 1e2, which an absolute per-element atol can't - cover. Relative to the gradient's maximum the error is 4e-7 in FP32 and 5e-3 in BF16, and the - PyTorch reference fails the same elements, plus more. `--seq 64` passes in both dtypes. +- **FP32 gtest at larger S.** H3 runs this op with BF16 activations and weights, and the BF16 + gtest passes at S = 257, 1024 and 4097. With `--dtype fp32` (every input FP32), the FP32 + reduction tolerance (atol = rtol = 1e-4) misses `d_temb` from S = 257 and `dW` from + S = 1024. Both gradients are FP32 sums over 3·S rows. Where large terms cancel, the + summation error exceeds the 1e-4 absolute floor. The PyTorch reference misses the same + gradients at the same S (`--batch 3`, seed 123, max abs error): + + | S | `d_temb` CUDA | `d_temb` PyTorch | `dW` CUDA | `dW` PyTorch | + | --- | --- | --- | --- | --- | + | 64 | 9.2e-5 | 1.5e-4 | 1.5e-4 | 9.2e-5 | + | 257 | 3.1e-4 ✗ | 5.9e-4 ✗ | 8.5e-4 | 2.4e-4 | + | 1024 | 1.5e-3 ✗ | 2.4e-3 ✗ | 2.0e-3 ✗ | 1.5e-3 ✗ | + | 4097 | 5.9e-3 ✗ | 9.3e-3 ✗ | 5.9e-3 ✗ | 3.9e-3 ✗ | + + ✗ fails the gtest. Unmarked values above 1e-4 pass through the rtol term. In a PyTorch + emulation, accumulating only the projection backward in FP64 clears S = 257 but not + S ≥ 1024, because the segment sums contribute as well. Making every reduction FP64 would + change the `h3_det_linear` kernels shared with `timestep_mlp_fp32` and + `adaln_projection_3mod`. The alternative is an FP32 gradient + tolerance that scales with reduction length. Which of the two to take is an open question + for the maintainers. - There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py b/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py index 94aaa28c1..b40924e12 100644 --- a/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py +++ b/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py @@ -58,9 +58,10 @@ def validate_h3_final_adaln_out( class NativeH3FinalAdaLNOutOp: """PyTorch reference: ``forward`` replays diffusers' ``norm_out``; - ``forward_fp32`` is the FP64 golden, returned in FP32. It rounds only at the - module boundaries the model declares (the SiLU cast and the BF16 table), - straight-through for the gradient.""" + ``forward_fp32`` is the FP64 golden, returned in FP32. It rounds only where + the model itself stores a value in its own dtype (the SiLU cast, the BF16 + table, the ``norm_out.norm`` output and ``1 + scale``), straight-through for + the gradient, and computes everything else in FP64.""" op_class = "reduction" @@ -89,7 +90,15 @@ def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps shift, scale = table.chunk(2, dim=-1) x64 = x.double() n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * norm_weight.double() - out = n * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( - 0, timestep_indices + # norm_out.norm is an x.dtype module and diffusers forms 1 + scale in + # x.dtype, so both are stored rounded. 1 + scale is shared by every + # position of a timestep: left unrounded, its error is systematic and + # grows with S in d_norm_weight; n's rounding enters d_scale (and so + # dW, d_temb) summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.index_select(0, timestep_indices) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() ) + out = n * one_plus_scale + shift.index_select(0, timestep_indices) return out.float() diff --git a/tests/h3/test_h3_final_adaln_out.py b/tests/h3/test_h3_final_adaln_out.py index 17a3843b8..7a1f5b6d8 100644 --- a/tests/h3/test_h3_final_adaln_out.py +++ b/tests/h3/test_h3_final_adaln_out.py @@ -169,3 +169,31 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("final_adaln_out", device="cuda") assert type(op).__name__ == "H3FinalAdaLNOutCudaOp" + + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_bf16_gradients_meet_gtest_contract(self, seq): + # The check_operator path; S=257 used to miss d_norm_weight before the + # golden stored norm(x) and 1 + scale in BF16 like the model does. + import argparse + + from rl_engine.kernels.gtest import run_operator_suite + from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case + + _cuda_op() + args = argparse.Namespace( + op="final_adaln_out", + candidate="cuda", + arch_key=None, + batch=3, + seq=seq, + normalized_dim=5376, + seed=123, + input_mode="random", + ) + report = run_operator_suite( + "final_adaln_out", + candidates=[make_candidate(args)], + cases=[make_operator_case(args, torch.bfloat16, torch.device("cuda"))], + check_grad=True, + ) + assert report.passed, report.candidates[0].cases[0].outputs From dcbd55bfd63cb067f407f3abdcc08fa674ea6638 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:57:30 +0000 Subject: [PATCH 36/89] docs(h3): regenerate final AdaLN evidence with the BF16 norm_out golden Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 7439e84a1..b7dd49d65 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -40,7 +40,7 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste summation tree versus cuBLAS. Rows are batch- and position-invariant. - **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into the projection backward without an extra BF16 rounding: - - `d_temb`, `dW` and `db` are 22–43× closer to FP64 in the recorded run than diffusers, + - `d_temb`, `dW` and `db` are 22–48× closer to FP64 in the recorded run than diffusers, whose error varies from run to run because it rounds the table gradient to BF16 and accumulates it with `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; @@ -62,9 +62,9 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.14 ms | 0.15 ms | 1.67 ms | 1.12 ms | -| 32768 | 0.40 ms | 0.78 ms | 2.97 ms | 5.39 ms | -| 131072 | 1.25 ms | 2.94 ms | 7.87 ms | 22.2 ms | +| 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | +| 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | +| 131072 | 1.24 ms | 2.94 ms | 7.75 ms | 20.6 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -77,11 +77,13 @@ There are two data files: - [`report.json`](../usage/evidence/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `372d6d7`, with FP64 leaves and upstream gradients and explicit BF16 forward boundaries. + tree at `001684d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + backward reference keeps only the SiLU and table roundings, so the errors include the BF16 + rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../usage/evidence/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `38d575f`. + tree at `001684d`. ## Tests From e880a103c8dc8e36eff468d22c811a2843a09ae7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:33:12 +0000 Subject: [PATCH 37/89] test(h3): assert adaln_row_gather batch invariance per logical row A position's six gathered rows are bitwise the same whichever other positions share the call (single positions, spans, the tail, a shuffled repacking), and a table row's gradient from the full packing equals the one from a packing of only its own positions, across tile boundaries. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/h3/test_h3_adaln_row_gather.py | 38 +++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/tests/h3/test_h3_adaln_row_gather.py b/tests/h3/test_h3_adaln_row_gather.py index 8f6558034..4bbe0dc10 100644 --- a/tests/h3/test_h3_adaln_row_gather.py +++ b/tests/h3/test_h3_adaln_row_gather.py @@ -8,7 +8,10 @@ * semantic indices fail closed (out-of-range tags or timesteps, RFC probes H2/H3), and flipping a tag or offsetting a timestep picks a different row; * backward is a deterministic FP32 segmented sum: repeat-bitwise and - correctly rounded, unlike the atomic BF16 ``index_select`` backward. + correctly rounded, unlike the atomic BF16 ``index_select`` backward; +* batch invariance: a position's output does not depend on which other + positions share the call, and a row's gradient does not depend on the + positions that reference other rows. """ from __future__ import annotations @@ -219,6 +222,39 @@ def test_repeat_bitwise(self): for _ in range(3): assert torch.equal(self._cuda_grad(rows, ti, tags, grads), first) + def test_forward_rows_invariant_to_batch_size_and_position(self): + rows = _rows(3, hidden=512, seed=5) + ti, tags = h3_packed_layout(4097, 3, seed=5) + full = _cuda_op()(rows, ti, tags) + perm = torch.randperm(4097, generator=torch.Generator().manual_seed(5)).cuda() + for pick in ( + torch.tensor([0]), + torch.tensor([2048]), + torch.arange(100, 357), + torch.arange(4000, 4097), + perm, + ): + pick = pick.cuda() + part = _cuda_op()(rows, ti[pick], tags[pick]) + for a, b in zip(part, full): + assert torch.equal(a, b[pick]) + + def test_backward_row_gradient_independent_of_other_rows_tokens(self): + # Each row's segment is tiled on its own, so a row's gradient from the + # full packing equals the one from a packing of only its own positions + # (same order), even though ~455 positions per row cross tile boundaries. + rows = _rows(3, hidden=512, seed=6) + ti, tags = h3_packed_layout(4097, 3, seed=6) + grads = self._grads(4097, 512, seed=6) + full = self._cuda_grad(rows, ti, tags, grads) + flat = ti * 3 + tags + for r in range(rows.shape[0]): + own = (flat == r).nonzero().squeeze(1) + assert own.numel() > 256 + alone = self._cuda_grad(rows, ti[own], tags[own], [g[own] for g in grads]) + assert torch.equal(alone[r], full[r]) + assert torch.count_nonzero(alone[torch.arange(rows.shape[0], device="cuda") != r]) == 0 + def test_unreferenced_rows_get_zero_grad(self): rows = _rows(3, hidden=16) ti = torch.zeros(7, dtype=torch.long, device="cuda") From 93c692f5479a6f2b54944ad108a2f29b8262af4a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:44:55 +0000 Subject: [PATCH 38/89] fix(h3): store norm(x) and 1 + scale in BF16 in the modulated RMSNorm golden forward_modulated_fp32 computed both in FP64, but the model stores them in x's dtype. 1 + scale is shared by every position of a row, so against the old golden the CUDA d_weight missed the BF16 reduction tolerance from S=1024 and d_scale at S=4097. With both declared straight-through every gradient passes at S=257/1024/4097; the new test fails on the old golden. No evidence changes: the report uses its own FP64 reference. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-rmsnorm.md | 2 +- rl_engine/kernels/ops/pytorch/h3/rmsnorm.py | 16 +++++++--- tests/h3/test_h3_rmsnorm.py | 35 +++++++++++++++++++++ 3 files changed, 48 insertions(+), 5 deletions(-) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 3c69706e7..21c6b5c6f 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -31,7 +31,7 @@ out = op.forward_modulated(x, weight, shift_msa, scale_msa, adaln_indices) | Backend | Wrapper | Native symbols | Status | | --- | --- | --- | --- | | CUDA (SM80+, validated on SM100) | `rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp` | `rl_engine._C.h3_rmsnorm_{forward,backward}` | Bitwise equal to `nn.RMSNorm` and to the diffusers modulation | -| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp` | n/a | `forward*`: provider path; `forward*_fp32`: FP64 golden | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp` | n/a | `forward*`: provider path; `forward*_fp32`: FP64 golden (the modulated one stores `norm(x)` and `1 + scale` in BF16, straight-through, as the model does) | | ROCm | n/a | n/a | Falls back to the PyTorch reference | ## Tensor Contract diff --git a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py index aeee51754..1727fd1d0 100644 --- a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py +++ b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py @@ -107,8 +107,9 @@ class NativeH3RMSNormOp: ``forward`` / ``forward_modulated`` are the provider path (``F.rms_norm`` and the eager modulation expression, with a deterministic shift/scale gradient; the raw diffusers path is ``h3_provider.provider_norm_modulate``). ``forward_fp32`` / - ``forward_modulated_fp32`` are the golden: the same math in FP64 with no - intermediate rounding, returned in FP32. + ``forward_modulated_fp32`` are the golden: the same math in FP64, returned in + FP32. The modulated golden rounds, straight-through, only the two values the + model stores in ``x``'s dtype: ``norm(x)`` and ``1 + scale``. """ op_class = "reduction" @@ -139,7 +140,14 @@ def forward_modulated_fp32(self, x, weight, shift, scale, index, eps: float = H3 validate_h3_rmsnorm(x, weight, eps) x64 = x.double() n = x64 * torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) * weight.double() - out = n * (1.0 + scale.double().index_select(0, index)) + shift.double().index_select( - 0, index + # The norm is an x.dtype module and diffusers forms 1 + scale in x.dtype. + # 1 + scale is shared by every position of a row, so leaving it unrounded + # skews d_norm_weight systematically with S; n's rounding enters d_scale + # summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.double().index_select(0, index) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() ) + out = n * one_plus_scale + shift.double().index_select(0, index) return out.float() diff --git a/tests/h3/test_h3_rmsnorm.py b/tests/h3/test_h3_rmsnorm.py index 53e46b16b..bf655e423 100644 --- a/tests/h3/test_h3_rmsnorm.py +++ b/tests/h3/test_h3_rmsnorm.py @@ -189,6 +189,41 @@ def golden(x_, w_, t_): again = self._grads(ours, (x, w, table), grad) assert all(torch.equal(a, b) for a, b in zip(got, again)) + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_modulated_gradients_meet_contract_against_golden(self, seq): + # Against forward_modulated_fp32 before it stored norm(x) and 1 + scale + # in BF16, d_weight missed from S=1024 and d_scale at S=4097. + from rl_engine.kernels.gtest.tolerance import load_contract, resolve_tolerance + + spec = resolve_tolerance( + load_contract(), + judgment="gradient_accuracy", + op_class="reduction", + dtype=torch.bfloat16, + ) + g = torch.Generator(device="cpu").manual_seed(seq) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.1).bfloat16().cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seq) + index = ti * 3 + tags + x = _x((1, seq, 5376), seed=seq) + grad = _x((1, seq, 5376), seed=seq + 1, scale=1.0) + op, native = _cuda_op(), NativeH3RMSNormOp() + + def run(fn, upstream): + def call(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return fn(x_, w_, shift, scale, index) + + return self._grads(call, (x, w, table), upstream) + + got = run(op.forward_modulated, grad) + ref = run(native.forward_modulated_fp32, grad.float()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + torch.testing.assert_close( + g_.float(), r.float(), atol=spec.atol, rtol=spec.rtol, msg=lambda m: f"{name}: {m}" + ) + def test_plain_backward_against_fp64_and_row_local_dx(self): w = (torch.rand(5376, generator=torch.Generator().manual_seed(7)) + 0.5).bfloat16().cuda() x = _x((4, 64, 5376), seed=7) From edca55999ce595595f48526d9f71f4cda61b4f1d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:43:58 +0000 Subject: [PATCH 39/89] feat(scripts): H3 prior-art runner: batch invariance, accuracy and latency vs existing implementations The RFC #420 reuse rule asks whether an existing implementation is batch-invariant and for accuracy and performance comparisons. scripts/h3_prior_art.py measures the operator's rl-kernel op next to diffusers, SGLang, Liger, Transformer Engine, and the vLLM / SGLang / Megatron batch-invariant modes, whichever import. Each global mode runs in its own subprocess with its environment set before CUDA starts. Batch invariance checks every row alone vs full batches, the full workload batch vs covering sub-batches, and a dense batch-size sweep. Only operators whose op is on the branch can be selected, so later stacked PRs reuse the same runner. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- scripts/h3_prior_art.py | 929 +++++++++++++++++++++++++++++++++++ scripts/plot_h3_prior_art.py | 126 +++++ 2 files changed, 1055 insertions(+) create mode 100644 scripts/h3_prior_art.py create mode 100644 scripts/plot_h3_prior_art.py diff --git a/scripts/h3_prior_art.py b/scripts/h3_prior_art.py new file mode 100644 index 000000000..c39adcf0a --- /dev/null +++ b/scripts/h3_prior_art.py @@ -0,0 +1,929 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Compare one RFC #420 operator with existing implementations: batch invariance, accuracy, latency. + +The RFC #420 reuse rule asks, for every operator, whether an existing implementation is +batch-invariant (then reuse it) and for accuracy and performance comparisons. This runner +measures the operator's rl-kernel CUDA op next to the existing implementations that import in +the current environment, and writes one JSON report per operator: + + python scripts/h3_prior_art.py --op adaln_row_gather \\ + --out docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json \\ + [--megatron-src /path/to/Megatron-LM] [--quick] + python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json + +Only operators whose rl-kernel op exists on the current branch can be selected. Optional +libraries (SGLang, Liger, Transformer Engine, vLLM, Megatron-LM from a source tree) are skipped +and recorded as unavailable when they do not import. Global batch-invariant modes (vLLM, +SGLang, Megatron) patch aten for the whole process, so each runs in its own subprocess, with +the environment variables it relies on set before CUDA initialises. + +Batch invariance is bitwise, for the forward, the per-row gradients, and repeatability of every +parameter and table gradient, with three checks: + +1. every row computed alone vs inside full batches of 64, 257 and 2048 rows, seeds 3-5; +2. the full workload-size batch vs sub-batches that together cover every row; +3. eight probe rows at the front, middle and back of batches of every size 1..9 and + 2^k - 1, 2^k, 2^k + 1, plus the whole batch reversed. + +Sparse probes (a few rows, a few batch sizes) miss row-specific or large-batch dependence, which +is why every row is checked. Accuracy is max|err| / max|ref| against the same computation in +FP64; latency is the median of 100 CUDA-event samples after 20 warm-ups. Run on a clean tree on +an otherwise idle GPU; the report records the commit, the tree state and every library version. +""" + +from __future__ import annotations + +import argparse +import contextlib +import importlib.metadata as md +import importlib.util +import json +import os +import statistics +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any, Callable + +REPO_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO_ROOT)) + +import torch # noqa: E402 +import torch.nn.functional as F # noqa: E402 + +H, D, F1, T3, EPS = 5376, 2688, 256, 3, 1e-5 + +#: operator -> rl-kernel module that must exist on the branch +OPS = { + "timestep_sinusoid": "rl_engine.kernels.ops.cuda.h3.timestep_sinusoid", + "timestep_mlp": "rl_engine.kernels.ops.cuda.h3.timestep_mlp", + "adaln_projection": "rl_engine.kernels.ops.cuda.h3.adaln_projection", + "adaln_row_gather": "rl_engine.kernels.ops.cuda.h3.adaln_row_gather", + "norm_modulate": "rl_engine.kernels.ops.cuda.h3.rmsnorm", + "gate_residual": "rl_engine.kernels.ops.cuda.h3.gate_residual", + "final_adaln_out": "rl_engine.kernels.ops.cuda.h3.final_adaln_out", +} +#: global batch-invariant modes compared for the GEMM operators, with their environment +MODES = { + "plain": {}, + "vllm": {"CUBLAS_WORKSPACE_CONFIG": ":16:8", "CUBLASLT_WORKSPACE_SIZE": "1"}, + "sglang": {}, + "sglang_ieee": {"TRITON_F32_DEFAULT": "ieee"}, + "megatron_te_native": {"CUBLASLT_WORKSPACE_SIZE": "0"}, + "megatron_triton": {}, + "megatron_triton_ieee": {"TRITON_F32_DEFAULT": "ieee"}, +} +GEMM_OPS = ("timestep_mlp", "adaln_projection") +LIBS = ("torch", "triton", "diffusers", "vllm", "flashinfer-python", "sglang", "liger-kernel") +LIBS += ("transformer_engine", "megatron-core") + +DEV = "cuda" + + +# --------------------------------------------------------------------------- # +# Batch-invariance checks +# --------------------------------------------------------------------------- # + + +class _Case: + """``fn(**inputs)`` with row arguments sliced together; fixed upstream gradients per seed.""" + + def __init__(self, fn, inputs, rows, grads, out_dims, seed): + self.fn, self.inputs, self.rows, self.grads, self.out_dims = ( + fn, + inputs, + rows, + grads, + out_dims, + ) + self.n = next(inputs[k].shape[d] for k, d in rows.items()) + outs = self._flat(fn(**inputs)) + gen = torch.Generator(device=DEV).manual_seed(seed) + self.ups = [torch.randn(o.shape, device=DEV, generator=gen).to(o.dtype) for o in outs] + + @staticmethod + def _flat(out): + return list(out) if isinstance(out, (tuple, list)) else [out] + + def run(self, idx=None): + args = {} + for k, v in self.inputs.items(): + if idx is not None and k in self.rows: + v = v.index_select(self.rows[k], idx) + if k in self.grads: + v = v.detach().clone().requires_grad_(True) + args[k] = v + outs = self._flat(self.fn(**args)) + ups = ( + self.ups + if idx is None + else [u.index_select(d, idx) for u, d in zip(self.ups, self.out_dims)] + ) + live = [(o, u) for o, u in zip(outs, ups) if o.requires_grad] + if live: + torch.autograd.backward([o for o, _ in live], [u for _, u in live]) + grads = {k: args[k].grad for k in self.grads} + return [o.detach() for o in outs], grads + + +def _sweep_sizes(n: int) -> list[int]: + sizes = set(range(1, min(n, 9) + 1)) + k = 16 + while k <= n: + sizes.update(v for v in (k - 1, k, k + 1) if v <= n) + k *= 2 + return sorted(sizes | {n}) + + +def batch_invariance(cand: dict[str, Any], quick: bool) -> dict[str, Any]: + fn, make, rows, grads = cand["fn"], cand["make"], cand["rows"], cand["grads"] + out_dims = cand.get("out_dims", (0,)) + params = [k for k in grads if k not in rows] + full_sizes, big, subs = cand["bi_sizes"] + if quick: + full_sizes, big, subs = (17, 64), 256, (1, 7, 64) + seeds = (3, 4, 5) + result: dict[str, Any] = {"all_rows": {}, "params_repeatable": None} + ok = True + + def same_row(a, b, d, i): + return torch.equal(a.select(d, 0), b.select(d, i)) + + for n in full_sizes: + fwd = grd = 0 + repeat = True + for seed in seeds: + case = _Case(fn, make(seed, n), rows, grads, out_dims, seed) + fo, fg = case.run() + for i in range(n): + o, g = case.run(torch.tensor([i], device=DEV)) + fwd += not all(same_row(a, b, d, i) for a, b, d in zip(o, fo, out_dims)) + grd += not all( + same_row(g[k], fg[k], rows[k], i) + for k in rows + if k in grads and g[k] is not None + ) + if params: + _, again = case.run() + repeat &= all(torch.equal(again[k], fg[k]) for k in params if fg[k] is not None) + result["all_rows"][str(n)] = { + "rows": len(seeds) * n, + "fwd_differ": fwd, + "rowgrad_differ": grd, + } + if params: + result["params_repeatable"] = repeat and result["params_repeatable"] is not False + ok &= fwd == 0 and grd == 0 and repeat + + fwd = grd = checked = 0 + for seed in seeds[:2]: + case = _Case(fn, make(seed, big), rows, grads, out_dims, seed) + fo, fg = case.run() + for sb in subs: + step = sb if sb >= 1024 else max(sb, big // 512) + for start in range(0, big, step): + idx = torch.arange(start, min(start + sb, big), device=DEV) + o, g = case.run(idx) + checked += 1 + fwd += not all( + torch.equal(a, b.index_select(d, idx)) for a, b, d in zip(o, fo, out_dims) + ) + grd += not all( + torch.equal(g[k], fg[k].index_select(rows[k], idx)) + for k in rows + if k in grads and g[k] is not None + ) + del case, fo, fg + torch.cuda.empty_cache() + result["full_vs_sub_batches"] = { + "full_batch": big, + "sub_batch_sizes": list(subs), + "sub_batches": checked, + "fwd_differ": fwd, + "rowgrad_differ": grd, + } + ok &= fwd == 0 and grd == 0 + + n = full_sizes[-1] + sizes = _sweep_sizes(n) + probes = sorted({0, n - 1, *list(range(n // 8 + 3, n - 1, n // 8))[:6]}) + compared = failed = 0 + first: list[dict[str, int]] = [] + for seed in seeds: + case = _Case(fn, make(seed, n), rows, grads, out_dims, seed) + alone = {r: case.run(torch.tensor([r], device=DEV)) for r in probes} + gen = torch.Generator().manual_seed(seed) + + def compare(idx, pos, r, seed=seed, case=case, alone=alone): + nonlocal compared, failed + o, g = case.run(idx.to(DEV)) + ao, ag = alone[r] + good = all( + torch.equal(a.select(d, pos), b.select(d, 0)) for a, b, d in zip(o, ao, out_dims) + ) + good &= all( + torch.equal(g[k].select(rows[k], pos), ag[k].select(rows[k], 0)) + for k in rows + if k in grads and g[k] is not None + ) + compared += 1 + if not good: + failed += 1 + if len(first) < 8: + first.append( + {"seed": seed, "batch": int(idx.numel()), "position": pos, "row": r} + ) + + for m in sizes: + for r in probes: + others = torch.randperm(n, generator=gen) + others = others[others != r][: m - 1] + for pos in sorted({0, (m - 1) // 2, m - 1}): + compare(torch.cat([others[:pos], torch.tensor([r]), others[pos:]]), pos, r) + for r in probes: + compare(torch.arange(n - 1, -1, -1), n - 1 - r, r) + result["batch_size_sweep"] = { + "batch_sizes": sizes, + "probe_rows": probes, + "comparisons": compared, + "differ": failed, + "first_failures": first, + } + ok &= failed == 0 + result["batch_invariant"] = bool(ok) + return result + + +# --------------------------------------------------------------------------- # +# Accuracy and latency +# --------------------------------------------------------------------------- # + + +def _time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 100) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(iters): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1e3) + return statistics.median(samples) + + +def _rel(a: torch.Tensor, ref: torch.Tensor) -> float: + return ((a.double() - ref).abs().max() / ref.abs().max().clamp_min(1e-300)).item() + + +def accuracy_latency(cand: dict[str, Any], mode_off, quick: bool) -> dict[str, Any]: + out: dict[str, Any] = {} + for size in cand["perf_sizes"][:2] if quick else cand["perf_sizes"]: + inputs = cand["make"](11, size) + names = list(inputs) + grads = [k for k in cand["grads"] if k in names] + leaves = { + k: (v.detach().clone().requires_grad_(True) if k in grads else v) + for k, v in inputs.items() + } + y = cand["fn"](**leaves) + y = torch.cat([t.reshape(-1) for t in y]) if isinstance(y, (tuple, list)) else y + with mode_off(): + ref_in = { + k: (v.detach().double().requires_grad_(k in grads) if v.is_floating_point() else v) + for k, v in inputs.items() + } + ref = cand["ref"](**ref_in) + ref = torch.cat([t.reshape(-1) for t in ref]) if isinstance(ref, (tuple, list)) else ref + row: dict[str, Any] = {"fwd_rel_err": _rel(y, ref.detach())} + if grads: + gen = torch.Generator(device=DEV).manual_seed(9) + up = torch.randn(y.shape, device=DEV, generator=gen).to(y.dtype) + y.backward(up) + with mode_off(): + ref.backward(up.double()) + row["grad_rel_err"] = {k: _rel(leaves[k].grad, ref_in[k].grad) for k in grads} + row["fwd_us"] = _time_us(lambda inputs=inputs: cand["fn"](**inputs)) + if grads: + + def step(inputs=inputs): + lv = { + k: (v.detach().clone().requires_grad_(True) if k in grads else v) + for k, v in inputs.items() + } + o = cand["fn"](**lv) + outs = list(o) if isinstance(o, (tuple, list)) else [o] + torch.autograd.backward(outs, [torch.ones_like(t) for t in outs]) + + row["fwd_bwd_us"] = _time_us(step) + out[str(size)] = row + del inputs, leaves, y, ref + torch.cuda.empty_cache() + return out + + +# --------------------------------------------------------------------------- # +# Candidates +# --------------------------------------------------------------------------- # + + +def _gen(seed: int, salt: int) -> torch.Generator: + return torch.Generator(device=DEV).manual_seed(seed * 1000 + salt) + + +def _rn(seed, salt, *shape, dtype=torch.float32, scale=1.0): + return (torch.randn(*shape, device=DEV, generator=_gen(seed, salt)) * scale).to(dtype) + + +def _layout(n: int, seed: int): + from rl_engine.testing.h3_cases import h3_packed_layout + + ti, tags = h3_packed_layout(n, T3, seed=seed) + return ti.contiguous(), tags.contiguous() + + +def _version(name: str) -> str | None: + try: + return md.version(name) + except md.PackageNotFoundError: + return None + + +T_BI = ((64, 257, 2048), 4096, (1, 7, 1024, 2048)) +S_BI = ((64, 257, 2048), 131072, (1, 7, 4097, 32768, 65536)) +T_PERF = (1, 3, 4, 64, 256, 2048) +S_PERF = (4097, 32768) + + +def _sinusoid_ref(t): + half = F1 // 2 + exponent = -torch.log(torch.tensor(10000.0, dtype=torch.float64, device=DEV)) + freqs = torch.exp(exponent * torch.arange(half, device=DEV, dtype=torch.float64) / half) + e = t.double()[:, None] * freqs[None] + return torch.cat([torch.cos(e), torch.sin(e)], -1) + + +def candidates(op: str, mode: str) -> list[dict[str, Any]]: + """Candidate dicts; a factory that raises (missing library) becomes an unavailable entry.""" + + from rl_engine.testing import h3_provider as P + + plain = mode == "plain" + found: list[tuple[str, Callable[[], dict[str, Any]]]] = [] + + if op == "timestep_sinusoid": + base = { + "make": lambda s, n: {"t": torch.rand(n, device=DEV, generator=_gen(s, 1))}, + "rows": {"t": 0}, + "grads": [], + "ref": _sinusoid_ref, + "bi_sizes": ((64, 257, 2048), 8192, (1, 7, 2048, 4097)), + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers get_timestep_embedding (op-for-op replay)", + "fn": lambda t: P.provider_time_proj(t), + } + + def sglang(): + from sglang.kernels.ops.diffusion.modulate.timestep_embedding_jit import ( + timestep_embedding, + ) + + return { + **base, + "source": f"SGLang {_version('sglang')} timestep_embedding", + "fn": lambda t: timestep_embedding( + t, F1, flip_sin_to_cos=True, downscale_freq_shift=0.0 + ), + } + + def ours(): + from rl_engine.kernels.ops.cuda.h3.timestep_sinusoid import H3TimestepSinusoidCudaOp + + op_ = H3TimestepSinusoidCudaOp() + return { + **base, + "source": "rl-kernel H3TimestepSinusoidCudaOp (check_range=False)", + "fn": lambda t: op_.forward(t, check_range=False), + } + + found = [("diffusers", diffusers), ("sglang", sglang), ("rl_kernel", ours)] + + elif op == "timestep_mlp": + + def make(s, n): + return { + "x": _rn(s, 2, n, F1), + "w1": _rn(1, 3, H, F1, scale=F1**-0.5), + "b1": _rn(1, 4, H, scale=0.02), + "w2": _rn(1, 5, D, H, scale=H**-0.5), + "b2": _rn(1, 6, D, scale=0.02), + } + + base = { + "make": make, + "rows": {"x": 0}, + "grads": ["x", "w1", "b1", "w2", "b2"], + "ref": lambda x, w1, b1, w2, b2: F.linear(F.silu(F.linear(x, w1, b1)), w2, b2), + "bi_sizes": T_BI, + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": f"diffusers TimestepEmbedding, FP32 F.linear [{mode}]", + "fn": lambda x, w1, b1, w2, b2: P.provider_time_embedder(x, w1, b1, w2, b2), + } + + def ours(): + from rl_engine.kernels.ops.cuda.h3.timestep_mlp import H3TimestepMLPCudaOp + + op_ = H3TimestepMLPCudaOp() + return { + **base, + "source": "rl-kernel H3TimestepMLPCudaOp", + "fn": lambda x, w1, b1, w2, b2: op_(x, w1, b1, w2, b2), + } + + found = [(f"diffusers[{mode}]", diffusers)] + ([("rl_kernel", ours)] if plain else []) + + elif op == "adaln_projection": + + def make(s, n): + return { + "temb": _rn(s, 7, n, D), + "w": _rn(1, 8, 18 * H, D, dtype=torch.bfloat16, scale=D**-0.5), + "b": _rn(1, 9, 18 * H, dtype=torch.bfloat16, scale=0.02), + } + + base = { + "make": make, + "rows": {"temb": 0}, + "grads": ["temb", "w", "b"], + "ref": lambda temb, w, b: F.linear(F.silu(temb), w, b), + "bi_sizes": T_BI, + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": f"diffusers AdaLN projection, BF16 F.linear [{mode}]", + "fn": lambda temb, w, b: F.linear(F.silu(temb).to(w.dtype), w, b), + } + + def ours(): + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + + op_ = H3AdaLNProjectionCudaOp() + return { + **base, + "source": "rl-kernel H3AdaLNProjectionCudaOp", + "fn": lambda temb, w, b: op_.forward_table(temb, w, b), + } + + found = [(f"diffusers[{mode}]", diffusers)] + ([("rl_kernel", ours)] if plain else []) + + elif op == "adaln_row_gather": + + def make(s, n): + ti, tags = _layout(n, s) + return { + "table": _rn(1, 14, 3 * T3, 6 * H, dtype=torch.bfloat16, scale=0.1), + "ti": ti, + "tags": tags, + } + + def ref(table, ti, tags): + return tuple(table.index_select(0, ti * 3 + tags).chunk(6, dim=1)) + + base = { + "make": make, + "rows": {"ti": 0, "tags": 0}, + "grads": ["table"], + "out_dims": (0,) * 6, + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return {**base, "source": "diffusers six index_select calls", "fn": ref} + + def ours(): + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import H3AdaLNRowGatherCudaOp + + op_ = H3AdaLNRowGatherCudaOp() + return { + **base, + "source": "rl-kernel H3AdaLNRowGatherCudaOp", + "fn": lambda table, ti, tags: op_(table, ti, tags), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + elif op == "norm_modulate": + + def make(s, n): + ti, tags = _layout(n, s) + w = (0.5 + torch.rand(H, device=DEV, generator=_gen(1, 11))).bfloat16() + return { + "x": _rn(s, 10, n, H, dtype=torch.bfloat16, scale=2.0), + "w": w, + "sh": _rn(1, 12, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "sc": _rn(1, 13, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "idx": ti * 3 + tags, + } + + def ref(x, w, sh, sc, idx): + n = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w + return n * (1 + sc.index_select(0, idx)) + sh.index_select(0, idx) + + base = { + "make": make, + "rows": {"x": 0, "idx": 0}, + "grads": ["x", "w", "sh", "sc"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def per_token(fn): + # Liger / SGLang take per-token shift/scale rows: gather them first (index_select). + return lambda x, w, sh, sc, idx: fn( + x, w, sh.index_select(0, idx), sc.index_select(0, idx) + ) + + def diffusers(): + return { + **base, + "source": "diffusers composition (F.rms_norm + index_select modulation)", + "fn": lambda x, w, sh, sc, idx: P.provider_norm_modulate(x, w, sh, sc, idx), + } + + def liger(): + from liger_kernel.ops.modulated_rms_norm import LigerModulatedRMSNormFunction + + return { + **base, + "source": f"Liger {_version('liger-kernel')} modulated RMSNorm + index_select", + "fn": per_token( + lambda x, w, shr, scr: LigerModulatedRMSNormFunction.apply( + x, w, scr, shr, EPS, 0.0, "llama", False + ) + ), + } + + def liger_rlk_gather(): + # Same Liger kernel, with the rows gathered by rl-kernel's deterministic row gather + # (adaln_row_gather) instead of index_select: separates Liger from the gather. + from liger_kernel.ops.modulated_rms_norm import LigerModulatedRMSNormFunction + + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import H3AdaLNRowGatherCudaOp + + gather = H3AdaLNRowGatherCudaOp() + + def fn(x, w, sh, sc, idx): + pad = torch.zeros_like(sh) + table = torch.cat([sh, sc, pad, pad, pad, pad], dim=1) + rows = gather(table, torch.div(idx, 3, rounding_mode="floor"), idx % 3) + return LigerModulatedRMSNormFunction.apply( + x, w, rows[1], rows[0], EPS, 0.0, "llama", False + ) + + version = _version("liger-kernel") + return { + **base, + "source": f"Liger {version} modulated RMSNorm + rl-kernel row gather", + "fn": fn, + } + + def sglang(): + from sglang.kernels.ops.diffusion.norm.scale_residual_norm_cutedsl import ( + fused_norm_scale_shift, + ) + + return { + **base, + "grads": [], + "source": f"SGLang {_version('sglang')} fused_norm_scale_shift (forward only)", + "fn": per_token( + lambda x, w, shr, scr: fused_norm_scale_shift( + x[None], w, None, scr[None], shr[None], "rms", EPS + )[0] + ), + } + + def ours(): + from rl_engine.kernels.ops.cuda.h3.rmsnorm import H3RMSNormCudaOp + + op_ = H3RMSNormCudaOp() + return { + **base, + "source": "rl-kernel H3RMSNormCudaOp.forward_modulated", + "fn": lambda x, w, sh, sc, idx: op_.forward_modulated(x, w, sh, sc, idx), + } + + def norm_only(make): + def m(s, n): + d = make(s, n) + return {"x": d["x"], "w": d["w"]} + + return m + + def norm_ref(x, w): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w + + def torch_norm(): + # diffusers' norm without the modulation, to show which half is batch-invariant + return { + **base, + "make": norm_only(make), + "rows": {"x": 0}, + "grads": ["x", "w"], + "ref": norm_ref, + "source": "torch F.rms_norm (no modulation)", + "fn": lambda x, w: F.rms_norm(x, (H,), w, EPS), + } + + def te_norm(): + import transformer_engine.pytorch as te + + mod = te.RMSNorm(H, eps=EPS, params_dtype=torch.bfloat16, device=DEV) + + def fn(x, w): + # TE's backward reads its own parameter, so copy w in (functional_call + # gives a wrong dx); dweight is then TE's parameter gradient, not checked. + with torch.no_grad(): + mod.weight.copy_(w) + return mod(x) + + return { + **base, + "make": norm_only(make), + "rows": {"x": 0}, + "grads": ["x"], + "ref": norm_ref, + "source": f"TE {_version('transformer_engine')} RMSNorm (no modulation)", + "fn": fn, + } + + found = [ + ("diffusers", diffusers), + ("torch_rms_norm", torch_norm), + ("transformer_engine", te_norm), + ("liger", liger), + ("liger_rlk_gather", liger_rlk_gather), + ("sglang", sglang), + ("rl_kernel", ours), + ] + + elif op == "gate_residual": + + def make(s, n): + ti, tags = _layout(n, s) + return { + "res": _rn(s, 15, n, H, dtype=torch.bfloat16), + "y": _rn(s, 16, n, H, dtype=torch.bfloat16), + "gate": _rn(1, 17, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "idx": ti * 3 + tags, + } + + def ref(res, y, gate, idx): + return res + gate.index_select(0, idx) * y + + base = { + "make": make, + "rows": {"res": 0, "y": 0, "idx": 0}, + "grads": ["res", "y", "gate"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers residual + gate.index_select(...) * y", + "fn": lambda res, y, gate, idx: P.provider_gate_residual(res, gate, idx, y), + } + + def ours(): + from rl_engine.kernels.ops.cuda.h3.gate_residual import H3GateResidualCudaOp + + op_ = H3GateResidualCudaOp() + return { + **base, + "source": "rl-kernel H3GateResidualCudaOp", + "fn": lambda res, y, gate, idx: op_(res, y, gate, idx), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + elif op == "final_adaln_out": + + def make(s, n): + ti, _ = _layout(n, s) + return { + "x": _rn(s, 18, n, H, dtype=torch.bfloat16, scale=2.0), + "nw": (0.5 + torch.rand(H, device=DEV, generator=_gen(1, 19))).bfloat16(), + "temb": _rn(1, 20, T3, D), + "w": _rn(1, 21, 2 * H, D, dtype=torch.bfloat16, scale=D**-0.5), + "b": _rn(1, 22, 2 * H, dtype=torch.bfloat16, scale=0.02), + "ti": ti, + } + + def ref(x, nw, temb, w, b, ti): + return P.provider_final_adaln_out(x, nw, temb, w, b, ti) + + base = { + "make": make, + "rows": {"x": 0, "ti": 0}, + "grads": ["x", "nw", "temb", "w", "b"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers MiniMaxH3AdaLayerNormOut (op-for-op replay)", + "fn": ref, + } + + def ours(): + from rl_engine.kernels.ops.cuda.h3.final_adaln_out import H3FinalAdaLNOutCudaOp + + op_ = H3FinalAdaLNOutCudaOp() + return { + **base, + "source": "rl-kernel H3FinalAdaLNOutCudaOp", + "fn": lambda x, nw, temb, w, b, ti: op_(x, nw, temb, w, b, ti), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + result = [] + for name, factory in found: + try: + result.append({"name": name, **factory()}) + except Exception as exc: # noqa: BLE001 - optional library missing or unusable + result.append({"name": name, "unavailable": f"{type(exc).__name__}: {exc}"[:300]}) + return result + + +# --------------------------------------------------------------------------- # +# Modes and driver +# --------------------------------------------------------------------------- # + + +def enable_mode(mode: str, megatron_src: str | None): + """Switch on a global batch-invariant mode; returns a context manager that turns it off.""" + + base = mode[:-5] if mode.endswith("_ieee") else mode + if base == "plain": + return contextlib.nullcontext + if base == "vllm": + from vllm.model_executor.determinism.batch_invariant import enable_batch_invariant_mode + + enable_batch_invariant_mode() + return contextlib.nullcontext + if base == "sglang": + from sglang.srt.batch_invariant_ops import batch_invariant_ops as mod + + mod.enable_batch_invariant_mode() + kwargs: dict[str, str] = {} + else: + if megatron_src: + sys.path.insert(0, megatron_src) + from megatron.core.transformer.custom_layers import batch_invariant_kernels as mod + + kwargs = {"backend": base.split("_", 1)[1]} + mod.enable_batch_invariant_mode(**kwargs) + if kwargs["backend"] != "triton": + return contextlib.nullcontext + + @contextlib.contextmanager + def off(): + # The Triton matmuls have no FP64 configuration: FP64 references run with the mode off. + mod.disable_batch_invariant_mode() + try: + yield + finally: + mod.enable_batch_invariant_mode(**kwargs) + + return off + + +def run_mode(op: str, mode: str, megatron_src: str | None, quick: bool) -> dict[str, Any]: + try: + mode_off = enable_mode(mode, megatron_src) + except Exception as exc: # noqa: BLE001 + return {"unavailable": f"{type(exc).__name__}: {exc}"[:300]} + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + out: dict[str, Any] = {} + for cand in candidates(op, mode): + name = cand.pop("name") + if "unavailable" in cand: + out[name] = cand + print(f" [{mode}] {name}: unavailable ({cand['unavailable'][:80]})", flush=True) + continue + entry: dict[str, Any] = {"source": cand["source"], "has_backward": bool(cand["grads"])} + entry["batch_invariance"] = batch_invariance(cand, quick) + entry["accuracy_latency"] = accuracy_latency(cand, mode_off, quick) + out[name] = entry + print( + f" [{mode}] {name}: batch_invariant={entry['batch_invariance']['batch_invariant']}", + flush=True, + ) + torch.cuda.empty_cache() + return out + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--megatron-src", default=None, help="Megatron-LM source tree (optional)") + parser.add_argument("--quick", action="store_true", help="small sizes, for a smoke test only") + parser.add_argument("--mode", default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + if importlib.util.find_spec(OPS[args.op]) is None: + raise SystemExit(f"{args.op}: {OPS[args.op]} is not on this branch") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + + if args.mode is not None: # child process: one global mode + args.out.write_text(json.dumps(run_mode(args.op, args.mode, args.megatron_src, args.quick))) + return + + from rl_engine.testing.h3_chain import environment, git_state + + modes = list(MODES) if args.op in GEMM_OPS else ["plain"] + results: dict[str, Any] = {} + for mode in modes: + print(f"[{args.op}] mode {mode}", flush=True) + with tempfile.TemporaryDirectory() as tmp: + part = Path(tmp) / "part.json" + cmd = [sys.executable, __file__, "--op", args.op, "--out", str(part), "--mode", mode] + if args.megatron_src: + cmd += ["--megatron-src", args.megatron_src] + if args.quick: + cmd.append("--quick") + env = {**os.environ, **MODES[mode]} + proc = subprocess.run(cmd, env=env) + results[mode] = ( + json.loads(part.read_text()) + if proc.returncode == 0 and part.exists() + else {"unavailable": f"subprocess exited with {proc.returncode}"} + ) + megatron_commit = None + if args.megatron_src: + megatron_commit = ( + subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + cwd=args.megatron_src, + capture_output=True, + text=True, + ).stdout.strip() + or None + ) + report = { + "kind": "h3_prior_art", + "rfc": "RL-Align/RL-Kernel#420", + "op": args.op, + **git_state(), + "environment": { + **environment(), + "libraries": {lib: _version(lib) for lib in LIBS}, + "megatron_source_commit": megatron_commit, + "mode_environment": {m: MODES[m] for m in modes}, + }, + "quick": args.quick, + "results": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_h3_prior_art.py b/scripts/plot_h3_prior_art.py new file mode 100644 index 000000000..f8e1b2cfe --- /dev/null +++ b/scripts/plot_h3_prior_art.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plot a ``scripts/h3_prior_art.py`` report: latency, accuracy and batch invariance. + + python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/.json + +Writes ``.png`` next to the report. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + + +def _entries(report): + for mode, cands in report["results"].items(): + if "unavailable" in cands: + continue + for name, entry in cands.items(): + if "unavailable" not in entry: + yield (name if name.endswith("]") or mode == "plain" else f"{name}[{mode}]"), entry + + +def _differ(bi) -> int: + rows = sum(v["fwd_differ"] + v["rowgrad_differ"] for v in bi["all_rows"].values()) + sub = bi["full_vs_sub_batches"] + return rows + sub["fwd_differ"] + sub["rowgrad_differ"] + bi["batch_size_sweep"]["differ"] + + +def main() -> None: + path = Path(sys.argv[1]) + report = json.loads(path.read_text()) + entries = list(_entries(report)) + names = [n for n, _ in entries] + fig, axes = plt.subplots(1, 3, figsize=(20, 6.5), layout="constrained") + fig.suptitle( + f"{report['op']} vs existing implementations — {report['environment']['gpu']}, " + f"commit {report['rl_kernel_commit'][:7]}", + fontsize=13, + ) + + ax = axes[0] + for name, entry in entries: + perf = entry["accuracy_latency"] + key = "fwd_bwd_us" if entry["has_backward"] else "fwd_us" + sizes = sorted(perf, key=int) + ax.plot( + [int(s) for s in sizes], + [perf[s][key] for s in sizes], + marker="o", + ls="-" if entry["has_backward"] else "--", + label=name + ("" if entry["has_backward"] else " (forward only)"), + ) + ax.set_xscale("log", base=2) + ax.set_yscale("log") + ax.set_xlabel("rows (timesteps or tokens)") + ax.set_ylabel("µs, median") + ax.set_title("forward + backward latency (dashed: forward only)") + ax.grid(True, which="both", alpha=0.3) + ax.legend(fontsize=7) + + ax = axes[1] + first = entries[0][1]["accuracy_latency"] + size = "3" if "3" in first else sorted(first, key=int)[-1] + ys = range(len(entries)) + fwd = [e["accuracy_latency"][size]["fwd_rel_err"] for _, e in entries] + grad = [ + max(e["accuracy_latency"][size].get("grad_rel_err", {}).values(), default=0) + for _, e in entries + ] + ax.barh([y - 0.2 for y in ys], fwd, 0.4, label="forward") + ax.barh([y + 0.2 for y in ys], [g or float("nan") for g in grad], 0.4, label="worst gradient") + for y, a, b in zip(ys, fwd, grad): + ax.text( + max(a, b or 0) * 1.3, + y, + f"{a:.1e} / {b:.1e}" if b else f"{a:.1e}", + va="center", + fontsize=7, + ) + ax.set_xscale("log") + ax.set_yticks(list(ys), names, fontsize=7) + ax.invert_yaxis() + ax.set_title(f"error vs FP64, max|err| / max|ref| (size {size})") + ax.legend(fontsize=7) + ax.grid(True, axis="x", alpha=0.3) + + ax = axes[2] + diff = [_differ(e["batch_invariance"]) for _, e in entries] + colors = [ + "#3aa676" if e["batch_invariance"]["batch_invariant"] else "#d14b4b" for _, e in entries + ] + ax.barh(list(ys), [max(d, 0.5) for d in diff], color=colors) + for y, (_, e), d in zip(ys, entries, diff): + bi = e["batch_invariance"] + rep = bi["params_repeatable"] + note = "" if rep is None else ("; params repeatable" if rep else "; params NOT repeatable") + verdict = "batch-invariant" if bi["batch_invariant"] else "NOT batch-invariant" + ax.text(max(d, 0.5) * 1.3, y, f"{d} differ — {verdict}{note}", va="center", fontsize=7) + ax.set_xscale("log") + ax.set_yticks(list(ys), names, fontsize=7) + ax.invert_yaxis() + ax.set_xlim(0.4, max(max(diff), 1) * 1e3) + ax.set_title("batch invariance: bitwise comparisons that differ") + ax.set_xlabel( + "every row alone vs full batches + full batch vs covering sub-batches + size sweep", + fontsize=7, + ) + ax.grid(True, axis="x", alpha=0.3) + + out = path.with_suffix(".png") + fig.savefig(out, dpi=120) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() From 0d3794de6cf10777c5e9e1cd1736bcb09c9d161e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:07:40 +0000 Subject: [PATCH 40/89] fix(scripts): keep h3_prior_art mode results next to the report, not in /tmp A cluster job's epilog can clear the user's /tmp while another of their jobs on the same node is still running, which dropped the result of a mode subprocess. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- scripts/h3_prior_art.py | 32 +++++++++++++++++--------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/scripts/h3_prior_art.py b/scripts/h3_prior_art.py index c39adcf0a..2907d16ab 100644 --- a/scripts/h3_prior_art.py +++ b/scripts/h3_prior_art.py @@ -45,7 +45,6 @@ import statistics import subprocess import sys -import tempfile from pathlib import Path from typing import Any, Callable @@ -878,20 +877,23 @@ def main() -> None: results: dict[str, Any] = {} for mode in modes: print(f"[{args.op}] mode {mode}", flush=True) - with tempfile.TemporaryDirectory() as tmp: - part = Path(tmp) / "part.json" - cmd = [sys.executable, __file__, "--op", args.op, "--out", str(part), "--mode", mode] - if args.megatron_src: - cmd += ["--megatron-src", args.megatron_src] - if args.quick: - cmd.append("--quick") - env = {**os.environ, **MODES[mode]} - proc = subprocess.run(cmd, env=env) - results[mode] = ( - json.loads(part.read_text()) - if proc.returncode == 0 and part.exists() - else {"unavailable": f"subprocess exited with {proc.returncode}"} - ) + # The part file sits next to the report, not in /tmp: cluster job epilogs can clear a + # user's /tmp while another of their jobs on the same node is still running. + part = args.out.with_name(f".{args.out.stem}.{mode}.part.json") + part.parent.mkdir(parents=True, exist_ok=True) + part.unlink(missing_ok=True) + cmd = [sys.executable, __file__, "--op", args.op, "--out", str(part), "--mode", mode] + if args.megatron_src: + cmd += ["--megatron-src", args.megatron_src] + if args.quick: + cmd.append("--quick") + proc = subprocess.run(cmd, env={**os.environ, **MODES[mode]}) + results[mode] = ( + json.loads(part.read_text()) + if proc.returncode == 0 and part.exists() + else {"unavailable": f"subprocess exited with {proc.returncode}"} + ) + part.unlink(missing_ok=True) megatron_commit = None if args.megatron_src: megatron_commit = ( From 4b01e4fd07b7823f7314d2c856d18bbe741107a8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:15:42 +0000 Subject: [PATCH 41/89] fix(scripts): name the failing check in the H3 prior-art figure Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- scripts/plot_h3_prior_art.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/scripts/plot_h3_prior_art.py b/scripts/plot_h3_prior_art.py index f8e1b2cfe..cd3bbe342 100644 --- a/scripts/plot_h3_prior_art.py +++ b/scripts/plot_h3_prior_art.py @@ -103,14 +103,28 @@ def main() -> None: for y, (_, e), d in zip(ys, entries, diff): bi = e["batch_invariance"] rep = bi["params_repeatable"] - note = "" if rep is None else ("; params repeatable" if rep else "; params NOT repeatable") + note = ( + "" + if rep is None + else ( + "; parameter/table grads repeatable" + if rep + else "; parameter/table grads NOT repeatable" + ) + ) verdict = "batch-invariant" if bi["batch_invariant"] else "NOT batch-invariant" - ax.text(max(d, 0.5) * 1.3, y, f"{d} differ — {verdict}{note}", va="center", fontsize=7) + ax.text( + max(d, 0.5) * 1.3, + y, + f"{verdict}: {d} row comparisons differ{note}", + va="center", + fontsize=7, + ) ax.set_xscale("log") ax.set_yticks(list(ys), names, fontsize=7) ax.invert_yaxis() ax.set_xlim(0.4, max(max(diff), 1) * 1e3) - ax.set_title("batch invariance: bitwise comparisons that differ") + ax.set_title("batch invariance: bitwise row comparisons that differ, and gradient repeats") ax.set_xlabel( "every row alone vs full batches + full batch vs covering sub-batches + size sweep", fontsize=7, From d10034c7baf15476eeecdd873056fc0c5837d9a0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:16:28 +0000 Subject: [PATCH 42/89] fix(scripts): shorter batch-invariance panel title in the H3 prior-art figure Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- scripts/plot_h3_prior_art.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/plot_h3_prior_art.py b/scripts/plot_h3_prior_art.py index cd3bbe342..d64d4fd56 100644 --- a/scripts/plot_h3_prior_art.py +++ b/scripts/plot_h3_prior_art.py @@ -124,7 +124,7 @@ def main() -> None: ax.set_yticks(list(ys), names, fontsize=7) ax.invert_yaxis() ax.set_xlim(0.4, max(max(diff), 1) * 1e3) - ax.set_title("batch invariance: bitwise row comparisons that differ, and gradient repeats") + ax.set_title("batch invariance (bitwise)") ax.set_xlabel( "every row alone vs full batches + full batch vs covering sub-batches + size sweep", fontsize=7, From 9bf060591972ac62235c3829eb899b4261222ffa Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:16:45 +0000 Subject: [PATCH 43/89] docs(h3): adaln_row_gather vs existing implementations (B200, 6c900ae) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 1fb114792..a40b1a03e 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,6 +131,29 @@ There are two data files: Both were written from a clean tree at commit `fa551c2`. +## Existing implementations (RFC #420 reuse rule) + +![adaln_row_gather vs existing implementations](../usage/evidence/h3-prior-art-b200/adaln_row_gather.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers six index_select calls | **no** (param/table grads not repeatable) | 0.0e+00 / 6.0e-02 / 1381 µs | 0.0e+00 / 2.1e-01 / 9753 µs | +| rl-kernel H3AdaLNRowGatherCudaOp | yes | 0.0e+00 / 2.5e-03 / 1106 µs | 0.0e+00 / 3.0e-03 / 2776 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](../usage/evidence/h3-prior-art-b200/adaln_row_gather.json) +was written from a clean tree at `6c900ae` by + +```bash +python scripts/h3_prior_art.py --op adaln_row_gather --out docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From dd35c14482d7a7a5693b78848230acccd0e555ca Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:16:48 +0000 Subject: [PATCH 44/89] docs(h3): gate_residual vs existing implementations (B200, 4010854) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 2ef31ba54..6b923d740 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -81,6 +81,29 @@ The data is in [`report.json`](../usage/evidence/h3-adaln-gate-residual-b200/rep written by `scripts/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. +## Existing implementations (RFC #420 reuse rule) + +![gate_residual vs existing implementations](../usage/evidence/h3-prior-art-b200/gate_residual.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers residual + gate.index_select(...) * y | **no** (param/table grads not repeatable) | 3.0e-03 / 4.3e-02 / 637 µs | 2.9e-03 / 1.4e-01 / 2502 µs | +| rl-kernel H3GateResidualCudaOp | yes | 3.0e-03 / 2.6e-03 / 1210 µs | 2.9e-03 / 2.6e-03 / 1596 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../usage/evidence/h3-prior-art-b200/gate_residual.json) +was written from a clean tree at `4010854` by + +```bash +python scripts/h3_prior_art.py --op gate_residual --out docs/usage/evidence/h3-prior-art-b200/gate_residual.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/gate_residual.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 3146af19f87d0f275753714f121bb6a274d0d53f Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:30:01 +0000 Subject: [PATCH 45/89] docs(h3): timestep_sinusoid vs existing implementations (B200, 7a82917) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-sinusoid.md | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 316b02132..aa1189202 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -94,6 +94,30 @@ The data is in [`report.json`](../usage/evidence/h3-timestep-sinusoid-b200/repor `scripts/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. +## Existing implementations (RFC #420 reuse rule) + +![timestep_sinusoid vs existing implementations](../usage/evidence/h3-prior-art-b200/timestep_sinusoid.png) + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers get_timestep_embedding (op-for-op replay) | yes | 6.7e-08 / — / 64 µs (fwd only) | 8.6e-08 / — / 64 µs (fwd only) | 9.9e-08 / — / 65 µs (fwd only) | +| SGLang 0.5.21 timestep_embedding | yes | 6.7e-08 / — / 13 µs (fwd only) | 8.6e-08 / — / 12 µs (fwd only) | 9.9e-08 / — / 13 µs (fwd only) | +| rl-kernel H3TimestepSinusoidCudaOp (check_range=False) | yes | 6.7e-08 / — / 16 µs (fwd only) | 8.6e-08 / — / 16 µs (fwd only) | 9.9e-08 / — / 16 µs (fwd only) | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../usage/evidence/h3-prior-art-b200/timestep_sinusoid.json) +was written from a clean tree at `7a82917` by + +```bash +python scripts/h3_prior_art.py --op timestep_sinusoid --out docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 4c040db9a30b8a7052145a7a5850110ab2e09740 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:30:05 +0000 Subject: [PATCH 46/89] docs(h3): timestep_mlp vs existing implementations (B200, ffd958a) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index a57f98033..237c7ae54 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -115,6 +115,35 @@ The data is in [`report.json`](../usage/evidence/h3-timestep-mlp-b200/report.jso by `scripts/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. +## Existing implementations (RFC #420 reuse rule) + +![timestep_mlp vs existing implementations](../usage/evidence/h3-prior-art-b200/timestep_mlp.png) + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers TimestepEmbedding, FP32 F.linear [plain] | **no** (14214 rows, 4108 sub-batches, 2232 sweep cases) | 5.7e-07 / 4.1e-07 / 777 µs | 5.7e-07 / 1.0e-06 / 832 µs | 3.7e-06 / 2.0e-06 / 3609 µs | +| rl-kernel H3TimestepMLPCudaOp | yes | 2.6e-07 / 4.2e-07 / 748 µs | 2.3e-07 / 7.6e-07 / 4595 µs | 2.6e-07 / 2.3e-06 / 42993 µs | +| diffusers TimestepEmbedding, FP32 F.linear [vllm] | **no** (14214 rows, 4096 sub-batches, 2232 sweep cases) | 1.8e-07 / 4.5e-07 / 752 µs | 3.7e-06 / 2.8e-06 / 921 µs | 3.7e-06 / 3.2e-06 / 3826 µs | +| diffusers TimestepEmbedding, FP32 F.linear [sglang] | yes | 1.6e-03 / 1.6e-03 / 1036 µs | 1.5e-03 / 1.5e-03 / 1140 µs | 1.5e-03 / 1.7e-03 / 2636 µs | +| diffusers TimestepEmbedding, FP32 F.linear [sglang_ieee] | yes | 3.3e-06 / 2.3e-06 / 2234 µs | 3.7e-06 / 2.8e-06 / 2428 µs | 3.7e-06 / 3.2e-06 / 7329 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_te_native] | **no** (14214 rows, 4108 sub-batches, 2232 sweep cases) | 1.8e-07 / 4.1e-07 / 729 µs | 3.7e-06 / 1.0e-06 / 814 µs | 3.7e-06 / 2.0e-06 / 3636 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_triton] | yes | 1.6e-03 / 1.6e-03 / 1027 µs | 1.5e-03 / 1.5e-03 / 1137 µs | 1.5e-03 / 1.7e-03 / 2634 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_triton_ieee] | yes | 3.3e-06 / 2.3e-06 / 2309 µs | 3.7e-06 / 2.8e-06 / 2435 µs | 3.7e-06 / 3.2e-06 / 7349 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../usage/evidence/h3-prior-art-b200/timestep_mlp.json) +was written from a clean tree at `ffd958a` by + +```bash +python scripts/h3_prior_art.py --op timestep_mlp --out docs/usage/evidence/h3-prior-art-b200/timestep_mlp.json --megatron-src +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/timestep_mlp.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 28cba8702a7e5219ee6360aef899e8b578a2ab6a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:30:08 +0000 Subject: [PATCH 47/89] docs(h3): adaln_projection vs existing implementations (B200, c88691c) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 29 +++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 0e27a5761..910165a3b 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -131,6 +131,35 @@ written by `scripts/h3_evidence.py` from a clean tree at commit `d06e120`. The r records that a timestep's 18 modulation rows are bitwise identical whether it runs alone or in a batch of 9. +## Existing implementations (RFC #420 reuse rule) + +![adaln_projection vs existing implementations](../usage/evidence/h3-prior-art-b200/adaln_projection.png) + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers AdaLN projection, BF16 F.linear [plain] | **no** (14214 rows, 3080 sub-batches, 2232 sweep cases) | 3.6e-03 / 3.9e-03 / 716 µs | 3.4e-03 / 3.2e-03 / 728 µs | 3.2e-03 / 3.5e-03 / 2346 µs | +| rl-kernel H3AdaLNProjectionCudaOp | yes | 3.6e-03 / 3.9e-03 / 2276 µs | 3.4e-03 / 3.2e-03 / 77446 µs | 3.2e-03 / 3.5e-03 / 626992 µs | +| diffusers AdaLN projection, BF16 F.linear [vllm] | yes | 3.6e-03 / 3.9e-03 / 681 µs | 3.4e-03 / 3.2e-03 / 722 µs | 3.2e-03 / 3.5e-03 / 2330 µs | +| diffusers AdaLN projection, BF16 F.linear [sglang] | yes | 3.6e-03 / 3.9e-03 / 3096 µs | 3.4e-03 / 3.2e-03 / 3690 µs | 3.2e-03 / 3.5e-03 / 14296 µs | +| diffusers AdaLN projection, BF16 F.linear [sglang_ieee] | yes | 3.6e-03 / 3.9e-03 / 3138 µs | 3.4e-03 / 3.2e-03 / 3740 µs | 3.2e-03 / 3.5e-03 / 14325 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_te_native] | **no** (7107 rows, 2056 sub-batches, 1104 sweep cases) | 3.6e-03 / 3.9e-03 / 702 µs | 3.4e-03 / 3.2e-03 / 706 µs | 3.2e-03 / 3.5e-03 / 2330 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_triton] | yes | 3.6e-03 / 3.9e-03 / 2532 µs | 3.4e-03 / 3.2e-03 / 3204 µs | 3.2e-03 / 3.5e-03 / 13771 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_triton_ieee] | yes | 3.6e-03 / 3.9e-03 / 2533 µs | 3.4e-03 / 3.2e-03 / 3206 µs | 3.2e-03 / 3.5e-03 / 13773 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../usage/evidence/h3-prior-art-b200/adaln_projection.json) +was written from a clean tree at `c88691c` by + +```bash +python scripts/h3_prior_art.py --op adaln_projection --out docs/usage/evidence/h3-prior-art-b200/adaln_projection.json --megatron-src +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/adaln_projection.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 656b64ba6aad69bff68f15ca9fa440b941f77764 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:30:12 +0000 Subject: [PATCH 48/89] docs(h3): norm_modulate vs existing implementations (B200, ee83dec) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-rmsnorm.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 21c6b5c6f..fe58b4b8b 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -99,6 +99,34 @@ The data is in [`report.json`](../usage/evidence/h3-rmsnorm-b200/report.json), w - bitwise equality with diffusers for the modulation; - row invariance. +## Existing implementations (RFC #420 reuse rule) + +![norm_modulate vs existing implementations](../usage/evidence/h3-prior-art-b200/norm_modulate.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers composition (F.rms_norm + index_select modulation) | **no** (param/table grads not repeatable) | 8.6e-03 / 4.5e-02 / 872 µs | 7.5e-03 / 1.9e-01 / 4434 µs | +| torch F.rms_norm (no modulation) | yes | 2.1e-03 / 2.1e-03 / 331 µs | 2.0e-03 / 2.8e-03 / 850 µs | +| TE 2.20.2 RMSNorm (no modulation) | yes | 2.1e-03 / 2.1e-03 / 471 µs | 2.0e-03 / 2.8e-03 / 1597 µs | +| Liger 0.8.4 modulated RMSNorm + index_select | **no** (param/table grads not repeatable) | 6.6e-03 / 4.7e-02 / 810 µs | 6.9e-03 / 2.0e-01 / 3911 µs | +| Liger 0.8.4 modulated RMSNorm + rl-kernel row gather | yes | 6.6e-03 / 6.5e-03 / 1824 µs | 6.9e-03 / 8.0e-03 / 4763 µs | +| SGLang 0.5.21 fused_norm_scale_shift (forward only) | yes | 4.4e-03 / — / 85 µs (fwd only) | 4.1e-03 / — / 436 µs (fwd only) | +| rl-kernel H3RMSNormCudaOp.forward_modulated | yes | 8.6e-03 / 4.2e-03 / 1093 µs | 7.5e-03 / 5.0e-03 / 2653 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](../usage/evidence/h3-prior-art-b200/norm_modulate.json) +was written from a clean tree at `ee83dec` by + +```bash +python scripts/h3_prior_art.py --op norm_modulate --out docs/usage/evidence/h3-prior-art-b200/norm_modulate.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/norm_modulate.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 88053296fef5e325c8ed84c89f3434dd68738b85 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:30:16 +0000 Subject: [PATCH 49/89] docs(h3): final_adaln_out vs existing implementations (B200, 03da729) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index b7dd49d65..e7bac8c59 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -85,6 +85,29 @@ There are two data files: T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean tree at `001684d`. +## Existing implementations (RFC #420 reuse rule) + +![final_adaln_out vs existing implementations](../usage/evidence/h3-prior-art-b200/final_adaln_out.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers MiniMaxH3AdaLayerNormOut (op-for-op replay) | **no** (param/table grads not repeatable) | 8.3e-03 / 7.3e-02 / 1021 µs | 8.9e-03 / 2.4e-01 / 4525 µs | +| rl-kernel H3FinalAdaLNOutCudaOp | yes | 8.3e-03 / 5.2e-03 / 1642 µs | 8.9e-03 / 5.7e-03 / 2959 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](../usage/evidence/h3-prior-art-b200/final_adaln_out.json) +was written from a clean tree at `03da729` by + +```bash +python scripts/h3_prior_art.py --op final_adaln_out --out docs/usage/evidence/h3-prior-art-b200/final_adaln_out.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/final_adaln_out.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 2e120f2a42b80c218c813a366b8363b57f9083d5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:45:32 +0000 Subject: [PATCH 50/89] fix(scripts): readable H3 prior-art figure when errors are close or labels are long Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- scripts/plot_h3_prior_art.py | 56 +++++++++++++++++------------------- 1 file changed, 27 insertions(+), 29 deletions(-) diff --git a/scripts/plot_h3_prior_art.py b/scripts/plot_h3_prior_art.py index d64d4fd56..4de1f19a0 100644 --- a/scripts/plot_h3_prior_art.py +++ b/scripts/plot_h3_prior_art.py @@ -19,6 +19,7 @@ matplotlib.use("Agg") import matplotlib.pyplot as plt # noqa: E402 +from matplotlib import transforms # noqa: E402 def _entries(report): @@ -79,57 +80,54 @@ def main() -> None: ] ax.barh([y - 0.2 for y in ys], fwd, 0.4, label="forward") ax.barh([y + 0.2 for y in ys], [g or float("nan") for g in grad], 0.4, label="worst gradient") + values = [v for v in fwd + grad if v] + if values and max(values) / min(values) > 10: + ax.set_xscale("log") + inside = transforms.blended_transform_factory(ax.transAxes, ax.transData) for y, a, b in zip(ys, fwd, grad): ax.text( - max(a, b or 0) * 1.3, + 0.98, y, - f"{a:.1e} / {b:.1e}" if b else f"{a:.1e}", + f"fwd {a:.1e} / grad {b:.1e}" if b else f"fwd {a:.1e}", + transform=inside, + ha="right", va="center", fontsize=7, + bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.8, "pad": 1}, ) - ax.set_xscale("log") ax.set_yticks(list(ys), names, fontsize=7) ax.invert_yaxis() ax.set_title(f"error vs FP64, max|err| / max|ref| (size {size})") - ax.legend(fontsize=7) + ax.legend(fontsize=7, loc="lower right") ax.grid(True, axis="x", alpha=0.3) ax = axes[2] - diff = [_differ(e["batch_invariance"]) for _, e in entries] - colors = [ - "#3aa676" if e["batch_invariance"]["batch_invariant"] else "#d14b4b" for _, e in entries - ] - ax.barh(list(ys), [max(d, 0.5) for d in diff], color=colors) - for y, (_, e), d in zip(ys, entries, diff): + for y, (_, e) in zip(ys, entries): bi = e["batch_invariance"] rep = bi["params_repeatable"] - note = ( - "" - if rep is None - else ( - "; parameter/table grads repeatable" - if rep - else "; parameter/table grads NOT repeatable" - ) - ) - verdict = "batch-invariant" if bi["batch_invariant"] else "NOT batch-invariant" + lines = ["batch-invariant" if bi["batch_invariant"] else "NOT batch-invariant"] + lines.append(f"{_differ(bi)} row comparisons differ") + if rep is not None: + lines.append("parameter/table grads " + ("repeatable" if rep else "NOT repeatable")) ax.text( - max(d, 0.5) * 1.3, + 0.02, y, - f"{verdict}: {d} row comparisons differ{note}", + " | ".join(lines), va="center", - fontsize=7, + fontsize=8, + color="#3aa676" if bi["batch_invariant"] else "#d14b4b", ) - ax.set_xscale("log") - ax.set_yticks(list(ys), names, fontsize=7) - ax.invert_yaxis() - ax.set_xlim(0.4, max(max(diff), 1) * 1e3) + ax.set_ylim(len(entries) - 0.5, -0.5) + ax.set_xlim(0, 1) + ax.axis("off") ax.set_title("batch invariance (bitwise)") - ax.set_xlabel( + ax.text( + 0.02, + len(entries) - 0.3, "every row alone vs full batches + full batch vs covering sub-batches + size sweep", fontsize=7, + color="#555555", ) - ax.grid(True, axis="x", alpha=0.3) out = path.with_suffix(".png") fig.savefig(out, dpi=120) From a9ed1267326f3efb50a9cf1d2a27188fdf52664a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:11:51 +0000 Subject: [PATCH 51/89] Fix H3 native validation, backward timing, and rank cleanup Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 4 +- csrc/cuda/h3/gate_residual.cu | 7 +++- csrc/cuda/h3/rmsnorm_modulate.cu | 23 +++++++---- docs/operators/h3-adaln-projection.md | 5 +++ docs/operators/h3-tp-adaln-3mod.md | 1 + rl_engine/testing/h3_report.py | 28 ++++++++++---- rl_engine/testing/h3_ws2.py | 55 ++++++++++++++++++++++----- tests/h3/test_h3_benchmark_timing.py | 50 +++++++++++++++++++++++- tests/h3/test_h3_native_validation.py | 51 ++++++++++++++++++++++++- tests/h3/test_h3_ws2_processes.py | 47 +++++++++++++++++++++++ 10 files changed, 239 insertions(+), 32 deletions(-) create mode 100644 tests/h3/test_h3_ws2_processes.py diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..9f3d00973 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -207,9 +207,9 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor TORCH_CHECK(grad.is_cuda() && grad.dim() == 3 && grad.is_contiguous(), "grad must be a contiguous (C, S, H) CUDA tensor"); for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { - TORCH_CHECK(t->is_cuda() && t->dim() == 1 && t->is_contiguous() && + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->dim() == 1 && t->is_contiguous() && t->scalar_type() == at::kLong, - "tile metadata must be contiguous 1-D int64 CUDA tensors"); + "tile metadata must be contiguous int64 tensors on ", grad.device()); } TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); const int64_t chunks = grad.size(0); diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index c2bd7921a..6e0b5a85b 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -50,6 +50,7 @@ struct Gate { int64_t stride; const int64_t* index; // (S,) int64_t seq; + int64_t num_rows; }; // One block per row (grid-strided); 16-byte vectors when the row allows it. @@ -60,7 +61,9 @@ __global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __re Gate gate, bool vectorized) { constexpr int kVecElems = 16 / sizeof(T); for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { - const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; const T* ar = a + r * n; const T* br = kDy ? nullptr : b + r * n; T* orow = out + r * n; @@ -148,7 +151,7 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); - return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; } void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index e2c05d2d0..19700d35a 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -106,6 +106,7 @@ struct Modulation { int64_t row_stride; const int64_t* index; // (S,) table row per sequence position int64_t seq; // S: row r of x uses index[r % S] + int64_t num_rows; }; template @@ -120,6 +121,7 @@ __global__ void __launch_bounds__(32 * kWarps) const T* scale = nullptr; if (mod.index != nullptr) { const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); shift = static_cast(mod.shift) + i * mod.row_stride; scale = static_cast(mod.scale) + i * mod.row_stride; } @@ -156,9 +158,12 @@ __global__ void __launch_bounds__(32 * kWarps) const int64_t r = blockIdx.x; const T* xr = x + r * n; const T* gr = g + r * n; - const T* scale = mod.index != nullptr - ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } const int numx = blockDim.x * blockDim.y; const int thrx = threadIdx.x + threadIdx.y * blockDim.x; float s = 0.0f; @@ -193,9 +198,12 @@ __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* const int64_t r1 = min(r0 + kRowTile, rows); float acc = 0.0f; for (int64_t r = r0; r < r1; ++r) { - const T* scale = mod.index != nullptr ? static_cast(mod.scale) + - mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); } partial[static_cast(blockIdx.x) * n + j] = acc; @@ -260,7 +268,7 @@ Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, int64_t n) { - Modulation mod{nullptr, nullptr, 0, nullptr, 1}; + Modulation mod{nullptr, nullptr, 0, nullptr, 1, 0}; if (!index.has_value()) { TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); return mod; @@ -287,6 +295,7 @@ Modulation make_modulation(const c10::optional& shift, mod.row_stride = sh.stride(0); mod.index = ix.data_ptr(); mod.seq = ix.size(0); + mod.num_rows = sh.size(0); return mod; } diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 910165a3b..2cd2cbe0d 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -135,6 +135,11 @@ or in a batch of 9. ![adaln_projection vs existing implementations](../usage/evidence/h3-prior-art-b200/adaln_projection.png) +The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` +streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional +weight passes and are outside its intended T ≤ 8 range. This accounts for the +steep increase in the reported CUDA forward + backward times (77.4 and 627 ms). + | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| | diffusers AdaLN projection, BF16 F.linear [plain] | **no** (14214 rows, 3080 sub-batches, 2232 sweep cases) | 3.6e-03 / 3.9e-03 / 716 µs | 3.4e-03 / 3.2e-03 / 728 µs | 3.2e-03 / 3.5e-03 / 2346 µs | diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 6f525efe6..7a33f6d26 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -83,6 +83,7 @@ byte-equal to WS1 computed on the same GPU. | 1 | forward + backward | 1.90 ms | 1.49 ms | 1.19 ms | 1.04 ms | | 3 | forward | 0.10 ms | 0.11 ms | 0.11 ms | 0.12 ms | | 3 | forward + backward | 2.25 ms | 1.77 ms | 1.48 ms | 1.32 ms | +| 4 | forward + backward | 2.09 ms | 1.83 ms | 1.57 ms | 1.45 ms | Times are for the slowest rank. The WS1 forward already streams the 520 MB weight at about 5 TB/s, so the TP forward is bounded by the all-gather. `DeterministicCollective` synchronises the diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 407906a88..3bd094314 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -590,9 +590,12 @@ def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: gate = _gate_view(table) grad = torch.randn_like(residual) - def backward(fn, residual=residual, y=y, table=table, grad=grad): + def prepare_backward(fn, residual=residual, y=y, table=table): leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] - fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + return fn(leaves[0], leaves[1], _gate_view(leaves[2])) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -606,10 +609,12 @@ def backward(fn, residual=residual, y=y, table=table, grad=grad): "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( r, g, i, y ), - "candidate_backward": lambda b=backward, i=index: b( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: provider_gate_residual(r, g, i, y_) ), } @@ -695,9 +700,12 @@ def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: x, nw, temb, w, b, ti = _final_inputs(seq) grad = torch.randn_like(x) - def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + def prepare_backward(fn, tensors=(x, nw, temb, w, b)): leaves = [t.detach().requires_grad_(True) for t in tensors] - fn(*leaves).backward(grad) + return fn(*leaves) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -708,8 +716,12 @@ def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), - "candidate_backward": lambda bw=backward, i=ti: bw(lambda *t: op(*t, i)), - "provider_backward": lambda bw=backward, i=ti: bw( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( lambda *t: provider_final_adaln_out(*t, i) ), } diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py index a0dc116f6..ad9041cf9 100644 --- a/rl_engine/testing/h3_ws2.py +++ b/rl_engine/testing/h3_ws2.py @@ -144,9 +144,9 @@ def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, inputs_path.with_name(f"rank{rank}.pt"), ) - queue.put({"rank": rank}) collective.close() dist.destroy_process_group() + queue.put({"rank": rank}) except Exception: # forwarded to the parent queue.put({"rank": rank, "error": traceback.format_exc()}) @@ -160,7 +160,9 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa """ import tempfile + import time from pathlib import Path + from queue import Empty import torch.multiprocessing as mp @@ -176,12 +178,45 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) for r in range(world) ] - for proc in procs: - proc.start() - status = sorted((queue.get(timeout=timeout) for _ in procs), key=lambda r: r["rank"]) - for proc in procs: - proc.join(timeout=120) - for result in status: - if "error" in result: - raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") - return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + started = [] + try: + for proc in procs: + proc.start() + started.append(proc) + deadline = time.monotonic() + timeout + pending = set(range(world)) + while pending: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"timed out waiting for ranks {sorted(pending)}") + try: + result = queue.get(timeout=min(0.1, remaining)) + except Empty: + for rank in pending: + if procs[rank].exitcode is not None: + raise RuntimeError( + f"rank {rank} exited with code {procs[rank].exitcode} " + "without reporting a result" + ) from None + continue + if "error" in result: + raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") + pending.remove(result["rank"]) + for rank, proc in enumerate(procs): + proc.join(timeout=max(0, deadline - time.monotonic())) + if proc.is_alive(): + raise TimeoutError(f"rank {rank} did not exit after reporting a result") + if proc.exitcode != 0: + raise RuntimeError(f"rank {rank} exited with code {proc.exitcode}") + return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + finally: + for proc in started: + if proc.is_alive(): + proc.terminate() + for proc in started: + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join() + queue.close() + queue.join_thread() diff --git a/tests/h3/test_h3_benchmark_timing.py b/tests/h3/test_h3_benchmark_timing.py index e6517bd70..47d5575b1 100644 --- a/tests/h3/test_h3_benchmark_timing.py +++ b/tests/h3/test_h3_benchmark_timing.py @@ -114,7 +114,7 @@ def run(key=key, duration=position): assert not any(key.endswith("_setup") for key in row) -@pytest.mark.parametrize("operator", ["norm", "gather"]) +@pytest.mark.parametrize("operator", ["norm", "gather", "gate", "final"]) def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( monkeypatch, cuda_clock, operator ): @@ -136,7 +136,9 @@ def backward_hook(grad): def track(fn, leaves, *args): if leaves[0].requires_grad: assert not cuda_clock.active - prepared_leaves.append(leaves) + prepared_leaves.append( + [leaf._base if leaf._base is not None else leaf for leaf in leaves] + ) cuda_clock.time += 1.0 outputs = fn(*leaves, *args) return track_backward(outputs) if leaves[0].requires_grad else outputs @@ -162,6 +164,50 @@ def forward(x, weight, shift, scale, index, **kwargs): monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) op = SimpleNamespace(forward_modulated=forward) build = h3_report._norm_perf + elif operator == "gate": + + def inputs(seq): + return ( + torch.randn(3, 6 * 4), + torch.randn(1, seq, 4), + torch.randn(1, seq, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_gate_residual + + def forward(residual, y, gate, index, **kwargs): + return track(lambda r, y_, g, i: provider(r, g, i, y_), [residual, y, gate], index) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_gate_inputs", inputs) + monkeypatch.setattr( + h3_report, "provider_gate_residual", lambda r, g, i, y: forward(r, y, g, i) + ) + op = SimpleNamespace(forward=forward) + build = h3_report._gate_perf + elif operator == "final": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 2), + torch.randn(8, 2), + torch.randn(8), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_final_adaln_out + + def forward(x, nw, temb, w, b, ti): + return track(provider, [x, nw, temb, w, b], ti) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_final_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", forward) + op = forward + build = h3_report._final_perf else: randn = torch.randn diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py index 8dd919028..8c71eba08 100644 --- a/tests/h3/test_h3_native_validation.py +++ b/tests/h3/test_h3_native_validation.py @@ -16,6 +16,7 @@ def native_extension(): from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE names = ( + "h3_adaln_row_gather_backward", "h3_rmsnorm_forward", "h3_rmsnorm_backward", "h3_gate_residual_forward", @@ -49,6 +50,10 @@ def native_case(): def _native_call(extension, operation, boundary, case): x, index = case["x"], case["index"] tiles = tuple(case[name] for name in ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile")) + if operation == "gather": + return extension.h3_adaln_row_gather_backward( + case["grad"].unsqueeze(0), *tiles, case["grad"].dtype + ) if operation == "rmsnorm": modulation = (case["shift"], case["scale"], index) if boundary == "forward": @@ -86,7 +91,7 @@ def test_native_rejects_invalid_row_index( _native_call(native_extension, operation, boundary, native_case) -@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual", "gather"]) @pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) @pytest.mark.parametrize("device", ["cpu", "other_cuda"]) def test_native_rejects_mixed_device_tiles( @@ -122,3 +127,47 @@ def test_native_accepts_same_device_inputs(native_extension, native_case, operat actual_grads = _native_call(native_extension, operation, "backward", native_case) for actual, expected in zip(actual_grads, expected_grads, strict=True): torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("boundary", ["forward", "backward"]) +@pytest.mark.parametrize("bad_index", [-1, 2]) +def test_native_rejects_out_of_range_index_in_subprocess( + native_extension, operation, boundary, bad_index +): + # A device assertion poisons the CUDA context, so isolate each invalid launch. + import os + import subprocess + import sys + + code = f""" +import torch +from rl_engine import _C +x = torch.ones(2, 8, device='cuda') +w = torch.ones(8, device='cuda') +index = torch.tensor([0, {bad_index}], device='cuda', dtype=torch.int64) +pos = torch.arange(2, device='cuda', dtype=torch.int64) +ends = pos + 1 +seg = torch.arange(3, device='cuda', dtype=torch.int64) +if {operation!r} == 'gate_residual': + if {boundary!r} == 'forward': + _C.h3_gate_residual_forward(x, x, x, index) + else: + _C.h3_gate_residual_backward(x, x, x, index, pos, pos, ends, seg) +else: + if {boundary!r} == 'forward': + _C.h3_rmsnorm_forward(x, w, 1e-5, x, x, index) + else: + rstd = torch.ones(2, device='cuda') + _C.h3_rmsnorm_backward(x, x, w, rstd, x, x, index, pos, pos, ends, seg) +torch.cuda.synchronize() +""" + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + timeout=60, + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + ) + assert result.returncode != 0 + assert "device-side assert" in result.stderr diff --git a/tests/h3/test_h3_ws2_processes.py b/tests/h3/test_h3_ws2_processes.py new file mode 100644 index 000000000..fb69430cd --- /dev/null +++ b/tests/h3/test_h3_ws2_processes.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU process lifecycle regressions for the multi-GPU evidence runner.""" + +import os +import time + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.testing import h3_ws2 + + +def _worker(rank, world, init, queue, path, target, kwargs): + mode = kwargs["mode"] + if mode == "success": + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + elif rank == 0 and mode == "error": + queue.put({"rank": rank, "error": "original rank traceback"}) + elif rank == 0 and mode == "crash": + os._exit(7) + else: + time.sleep(120) + + +@pytest.mark.parametrize("mode", ["success", "error", "crash", "timeout"]) +def test_run_world_reaps_workers(monkeypatch, mode): + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + if mode == "success": + ranks = h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert [r["rank"] for r in ranks] == [0, 1] + assert [r["value"].item() for r in ranks] == [0, 1] + elif mode == "timeout": + with pytest.raises(TimeoutError, match="timed out waiting for ranks"): + h3_ws2.run_world(2, None, {}, timeout=0.2, mode=mode) + else: + message = "original rank traceback" if mode == "error" else "rank 0 exited with code 7" + with pytest.raises(RuntimeError, match=message): + h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert time.monotonic() - start < 25 + assert {p.pid for p in mp.active_children()} == before From 72ed9702872be0ae4a5518ebe59579c264f54ead Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:14:09 +0000 Subject: [PATCH 52/89] Regenerate gate and final AdaLN backward-only B200 evidence Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 10 +++++----- docs/operators/h3-final-adaln-out.md | 13 +++++++------ 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6b923d740..8d3b4b380 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -61,13 +61,13 @@ hidden = op(residual, attn_output, gate_msa, adaln_indices) python benchmarks/benchmark_h3_conditioning.py --op adaln_gate_residual ``` -B200, B = 1, H = 5376, BF16: +B200, B = 1, H = 5376, BF16. Backward timings exclude input and forward-graph setup: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | -| 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | -| 131072 | 0.92 ms | 1.59 ms | 3.29 ms | 10.6 ms | +| 4097 | 0.07 ms | 0.09 ms | 1.12 ms | 0.58 ms | +| 32768 | 0.24 ms | 0.42 ms | 1.12 ms | 2.24 ms | +| 131072 | 0.88 ms | 1.61 ms | 2.22 ms | 8.92 ms | The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and @@ -78,7 +78,7 @@ tile setup. ![adaln_gate_residual on B200: latency and backward accuracy](../usage/evidence/h3-adaln-gate-residual-b200/figure.png) The data is in [`report.json`](../usage/evidence/h3-adaln-gate-residual-b200/report.json), -written by `scripts/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +written by `scripts/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index e7bac8c59..2c2e215c6 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -40,7 +40,7 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste summation tree versus cuBLAS. Rows are batch- and position-invariant. - **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into the projection backward without an extra BF16 rounding: - - `d_temb`, `dW` and `db` are 22–48× closer to FP64 in the recorded run than diffusers, + - `d_temb`, `dW` and `db` are 21–37× closer to FP64 in the recorded run than diffusers, whose error varies from run to run because it rounds the table gradient to BF16 and accumulates it with `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; @@ -58,13 +58,14 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste python benchmarks/benchmark_h3_conditioning.py --op final_adaln_out ``` -B200, B = 1, T = 3, pinned `norm_out` weights: +B200, B = 1, T = 3, pinned `norm_out` weights. Backward timings exclude input +and forward-graph setup: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | -| 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | -| 131072 | 1.24 ms | 2.94 ms | 7.75 ms | 20.6 ms | +| 4097 | 0.19 ms | 0.18 ms | 1.40 ms | 0.85 ms | +| 32768 | 0.46 ms | 0.80 ms | 2.59 ms | 4.49 ms | +| 131072 | 1.32 ms | 3.02 ms | 7.05 ms | 17.69 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -77,7 +78,7 @@ There are two data files: - [`report.json`](../usage/evidence/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `001684d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + tree at `bde1d8d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../usage/evidence/h3-final-adaln-out-b200/chain_replay.json): the From d9600066ccd459f6f71e2daf4e071aee6d28854a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:14:40 +0000 Subject: [PATCH 53/89] test(h3): consolidate native validation tests for review file limit Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/h3/test_h3_native_validation.py | 153 +++++++++++++++++ tests/h3/test_h3_rmsnorm_native_validation.py | 158 ------------------ 2 files changed, 153 insertions(+), 158 deletions(-) delete mode 100644 tests/h3/test_h3_rmsnorm_native_validation.py diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py index 8dd919028..9f2a06a07 100644 --- a/tests/h3/test_h3_native_validation.py +++ b/tests/h3/test_h3_native_validation.py @@ -122,3 +122,156 @@ def test_native_accepts_same_device_inputs(native_extension, native_case, operat actual_grads = _native_call(native_extension, operation, "backward", native_case) for actual, expected in zip(actual_grads, expected_grads, strict=True): torch.testing.assert_close(actual, expected) + + +TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") + + +@pytest.fixture +def native(): + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not all( + hasattr(_C, name) for name in ("h3_rmsnorm_forward", "h3_rmsnorm_backward") + ): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return _C + + +@pytest.fixture +def inputs(native): + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device=x.device) + _, rstd = native.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device=x.device).chunk(2, dim=1) + return { + "grad": torch.ones_like(x), + "x": x, + "weight": weight, + "rstd": rstd, + "shift": shift, + "scale": scale, + "index": torch.tensor([0, 1], dtype=torch.int64, device=x.device), + "sorted_pos": torch.tensor([0, 2, 1, 3], dtype=torch.int64, device=x.device), + "tile_begin": torch.tensor([0, 2], dtype=torch.int64, device=x.device), + "tile_end": torch.tensor([2, 4], dtype=torch.int64, device=x.device), + "seg_first_tile": torch.tensor([0, 1, 2], dtype=torch.int64, device=x.device), + } + + +def _call(native, inputs, entrypoint): + if entrypoint == "forward": + return native.h3_rmsnorm_forward( + inputs["x"], + inputs["weight"], + 1e-5, + inputs["shift"], + inputs["scale"], + inputs["index"], + ) + return native.h3_rmsnorm_backward(**inputs) + + +def _malformed(tensor, defect): + if defect == "cpu": + return tensor.cpu() + if defect == "dtype": + dtype = torch.float16 if tensor.dtype == torch.float32 else torch.int32 + return tensor.to(dtype) + if defect == "rank": + return tensor.unsqueeze(0) + if defect == "stride": + return torch.stack((tensor, tensor), dim=1)[:, 0] + if defect == "length": + return tensor[:-1] + raise ValueError(f"unknown defect: {defect}") + + +@pytest.mark.cuda_only +def test_valid_plain_and_modulated_backward(native, inputs): + plain = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + assert len(native.h3_rmsnorm_backward(**plain)) == 2 + dx, dweight, dshift, dscale = native.h3_rmsnorm_backward(**inputs) + assert dx.shape == inputs["x"].shape + assert dweight.shape == inputs["weight"].shape + assert dscale.shape == inputs["scale"].shape + torch.testing.assert_close(dshift, torch.full_like(inputs["shift"], 2.0)) + torch.cuda.synchronize(inputs["x"].device) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_empty_row_index(native, inputs, entrypoint): + inputs["index"] = inputs["index"][:0] + with pytest.raises(RuntimeError, match="row index must be a non-empty"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_zero_hidden_size(native, inputs, entrypoint): + inputs["x"] = inputs["x"][:, :0].contiguous() + inputs["weight"] = inputs["weight"][:0] + with pytest.raises(RuntimeError, match="at least one column"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +@pytest.mark.parametrize("mode", ["plain", "modulated"]) +def test_rejects_invalid_rstd(native, inputs, defect, mode): + inputs["rstd"] = _malformed(inputs["rstd"], defect) + if mode == "plain": + inputs = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +def test_rejects_invalid_segment_tiles(native, inputs, name, defect): + inputs[name] = _malformed(inputs[name], defect) + message = "tile_begin and tile_end" if defect == "length" and name.startswith("tile_") else name + with pytest.raises(RuntimeError, match=message): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_rejects_missing_segment_tiles(native, inputs, name): + inputs[name] = None + with pytest.raises(RuntimeError, match="needs the sorted segment tiles"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_plain_backward_rejects_segment_tiles(native, inputs, name): + plain = {key: inputs[key] for key in ("grad", "x", "weight", "rstd", name)} + with pytest.raises(RuntimeError, match="sorted segment tiles require modulation"): + native.h3_rmsnorm_backward(**plain) + + +@pytest.mark.cuda_only +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize( + ("entrypoint", "name"), + [ + ("forward", "weight"), + ("forward", "index"), + ("backward", "weight"), + ("backward", "index"), + ("backward", "rstd"), + ("backward", "sorted_pos"), + ("backward", "tile_begin"), + ("backward", "tile_end"), + ("backward", "seg_first_tile"), + ], +) +def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name): + other_device = (inputs["x"].device.index + 1) % torch.cuda.device_count() + inputs[name] = inputs[name].to(f"cuda:{other_device}") + message = "row index" if name == "index" else name + with pytest.raises(RuntimeError, match=message): + _call(native, inputs, entrypoint) diff --git a/tests/h3/test_h3_rmsnorm_native_validation.py b/tests/h3/test_h3_rmsnorm_native_validation.py deleted file mode 100644 index d6f2a0d81..000000000 --- a/tests/h3/test_h3_rmsnorm_native_validation.py +++ /dev/null @@ -1,158 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""Direct native calls reject malformed pointers and layouts before launching CUDA work.""" - -from __future__ import annotations - -import pytest -import torch - -pytestmark = [ - pytest.mark.cuda_only, - pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), -] - -TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") - - -@pytest.fixture -def native(): - from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE - - if not _EXT_AVAILABLE or not all( - hasattr(_C, name) for name in ("h3_rmsnorm_forward", "h3_rmsnorm_backward") - ): - pytest.skip("rl_engine._C lacks h3_rmsnorm_*") - return _C - - -@pytest.fixture -def inputs(native): - x = torch.randn(4, 8, device="cuda") - weight = torch.ones(8, device=x.device) - _, rstd = native.h3_rmsnorm_forward(x, weight, 1e-5) - shift, scale = torch.zeros(2, 16, device=x.device).chunk(2, dim=1) - return { - "grad": torch.ones_like(x), - "x": x, - "weight": weight, - "rstd": rstd, - "shift": shift, - "scale": scale, - "index": torch.tensor([0, 1], dtype=torch.int64, device=x.device), - "sorted_pos": torch.tensor([0, 2, 1, 3], dtype=torch.int64, device=x.device), - "tile_begin": torch.tensor([0, 2], dtype=torch.int64, device=x.device), - "tile_end": torch.tensor([2, 4], dtype=torch.int64, device=x.device), - "seg_first_tile": torch.tensor([0, 1, 2], dtype=torch.int64, device=x.device), - } - - -def _call(native, inputs, entrypoint): - if entrypoint == "forward": - return native.h3_rmsnorm_forward( - inputs["x"], - inputs["weight"], - 1e-5, - inputs["shift"], - inputs["scale"], - inputs["index"], - ) - return native.h3_rmsnorm_backward(**inputs) - - -def _malformed(tensor, defect): - if defect == "cpu": - return tensor.cpu() - if defect == "dtype": - dtype = torch.float16 if tensor.dtype == torch.float32 else torch.int32 - return tensor.to(dtype) - if defect == "rank": - return tensor.unsqueeze(0) - if defect == "stride": - return torch.stack((tensor, tensor), dim=1)[:, 0] - if defect == "length": - return tensor[:-1] - raise ValueError(f"unknown defect: {defect}") - - -def test_valid_plain_and_modulated_backward(native, inputs): - plain = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} - assert len(native.h3_rmsnorm_backward(**plain)) == 2 - dx, dweight, dshift, dscale = native.h3_rmsnorm_backward(**inputs) - assert dx.shape == inputs["x"].shape - assert dweight.shape == inputs["weight"].shape - assert dscale.shape == inputs["scale"].shape - torch.testing.assert_close(dshift, torch.full_like(inputs["shift"], 2.0)) - torch.cuda.synchronize(inputs["x"].device) - - -@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) -def test_rejects_empty_row_index(native, inputs, entrypoint): - inputs["index"] = inputs["index"][:0] - with pytest.raises(RuntimeError, match="row index must be a non-empty"): - _call(native, inputs, entrypoint) - - -@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) -def test_rejects_zero_hidden_size(native, inputs, entrypoint): - inputs["x"] = inputs["x"][:, :0].contiguous() - inputs["weight"] = inputs["weight"][:0] - with pytest.raises(RuntimeError, match="at least one column"): - _call(native, inputs, entrypoint) - - -@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) -@pytest.mark.parametrize("mode", ["plain", "modulated"]) -def test_rejects_invalid_rstd(native, inputs, defect, mode): - inputs["rstd"] = _malformed(inputs["rstd"], defect) - if mode == "plain": - inputs = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} - with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): - native.h3_rmsnorm_backward(**inputs) - - -@pytest.mark.parametrize("name", TILE_NAMES) -@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) -def test_rejects_invalid_segment_tiles(native, inputs, name, defect): - inputs[name] = _malformed(inputs[name], defect) - message = "tile_begin and tile_end" if defect == "length" and name.startswith("tile_") else name - with pytest.raises(RuntimeError, match=message): - native.h3_rmsnorm_backward(**inputs) - - -@pytest.mark.parametrize("name", TILE_NAMES) -def test_rejects_missing_segment_tiles(native, inputs, name): - inputs[name] = None - with pytest.raises(RuntimeError, match="needs the sorted segment tiles"): - native.h3_rmsnorm_backward(**inputs) - - -@pytest.mark.parametrize("name", TILE_NAMES) -def test_plain_backward_rejects_segment_tiles(native, inputs, name): - plain = {key: inputs[key] for key in ("grad", "x", "weight", "rstd", name)} - with pytest.raises(RuntimeError, match="sorted segment tiles require modulation"): - native.h3_rmsnorm_backward(**plain) - - -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") -@pytest.mark.parametrize( - ("entrypoint", "name"), - [ - ("forward", "weight"), - ("forward", "index"), - ("backward", "weight"), - ("backward", "index"), - ("backward", "rstd"), - ("backward", "sorted_pos"), - ("backward", "tile_begin"), - ("backward", "tile_end"), - ("backward", "seg_first_tile"), - ], -) -def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name): - other_device = (inputs["x"].device.index + 1) % torch.cuda.device_count() - inputs[name] = inputs[name].to(f"cuda:{other_device}") - message = "row index" if name == "index" else name - with pytest.raises(RuntimeError, match=message): - _call(native, inputs, entrypoint) From 86c168f2f5091b8f3b6b947dd3ef5a15ce30b19d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:42:12 +0000 Subject: [PATCH 54/89] fix(h3): address automated review bounds, device, and benchmark findings Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 20 +++- csrc/cuda/h3/gate_residual.cu | 7 +- csrc/cuda/h3/rmsnorm_modulate.cu | 23 ++-- docs/operators/h3-adaln-gate-residual.md | 7 +- docs/operators/h3-final-adaln-out.md | 7 +- docs/operators/h3-sp-norm-adaln.md | 2 +- docs/operators/h3-tp-adaln-3mod.md | 6 +- .../kernels/ops/cuda/h3/sp_norm_adaln.py | 27 +++-- .../ops/cuda/h3/tp_adaln_projection.py | 2 + rl_engine/testing/h3_chain.py | 4 +- rl_engine/testing/h3_report.py | 30 +++-- rl_engine/testing/h3_ws2.py | 62 +++++++++-- scripts/plot_h3_evidence.py | 14 ++- tests/h3/test_h3_benchmark_timing.py | 52 ++++++++- tests/h3/test_h3_native_validation.py | 103 ++++++++++++++++++ tests/h3/test_h3_sp_norm_adaln.py | 103 ++++++++++++++++++ tests/h3/test_h3_tp_adaln.py | 8 +- 17 files changed, 423 insertions(+), 54 deletions(-) diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..0de2d98d5 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -110,7 +110,10 @@ __global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, const int64_t h = j - c * hidden; const g_t* g = grad + c * seq * hidden + h; float acc = 0.0f; - for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t begin = tile_begin[tile], end = tile_end[tile]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= seq); + for (int64_t p = begin; p < end; ++p) { + CUDA_KERNEL_ASSERT(sorted_pos[p] >= 0 && sorted_pos[p] < seq); acc += to_float(g[sorted_pos[p] * hidden]); } partial[tile * width + j] = acc; @@ -119,12 +122,14 @@ __global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, template __global__ void adaln_row_gather_fold_kernel(const float* __restrict__ partial, const int64_t* __restrict__ seg_first_tile, - out_t* __restrict__ out, int64_t width) { + out_t* __restrict__ out, int64_t width, int64_t tiles) { const int64_t r = blockIdx.x; const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; if (j >= width) return; float acc = 0.0f; - for (int64_t tile = seg_first_tile[r]; tile < seg_first_tile[r + 1]; ++tile) { + const int64_t begin = seg_first_tile[r], end = seg_first_tile[r + 1]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= tiles); + for (int64_t tile = begin; tile < end; ++tile) { acc += partial[tile * width + j]; } out[r * width + j] = from_float(acc); @@ -207,9 +212,10 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor TORCH_CHECK(grad.is_cuda() && grad.dim() == 3 && grad.is_contiguous(), "grad must be a contiguous (C, S, H) CUDA tensor"); for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { - TORCH_CHECK(t->is_cuda() && t->dim() == 1 && t->is_contiguous() && + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && + t->dim() == 1 && t->is_contiguous() && t->scalar_type() == at::kLong, - "tile metadata must be contiguous 1-D int64 CUDA tensors"); + "tile metadata must be contiguous 1-D int64 tensors on grad device"); } TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); const int64_t chunks = grad.size(0); @@ -218,6 +224,8 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor const int64_t width = chunks * hidden; const int64_t num_rows = seg_first_tile.numel() - 1; const int64_t tiles = tile_begin.numel(); + TORCH_CHECK(chunks > 0 && seq > 0 && hidden > 0 && num_rows > 0, + "gradient dimensions and table row count must be positive"); TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries"); const c10::cuda::CUDAGuard device_guard(grad.device()); auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype)); @@ -244,7 +252,7 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor adaln_row_gather_fold_kernel <<(num_rows), col_blocks), threads, 0, stream>>>( partial.data_ptr(), seg_first_tile.data_ptr(), - reinterpret_cast(out.data_ptr()), width); + reinterpret_cast(out.data_ptr()), width, tiles); }); C10_CUDA_KERNEL_LAUNCH_CHECK(); return out; diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 2d8aea209..eb0b03f50 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -52,6 +52,7 @@ struct Gate { int64_t stride; const int64_t* index; // (S,) int64_t seq; + int64_t num_rows; }; // One block per row (grid-strided); 16-byte vectors when the row allows it. @@ -62,7 +63,9 @@ __global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __re Gate gate, bool vectorized) { constexpr int kVecElems = 16 / sizeof(T); for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { - const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; const T* ar = a + r * n; const T* br = kDy ? nullptr : b + r * n; T* orow = out + r * n; @@ -150,7 +153,7 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); - return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; } void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 42e3ee18a..f9304ab13 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -108,6 +108,7 @@ struct Modulation { const void* shift; // (R, N) rows with stride `row_stride` (elements), or nullptr const void* scale; int64_t row_stride; + int64_t num_rows; const int64_t* index; // (S,) table row per sequence position int64_t seq; // S: row r of x uses index[r % S] }; @@ -124,6 +125,7 @@ __global__ void __launch_bounds__(32 * kWarps) const T* scale = nullptr; if (mod.index != nullptr) { const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); shift = static_cast(mod.shift) + i * mod.row_stride; scale = static_cast(mod.scale) + i * mod.row_stride; } @@ -160,9 +162,12 @@ __global__ void __launch_bounds__(32 * kWarps) const int64_t r = blockIdx.x; const T* xr = x + r * n; const T* gr = g + r * n; - const T* scale = mod.index != nullptr - ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } const int numx = blockDim.x * blockDim.y; const int thrx = threadIdx.x + threadIdx.y * blockDim.x; float s = 0.0f; @@ -204,9 +209,12 @@ __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* float acc = 0.0f; for (int64_t p = p0; p < p1; ++p) { const int64_t r = row_list != nullptr ? row_list[p] : p; - const T* scale = mod.index != nullptr ? static_cast(mod.scale) + - mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); } partial[static_cast(blockIdx.x) * n + j] = acc; @@ -271,7 +279,7 @@ Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, int64_t n) { - Modulation mod{nullptr, nullptr, 0, nullptr, 1}; + Modulation mod{nullptr, nullptr, 0, 0, nullptr, 1}; if (!index.has_value()) { TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); return mod; @@ -296,6 +304,7 @@ Modulation make_modulation(const c10::optional& shift, mod.shift = sh.data_ptr(); mod.scale = sc.data_ptr(); mod.row_stride = sh.stride(0); + mod.num_rows = sh.size(0); mod.index = ix.data_ptr(); mod.seq = ix.size(0); return mod; diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6b923d740..9d5ee7e09 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -63,7 +63,7 @@ python benchmarks/benchmark_h3_conditioning.py --op adaln_gate_residual B200, B = 1, H = 5376, BF16: -| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| S | CUDA fwd | diffusers fwd | CUDA fwd+bwd | diffusers fwd+bwd | | --- | --- | --- | --- | --- | | 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | | 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | @@ -73,6 +73,11 @@ The forward is one pass over `residual` and `sublayer_output`, using 16-byte vec block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and tile setup. +The stored timing evidence includes forward graph setup in the columns labeled +`fwd+bwd`; those values are not backward-only measurements. The current benchmark +prepares the graph outside the timed backward region. Historical timings are retained +with their original commit and corrected scope. + ## Evidence ![adaln_gate_residual on B200: latency and backward accuracy](../usage/evidence/h3-adaln-gate-residual-b200/figure.png) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index e7bac8c59..6022af96a 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -60,7 +60,7 @@ python benchmarks/benchmark_h3_conditioning.py --op final_adaln_out B200, B = 1, T = 3, pinned `norm_out` weights: -| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| S | CUDA fwd | diffusers fwd | CUDA fwd+bwd | diffusers fwd+bwd | | --- | --- | --- | --- | --- | | 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | | 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | @@ -69,6 +69,11 @@ B200, B = 1, T = 3, pinned `norm_out` weights: At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. +The stored timing evidence includes forward graph setup in the columns labeled +`fwd+bwd`; those values are not backward-only measurements. The current benchmark +prepares the graph outside the timed backward region. Historical timings are retained +with their original commit and corrected scope. + ## Evidence ![final_adaln_out on B200: latency and backward accuracy](../usage/evidence/h3-final-adaln-out-b200/figure.png) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 8506060ea..65b02d021 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -65,7 +65,7 @@ How many rows move depends on how the table rows interleave along the sequence: | Packing | Rows sent per rank (S = 32768, SP8, 4096 rows per rank) | | --- | --- | -| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 174 | +| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 224 | | interleaved (modality and timestep random per position; stress case) | 0 – 1536 | ## Evidence diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 6f525efe6..0f6a64a35 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -28,8 +28,10 @@ op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the `N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. `N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard -never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 qualify. Anything else fails closed when -the op is built. +never splits a WS1 partial sum. The shard constructor accepts positive TP sizes satisfying +this alignment, including 1, 2, 3, 4, 6, 7, 8 and 9 (TP=9 gives 10752 columns per rank). +The `DeterministicCollective` backend supports only 1, 2, 4 and 8 ranks; shard alignment +alone does not imply backend support. | TP | Columns per rank | Rank 0 owns | | --- | --- | --- | diff --git a/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py index 19a867b3e..7b5eb20d3 100644 --- a/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py +++ b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py @@ -211,16 +211,18 @@ class SPPlan: num_rows: int | None -def sp_plan(layout: SPRowLayout, index_full=None, num_rows: int | None = None) -> SPPlan: +def sp_plan( + layout: SPRowLayout, index_full=None, num_rows: int | None = None, *, device=None +) -> SPPlan: """The reduction plan for one row index (``None``: the plain norm's ``dweight`` only).""" - device = index_full.device if index_full is not None else torch.device("cuda") + device = index_full.device if index_full is not None else torch.device(device or "cuda") families = [_dweight_tiles(layout.batch * layout.seq_len, device)] if index_full is None: return SPPlan(layout, _Plan(layout, families, device), None, None, None, None) *seg_tiles, seg_first_tile = _segment_tiles(index_full.repeat(layout.batch), num_rows) plan = _Plan(layout, [*families, tuple(seg_tiles)], device) - local_index = index_full[layout.lo : layout.hi].contiguous() + local_index = index_full[layout.lo : layout.hi].clone(memory_format=torch.contiguous_format) received = index_full[plan.recv_global % layout.seq_len] index_buf = torch.cat([local_index.repeat(layout.batch), received]).contiguous() return SPPlan(layout, plan, local_index, index_buf, seg_first_tile, num_rows) @@ -334,18 +336,21 @@ def __init__(self, collective, seq_len: int, batch: int) -> None: raise RuntimeError("rl_engine._C lacks the H3 SP symbols; rebuild the CUDA extension") self.collective = collective self.layout = sp_row_layout(seq_len, batch, collective.world_size, collective.rank) - self._plans: list[tuple[object, int, int | None, SPPlan]] = [] + self._plans: list[tuple[object, int, int | None, torch.device, SPPlan]] = [] - def plan(self, index: torch.Tensor | None, num_rows: int | None) -> SPPlan: + def plan(self, index: torch.Tensor | None, num_rows: int | None, *, device=None) -> SPPlan: """The cached plan for this index object (rebuilt if it was modified in place).""" + device = index.device if index is not None else torch.device(device or "cuda") + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) version = -1 if index is None else index._version - for obj, ver, rows, plan in self._plans: - if obj is index and ver == version and rows == num_rows: + for obj, ver, rows, cached_device, plan in self._plans: + if obj is index and ver == version and rows == num_rows and cached_device == device: return plan - plan = sp_plan(self.layout, index, num_rows) + plan = sp_plan(self.layout, index, num_rows, device=device) # Holding the index keeps its storage alive, so a new tensor can never alias it. - self._plans = [(index, version, num_rows, plan), *self._plans[:3]] + self._plans = [(index, version, num_rows, device, plan), *self._plans[:3]] return plan def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | None) -> None: @@ -368,7 +373,9 @@ def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | No def norm(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: self._check(x, None, None) validate_h3_rmsnorm(x, weight, eps) - return _SPRMSNorm.apply(x, weight, None, None, self.plan(None, None), self.collective, eps) + return _SPRMSNorm.apply( + x, weight, None, None, self.plan(None, None, device=x.device), self.collective, eps + ) def norm_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): self._check(x, index, shift.shape[0]) diff --git a/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py index 0eb5a4709..f5ec080a3 100644 --- a/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py +++ b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py @@ -84,6 +84,8 @@ def slots(self) -> list[tuple[int, int, int, int]]: def adaln_column_shard(n_total: int, tp: int, rank: int) -> AdaLNColumnShard: if tp < 1 or not 0 <= rank < tp: raise ValueError(f"need 0 <= rank < tp, got rank={rank}, tp={tp}") + if n_total <= 0: + raise ValueError(f"N must be positive, got {n_total}") if n_total % (H3_MODALITY_NUM * H3_ADALN_CHUNKS): raise ValueError(f"N={n_total} is not 3 modalities x 6 chunks x H") if n_total % (tp * DINPUT_CHUNK): diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 3af6ffd38..8f41be297 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -378,7 +378,9 @@ def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torc else: value = stage.golden(golden_op(registry, stage.op_type), run_ctx, value) outputs = list(value) - torch.autograd.backward(outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs)]) + torch.autograd.backward( + outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs, strict=True)] + ) return [leaves[name].grad for name in GRAD_LEAVES] diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 407906a88..877a2534d 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -590,9 +590,13 @@ def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: gate = _gate_view(table) grad = torch.randn_like(residual) - def backward(fn, residual=residual, y=y, table=table, grad=grad): + def prepare_backward(fn, residual=residual, y=y, table=table, grad=grad): leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] - fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + return fn(leaves[0], leaves[1], _gate_view(leaves[2])), grad + + def backward(state): + output, grad = state + output.backward(grad) cases.append( { @@ -606,10 +610,12 @@ def backward(fn, residual=residual, y=y, table=table, grad=grad): "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( r, g, i, y ), - "candidate_backward": lambda b=backward, i=index: b( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: provider_gate_residual(r, g, i, y_) ), } @@ -695,9 +701,13 @@ def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: x, nw, temb, w, b, ti = _final_inputs(seq) grad = torch.randn_like(x) - def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + def prepare_backward(fn, tensors=(x, nw, temb, w, b), grad=grad): leaves = [t.detach().requires_grad_(True) for t in tensors] - fn(*leaves).backward(grad) + return fn(*leaves), grad + + def backward(state): + output, grad = state + output.backward(grad) cases.append( { @@ -708,8 +718,12 @@ def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), - "candidate_backward": lambda bw=backward, i=ti: bw(lambda *t: op(*t, i)), - "provider_backward": lambda bw=backward, i=ti: bw( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( lambda *t: provider_final_adaln_out(*t, i) ), } diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py index b2e3b448f..13174f7f6 100644 --- a/rl_engine/testing/h3_ws2.py +++ b/rl_engine/testing/h3_ws2.py @@ -325,9 +325,9 @@ def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, inputs_path.with_name(f"rank{rank}.pt"), ) - queue.put({"rank": rank}) collective.close() dist.destroy_process_group() + queue.put({"rank": rank}) except Exception: # forwarded to the parent queue.put({"rank": rank, "error": traceback.format_exc()}) @@ -341,7 +341,9 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa """ import tempfile + import time from pathlib import Path + from queue import Empty import torch.multiprocessing as mp @@ -357,12 +359,52 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) for r in range(world) ] - for proc in procs: - proc.start() - status = sorted((queue.get(timeout=timeout) for _ in procs), key=lambda r: r["rank"]) - for proc in procs: - proc.join(timeout=120) - for result in status: - if "error" in result: - raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") - return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + started = [] + deadline = time.monotonic() + timeout + status = {} + try: + for proc in procs: + proc.start() + started.append(proc) + while len(status) < world or any(proc.is_alive() for proc in procs): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"rank results timed out after {timeout}s") + try: + result = queue.get(timeout=min(0.1, remaining)) + except Empty: + exited = [ + (rank, proc.exitcode) + for rank, proc in enumerate(procs) + if proc.exitcode is not None and (proc.exitcode != 0 or rank not in status) + ] + if not exited: + continue + # A worker may have flushed its final message between the + # timed read and the exit-code check. + try: + result = queue.get_nowait() + except Empty: + rank, code = exited[0] + raise RuntimeError( + f"rank {rank} exited with code {code} without completing" + ) from None + if "error" in result: + raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") + status[result["rank"]] = result + for rank, proc in enumerate(procs): + proc.join() + if proc.exitcode != 0: + raise RuntimeError(f"rank {rank} exited with code {proc.exitcode}") + return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + finally: + for proc in started: + if proc.is_alive(): + proc.terminate() + for proc in started: + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join() + queue.close() + queue.join_thread() diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 7328207ad..bd2ada020 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -440,6 +440,8 @@ def repeat(mode: str) -> str: def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: + combined = all(row.get("backward_timing_scope") == "forward_and_backward" for row in perf) + backward_label = "fwd+bwd" if combined else "bwd" _style(ax, "Forward and backward time (log, lower is better)", "ms") _bars( ax, @@ -447,9 +449,14 @@ def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: [ ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), (f"{provider} fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), - ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), ( - f"{provider} bwd", + f"CUDA {backward_label}", + RL_KERNEL, + [row["candidate_backward_us"] / 1e3 for row in perf], + "////", + ), + ( + f"{provider} {backward_label}", PROVIDER, [row["provider_backward_us"] / 1e3 for row in perf], "////", @@ -459,6 +466,8 @@ def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: log=True, ) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + def _backward_errors(ax, backward: dict[str, Any], keys: list[str], note: str) -> None: _style(ax, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") @@ -490,6 +499,7 @@ def label(name: str, mode: str) -> str: for text in ax.texts: if text.get_text() == f"{floor:.1e}": text.set_text("exact") + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=1) ax.set_ylim(floor / 3, 1e4) ax.text(0.98, 0.80, note, transform=ax.transAxes, ha="right", fontsize=8, color=INK) diff --git a/tests/h3/test_h3_benchmark_timing.py b/tests/h3/test_h3_benchmark_timing.py index e6517bd70..d709468b3 100644 --- a/tests/h3/test_h3_benchmark_timing.py +++ b/tests/h3/test_h3_benchmark_timing.py @@ -114,7 +114,7 @@ def run(key=key, duration=position): assert not any(key.endswith("_setup") for key in row) -@pytest.mark.parametrize("operator", ["norm", "gather"]) +@pytest.mark.parametrize("operator", ["norm", "gather", "gate", "final"]) def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( monkeypatch, cuda_clock, operator ): @@ -136,7 +136,7 @@ def backward_hook(grad): def track(fn, leaves, *args): if leaves[0].requires_grad: assert not cuda_clock.active - prepared_leaves.append(leaves) + prepared_leaves.append([leaf if leaf.is_leaf else leaf._base for leaf in leaves]) cuda_clock.time += 1.0 outputs = fn(*leaves, *args) return track_backward(outputs) if leaves[0].requires_grad else outputs @@ -162,6 +162,54 @@ def forward(x, weight, shift, scale, index, **kwargs): monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) op = SimpleNamespace(forward_modulated=forward) build = h3_report._norm_perf + elif operator == "gate": + + def inputs(seq): + return ( + torch.randn(3, 24), + torch.randn(1, seq, 4), + torch.randn(1, seq, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_gate_residual + + def forward(residual, y, gate, index, **kwargs): + return track(lambda r, y_, g: provider(r, g, index, y_), [residual, y, gate]) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_gate_inputs", inputs) + monkeypatch.setattr( + h3_report, "provider_gate_residual", lambda r, g, i, y: forward(r, y, g, i) + ) + op = SimpleNamespace(forward=forward) + build = h3_report._gate_perf + elif operator == "final": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 2), + torch.randn(8, 2), + torch.randn(8), + torch.arange(seq) % 3, + ) + + def forward(x, nw, temb, w, b, ti): + def compute(x, nw, temb, w, b): + rows = torch.nn.functional.linear(torch.nn.functional.silu(temb), w, b) + shift, scale = rows.chunk(2, dim=-1) + norm = torch.nn.functional.rms_norm(x, (4,), nw) + return norm * (1 + scale[ti]) + shift[ti] + + return track(compute, [x, nw, temb, w, b]) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_final_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", forward) + op = forward + build = h3_report._final_perf else: randn = torch.randn diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py index 9f2a06a07..4e59cfcce 100644 --- a/tests/h3/test_h3_native_validation.py +++ b/tests/h3/test_h3_native_validation.py @@ -275,3 +275,106 @@ def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name) message = "row index" if name == "index" else name with pytest.raises(RuntimeError, match=message): _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize( + "operation", ["gate_forward", "gate_dy", "norm_forward", "norm_dx", "norm_partials"] +) +@pytest.mark.parametrize("bad_index", [-1, 2]) +def test_native_table_bounds_fail_in_isolated_process(native_extension, operation, bad_index): + import subprocess + import sys + import textwrap + + script = textwrap.dedent(""" + import sys + import torch + from rl_engine import _C + op, bad = sys.argv[1], int(sys.argv[2]) + x = torch.ones(2, 8, device="cuda") + w = torch.ones(8, device="cuda") + table = torch.ones(2, 8, device="cuda") + index = torch.tensor([bad, 0], device="cuda") + rstd = torch.ones(2, device="cuda") + rows = torch.arange(2, device="cuda") + begin = torch.tensor([0], device="cuda") + end = torch.tensor([2], device="cuda") + try: + if op == "gate_forward": + _C.h3_gate_residual_forward(x, x, table, index) + elif op == "gate_dy": + _C.h3_gate_residual_backward_dy(x, table, index) + elif op == "norm_forward": + _C.h3_rmsnorm_forward(x, w, 1e-5, table, table, index) + elif op == "norm_dx": + _C.h3_rmsnorm_backward_dx(x, x, w, rstd, table, table, index) + else: + _C.h3_rmsnorm_backward_partials( + x, x, w, rstd, table, table, index, rows, begin, end, rows, begin, end) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid index reached native table lookup without an assertion") + """) + result = subprocess.run( + [sys.executable, "-c", script, operation, str(bad_index)], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +def test_row_gather_rejects_metadata_on_other_device(native_extension, metadata): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + grad = torch.ones(6, 2, 8, device="cuda:0") + tensors = dict( + sorted_pos=torch.tensor([0, 1], device="cuda:0"), + tile_begin=torch.tensor([0], device="cuda:0"), + tile_end=torch.tensor([2], device="cuda:0"), + seg_first_tile=torch.tensor([0, 1], device="cuda:0"), + ) + tensors[metadata] = tensors[metadata].to("cuda:1") + with pytest.raises(RuntimeError, match="grad device"): + native_extension.h3_adaln_row_gather_backward(grad, **tensors, out_dtype=torch.float32) + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +@pytest.mark.parametrize("bound", ["negative", "past_end"]) +def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metadata, bound): + import subprocess + import sys + import textwrap + + script = textwrap.dedent(""" + import sys + import torch + from rl_engine import _C + name, bound = sys.argv[1:] + grad = torch.ones(6, 2, 8, device="cuda") + metadata = dict(sorted_pos=torch.tensor([0, 1], device="cuda"), + tile_begin=torch.tensor([0], device="cuda"), + tile_end=torch.tensor([2], device="cuda"), + seg_first_tile=torch.tensor([0, 1], device="cuda")) + metadata[name][0] = -1 if bound == "negative" else 3 + try: + _C.h3_adaln_row_gather_backward(grad, **metadata, out_dtype=torch.float32) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid metadata was not rejected") + """) + result = subprocess.run( + [sys.executable, "-c", script, metadata, bound], capture_output=True, text=True, timeout=60 + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout diff --git a/tests/h3/test_h3_sp_norm_adaln.py b/tests/h3/test_h3_sp_norm_adaln.py index 04d58395c..2b91aad64 100644 --- a/tests/h3/test_h3_sp_norm_adaln.py +++ b/tests/h3/test_h3_sp_norm_adaln.py @@ -208,3 +208,106 @@ def test_nccl_region_byte_equal_to_ws1(world): ranks = run_world(world, sp_region_rank, {k: v.cpu() for k, v in case.items()}) assert sp_matches_ws1(ws1, ranks) == ALL_TRUE assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@requires_cuda +def test_plan_snapshots_caller_index(): + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import sp_plan + + index = torch.tensor([0, 1, 0, 1], device="cuda") + plan = sp_plan(sp_row_layout(4, 1, 1, 0), index, 2) + local, exchanged = plan.local_index.clone(), plan.index_buf.clone() + index.fill_(99) + assert torch.equal(plan.local_index, local) + assert torch.equal(plan.index_buf, exchanged) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_plain_norm_uses_activation_device_and_separate_cache(): + _require_ext() + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + collective = SimpleNamespace(rank=0, world_size=1, all_gather=lambda x: x) + op = H3SPNormAdaLNCudaOp(collective, 5, 1) + with torch.cuda.device(0): + for device in ("cuda:1", "cuda:0"): + x = torch.randn(1, 5, 8, device=device, requires_grad=True) + weight = torch.randn(8, device=device, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + ref_w = weight.detach().clone().requires_grad_(True) + result = op.norm(x, weight) + expected = torch.nn.functional.rms_norm(ref_x, (8,), ref_w, 1e-5) + result.sum().backward() + expected.sum().backward() + torch.testing.assert_close(result, expected) + torch.testing.assert_close(x.grad, ref_x.grad) + torch.testing.assert_close(weight.grad, ref_w.grad) + assert len(op._plans) == 2 + assert op.plan(None, None, device="cuda:1").plan.send_local.device.index == 1 + assert op.plan(None, None, device="cuda:0").plan.send_local.device.index == 0 + + +def _worker_exit_without_result(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + os._exit(kwargs["exit_code"]) + time.sleep(60) + + +def _worker_reports_error(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + queue.put({"rank": rank, "error": "intentional worker failure"}) + time.sleep(60) + + +def _worker_never_reports(rank, world, init, queue, path, target, kwargs): + import time + + time.sleep(60) + + +@pytest.mark.parametrize("mode", ["clean_exit", "crash", "reported_error", "timeout"]) +def test_run_world_cleans_up_failed_workers(monkeypatch, mode): + import time + import torch.multiprocessing as mp + from rl_engine.testing import h3_ws2 + + workers = { + "clean_exit": _worker_exit_without_result, + "crash": _worker_exit_without_result, + "reported_error": _worker_reports_error, + "timeout": _worker_never_reports, + } + monkeypatch.setattr(h3_ws2, "_bootstrap", workers[mode]) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + expected = TimeoutError if mode == "timeout" else RuntimeError + with pytest.raises(expected): + run_world( + 2, + None, + {}, + timeout=1 if mode == "timeout" else 30, + exit_code=7 if mode == "crash" else 0, + ) + assert time.monotonic() - start < 20 + assert {p.pid for p in mp.active_children()} <= before + + +def _worker_reports_result(rank, world, init, queue, path, target, kwargs): + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + + +def test_run_world_returns_rank_ordered_results(monkeypatch): + from rl_engine.testing import h3_ws2 + + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker_reports_result) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + results = run_world(2, None, {}, timeout=30) + assert [result["rank"] for result in results] == [0, 1] + assert [result["value"].item() for result in results] == [0, 1] diff --git a/tests/h3/test_h3_tp_adaln.py b/tests/h3/test_h3_tp_adaln.py index 8b341c8d3..91079b691 100644 --- a/tests/h3/test_h3_tp_adaln.py +++ b/tests/h3/test_h3_tp_adaln.py @@ -48,7 +48,7 @@ def _inputs(num_t, n_out, k, dtype=torch.bfloat16, seed=0): class TestOwnership: - @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8]) + @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8, 9]) def test_shards_tile_the_table_once(self, tp): shards = [adaln_column_shard(N_H3, tp, r) for r in range(tp)] assert [s.begin for s in shards] == [0, *[s.end for s in shards[:-1]]] @@ -146,3 +146,9 @@ def test_nccl_byte_equal_to_ws1(h3_weights_cpu, world): assert tp_matches_ws1(ws1, ranks) == ALL_TRUE assert [r["readback"]["rank"] for r in ranks] == list(range(world)) assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@pytest.mark.parametrize("n_total", [0, -1152]) +def test_rejects_nonpositive_projection(n_total): + with pytest.raises(ValueError, match="N must be positive"): + adaln_column_shard(n_total, 1, 0) From 78b8942a5c45d603d663f97323fe587b9a5c8978 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:39:10 +0000 Subject: [PATCH 55/89] docs(h3): point row-gather validation to migrated tests Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 2813e8934..e8a5adf5c 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -62,7 +62,7 @@ read-back; pass `check_range=False` when the indices are already validated. the gradient. `index_select`'s own backward is an atomic scatter-add in the gradient dtype (BF16), which is neither deterministic nor accurate. - **Tolerance class.** The VJP is a segmented reduction, so gtest judges the op as - `reduction`. The forward is asserted bitwise in `tests/h3`. + `reduction`. The forward is asserted bitwise in `tests/models/minimax_h3`. Measured on a B200 (torch 2.13.0+cu130): From bf902e4a3314e43891d5e40ca32706d6a1e38c85 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:49:17 +0000 Subject: [PATCH 56/89] fix(h3): keep timestep_sinusoid_h3 out of fast-math builds --use_fast_math swaps expf/sinf/cosf for the __expf/__sinf/__cosf approximations and enables FTZ, which breaks the kernel's bitwise parity with diffusers. nvcc 13 defines no macro for it, so the build drops the source and gates the binding on KERNEL_ALIGN_WITH_H3_TIMESTEP; dispatch then falls back to the PyTorch reference. Regression checks: build-config tests for both modes, a registry fallback test, a dense FP32 sweep of [0, 1] (a fast-math compile of the same source mismatches 146.8M of 273.5M elements), and a direct bitwise check against SGLang's timestep_embedding. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- build_tools/extensions.py | 13 ++++++++-- csrc/bindings/ops.cpp | 7 +++-- csrc/cuda/h3/timestep_sinusoid.cu | 3 +++ .../minimax_h3/timestep_sinusoid.py | 5 ++-- .../build/test_build_platform_collectives.py | 26 +++++++++++++++++++ .../minimax_h3/test_h3_timestep_sinusoid.py | 21 +++++++++++++++ tests/runtime/test_dispatch.py | 14 ++++++++++ 7 files changed, 81 insertions(+), 8 deletions(-) diff --git a/build_tools/extensions.py b/build_tools/extensions.py index caba51327..781484c3c 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -214,6 +214,7 @@ def get_extensions(): ) if is_rocm or torch.cuda.is_available(): + use_fast_math = envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH) cuda_sources = [ "csrc/bindings/ops.cpp", "csrc/cuda/logprob/fused_logp_kernel.cu", @@ -236,13 +237,18 @@ def get_extensions(): # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. cuda_sources.append("csrc/cuda/collectives/deterministic_collective.cu") - cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") + # The H3 timestep kernel promises bitwise parity with the precise + # libdevice expf/sinf/cosf; --use_fast_math swaps in the __expf/ + # __sinf/__cosf approximations and flushes denormals. Leave it out + # of fast-math builds so the registry falls back to the reference. + if not use_fast_math: + cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + if use_fast_math: nvcc_flags.append("--use_fast_math") if not is_rocm: cc_major, cc_minor = torch.cuda.get_device_capability() @@ -329,6 +335,9 @@ def get_extensions(): if os.name != "nt" and not is_rocm: # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). extra_link_args.append("-lcuda") + if "csrc/cuda/h3/timestep_sinusoid.cu" in cuda_sources: + cxx_flags.append("-DKERNEL_ALIGN_WITH_H3_TIMESTEP") + nvcc_flags.append("-DKERNEL_ALIGN_WITH_H3_TIMESTEP") if not is_rocm: sm90_srcs = [ diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index a546e0dee..734bf4eb7 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -501,8 +501,8 @@ at::Tensor prefix_shared_attention( #endif // MiniMax-H3 (RFC #420) conditioning-path declarations. CUDA only. -#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \ - (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) +// Fast-math builds omit the source (bitwise contract), so gate on the build define. +#if defined(KERNEL_ALIGN_WITH_H3_TIMESTEP) torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, double max_period, bool check_range); #endif @@ -766,8 +766,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { #endif #endif -#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \ - (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) +#if defined(KERNEL_ALIGN_WITH_H3_TIMESTEP) m.def("h3_timestep_sinusoid_forward", torch::wrap_pybind_function(h3_timestep_sinusoid_forward), "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0, diff --git a/csrc/cuda/h3/timestep_sinusoid.cu b/csrc/cuda/h3/timestep_sinusoid.cu index 9322920a4..45177200b 100644 --- a/csrc/cuda/h3/timestep_sinusoid.cu +++ b/csrc/cuda/h3/timestep_sinusoid.cu @@ -13,6 +13,9 @@ // to be bitwise equal to that path. Each element depends only on (t, k): // batch-, position- and repeat-invariant by construction. Build without // --use_fast_math: expf/sinf/cosf must stay the precise libdevice functions. +// nvcc defines no macro for fast math, so build_tools/extensions.py leaves this +// file out of KERNEL_ALIGN_USE_FAST_MATH=1 builds and the op falls back to the +// PyTorch reference (tests/build/test_build_platform_collectives.py). #include #include diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py index 70f53c33a..b792d3f73 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py @@ -7,8 +7,8 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3 import H3_FREQ_DIM, H3_MAX_PERIOD from rl_engine.reference.minimax_h3.fixed_order import tree_sum_lastdim_fp32 from rl_engine.reference.minimax_h3.timestep_sinusoid import ( @@ -73,7 +73,8 @@ def __init__(self) -> None: if not h3_sinusoid_cuda_available(): raise RuntimeError( "rl_engine._C.h3_timestep_sinusoid_forward is unavailable; rebuild the " - "CUDA extension with csrc/cuda/h3/timestep_sinusoid.cu" + "CUDA extension with csrc/cuda/h3/timestep_sinusoid.cu (it is omitted " + "from KERNEL_ALIGN_USE_FAST_MATH=1 builds to keep bitwise parity)" ) def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index b813506b9..c3ec15389 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -83,3 +83,29 @@ def test_package_data_real_json_is_not_reinterpreted(tmp_path) -> None: assert _resolve_package_data_source(source, project_root=tmp_path) == source assert "build_py" in get_cmdclass() + + +H3_TIMESTEP_SOURCE = "csrc/cuda/h3/timestep_sinusoid.cu" +H3_TIMESTEP_DEFINE = "-DKERNEL_ALIGN_WITH_H3_TIMESTEP" + + +def test_cuda_build_compiles_h3_timestep_without_fast_math(monkeypatch) -> None: + monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) + extension = _load_extension_config(monkeypatch, hip=None) + + assert H3_TIMESTEP_SOURCE in extension["sources"] + assert H3_TIMESTEP_DEFINE in extension["extra_compile_args"]["cxx"] + assert "--use_fast_math" not in extension["extra_compile_args"]["nvcc"] + + +def test_fast_math_build_excludes_bitwise_h3_timestep_kernel(monkeypatch) -> None: + # --use_fast_math swaps expf/sinf/cosf for approximations and enables FTZ, + # which breaks the kernel's bitwise parity promise; nvcc exposes no macro to + # reject it in-source, so the build must drop the kernel and its binding. + monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", "1") + extension = _load_extension_config(monkeypatch, hip=None) + + assert "--use_fast_math" in extension["extra_compile_args"]["nvcc"] + assert H3_TIMESTEP_SOURCE not in extension["sources"] + assert H3_TIMESTEP_DEFINE not in extension["extra_compile_args"]["cxx"] + assert H3_TIMESTEP_DEFINE not in extension["extra_compile_args"]["nvcc"] diff --git a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py index 2e1232feb..5f1ade3ee 100644 --- a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py +++ b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py @@ -38,6 +38,12 @@ def _cuda_op(): return H3TimestepSinusoidCudaOp() +def _dense_unit_sweep() -> torch.Tensor: + bits = torch.arange(0, 0x3F800001, 997, dtype=torch.int32, device="cuda") + edges = torch.tensor([0x00000001, 0x007FFFFF, 0x00800000, 0x3F800000], dtype=torch.int32) + return torch.cat([bits, edges.cuda()]).view(torch.float32) + + class TestReference: def test_layout_is_cos_then_sin(self): t = torch.tensor([0.0, 1.0]) @@ -89,6 +95,21 @@ def test_bitwise_equal_to_diffusers_cuda_path(self, num): assert torch.equal(out, provider_time_proj(t, 256)) assert torch.equal(out, NativeH3TimestepSinusoidOp().forward(t)) + def test_bitwise_over_dense_fp32_sweep_of_unit_interval(self): + # Every 997th FP32 bit pattern in [0, 1] plus the subnormals' edge: a + # fast-math build (approximate __expf/__sinf/__cosf, FTZ) fails here. + t = _dense_unit_sweep() + assert torch.equal(_cuda_op().forward(t, check_range=False), provider_time_proj(t, 256)) + + def test_sglang_timestep_embedding_is_bitwise_equal(self): + # RFC #420 reuse rule: parity with SGLang is checked bit for bit, not + # through matching FP64 error summaries. + jit = pytest.importorskip("sglang.kernels.ops.diffusion.modulate.timestep_embedding_jit") + t = torch.cat([_dense_unit_sweep(), h3_timesteps(4097, seed=3)]) + theirs = jit.timestep_embedding(t, 256, flip_sin_to_cos=True, downscale_freq_shift=0.0) + assert torch.equal(theirs, provider_time_proj(t, 256)) + assert torch.equal(theirs, _cuda_op().forward(t, check_range=False)) + @pytest.mark.parametrize("num_channels", [2, 8, 96, 256, 320]) def test_other_channel_counts_match_reference(self, num_channels): t = h3_timesteps(5) diff --git a/tests/runtime/test_dispatch.py b/tests/runtime/test_dispatch.py index 389bf0eb6..2d095aa14 100644 --- a/tests/runtime/test_dispatch.py +++ b/tests/runtime/test_dispatch.py @@ -246,6 +246,20 @@ def test_h3_ops_prefer_cuda_and_fall_back_to_pytorch_reference(op_name): assert [backend.name for backend in chain] == [cuda_chain[-1].name] +def test_h3_timestep_falls_back_when_fast_math_build_omits_kernel(monkeypatch): + # KERNEL_ALIGN_USE_FAST_MATH=1 builds leave out the bitwise timestep kernel + # and its binding; dispatch must then pick the PyTorch reference. + from rl_engine.backends.cuda.model_specific.minimax_h3 import timestep_sinusoid + from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp + + monkeypatch.setattr(timestep_sinusoid, "h3_sinusoid_cuda_available", lambda: False) + monkeypatch.setattr(torch.version, "hip", None, raising=False) + + op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda:0") + + assert isinstance(op, NativeH3TimestepSinusoidOp) + + def test_npu_available_handles_runtime_failure(monkeypatch): class BrokenNPU: @staticmethod From b4b58b8cd9ce426449afc4f95bd7a5d598e12061 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 10:54:24 +0000 Subject: [PATCH 57/89] build(h3): gate every H3 bitwise kernel on one fast-math guard Generalise the timestep_sinusoid fast-math guard so later H3 kernels join it by name: build_tools/extensions.py lists them in H3_BITWISE_SOURCES, drops the whole list from KERNEL_ALIGN_USE_FAST_MATH=1 builds, and defines KERNEL_ALIGN_WITH_H3_KERNELS (replacing KERNEL_ALIGN_WITH_H3_TIMESTEP) only when they are compiled. The build test now asserts that no csrc/cuda/h3 source survives a fast-math build. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- build_tools/extensions.py | 20 ++++++----- csrc/bindings/ops.cpp | 6 ++-- csrc/cuda/h3/timestep_sinusoid.cu | 6 ++-- .../build/test_build_platform_collectives.py | 34 +++++++++++-------- 4 files changed, 37 insertions(+), 29 deletions(-) diff --git a/build_tools/extensions.py b/build_tools/extensions.py index 781484c3c..f1493c45f 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -14,6 +14,14 @@ _PROJECT_ROOT = Path(__file__).resolve().parents[1] +# MiniMax-H3 (RFC #420) kernels promise bitwise parity with diffusers' eager +# FP32 math. --use_fast_math swaps expf/sinf/cosf/rsqrtf and division for +# approximations and flushes denormals, so these sources are left out of +# KERNEL_ALIGN_USE_FAST_MATH=1 builds; their bindings are gated on +# KERNEL_ALIGN_WITH_H3_KERNELS and dispatch falls back to the PyTorch reference. +H3_BITWISE_SOURCES = ("csrc/cuda/h3/timestep_sinusoid.cu",) +H3_KERNELS_DEFINE = "-DKERNEL_ALIGN_WITH_H3_KERNELS" + def _resolve_package_data_source( filename: str | os.PathLike[str], @@ -237,12 +245,8 @@ def get_extensions(): # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. cuda_sources.append("csrc/cuda/collectives/deterministic_collective.cu") - # The H3 timestep kernel promises bitwise parity with the precise - # libdevice expf/sinf/cosf; --use_fast_math swaps in the __expf/ - # __sinf/__cosf approximations and flushes denormals. Leave it out - # of fast-math builds so the registry falls back to the reference. if not use_fast_math: - cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") + cuda_sources.extend(H3_BITWISE_SOURCES) # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") @@ -335,9 +339,9 @@ def get_extensions(): if os.name != "nt" and not is_rocm: # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). extra_link_args.append("-lcuda") - if "csrc/cuda/h3/timestep_sinusoid.cu" in cuda_sources: - cxx_flags.append("-DKERNEL_ALIGN_WITH_H3_TIMESTEP") - nvcc_flags.append("-DKERNEL_ALIGN_WITH_H3_TIMESTEP") + if not is_rocm and not use_fast_math: + cxx_flags.append(H3_KERNELS_DEFINE) + nvcc_flags.append(H3_KERNELS_DEFINE) if not is_rocm: sm90_srcs = [ diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 734bf4eb7..265937e10 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -501,8 +501,8 @@ at::Tensor prefix_shared_attention( #endif // MiniMax-H3 (RFC #420) conditioning-path declarations. CUDA only. -// Fast-math builds omit the source (bitwise contract), so gate on the build define. -#if defined(KERNEL_ALIGN_WITH_H3_TIMESTEP) +// Fast-math builds omit the H3 sources (bitwise contract), so gate on the build define. +#if defined(KERNEL_ALIGN_WITH_H3_KERNELS) torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, double max_period, bool check_range); #endif @@ -766,7 +766,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { #endif #endif -#if defined(KERNEL_ALIGN_WITH_H3_TIMESTEP) +#if defined(KERNEL_ALIGN_WITH_H3_KERNELS) m.def("h3_timestep_sinusoid_forward", torch::wrap_pybind_function(h3_timestep_sinusoid_forward), "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0, diff --git a/csrc/cuda/h3/timestep_sinusoid.cu b/csrc/cuda/h3/timestep_sinusoid.cu index 45177200b..2acfd516c 100644 --- a/csrc/cuda/h3/timestep_sinusoid.cu +++ b/csrc/cuda/h3/timestep_sinusoid.cu @@ -13,9 +13,9 @@ // to be bitwise equal to that path. Each element depends only on (t, k): // batch-, position- and repeat-invariant by construction. Build without // --use_fast_math: expf/sinf/cosf must stay the precise libdevice functions. -// nvcc defines no macro for fast math, so build_tools/extensions.py leaves this -// file out of KERNEL_ALIGN_USE_FAST_MATH=1 builds and the op falls back to the -// PyTorch reference (tests/build/test_build_platform_collectives.py). +// nvcc defines no macro for fast math, so build_tools/extensions.py lists this +// file in H3_BITWISE_SOURCES, which KERNEL_ALIGN_USE_FAST_MATH=1 builds leave +// out; the op then falls back to the PyTorch reference. #include #include diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index c3ec15389..820347673 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -10,7 +10,12 @@ import torch from torch.utils import cpp_extension -from build_tools.extensions import _resolve_package_data_source, get_cmdclass +from build_tools.extensions import ( + H3_BITWISE_SOURCES, + H3_KERNELS_DEFINE, + _resolve_package_data_source, + get_cmdclass, +) def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: @@ -85,27 +90,26 @@ def test_package_data_real_json_is_not_reinterpreted(tmp_path) -> None: assert "build_py" in get_cmdclass() -H3_TIMESTEP_SOURCE = "csrc/cuda/h3/timestep_sinusoid.cu" -H3_TIMESTEP_DEFINE = "-DKERNEL_ALIGN_WITH_H3_TIMESTEP" - - -def test_cuda_build_compiles_h3_timestep_without_fast_math(monkeypatch) -> None: +def test_cuda_build_compiles_h3_kernels_without_fast_math(monkeypatch) -> None: monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) extension = _load_extension_config(monkeypatch, hip=None) - assert H3_TIMESTEP_SOURCE in extension["sources"] - assert H3_TIMESTEP_DEFINE in extension["extra_compile_args"]["cxx"] + assert set(H3_BITWISE_SOURCES) <= set(extension["sources"]) + assert H3_KERNELS_DEFINE in extension["extra_compile_args"]["cxx"] + assert H3_KERNELS_DEFINE in extension["extra_compile_args"]["nvcc"] assert "--use_fast_math" not in extension["extra_compile_args"]["nvcc"] -def test_fast_math_build_excludes_bitwise_h3_timestep_kernel(monkeypatch) -> None: - # --use_fast_math swaps expf/sinf/cosf for approximations and enables FTZ, - # which breaks the kernel's bitwise parity promise; nvcc exposes no macro to - # reject it in-source, so the build must drop the kernel and its binding. +def test_fast_math_build_excludes_bitwise_h3_kernels(monkeypatch) -> None: + # --use_fast_math swaps expf/sinf/cosf/rsqrtf and division for approximations + # and enables FTZ, which breaks the H3 kernels' bitwise parity promise; nvcc + # exposes no macro to reject it in-source, so the build drops the kernels + # and their bindings. monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", "1") extension = _load_extension_config(monkeypatch, hip=None) assert "--use_fast_math" in extension["extra_compile_args"]["nvcc"] - assert H3_TIMESTEP_SOURCE not in extension["sources"] - assert H3_TIMESTEP_DEFINE not in extension["extra_compile_args"]["cxx"] - assert H3_TIMESTEP_DEFINE not in extension["extra_compile_args"]["nvcc"] + assert not set(H3_BITWISE_SOURCES) & set(extension["sources"]) + assert not [s for s in extension["sources"] if s.startswith("csrc/cuda/h3/")] + assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["cxx"] + assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["nvcc"] From 231de1607f407fb74b138b511afc0566d2d2a501 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:01:40 +0000 Subject: [PATCH 58/89] test(h3): skip the CUDA chain e2e when the H3 kernels are not built KERNEL_ALIGN_USE_FAST_MATH=1 builds omit every H3 kernel, so the registry dispatches the PyTorch references and test_chain_forward's 'backend is a CudaOp' promise does not apply. Skip it there via a shared require_h3_kernels fixture instead of failing. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/conftest.py | 12 ++++++++++++ tests/models/minimax_h3/test_h3_conditioning_e2e.py | 2 +- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/conftest.py b/tests/models/minimax_h3/conftest.py index f8678d581..c1c9296a1 100644 --- a/tests/models/minimax_h3/conftest.py +++ b/tests/models/minimax_h3/conftest.py @@ -12,6 +12,7 @@ import pytest import torch +from rl_engine.backends.extension import _C, _EXT_AVAILABLE from rl_engine.validation.models.h3_weights import ( WEIGHTS_ENV, h3_weights_dir, @@ -27,3 +28,14 @@ def h3_weights_cpu() -> dict[str, torch.Tensor]: return load_h3_conditioning_weights("cpu") except FileNotFoundError as exc: pytest.skip(str(exc)) + + +def h3_kernels_built() -> bool: + """False for KERNEL_ALIGN_USE_FAST_MATH=1 builds, which omit every H3 kernel.""" + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) + + +@pytest.fixture +def require_h3_kernels() -> None: + if not h3_kernels_built(): + pytest.skip("rl_engine._C has no H3 kernels (not built, or a fast-math build)") diff --git a/tests/models/minimax_h3/test_h3_conditioning_e2e.py b/tests/models/minimax_h3/test_h3_conditioning_e2e.py index d4728fd48..104090d10 100644 --- a/tests/models/minimax_h3/test_h3_conditioning_e2e.py +++ b/tests/models/minimax_h3/test_h3_conditioning_e2e.py @@ -39,7 +39,7 @@ def registry(): @pytest.mark.parametrize("num_timesteps, seq_len", CASES) -def test_chain_forward(registry, chain_weights, num_timesteps, seq_len): +def test_chain_forward(require_h3_kernels, registry, chain_weights, num_timesteps, seq_len): report = h3_chain.run_case( registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len ) From b52916ac9517ba5156071b3dc4c783c54573f48c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:14:32 +0000 Subject: [PATCH 59/89] test(h3): skip the CUDA chain backward e2e when the H3 kernels are not built test_chain_backward runs H3AdaLNModulationCudaOp directly, which a KERNEL_ALIGN_USE_FAST_MATH=1 build does not provide; use the shared require_h3_kernels fixture like test_chain_forward. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_conditioning_e2e.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_conditioning_e2e.py b/tests/models/minimax_h3/test_h3_conditioning_e2e.py index 3372ff62b..07cd49776 100644 --- a/tests/models/minimax_h3/test_h3_conditioning_e2e.py +++ b/tests/models/minimax_h3/test_h3_conditioning_e2e.py @@ -65,7 +65,7 @@ def test_chain_forward(require_h3_kernels, registry, chain_weights, num_timestep @pytest.mark.parametrize("num_timesteps, seq_len", [(1, 257), (3, 4097)]) -def test_chain_backward(registry, chain_weights, num_timesteps, seq_len): +def test_chain_backward(require_h3_kernels, registry, chain_weights, num_timesteps, seq_len): """Parameter gradients of the whole chain: deterministic, and FP32-accurate when fused.""" report = h3_chain.run_backward_case( From be6a94ae0b10a30a6d39f52a916583f52a6b74d7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:26:14 +0000 Subject: [PATCH 60/89] tools(h3): direct bitwise comparison of reuse candidates h3_prior_art.py compares implementations through batch invariance and FP64 error summaries, which do not show that two implementations produce the same bits. h3_bitwise_reuse.py runs every candidate that h3_prior_art.candidates() defines on identical inputs (each global batch-invariant mode in its own process) and compares every forward output and leaf gradient byte for byte with diffusers' path and with the rl-kernel op. For timestep_sinusoid it uses three timestep sets, including a dense FP32 sweep of [0, 1], and --fast-math-variant also compares a --use_fast_math build of the kernel. plot_h3_bitwise_reuse.py draws one heatmap per report. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tools/validation/models/h3_bitwise_reuse.py | 250 ++++++++++++++++++ .../models/plot_h3_bitwise_reuse.py | 111 ++++++++ 2 files changed, 361 insertions(+) create mode 100644 tools/validation/models/h3_bitwise_reuse.py create mode 100644 tools/validation/models/plot_h3_bitwise_reuse.py diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py new file mode 100644 index 000000000..d43e5c70e --- /dev/null +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Direct bitwise comparison of one RFC #420 operator with its reuse candidates. + +``h3_prior_art.py`` reports each implementation's batch invariance and its error against FP64. +Matching error summaries do not show that two implementations produce the same bits, so before +a candidate is adopted this runner compares outputs directly. Every candidate that +``h3_prior_art.candidates()`` defines runs on the same inputs (each global batch-invariant mode +in its own process, as in ``h3_prior_art.py``), and every forward output and leaf gradient is +compared byte for byte with diffusers' path and with the rl-kernel op: + + python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \\ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json \\ + [--megatron-src /path/to/Megatron-LM] [--fast-math-variant] + python tools/validation/models/plot_h3_bitwise_reuse.py \\ + reports/experiments/h3-bitwise-reuse-b200/*.json --out figure.png + +For ``timestep_sinusoid`` the inputs are three timestep sets: 2^20 uniform draws, every 997th +FP32 bit pattern in [0, 1] (a dense sweep, mostly tiny values), and a 100001-point grid. +``--fast-math-variant`` also compiles ``csrc/cuda/h3/timestep_sinusoid.cu`` with +``--use_fast_math`` and compares it, to show what the fast-math build guard prevents. Other +operators use each candidate's own input generator at four sizes. + +Intermediate tensors are written next to the report (not to /tmp, which cluster job epilogs +can clear) and removed afterwards. Run on a clean tree; the report records the commit, the tree +state and every library version. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +import h3_prior_art as PA +import torch + +SEED = 7 +SIZES = { + "timestep_mlp": (1, 4, 64, 2048), + "adaln_projection": (1, 4, 64, 2048), + "adaln_row_gather": (1, 257, 4097, 32768), + "norm_modulate": (1, 257, 4097, 32768), + "gate_residual": (1, 257, 4097, 32768), + "final_adaln_out": (1, 257, 4097, 32768), +} +SINUSOID_SETS = ("random_2^20", "dense_fp32_sweep", "grid_100001") +_FAST_MATH_BINDING = """ +#include +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor, int64_t, double, bool); +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("forward", &h3_timestep_sinusoid_forward); } +""" + + +def _sinusoid_inputs(name: str) -> dict[str, torch.Tensor]: + if name == "random_2^20": + t = torch.rand(1 << 20, device="cuda", generator=torch.Generator("cuda").manual_seed(0)) + elif name == "dense_fp32_sweep": + bits = torch.arange(0, 0x3F800001, 997, dtype=torch.int32, device="cuda") + t = bits.view(torch.float32) + else: + t = torch.linspace(0, 1, 100001, device="cuda") + return {"t": t} + + +def _fast_math_sinusoid(build_dir: Path): + from torch.utils.cpp_extension import load + + build_dir.mkdir(parents=True, exist_ok=True) + binding = build_dir / "binding.cpp" + binding.write_text(_FAST_MATH_BINDING) + root = Path(__file__).resolve().parents[3] + mod = load( + name="h3_timestep_sinusoid_fast_math", + sources=[str(root / "csrc/cuda/h3/timestep_sinusoid.cu"), str(binding)], + build_directory=str(build_dir), + extra_cuda_cflags=["-O3", "--use_fast_math"], + ) + return { + "source": "rl-kernel timestep_sinusoid.cu compiled with --use_fast_math", + "grads": [], + "fn": lambda t: mod.forward(t, PA.F1, 10000.0, False), + } + + +def _run(cand: dict[str, Any], inputs: dict[str, torch.Tensor]): + args = {} + for k, v in inputs.items(): + args[k] = v.detach().clone().requires_grad_(True) if k in cand["grads"] else v + out = cand["fn"](**args) + outs = list(out) if isinstance(out, (tuple, list)) else [out] + live = [o for o in outs if o.requires_grad] + if live: + gen = torch.Generator("cuda").manual_seed(SEED) + ups = [torch.randn(o.shape, device="cuda", generator=gen).to(o.dtype) for o in live] + torch.autograd.backward(live, ups) + grads = {k: args[k].grad.cpu() if args[k].grad is not None else None for k in cand["grads"]} + return [o.detach().cpu() for o in outs], grads + + +def dump(op: str, mode: str, parts: Path, megatron_src: str | None, fast_math: bool) -> None: + """Child process: run every candidate of one mode and save its outputs.""" + + try: + PA.enable_mode(mode, megatron_src) + except Exception as exc: # noqa: BLE001 + (parts / f"{mode}.unavailable").write_text(f"{type(exc).__name__}: {exc}"[:300]) + return + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + cands = PA.candidates(op, mode) + if fast_math and op == "timestep_sinusoid" and mode == "plain": + cands.append({"name": "rl_kernel_fast_math", **_fast_math_sinusoid(parts / "build")}) + for cand in cands: + name = cand.pop("name") + if "unavailable" in cand: + (parts / f"{name}.unavailable").write_text(cand["unavailable"][:300]) + continue + cases = SINUSOID_SETS if op == "timestep_sinusoid" else SIZES[op] + for case in cases: + inputs = ( + _sinusoid_inputs(case) if op == "timestep_sinusoid" else cand["make"](SEED, case) + ) + outs, grads = _run(cand, inputs) + torch.save( + {"source": cand["source"], "out": outs, "grads": grads}, + parts / f"{name}__{case}.pt", + ) + print(f" [{mode}] {name}: done", flush=True) + torch.cuda.empty_cache() + + +def _same(a: torch.Tensor | None, b: torch.Tensor | None) -> dict[str, Any] | None: + if a is None or b is None: + return None + if a.shape != b.shape or a.dtype != b.dtype: + return {"equal": False, "note": f"{a.dtype}{tuple(a.shape)} vs {b.dtype}{tuple(b.shape)}"} + width = a.element_size() + differ = ( + a.contiguous().view(torch.uint8).view(-1, width) + != b.contiguous().view(torch.uint8).view(-1, width) + ).any(1) + return { + "equal": not bool(differ.any()), + "differ": int(differ.sum()), + "numel": a.numel(), + "max_abs": float((a.double() - b.double()).abs().max()), + } + + +def compare(parts: Path) -> dict[str, Any]: + runs: dict[str, dict[str, Any]] = {} + for f in sorted(parts.glob("*.pt")): + name, case = f.stem.rsplit("__", 1) + runs.setdefault(name, {})[case] = torch.load(f) + refs = [r for r in ("diffusers[plain]", "diffusers", "rl_kernel") if r in runs] + out: dict[str, Any] = {} + for name, by_case in runs.items(): + entry: dict[str, Any] = {"source": next(iter(by_case.values()))["source"]} + for ref in refs: + if ref == name: + continue + entry[f"vs {ref}"] = { + case: { + "out": [_same(a, b) for a, b in zip(run["out"], runs[ref][case]["out"])], + "grads": { + g: _same(run["grads"].get(g), runs[ref][case]["grads"].get(g)) + for g in run["grads"] + }, + } + for case, run in by_case.items() + } + out[name] = entry + for f in parts.glob("*.unavailable"): + out[f.stem] = {"unavailable": f.read_text()} + return out + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--op", required=True, choices=sorted(PA.OPS)) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--megatron-src", default=None, help="Megatron-LM source tree (optional)") + parser.add_argument( + "--fast-math-variant", + action="store_true", + help="timestep_sinusoid only: also compare a --use_fast_math build of the kernel", + ) + parser.add_argument("--mode", default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + if importlib.util.find_spec(PA.OPS[args.op]) is None: + raise SystemExit(f"{args.op}: {PA.OPS[args.op]} is not on this branch") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + parts = args.out.with_name(f".{args.out.stem}.parts") + + if args.mode is not None: # child process: one global mode + dump(args.op, args.mode, parts, args.megatron_src, args.fast_math_variant) + return + + from rl_engine.validation.models.h3_chain import environment, git_state + + shutil.rmtree(parts, ignore_errors=True) + parts.mkdir(parents=True) + modes = list(PA.MODES) if args.op in PA.GEMM_OPS else ["plain"] + for mode in modes: + print(f"[{args.op}] mode {mode}", flush=True) + cmd = [sys.executable, __file__, "--op", args.op, "--out", str(args.out), "--mode", mode] + if args.megatron_src: + cmd += ["--megatron-src", args.megatron_src] + if args.fast_math_variant: + cmd.append("--fast-math-variant") + proc = subprocess.run(cmd, env={**os.environ, **PA.MODES[mode]}) + if proc.returncode != 0: + (parts / f"{mode}.unavailable").write_text(f"subprocess exited with {proc.returncode}") + results = compare(parts) + shutil.rmtree(parts, ignore_errors=True) + report = { + "kind": "h3_bitwise_reuse", + "rfc": "RL-Align/RL-Kernel#420", + "op": args.op, + **git_state(), + "environment": { + **environment(), + "libraries": {lib: PA._version(lib) for lib in PA.LIBS}, + "mode_environment": {m: PA.MODES[m] for m in modes}, + }, + "cases": list(SINUSOID_SETS if args.op == "timestep_sinusoid" else SIZES[args.op]), + "results": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_bitwise_reuse.py b/tools/validation/models/plot_h3_bitwise_reuse.py new file mode 100644 index 000000000..85983fc0a --- /dev/null +++ b/tools/validation/models/plot_h3_bitwise_reuse.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plot ``h3_bitwise_reuse.py`` reports: share of forward outputs that are not bitwise equal. + +One heatmap per report. Rows are implementations, columns are the cases compared against +diffusers' path and against the rl-kernel op; 0 means bitwise identical. + + python tools/validation/models/plot_h3_bitwise_reuse.py \\ + reports/experiments/h3-bitwise-reuse-b200/*.json --out figure.png +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + +_REFS = ("diffusers[plain]", "diffusers", "rl_kernel") + + +def _cell(entry: dict | None) -> float | None: + if not entry: + return None + out = entry["out"][0] + if out is None or "differ" not in out: + return None + return 0.0 if out["equal"] else out["differ"] / out["numel"] + + +def _text(frac: float | None) -> str: + if frac is None: + return "—" + if frac == 0: + return "0" + return f"{frac:.1%}" if frac >= 0.001 else f"{frac:.2%}" + + +def _panel(ax, report: dict) -> None: + results = {k: v for k, v in report["results"].items() if "unavailable" not in v} + cases = report["cases"] + diffusers = next((r for r in _REFS[:2] if r in results), None) + refs = [r for r in (diffusers, "rl_kernel") if r] + rows = ["rl_kernel"] + sorted(r for r in results if r not in _REFS) + cols = [(ref, case) for ref in refs for case in cases] + grid = np.full((len(rows), len(cols)), np.nan) + for i, row in enumerate(rows): + for j, (ref, case) in enumerate(cols): + frac = _cell(results.get(row, {}).get(f"vs {ref}", {}).get(str(case))) + if frac is not None: + grid[i, j] = frac + ax.text( + j, + i, + "—" if row == ref else _text(frac), + ha="center", + va="center", + fontsize=7.5, + color="white" if frac is not None and frac > 0.3 else "black", + fontweight="bold" if frac == 0 else None, + ) + ax.imshow(np.sqrt(np.nan_to_num(grid)), cmap="Reds", vmin=0, vmax=1, aspect="auto") + ax.set_yticks(range(len(rows)), rows, fontsize=8) + ax.set_xticks(range(len(cols)), [str(c) for _, c in cols], fontsize=7, rotation=45) + transform = ax.get_xaxis_transform() + for k, ref in enumerate(refs): + if k: + ax.axvline(k * len(cases) - 0.5, color="k", lw=1.5) + label = "vs rl-kernel" if ref == "rl_kernel" else "vs diffusers" + ax.text( + k * len(cases) + (len(cases) - 1) / 2, + 1.01, + label, + ha="center", + va="bottom", + fontsize=8.5, + transform=transform, + ) + ax.set_title(f"{report['op']} @ {report['rl_kernel_commit'][:8]}", fontsize=10, pad=18) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("reports", nargs="+", type=Path) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + reports = [json.loads(p.read_text()) for p in args.reports] + fig, axes = plt.subplots(1, len(reports), figsize=(6.4 * len(reports), 5.8), squeeze=False) + for ax, report in zip(axes[0], reports): + _panel(ax, report) + gpu = reports[0]["environment"].get("gpu", "") + fig.suptitle( + "Forward outputs that are not bitwise equal\n" + f"(share of elements; 0 = bitwise identical), {gpu}", + fontsize=10, + ) + fig.tight_layout() + args.out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(args.out, dpi=125) + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() From 367883ad5f837aa585a3b2c5e492d5b55920fd30 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:34:04 +0000 Subject: [PATCH 61/89] evidence(h3): timestep_sinusoid direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py --fast-math-variant from a clean tree at 0eb5ddd5. SGLang 0.5.21 timestep_embedding, diffusers' path and this kernel are bitwise identical on 2.22M timesteps; the --use_fast_math build differs in 53.7-82.7% of outputs. The operator doc gains the section, the command and the fast-math limitation. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-sinusoid.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 7bf0f13b7..c5f69a15b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -118,6 +118,29 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs: +2^20 uniform timesteps, every 997th FP32 bit pattern in [0, 1], and a 100001-point grid. It +compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this +kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). + +![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) + +SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all +2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. +[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +was written from a clean tree at `0eb5ddd5` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json +python tools/validation/models/plot_h3_bitwise_reuse.py \ + reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png +``` + ## Tests ```bash @@ -135,3 +158,6 @@ python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-timest - There is no ROCm kernel. ROCm dispatches the PyTorch reference. - `max_period` is fixed at 10000, the only value H3 uses. +- `KERNEL_ALIGN_USE_FAST_MATH=1` builds leave out every H3 kernel (`H3_BITWISE_SOURCES` in + `build_tools/extensions.py`): fast math replaces `expf`/`sinf`/`cosf` with approximations and + breaks bitwise parity. Such builds dispatch the PyTorch reference. From 1a21cd6ae3189ac7122cd89054beab00bf77fc83 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:38:43 +0000 Subject: [PATCH 62/89] evidence(h3): timestep_mlp direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 0df83a5a. No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 7f2d477e3..8049a87df 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -144,6 +144,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) + +No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. +[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png +``` + ## Tests ```bash From c4f6145ffa65a114b6a9e2b2977954626eeb1ef9 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:46:47 +0000 Subject: [PATCH 63/89] evidence(h3): adaln_projection direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 0e476483. This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index e93cdee6b..92dca1376 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -160,6 +160,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) + +This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. +[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png +``` + ## Tests ```bash From 62f8009e7ae70381d99fa722c7ddbbe1719ae4dc Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:23:35 +0000 Subject: [PATCH 64/89] test(h3): fix the flaky FP32 table-gradient check in adaln_modulation test_backward_keeps_table_gradient_fp32 drew temb from an unseeded torch.randn and failed about 1 run in 30 with a dW correctly-rounded fraction near 0.9997. The kernel is deterministic (150 repeats bitwise equal); the golden was inconsistent: it rounded an FP64 SiLU to BF16, while the kernel and diffusers round the FP32 SiLU. Near a BF16 tie the two pick different neighbours for 1-2 of the 8064 activations, which shifts a whole dW column (10 of 300 draws fell below 0.9999). Round the golden's activation from F.silu(temb.float()) and seed temb. Over 300 temb draws the fraction is then 0.999981-0.999982 and the d_temb error at most 2.0e-6 (limits 0.9999 and 1e-5); 30 repeated runs pass. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_adaln_modulation.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py index cd6647739..066445fc0 100644 --- a/tests/models/minimax_h3/test_h3_adaln_modulation.py +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -70,7 +70,8 @@ def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): fused, proj, gather = _ops() weight = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.weight"].cuda() bias = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.bias"].cuda() - temb = torch.randn(3, 2688, device="cuda") * 2 + g = torch.Generator(device="cuda").manual_seed(0) + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 ti, tags = h3_packed_layout(4097, 3, seed=1) grads = _grads(4097, 5376, seed=1) ours = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) @@ -84,7 +85,11 @@ def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] act = ref[0] * torch.sigmoid(ref[0]) - act = act + (act.to(torch.bfloat16).double() - act).detach() # identity-VJP cast + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column (about 1 temb draw in 30 dropped below 0.9999 that way). + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # identity-VJP cast table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376) index = ti * 3 + tags torch.autograd.backward( From 7aeb6ba16b34c9b14de3cf5e2cc6c420b84ec505 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:25:43 +0000 Subject: [PATCH 65/89] evidence(h3): adaln_row_gather direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 46347118. There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index e8a5adf5c..1d0387183 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -154,6 +154,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![adaln_row_gather bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) + +There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. +[`adaln_row_gather.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png +``` + ## Tests ```bash From cad31b95545347344cf11a6fef6faf949cd72b17 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:29:03 +0000 Subject: [PATCH 66/89] evidence(h3): norm_modulate direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 228e50b9. This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-rmsnorm.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 13b585606..7a16f4a57 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -127,6 +127,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![norm_modulate bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) + +This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. +[`norm_modulate.json`](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json \ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png +``` + ## Tests ```bash From 3bbcca42f7e8a5a8a7b4d8753a46c108cd1bd653 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:30:51 +0000 Subject: [PATCH 67/89] evidence(h3): gate_residual direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 910d614a. There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index e69de13cd..e003fd5b8 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -104,6 +104,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) + +There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. +[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/gate_residual.json \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.png +``` + ## Tests ```bash From 28099904ff6ed0197d9389c8ca456598bbbb9045 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:32:41 +0000 Subject: [PATCH 68/89] evidence(h3): final_adaln_out direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 4430ef16. There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index d3c938d55..3b571efec 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -108,6 +108,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![final_adaln_out bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) + +There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. +[`final_adaln_out.json`](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png +``` + ## Tests ```bash From 49fa04a0d191092c8dd1dda2fea3671f64423d3d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:46:51 +0000 Subject: [PATCH 69/89] test(h3): round the projection golden's activation from the FP32 SiLU test_backward_against_golden rounded an FP64 SiLU to BF16 for its declared cast, while the kernel and diffusers round the FP32 SiLU. Near a BF16 tie the two pick different neighbours and shift a whole dW column. The test was seeded (seed 5 scores 1.0), so it never failed, but over 300 temb draws 8 fell below 0.9999 (worst 0.999646). With the FP32-SiLU golden every dW/db element is correctly rounded for all 300 draws; d_temb is unchanged. Same fix as test_h3_adaln_modulation.py. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_adaln_projection.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 6f9323ed5..91732f6d2 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -194,7 +194,11 @@ def test_backward_against_golden(self, block0): ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] act = ref[0] * torch.sigmoid(ref[0]) - act = act + (act.to(torch.bfloat16).double() - act).detach() # declared cast, identity VJP + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column. + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # declared cast, identity VJP table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376).chunk(6, dim=-1) torch.autograd.backward(list(table), [gr.double() for gr in grads]) # tolerance_contract.json gradient_accuracy / reduction: FP32 temb grad From 55370d6c2aaa79ae477634932c1856f452f914cb Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:20:09 +0000 Subject: [PATCH 70/89] test(h3): require exact dW/db equality in the projection golden check With the FP32-SiLU golden (8bb2e092), the BF16 weight and bias gradients equal the correctly rounded golden at every element for all 500 temb draws tried (T = 2, two products per element). Replace the '> 0.999 of elements' threshold with torch.equal; the FP64 closeness check stays. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_adaln_projection.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 91732f6d2..5f0e7713e 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -206,7 +206,10 @@ def test_backward_against_golden(self, block0): torch.testing.assert_close(leaves[0].grad.double(), ref[0].grad, atol=1e-4, rtol=1e-4) for leaf, r in zip(leaves[1:], ref[1:], strict=True): torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) - assert (leaf.grad == r.grad.to(torch.bfloat16)).float().mean() > 0.999 + # With T = 2 each dW/db element sums two products, and the one BF16 cast + # of the FP32 sum is the correctly rounded value of every element + # (checked on 500 of 500 temb draws). + assert torch.equal(leaf.grad, r.grad.to(torch.bfloat16)) @requires_cuda From 0d433f6f46072ab73818ed40bcfb63b235015f3d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:36:32 +0000 Subject: [PATCH 71/89] test(h3): round the projection forward golden from the FP32 SiLU, require 0.9995 The two forward checks compared the kernel with forward_fp32, which rounds an FP64 SiLU to BF16 while the kernel and diffusers round the FP32 SiLU. Near a BF16 tie that shifts a whole output column: over 300 temb draws the old golden matched as little as 95.7% of the outputs, so the '> 0.999' checks only passed for their fixed seeds. A test-local _declared_golden now rounds F.silu(temb.float()). Exact equality is not reachable: the FP32 tensor-core accumulation over K = 2688 leaves 32-104 outputs off on every draw (ties, and a few ULPs where near-zero sums cancel). Over 300 draws per T = 1..4 the worst match is 0.999618-0.999731, so both checks now require >= 0.9995. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/test_h3_adaln_projection.py | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 5f0e7713e..29269b061 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -20,7 +20,10 @@ import torch.nn.functional as F from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM -from rl_engine.reference.minimax_h3.adaln_projection import NativeH3AdaLNProjectionOp +from rl_engine.reference.minimax_h3.adaln_projection import ( + NativeH3AdaLNProjectionOp, + split_adaln_table, +) from rl_engine.validation.models.h3_provider import provider_adaln_modulation requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") @@ -64,6 +67,19 @@ def _flat(outputs): return torch.cat([out.reshape(-1) for out in outputs]) +def _declared_golden(temb, weight, bias): + """FP64 projection of the declared BF16 activation, rounded from the FP32 SiLU. + + The kernel and diffusers round the FP32 SiLU; ``forward_fp32`` rounds an FP64 + SiLU, which can pick the other BF16 neighbour near a tie and shift a whole + output column (one input in 300 then matched only 95.7% of the outputs). + """ + + act = F.silu(temb.float()).to(weight.dtype).double() + table = F.linear(act, weight.double(), bias.double()) + return tuple(chunk.float() for chunk in split_adaln_table(table, weight.shape[0] // 18)) + + class TestReference: def test_provider_layout_matches_diffusers(self): """Match the provider values and timestep-major modality layout on CPU.""" @@ -151,9 +167,12 @@ def test_forward_is_correctly_rounded_declared_golden(self, block0, num): # tolerance_contract.json forward_accuracy / reduction / bfloat16. torch.testing.assert_close(out.float(), gold, atol=5e-2, rtol=2e-2) # Beyond the contract: almost every element is the correctly rounded - # golden value; the rest are 1-ULP ties of the FP32 accumulation. - match = (_flat(outputs) == _flat(golden).bfloat16()).float().mean().item() - assert match > 0.999 + # golden value. The rest come from the FP32 tensor-core accumulation over + # K = 2688 (ties, and a few ULPs where near-zero sums cancel), so exact + # equality is not reachable: over 300 temb draws per T the worst match was + # 0.999618 (T = 1), 0.999690, 0.999728 and 0.999731 (T = 4). + declared = _flat(_declared_golden(temb, weight, bias)).bfloat16() + assert (_flat(outputs) == declared).float().mean().item() >= 0.9995 def test_early_cast_is_numerically_detectable(self, block0): """Probe H7: a BF16 cast before the SiLU moves about half the outputs.""" @@ -162,9 +181,9 @@ def test_early_cast_is_numerically_detectable(self, block0): temb = _temb(3, seed=7) op = NativeH3AdaLNProjectionOp() ours = _flat(_cuda_op()(temb, weight, bias)) - declared = _flat(op.forward_fp32(temb, weight, bias)).bfloat16() + declared = _flat(_declared_golden(temb, weight, bias)).bfloat16() early = _flat(op.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() - assert (ours == declared).float().mean() > 0.999 + assert (ours == declared).float().mean() >= 0.9995 assert (ours == early).float().mean() < 0.8 def test_rows_are_batch_and_position_invariant(self, block0): From 579775f1306a3afb4103de10925db79a67cdc503 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:50:58 +0000 Subject: [PATCH 72/89] fix(h3): round the adaln_projection golden's activation from the FP32 SiLU NativeH3AdaLNProjectionOp.forward_fp32 rounded an FP64 SiLU to the weight dtype for the declared cast, while the model and the kernel round the FP32 SiLU. Near a BF16 tie the two pick different neighbours and shift a whole output column: over 300 temb draws the old golden matched the kernel's BF16 outputs as little as 95.7% of the time. The golden now rounds F.silu(temb.float()); the VJP stays the FP64 SiLU derivative (identity through the cast). The test-local golden added in 42943618 is no longer needed and the forward checks use forward_fp32 again. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../reference/minimax_h3/adaln_projection.py | 11 ++++++---- .../minimax_h3/test_h3_adaln_projection.py | 22 +++---------------- 2 files changed, 10 insertions(+), 23 deletions(-) diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py index 03c1ad683..be83239c5 100644 --- a/rl_engine/reference/minimax_h3/adaln_projection.py +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -112,9 +112,12 @@ def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: hidden = validate_h3_adaln_projection(temb, weight, bias) t64 = temb.double() act = t64 * torch.sigmoid(t64) - # The declared cast rounds the value; its VJP is the identity. Written - # straight-through so autograd does not round the FP64 gradient to the - # weight dtype on its way back (``.to(bf16)`` would). - act = act + (act.to(weight.dtype).double() - act).detach() + # The declared cast rounds the FP32 SiLU, as the model (and the kernel) + # computes it: an FP64 SiLU can round to the other weight-dtype neighbour + # near a tie and shift a whole output column. Its VJP is the identity, + # written straight-through so autograd does not round the FP64 gradient + # to the weight dtype on its way back (``.to(bf16)`` would). + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() table = F.linear(act, weight.double(), bias.double()) return tuple(chunk.float() for chunk in split_adaln_table(table, hidden)) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 29269b061..5fb10a73f 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -20,10 +20,7 @@ import torch.nn.functional as F from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM -from rl_engine.reference.minimax_h3.adaln_projection import ( - NativeH3AdaLNProjectionOp, - split_adaln_table, -) +from rl_engine.reference.minimax_h3.adaln_projection import NativeH3AdaLNProjectionOp from rl_engine.validation.models.h3_provider import provider_adaln_modulation requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") @@ -67,19 +64,6 @@ def _flat(outputs): return torch.cat([out.reshape(-1) for out in outputs]) -def _declared_golden(temb, weight, bias): - """FP64 projection of the declared BF16 activation, rounded from the FP32 SiLU. - - The kernel and diffusers round the FP32 SiLU; ``forward_fp32`` rounds an FP64 - SiLU, which can pick the other BF16 neighbour near a tie and shift a whole - output column (one input in 300 then matched only 95.7% of the outputs). - """ - - act = F.silu(temb.float()).to(weight.dtype).double() - table = F.linear(act, weight.double(), bias.double()) - return tuple(chunk.float() for chunk in split_adaln_table(table, weight.shape[0] // 18)) - - class TestReference: def test_provider_layout_matches_diffusers(self): """Match the provider values and timestep-major modality layout on CPU.""" @@ -171,7 +155,7 @@ def test_forward_is_correctly_rounded_declared_golden(self, block0, num): # K = 2688 (ties, and a few ULPs where near-zero sums cancel), so exact # equality is not reachable: over 300 temb draws per T the worst match was # 0.999618 (T = 1), 0.999690, 0.999728 and 0.999731 (T = 4). - declared = _flat(_declared_golden(temb, weight, bias)).bfloat16() + declared = _flat(golden).bfloat16() assert (_flat(outputs) == declared).float().mean().item() >= 0.9995 def test_early_cast_is_numerically_detectable(self, block0): @@ -181,7 +165,7 @@ def test_early_cast_is_numerically_detectable(self, block0): temb = _temb(3, seed=7) op = NativeH3AdaLNProjectionOp() ours = _flat(_cuda_op()(temb, weight, bias)) - declared = _flat(_declared_golden(temb, weight, bias)).bfloat16() + declared = _flat(op.forward_fp32(temb, weight, bias)).bfloat16() early = _flat(op.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() assert (ours == declared).float().mean() >= 0.9995 assert (ours == early).float().mean() < 0.8 From a4bc99f958003230273f05de48cd03ac796a4474 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:54:31 +0000 Subject: [PATCH 73/89] fix(h3): round the final_adaln_out goldens' activation from the FP32 SiLU Same mismatch as a0834552 for adaln_projection: forward_fp32, the test's _golden and the evidence report's golden_backward rounded an FP64 SiLU to BF16 for the declared cast, while the model and the kernel round the FP32 SiLU. Near a BF16 tie that picks the other neighbour and shifts a whole column of the norm_out table. All three now round F.silu(temb.float()); the VJP is unchanged. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/reference/minimax_h3/final_adaln_out.py | 5 ++++- rl_engine/validation/models/h3_report.py | 4 +++- tests/models/minimax_h3/test_h3_final_adaln_out.py | 4 +++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/rl_engine/reference/minimax_h3/final_adaln_out.py b/rl_engine/reference/minimax_h3/final_adaln_out.py index 0b19aaee3..a21a5cbdb 100644 --- a/rl_engine/reference/minimax_h3/final_adaln_out.py +++ b/rl_engine/reference/minimax_h3/final_adaln_out.py @@ -81,7 +81,10 @@ def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps ) t64 = temb.double() act = t64 * torch.sigmoid(t64) - act = act + (act.to(weight.dtype).double() - act).detach() # declared cast, identity VJP + # Declared cast of the FP32 SiLU, as the model computes it (an FP64 SiLU can + # round to the other neighbour near a tie); identity VJP. + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() table = F.linear(act, weight.double(), bias.double()) # norm_out.linear is a BF16 module: its (T, 2H) output is rounded to the # weight dtype, and every position of a timestep shares that row, so the diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 51407a030..107ee2ac2 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -742,7 +742,9 @@ def grads(fn, dtype=None): def golden_backward(x_, nw_, t_, w_, b_): # Keep the declared BF16 boundaries with FP64 leaves and identity VJPs. act = t_ * torch.sigmoid(t_) - act = act + (act.to(torch.bfloat16).double() - act).detach() + # The declared cast rounds the FP32 SiLU, as the model computes it. + declared = torch.nn.functional.silu(t_.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() table = torch.nn.functional.linear(act, w_, b_) table = table + (table.to(torch.bfloat16).double() - table).detach() shift, scale = table.chunk(2, dim=-1) diff --git a/tests/models/minimax_h3/test_h3_final_adaln_out.py b/tests/models/minimax_h3/test_h3_final_adaln_out.py index 5095d2c44..5fb964677 100644 --- a/tests/models/minimax_h3/test_h3_final_adaln_out.py +++ b/tests/models/minimax_h3/test_h3_final_adaln_out.py @@ -53,7 +53,9 @@ def _inputs(seq, num_timesteps=3, seed=0, batch=1): def _golden(x, norm_weight, temb, weight, bias, timestep_indices): act = temb * torch.sigmoid(temb) - act = act + (act.to(torch.bfloat16).double() - act).detach() + # Declared cast of the FP32 SiLU, as the model computes it; identity VJP. + declared = F.silu(temb.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() table = F.linear(act, weight, bias) table = table + (table.to(torch.bfloat16).double() - table).detach() shift, scale = table.chunk(2, dim=-1) From 068fcaa378ff05dc11b20f33b947912317079304 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:19:59 +0000 Subject: [PATCH 74/89] evidence(h3): regenerate adaln_projection_3mod report after the golden fix Written by tools/validation/models/h3_evidence.py from a clean tree at a0834552 on an idle B200. The accuracy fields are identical to the previous report (none of its 20 draws has an activation on a BF16 tie); the op latencies were re-measured: CUDA 100/100/101 us and provider 112/100/101 us for T = 1/2/4. The doc's table and evidence pointer follow. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 92dca1376..a76cc0417 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -108,14 +108,14 @@ B200, pinned BF16 weights (520 MB per block, streamed once per call): | T | CUDA op | provider (SiLU + cast + cuBLAS) | CUDA GEMV kernel | cuBLAS kernel | | --- | --- | --- | --- | --- | -| 1 | 105 µs | 115 µs | 78.3 µs (6.6 TB/s) | 93 µs | -| 2 | 105 µs | 103 µs | 78.4 µs | 79 µs | -| 4 | 106 µs | 104 µs | 79.0 µs | 79 µs | +| 1 | 100 µs | 112 µs | 78.3 µs (6.6 TB/s) | 93 µs | +| 2 | 100 µs | 100 µs | 78.4 µs | 79 µs | +| 4 | 101 µs | 101 µs | 79.0 µs | 79 µs | The kernel takes the same time for every T <= 8, because it always computes 8 rows. It -matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 24 µs of the op time is +matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 22 µs of the op time is the SiLU and cast kernels plus the Python wrapper. With those included, the op is within -about 2% of the provider for T >= 2. +about 1% of the provider for T >= 2. An earlier FMA-only BF16 kernel ran at 81, 98 and 166 µs for T = 1, 2 and 4. Its time grew with T because every extra row adds FMAs per weight element. Packed FP32 FMA @@ -127,7 +127,9 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `d06e120`. The report also +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 +SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone or in a batch of 9. From ecb30e25a04177c7f42f6e2fe08a9f4dee4a5216 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:20:24 +0000 Subject: [PATCH 75/89] evidence(h3): regenerate adaln_row_gather report and chain replay after the golden fix Written by tools/validation/models/h3_evidence.py and h3_chain_replay.py (--timesteps 1,2,3,4 --seq-lens 3,257,4097,32768 --backward) from a clean tree at d3b37fe1 on an idle B200. Every CUDA field and pass/fail flag matches the previous files; only timings and the diffusers backward errors (not repeatable run to run) changed. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 1d0387183..c870b6983 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -129,7 +129,10 @@ There are two data files: - [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -Both were written from a clean tree at commit `fa551c2`. +Both were written from a clean tree at commit `d3b37fe1` on an otherwise idle B200, regenerated +after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU (as the model +does). Every CUDA field and every pass/fail flag is unchanged; only timings and the diffusers +backward errors, which vary from run to run, moved. ## Existing implementations (RFC #420 reuse rule) From a1e37ef3894a5dc627d70a85c99e79d89ceefea8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:24:25 +0000 Subject: [PATCH 76/89] evidence(h3): regenerate the row-gather report on the newer runner after the golden fix This branch's adaln_row_gather report.json uses the newer evidence runner (backward-only timings, alternating order). Regenerated by tools/validation/models/h3_evidence.py from a clean tree at 14448146 on an idle B200; chain_replay.json comes from the #486 regeneration at d3b37fe1. Every CUDA pass/fail flag is unchanged; the op-backward correctly rounded fraction (unseeded draw) is 99.99931% vs 99.99965% before, and the doc quotes the new value. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 3cfd6adf6..4ffbb0b90 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -70,7 +70,7 @@ Measured on a B200 (torch 2.13.0+cu130): | --- | --- | --- | | Forward vs `index_select` | bitwise equal | | | `d_rows` repeat-bitwise | yes | no | -| `d_rows` correctly rounded from the FP64 sum | 99.99965% | 7.8% | +| `d_rows` correctly rounded from the FP64 sum | 99.99931% | 7.9% | | `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | ### Fused modulation: `H3AdaLNModulationCudaOp` @@ -131,9 +131,12 @@ There are two data files: - [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `80e4609` with backward-only +`report.json` was regenerated from a clean tree at commit `14448146` with backward-only timings and alternating execution order. `chain_replay.json` was written from a clean tree -at commit `fa551c2`. +at commit `d3b37fe1`. Both were regenerated after the AdaLN projection golden's declared +cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is +unchanged. The op-backward check draws unseeded rows and gradients, so its correctly +rounded fraction moves slightly between runs (99.99965% in the previous report). ## Existing implementations (RFC #420 reuse rule) From 7d716f227333a5265756f8f6fe3680ec5f38cf5e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:25:03 +0000 Subject: [PATCH 77/89] evidence(h3): regenerate final_adaln_out report and chain replay after the golden fix Written by tools/validation/models/h3_evidence.py and h3_chain_replay.py (--timesteps 1,2,3,4 --seq-lens 3,257,4097,32768 --backward) from a clean tree at 50e913fa on an idle B200. The CUDA accuracy and every pass/fail flag are unchanged; the h3_rmsnorm stage's chained error against the golden drops in all 16 cases (worst 0.0796 -> 0.0671). The latency table in the doc follows the new report. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 3b571efec..44e08e507 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -62,9 +62,9 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | -| 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | -| 131072 | 1.24 ms | 2.94 ms | 7.75 ms | 20.6 ms | +| 4097 | 0.18 ms | 0.19 ms | 1.50 ms | 1.07 ms | +| 32768 | 0.44 ms | 0.84 ms | 2.75 ms | 5.34 ms | +| 131072 | 1.29 ms | 2.99 ms | 7.76 ms | 20.8 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -77,13 +77,16 @@ There are two data files: - [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `001684d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + tree at `50e913fa` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `001684d`. + tree at `50e913fa`. Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` + stage's chained error against the golden is lower in all 16 cases (worst 0.0796 before, + 0.0671 now). ## Existing implementations (RFC #420 reuse rule) From 76594af5cb43cb73698faa7b1698cdb85ab9aff8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:28:22 +0000 Subject: [PATCH 78/89] evidence(h3): regenerate this branch's final_adaln_out report after the golden fix This branch carries a newer final_adaln_out report.json than #489. Regenerated by tools/validation/models/h3_evidence.py from a clean tree at 0a2d5589 on an idle B200; chain_replay.json comes from #489's regeneration at 50e913fa. CUDA accuracy and every pass/fail flag are unchanged; the doc's latency table follows the new report. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index b95c906b0..3e884dcf2 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -63,9 +63,9 @@ and forward-graph setup: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.19 ms | 0.18 ms | 1.40 ms | 0.85 ms | -| 32768 | 0.46 ms | 0.80 ms | 2.59 ms | 4.49 ms | -| 131072 | 1.32 ms | 3.02 ms | 7.05 ms | 17.69 ms | +| 4097 | 0.19 ms | 0.21 ms | 1.37 ms | 0.83 ms | +| 32768 | 0.42 ms | 0.82 ms | 2.53 ms | 4.45 ms | +| 131072 | 1.28 ms | 2.97 ms | 7.02 ms | 17.61 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -78,13 +78,16 @@ There are two data files: - [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `bde1d8d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `001684d`. + tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 + before, 0.0671 now). ## Existing implementations (RFC #420 reuse rule) From a6c85451abb1c818b244a3a9dde16f8d283cdcc3 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:26:07 +0000 Subject: [PATCH 79/89] fix(h3): seed the row-gather op-backward check in the evidence report _gather_accuracy drew its table rows and upstream gradients from the global CUDA generator, so the CUDA op-backward correctly rounded fraction moved between runs (99.99965% and 99.99931% in two reports). Draw both from a seeded generator; three runs now give the same CUDA fraction. The diffusers fraction still varies run to run (atomic scatter-add). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 8b17494a2..505a49e56 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -351,7 +351,9 @@ def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("adaln_row_gather", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) - rows = torch.randn(9, 6 * 5376, device="cuda").bfloat16() + # Seeded so the op-backward fractions below are reproducible from this commit. + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(9, 6 * 5376, device="cuda", generator=gen).bfloat16() forward_bitwise = {} for seq in (1, 257, 4097, 32768): ti, tags = h3_packed_layout(seq, 3, seed=seq) @@ -361,7 +363,7 @@ def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. ti, tags = h3_packed_layout(4097, 3, seed=1) - grads = [torch.randn(4097, 5376, device="cuda").bfloat16() for _ in range(6)] + grads = [torch.randn(4097, 5376, device="cuda", generator=gen).bfloat16() for _ in range(6)] def grad_of(fn): leaf = rows.detach().clone().requires_grad_(True) From cbbfaedc8702bb70a6b74d18c8b1c7b0c5c48346 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:29:33 +0000 Subject: [PATCH 80/89] evidence(h3): regenerate the row-gather report with the seeded op-backward check Written by tools/validation/models/h3_evidence.py from a clean tree at 40f4400c on an idle B200. With rows and gradients seeded, the CUDA op-backward correctly rounded fraction is 99.99931% (99.99999% before, which came from the default CUDA generator's state after the runner's earlier draws). Every other CUDA field and pass/fail flag is unchanged. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index c870b6983..9fa8a690e 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -70,7 +70,7 @@ Measured on a B200 (torch 2.13.0+cu130): | --- | --- | --- | | Forward vs `index_select` | bitwise equal | | | `d_rows` repeat-bitwise | yes | no | -| `d_rows` correctly rounded from the FP64 sum | 99.99999% | 7.8% | +| `d_rows` correctly rounded from the FP64 sum | 99.99931% | 7.8% | | `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | ### Fused modulation: `H3AdaLNModulationCudaOp` @@ -129,10 +129,14 @@ There are two data files: - [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -Both were written from a clean tree at commit `d3b37fe1` on an otherwise idle B200, regenerated +`report.json` was rewritten from a clean tree at commit `40f4400c`, after the op-backward check +was seeded (it previously took its rows and gradients from the process's default CUDA +generator, so its fraction depended on how many draws the runner made before it; 99.99999% in +the previous report). `chain_replay.json` was written from a clean tree at commit `d3b37fe1`. +Both ran on an otherwise idle B200 and were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU (as the model -does). Every CUDA field and every pass/fail flag is unchanged; only timings and the diffusers -backward errors, which vary from run to run, moved. +does). Apart from the seeded op-backward fraction, every CUDA field and every pass/fail flag +is unchanged; only timings and the diffusers backward errors, which vary from run to run, moved. ## Existing implementations (RFC #420 reuse rule) From fe0350c1cf8015aa4b6b05b6c865adfed9a492af Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:30:04 +0000 Subject: [PATCH 81/89] evidence(h3): regenerate this branch's row-gather report with the seeded check Written by tools/validation/models/h3_evidence.py from a clean tree at f020f7d5 (same code as this merge's parent tree) on an idle B200. The seeded CUDA op-backward fraction is 99.99931%, matching #486's seeded run; every other CUDA field and pass/fail flag is unchanged. The doc no longer says the fraction varies between runs: the old draw was deterministic per runner version, not per run. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 4ffbb0b90..88fee5f3d 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,12 +131,13 @@ There are two data files: - [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `14448146` with backward-only -timings and alternating execution order. `chain_replay.json` was written from a clean tree -at commit `d3b37fe1`. Both were regenerated after the AdaLN projection golden's declared -cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is -unchanged. The op-backward check draws unseeded rows and gradients, so its correctly -rounded fraction moves slightly between runs (99.99965% in the previous report). +`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +timings and alternating execution order, and with the op-backward check's rows and gradients +drawn from a seeded generator. Before, they came from the process's default CUDA generator, +so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). +`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the +model does; every CUDA pass/fail flag is unchanged. ## Existing implementations (RFC #420 reuse rule) From fb1bd8aa5035051d3a4f98058b19f88d9449cac7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:55:22 +0000 Subject: [PATCH 82/89] fix(h3): seed the timestep_mlp timing inputs in the evidence report _mlp_perf drew its timestep features from the process's default CUDA generator, so the inputs depended on how many draws ran first. Seed them per T. Only timing inputs change; no reported accuracy field depends on them, so the committed report needs no regeneration. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index e42382ca8..f7bd01520 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -139,7 +139,8 @@ def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: weight_bytes = sum(p.numel() * p.element_size() for p in params) cases = [] for num in (1, 2, 3, 4): - x = torch.rand(num, 256, device="cuda") * 2 - 1 + gen = torch.Generator(device="cuda").manual_seed(num) + x = torch.rand(num, 256, device="cuda", generator=gen) * 2 - 1 cases.append( { "op": "timestep_mlp_fp32", From 7f2fa858bcfe42e51130f20a95352f002a66c931 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:56:42 +0000 Subject: [PATCH 83/89] fix(h3): seed the adaln_projection timing and invariance inputs in the report _projection_perf and the batch-invariance part of _projection_accuracy drew temb from the process's default CUDA generator. Seed both. The per-draw accuracy inputs were already seeded; the full accuracy section from this commit is identical to the committed report's, so it needs no regeneration. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 92e8bd632..2dff72a8a 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -246,7 +246,8 @@ def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: weight, bias = _adaln_params() cases = [] for num in (1, 2, 3, 4): - temb = torch.randn(num, 2688, device="cuda") + gen = torch.Generator(device="cuda").manual_seed(num) + temb = torch.randn(num, 2688, device="cuda", generator=gen) cases.append( { "op": "adaln_projection_3mod", @@ -285,7 +286,8 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, cuda_frac.append(float((ours == gold).float().mean())) provider_frac.append(float((theirs == gold).float().mean())) early_frac.append(float((ours == early).float().mean())) - temb = torch.randn(9, 2688, device="cuda") + gen = torch.Generator(device="cuda").manual_seed(draws) + temb = torch.randn(9, 2688, device="cuda", generator=gen) full = op(temb, weight, bias) invariant = all( all( From 2553881ba979522645a592706a77dc374c870a79 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:57:53 +0000 Subject: [PATCH 84/89] fix(h3): seed the row-gather timing inputs in the evidence report _gather_perf drew its table rows and upstream gradients from the process's default CUDA generator. Seed them. Only timing inputs change; no reported accuracy field depends on them. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 4c5eaba33..347f3f856 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -317,12 +317,15 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: op = registry.get_op("adaln_row_gather", device="cuda") num_timesteps, hidden = 3, 5376 - rows = torch.randn(3 * num_timesteps, 6 * hidden, device="cuda").bfloat16() + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(3 * num_timesteps, 6 * hidden, device="cuda", generator=gen).bfloat16() chunks = rows.chunk(6, dim=-1) cases = [] for seq in GATHER_SEQ_LENS: ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) - grads = [torch.randn(seq, hidden, device="cuda").bfloat16() for _ in range(6)] + grads = [ + torch.randn(seq, hidden, device="cuda", generator=gen).bfloat16() for _ in range(6) + ] def backward(fn, ti=ti, tags=tags, grads=grads): leaf = rows.detach().requires_grad_(True) From 95186cb701161fab13b592c3b528d42dcb2ad274 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 20:00:41 +0000 Subject: [PATCH 85/89] fix(h3): seed the plain-RMSNorm inputs of the norm evidence check _norm_accuracy drew the inputs of its four plain-norm bitwise checks from the process's default CUDA generator. Draw them from one seeded generator. All four checks stay bitwise equal to F.rms_norm, and the rest of the accuracy section matches the committed report (the CUDA dscale error differs only in the 16th digit, from the FP64 reference's atomic index_select backward), so it needs no regeneration. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index d42942492..1102aa396 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -490,6 +490,7 @@ def backward(output, grad=grad): def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("h3_rmsnorm", device="cuda") weights = load_h3_conditioning_weights("cuda") + gen = torch.Generator(device="cuda").manual_seed(0) plain = {} for name in ( "transformer_blocks.0.norm1.weight", @@ -497,7 +498,7 @@ def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: "token_refiner.final_norm.weight", "norm_out.norm.weight", ): - x = torch.randn(2, 777, 5376, device="cuda").bfloat16() + x = torch.randn(2, 777, 5376, device="cuda", generator=gen).bfloat16() ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) plain[name] = bool(torch.equal(op(x, weights[name]), ref)) x, weight, shift, scale, index = _norm_inputs(4097, seed=1) From 8725b3d25e2b55d7304cfc8758827f48b13c5a82 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 20:11:43 +0000 Subject: [PATCH 86/89] test(h3): let the CPU timing shim accept seeded CUDA generators The gather case of test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events redirects torch.randn to the CPU. _gather_perf now draws from a seeded CUDA generator, which cannot produce a CPU tensor, so the shim swaps it for a CPU generator with the same seed (one per CUDA generator, so the draw sequence continues). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_benchmark_timing.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/models/minimax_h3/test_h3_benchmark_timing.py b/tests/models/minimax_h3/test_h3_benchmark_timing.py index 7dba40be6..95af1ff18 100644 --- a/tests/models/minimax_h3/test_h3_benchmark_timing.py +++ b/tests/models/minimax_h3/test_h3_benchmark_timing.py @@ -164,9 +164,17 @@ def forward(x, weight, shift, scale, index, **kwargs): build = h3_report._norm_perf else: randn = torch.randn + cpu_generators = {} def cpu_randn(*args, **kwargs): kwargs["device"] = "cpu" + # A seeded CUDA generator cannot draw a CPU tensor: continue the same seed on + # one CPU generator per CUDA generator. + gen = kwargs.get("generator") + if gen is not None and gen.device.type != "cpu": + if id(gen) not in cpu_generators: + cpu_generators[id(gen)] = torch.Generator().manual_seed(gen.initial_seed()) + kwargs["generator"] = cpu_generators[id(gen)] return randn(*args, **kwargs) provider = h3_report.provider_adaln_row_gather From 0033b3a5a513604c72987abbc2cc08bd22d5a67d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:22 +0000 Subject: [PATCH 87/89] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 12 ++++++------ docs/operators/h3-adaln-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-final-adaln-out.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-sp-norm-adaln.md | 2 +- docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ docs/operators/h3-tp-adaln-3mod.md | 2 +- 9 files changed, 46 insertions(+), 46 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6c972994d..9a5d815a5 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -75,15 +75,15 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](../../reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-gate-residual-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) -![gate_residual vs existing implementations](../../reports/experiments/h3-prior-art-b200/gate_residual.png) +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -94,7 +94,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../../reports/experiments/h3-prior-art-b200/gate_residual.json) +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) was written from a clean tree at `4010854` by ```bash @@ -111,10 +111,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a4cb57dd8..e8fc09537 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -124,9 +124,9 @@ dependence on T. ## Evidence -![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also @@ -135,7 +135,7 @@ or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![adaln_projection vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_projection.png) +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional @@ -157,7 +157,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../../reports/experiments/h3-prior-art-b200/adaln_projection.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) was written from a clean tree at `c88691c` by ```bash @@ -174,10 +174,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 88fee5f3d..94e32b711 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -122,13 +122,13 @@ alternates each iteration and is recorded in the report. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../../reports/experiments/h3-adaln-row-gather-b200/figure.png) +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward bitwise checks, op-level backward, and the chain-backward cases plotted above. -- [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. `report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only @@ -141,7 +141,7 @@ model does; every CUDA pass/fail flag is unchanged. ## Existing implementations (RFC #420 reuse rule) -![adaln_row_gather vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.png) +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) was written from a clean tree at `6c900ae` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_row_gather bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 93143b858..d40a49f5f 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -77,17 +77,17 @@ with their original commit and corrected scope. ## Evidence -![final_adaln_out on B200: latency and backward accuracy](../../reports/experiments/h3-final-adaln-out-b200/figure.png) +![final_adaln_out on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. -- [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the @@ -96,7 +96,7 @@ There are two data files: ## Existing implementations (RFC #420 reuse rule) -![final_adaln_out vs existing implementations](../../reports/experiments/h3-prior-art-b200/final_adaln_out.png) +![final_adaln_out vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -107,7 +107,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](../../reports/experiments/h3-prior-art-b200/final_adaln_out.json) +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.json) was written from a clean tree at `03da729` by ```bash @@ -124,10 +124,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![final_adaln_out bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) +![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 740f16232..75dd3d477 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -92,9 +92,9 @@ each iteration and is recorded in the report. ## Evidence -![h3_rmsnorm on B200: latency and backward accuracy](../../reports/experiments/h3-rmsnorm-b200/figure.png) +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report.json), written by +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; @@ -103,7 +103,7 @@ The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report. ## Existing implementations (RFC #420 reuse rule) -![norm_modulate vs existing implementations](../../reports/experiments/h3-prior-art-b200/norm_modulate.png) +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -119,7 +119,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](../../reports/experiments/h3-prior-art-b200/norm_modulate.json) +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) was written from a clean tree at `ee83dec` by ```bash @@ -136,10 +136,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![norm_modulate bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 1f19d0444..8eb91cc38 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -70,7 +70,7 @@ How many rows move depends on how the table rows interleave along the sequence: ## Evidence -![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](../../reports/experiments/h3-sp-norm-adaln-b200/figure.png) +![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr494/reports/experiments/h3-sp-norm-adaln-b200/figure.png) The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `6c0d380`, with real NCCL processes on one 8 x B200 node, one GPU per rank. The region is diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 8049a87df..cc5809cee 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -109,15 +109,15 @@ kernel time. ## Evidence -![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../../reports/experiments/h3-timestep-mlp-b200/figure.png) +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-mlp-b200/report.json), written +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![timestep_mlp vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_mlp.png) +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -134,7 +134,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../../reports/experiments/h3-prior-art-b200/timestep_mlp.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) was written from a clean tree at `ffd958a` by ```bash @@ -151,10 +151,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index c5f69a15b..75edf172b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -88,15 +88,15 @@ about 1 KB per timestep. ## Evidence -![timestep_sinusoid_h3 on B200: latency and error vs FP64](../../reports/experiments/h3-timestep-sinusoid-b200/figure.png) +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-sinusoid-b200/report.json). +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). `tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. ## Existing implementations (RFC #420 reuse rule) -![timestep_sinusoid vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -108,7 +108,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) was written from a clean tree at `7a82917` by ```bash @@ -126,11 +126,11 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). -![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. -[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) was written from a clean tree at `0eb5ddd5` by ```bash diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index c348bce35..318dbd2d0 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -72,7 +72,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ## Evidence -![tp_adaln_3mod on 8 x B200: byte equality and time](../../reports/experiments/h3-tp-adaln-3mod-b200/figure.png) +![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For From 4ef511c9e05e26073dc4baff6ec0075e5048be24 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:06:03 +0000 Subject: [PATCH 88/89] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 2 +- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-final-adaln-out.md | 6 +++--- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-sp-norm-adaln.md | 2 +- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- docs/operators/h3-tp-adaln-3mod.md | 2 +- 9 files changed, 17 insertions(+), 17 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 9a5d815a5..ffd200b0a 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -114,7 +114,7 @@ gradient with diffusers' path and with this kernel. ![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `a746bff0` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index e8fc09537..526e2b4cf 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -127,7 +127,7 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `579775f1` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone @@ -177,7 +177,7 @@ gradient with diffusers' path and with this kernel. ![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `50d1e924` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 94e32b711..923ad2f77 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,11 +131,11 @@ There are two data files: - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +`report.json` was regenerated from a clean tree at commit `29b07c5d` with backward-only timings and alternating execution order, and with the op-backward check's rows and gradients drawn from a seeded generator. Before, they came from the process's default CUDA generator, -so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). -`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +so the fraction depended on how many draws the runner made first (99.99965% at `d2a6546`). +`chain_replay.json` was written from a clean tree at commit `61d23ade`. Both were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is unchanged. @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `62f8009e` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index d40a49f5f..a46954f06 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -83,14 +83,14 @@ There are two data files: - [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + tree at `ddca8f95` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + tree at `a4bc99f9` (#489). Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 before, 0.0671 now). @@ -127,7 +127,7 @@ gradient with diffusers' path and with this kernel. ![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `8cb4abed` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 75dd3d477..ac625ae35 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -95,7 +95,7 @@ each iteration and is recorded in the report. ![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by -`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -139,7 +139,7 @@ gradient with diffusers' path and with this kernel. ![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `83797833` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 8eb91cc38..eeb90f72f 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -72,7 +72,7 @@ How many rows move depends on how the table rows interleave along the sequence: ![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr494/reports/experiments/h3-sp-norm-adaln-b200/figure.png) -The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `6c0d380`, with +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `ecde5bb`, with real NCCL processes on one 8 x B200 node, one GPU per rank. The region is `norm2(residual + gate_msa[row] * y)` with `shift_mlp`/`scale_mlp` modulation, H = 5376, BF16. It covers six cases: S = 4097 (block, interleaved, B = 2), 32768 (block, interleaved) and 131072 diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index cc5809cee..b65b19fca 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -112,7 +112,7 @@ kernel time. ![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written -by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `418998c`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) @@ -154,7 +154,7 @@ gradient with diffusers' path and with this kernel. ![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `73939f2d` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 75edf172b..ff6579cd8 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -131,7 +131,7 @@ kernel compiled with `--use_fast_math`, which the build excludes (see Known Limi SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) -was written from a clean tree at `0eb5ddd5` by +was written from a clean tree at `be6a94ae` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 318dbd2d0..868df445a 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -74,7 +74,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) -The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `507c9ad`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are byte-equal to WS1 computed on the same GPU. From 193f079ead3e51d8834c2f8cdd7ab222a04a7bd1 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:16:26 +0000 Subject: [PATCH 89/89] [WS1][CUDA][MiniMax-H3] ws1_one_h3_block: block 0 node-by-node replay and diagnostics (RFC #420, #539 split 1/3) Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .github/workflows/ci.yml | 7 +- .../benchmark_deterministic_attention.py | 175 ++++- build_tools/envs.py | 26 + build_tools/extensions.py | 27 + csrc/bindings/h3_ops.cpp | 35 + .../cuda/attention/deterministic_attention.cu | 341 ++++++++- docs/operators/README.md | 1 + docs/operators/attention.md | 57 ++ docs/operators/h3-one-block.md | 183 +++++ .../models/minimax_h3/diagnose-block.sbatch | 41 + .../cuda/attention/h3_full_attention.py | 150 ++++ .../minimax_h3/adaln_modulation.py | 4 +- .../minimax_h3/adaln_projection.py | 2 +- .../minimax_h3/adaln_row_gather.py | 2 +- .../minimax_h3/final_adaln_out.py | 4 +- .../minimax_h3/gate_residual.py | 4 +- .../cuda/model_specific/minimax_h3/rmsnorm.py | 4 +- .../minimax_h3/sp_norm_adaln.py | 8 +- .../model_specific/minimax_h3/timestep_mlp.py | 2 +- .../minimax_h3/tp_adaln_projection.py | 2 +- .../reference/attention/standard_attn.py | 35 + rl_engine/reference/minimax_h3/block.py | 98 +++ rl_engine/runtime/registry.py | 32 + rl_engine/validation/models/h3_block.py | 718 ++++++++++++++++++ rl_engine/validation/models/h3_cases.py | 124 +++ rl_engine/validation/models/h3_chain.py | 8 +- rl_engine/validation/models/h3_manifest.json | 34 + rl_engine/validation/models/h3_provider.py | 72 ++ rl_engine/validation/models/h3_report.py | 26 +- rl_engine/validation/models/h3_weights.py | 41 +- rl_engine/validation/models/h3_ws2.py | 2 +- .../validation/operators/operator_inputs.py | 10 + .../validation/operators/operator_specs.py | 10 + tests/backends/cuda/test_h3_full_attention.py | 344 +++++++++ tests/models/minimax_h3/conftest.py | 15 + .../minimax_h3/test_h3_adaln_row_gather.py | 6 +- tests/models/minimax_h3/test_h3_det_linear.py | 2 +- .../minimax_h3/test_h3_native_validation.py | 12 +- tests/models/minimax_h3/test_h3_one_block.py | 139 ++++ .../minimax_h3/test_h3_sp_norm_adaln.py | 2 + tools/migration/layout.json | 2 + .../validation/models/h3_block_diagnostic.py | 148 ++++ tools/validation/models/h3_block_replay.py | 114 +++ tools/validation/models/h3_block_ulp_sweep.py | 96 +++ .../models/h3_evidence_integrity.py | 34 + .../models/h3_first_step_diagnostic.py | 117 +++ tools/validation/models/plot_h3_block.py | 172 +++++ tools/weights/prepare_h3_weights.py | 66 +- 48 files changed, 3455 insertions(+), 99 deletions(-) create mode 100644 csrc/bindings/h3_ops.cpp create mode 100644 docs/operators/h3-one-block.md create mode 100644 examples/models/minimax_h3/diagnose-block.sbatch create mode 100644 rl_engine/backends/cuda/attention/h3_full_attention.py create mode 100644 rl_engine/reference/minimax_h3/block.py create mode 100644 rl_engine/validation/models/h3_block.py create mode 100644 tests/backends/cuda/test_h3_full_attention.py create mode 100644 tests/models/minimax_h3/test_h3_one_block.py create mode 100644 tools/validation/models/h3_block_diagnostic.py create mode 100644 tools/validation/models/h3_block_replay.py create mode 100644 tools/validation/models/h3_block_ulp_sweep.py create mode 100644 tools/validation/models/h3_evidence_integrity.py create mode 100644 tools/validation/models/h3_first_step_diagnostic.py create mode 100644 tools/validation/models/plot_h3_block.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4634569cd..58d2dd175 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,9 +5,9 @@ name: CI-Pipeline on: push: - branches: [ main, test ] + branches: [ main, test, test-h3 ] pull_request: - branches: [ main, test ] + branches: [ main, test, test-h3 ] permissions: contents: read @@ -71,6 +71,9 @@ jobs: - name: Check refactor packaging and compatibility run: python ci/run.py layout -q + - name: Run full H3 binding contracts + run: python -m pytest -q tests/models/minimax_h3/test_h3_full_model.py + - name: Run Mocked Hardware Discovery Tests run: | python -m pytest tests/runtime/test_dispatch.py -v diff --git a/benchmarks/operators/attention/benchmark_deterministic_attention.py b/benchmarks/operators/attention/benchmark_deterministic_attention.py index b41f98a00..b830c6393 100644 --- a/benchmarks/operators/attention/benchmark_deterministic_attention.py +++ b/benchmarks/operators/attention/benchmark_deterministic_attention.py @@ -5,12 +5,17 @@ Reports latency and peak memory for Qwen3-8B representative shapes, including the cost of full scores/P materialization. + +Use the ``h3`` subcommand for the strict-native/CUDA all-gradient matrix. """ import argparse +import math +import sys import time import torch +from torch.autograd.function import once_differentiable def benchmark_attention( @@ -93,5 +98,173 @@ def main(): torch.cuda.empty_cache() +def h3_value_grad(op, xs, dy): + out = op(*xs) + return [out.detach(), *[g.detach() for g in torch.autograd.grad(out, xs, dy)]] + + +def h3_bits_equal(a, b): + return torch.equal(a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8)) + + +class H3StrictNative(torch.autograd.Function): + """Unoptimized native baseline, compiled with the same strict arithmetic flags.""" + + @staticmethod + def forward(ctx, q, k, v): + from rl_engine.backends.cuda.attention.h3_full_attention import _extension + + q, k, v = (x.transpose(1, 2).contiguous() for x in (q, k, v)) + out, lse, p = _extension().deterministic_attention_forward( + q, k, v, False, 1.0 / math.sqrt(128), None + ) + ctx.save_for_backward(q, k, v, p) + lse = lse.transpose(1, 2) + ctx.mark_non_differentiable(lse) + return out.transpose(1, 2), lse + + @staticmethod + @once_differentiable + def backward(ctx, grad, unused_lse): + from rl_engine.backends.cuda.attention.h3_full_attention import _extension + + grads = _extension().deterministic_attention_backward( + grad.transpose(1, 2).contiguous(), *ctx.saved_tensors, False, 1.0 / math.sqrt(128), None + ) + return tuple(g.transpose(1, 2) for g in grads) + + +def h3_main(argv=None): + """Paired strict native/CUDA forward and all-QKV-gradient matrix.""" + import hashlib + import json + import statistics + from pathlib import Path + + ROOT = Path(__file__).resolve().parents[3] + sys.path.insert(0, str(ROOT)) + from rl_engine.backends.cuda.attention.h3_full_attention import H3FullAttentionOp, _extension + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument( + "--shapes", + nargs="+", + default=["1x129", "1x512", "1x1025", "1x4096", "2x1024", "4x1024"], + ) + parser.add_argument("--chunks", nargs="+", type=int, default=[64, 128, 256]) + parser.add_argument("--samples", type=int, default=6) + parser.add_argument("--warmups", type=int, default=2) + args = parser.parse_args(argv) + try: + shapes = [tuple(map(int, x.split("x"))) for x in args.shapes] + except ValueError: + parser.error("shapes must be positive BxS pairs") + if any(len(s) != 2 or min(s) <= 0 for s in shapes): + parser.error("shapes must be positive BxS pairs") + if args.samples < 3 or args.warmups < 1 or min(args.chunks) <= 0: + parser.error("use >=3 samples, >=1 warmup and positive chunks") + args.output.mkdir(parents=True, exist_ok=True) + ops = {"strict_native": lambda *xs: H3StrictNative.apply(*xs)[0]} + for chunk in dict.fromkeys(args.chunks): + ops[f"cuda-c{chunk}"] = H3FullAttentionOp(chunk_rows=chunk) + sources = [Path(__file__), ROOT / "csrc/cuda/attention/deterministic_attention.cu"] + sources += [ + ROOT / "rl_engine/backends/cuda/attention/h3_full_attention.py", + ROOT / "rl_engine/reference/attention/standard_attn.py", + ] + report = { + "status": "running", + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "extension_sha256": hashlib.sha256(Path(_extension().__file__).read_bytes()).hexdigest(), + "source_sha256": { + str(p.relative_to(ROOT)): hashlib.sha256(p.read_bytes()).hexdigest() for p in sources + }, + "provenance": {k: v.provenance() for k, v in ops.items() if k != "strict_native"}, + "measurement": ( + "CUDA events; eager forward + backward, including layout copies and allocations" + ), + "accuracy": "output/dQ/dK/dV raw-bit equality against strict native on identical inputs", + "warmups": args.warmups, + "samples": args.samples, + "cases": [], + } + + def save(): + temporary = args.output / "report.tmp" + temporary.write_text(json.dumps(report, indent=2) + "\n") + temporary.replace(args.output / "report.json") + + save() + for batch, seq in shapes: + # The native baseline materializes quadratic FP32 intermediates. + free, _ = torch.cuda.mem_get_info() + estimate = batch * 56 * seq * seq * 4 * 3 + batch * seq * 56 * 128 * 2 * 20 + if estimate > free * 0.75: + raise RuntimeError(f"{batch}x{seq}: insufficient free memory for safe paired baseline") + torch.manual_seed(42) + xs = [ + torch.randn(batch, seq, 56, 128, device="cuda", dtype=torch.bfloat16).requires_grad_() + for _ in range(3) + ] + dy = torch.randn_like(xs[0]) + golden = h3_value_grad(ops["strict_native"], xs, dy) + for name, op in ops.items(): + actual = h3_value_grad(op, xs, dy) + assert all(h3_bits_equal(a, g) for a, g in zip(actual, golden)), (batch, seq, name) + del actual + del golden + for op in ops.values(): + for _ in range(args.warmups): + h3_value_grad(op, xs, dy) + torch.cuda.synchronize() + rows = {name: [] for name in ops} + names = list(ops) + for sample in range(args.samples): + order = names[sample % len(names) :] + names[: sample % len(names)] + if sample % 2: + order = list(reversed(order)) + for name in order: + resident = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = h3_value_grad(ops[name], xs, dy) + end.record() + end.synchronize() + rows[name].append( + { + "ms": start.elapsed_time(end), + "extra_peak_bytes": torch.cuda.max_memory_allocated() - resident, + } + ) + del result + case = dict( + shape=[batch, seq, 56, 128], + seed=42, + backends={ + name: dict( + median_ms=statistics.median(s["ms"] for s in samples), + peak_extra_bytes=max(s["extra_peak_bytes"] for s in samples), + raw_samples=samples, + ) + for name, samples in rows.items() + }, + ) + report["cases"].append(case) + save() + print(json.dumps(case), flush=True) + del xs, dy + report["status"] = "passed" + save() + print("PAIRED_MATRIX_PASSED", flush=True) + + if __name__ == "__main__": - main() + if sys.argv[1:2] == ["h3"]: + h3_main(sys.argv[2:]) + else: + main() diff --git a/build_tools/envs.py b/build_tools/envs.py index 77dcb8e35..b4c901fd1 100644 --- a/build_tools/envs.py +++ b/build_tools/envs.py @@ -8,6 +8,7 @@ """ import os +import shlex def env_flag(name: str, default: bool = False) -> bool: @@ -27,6 +28,31 @@ def env_flag(name: str, default: bool = False) -> bool: KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC = "KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC" KERNEL_ALIGN_FORCE_SM90 = "KERNEL_ALIGN_FORCE_SM90" RL_KERNEL_REQUIRE_EXT = "RL_KERNEL_REQUIRE_EXT" +RL_KERNEL_H3_STRICT = "RL_KERNEL_H3_STRICT" + + +def h3_strict_cuda_flags() -> list[str]: + """Opt-in, cache-visible arithmetic profile for the H3 CUDA entry points. + + External nvcc flags are rejected in this profile: append flags and response + files can silently override these options after build_tools/extensions.py constructs them. + Other operators keep the existing build defaults when this opt-in is absent. + """ + if not env_flag(RL_KERNEL_H3_STRICT): + return [] + if env_flag(KERNEL_ALIGN_USE_FAST_MATH): + raise ValueError("RL_KERNEL_H3_STRICT is incompatible with KERNEL_ALIGN_USE_FAST_MATH") + for name in ("NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS"): + if shlex.split(os.environ.get(name, "")): + raise ValueError(f"Unset {name} for the audited RL_KERNEL_H3_STRICT build") + return [ + "-DRL_KERNEL_H3_STRICT=1", + "--fmad=false", + "--ftz=false", + "--prec-div=true", + "--prec-sqrt=true", + ] + KERNEL_ALIGN_FORCE_ASCEND = "KERNEL_ALIGN_FORCE_ASCEND" KERNEL_ALIGN_ASCEND_ARCH = "KERNEL_ALIGN_ASCEND_ARCH" diff --git a/build_tools/extensions.py b/build_tools/extensions.py index dd135beff..dbddf0c0a 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -119,6 +119,7 @@ def _native_extension_required() -> bool: """Whether the caller explicitly requested a native extension build.""" return ( envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or envs.env_flag(envs.RL_KERNEL_H3_STRICT) or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) or envs.env_flag("FORCE_CUDA") @@ -172,6 +173,7 @@ def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: def get_extensions(): + h3_flags = envs.h3_strict_cuda_flags() torch, _, GpuExtension = _load_torch_extension_tools() if torch is None: message = ( @@ -196,6 +198,9 @@ def get_extensions(): torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") is_rocm = getattr(torch.version, "hip", None) is not None + if h3_flags and (is_rocm or _musa_build_available(torch)): + raise RuntimeError("RL_KERNEL_H3_STRICT currently qualifies NVIDIA CUDA only") + if _musa_build_available(torch): extensions.append( GpuExtension( @@ -398,6 +403,28 @@ def get_extensions(): ) ) + if h3_flags: + # Separate DSO: strict arithmetic never changes other operators. + extensions.append( + GpuExtension( + name="rl_engine._H3_C", + sources=[ + "csrc/bindings/h3_ops.cpp", + "csrc/cuda/attention/deterministic_attention.cu", + ], + extra_compile_args={ + "cxx": ["-O3"], + "nvcc": [ + "-O3", + *h3_flags, + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + ], + }, + extra_link_args=torch_rpath, + ) + ) + extensions.extend(_ascend_extensions()) if _native_extension_required() and not extensions: diff --git a/csrc/bindings/h3_ops.cpp b/csrc/bindings/h3_ops.cpp new file mode 100644 index 000000000..f26859760 --- /dev/null +++ b/csrc/bindings/h3_ops.cpp @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +#include + +#include +#include +#include + +std::map h3_attention_build_info(); +std::vector h3_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v); +torch::Tensor h3_attention_probabilities( + torch::Tensor q, torch::Tensor k, torch::Tensor v); +std::vector h3_attention_backward_accumulate( + torch::Tensor grad_output, torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor probabilities, + torch::Tensor dk, torch::Tensor dv); + +// Expose the original core in the same strict build for numerical comparisons. +std::vector deterministic_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, + torch::optional key_padding_mask); +std::vector deterministic_attention_backward( + torch::Tensor grad_output, torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor probabilities, bool causal, double scale, + torch::optional key_padding_mask); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("h3_attention_build_info", &h3_attention_build_info); + m.def("h3_attention_forward", &h3_attention_forward); + m.def("h3_attention_probabilities", &h3_attention_probabilities); + m.def("h3_attention_backward_accumulate", &h3_attention_backward_accumulate); + m.def("deterministic_attention_forward", &deterministic_attention_forward); + m.def("deterministic_attention_backward", &deterministic_attention_backward); +} diff --git a/csrc/cuda/attention/deterministic_attention.cu b/csrc/cuda/attention/deterministic_attention.cu index 82ec824db..a3c7974dc 100644 --- a/csrc/cuda/attention/deterministic_attention.cu +++ b/csrc/cuda/attention/deterministic_attention.cu @@ -6,13 +6,17 @@ // All reductions use fixed order; no split-KV or dynamic dispatch. #include +#include #include #include #include #include #include +#include #include +#include +#include #include namespace { @@ -60,6 +64,117 @@ __global__ void qk_kernel( scores[out_idx] = scale * acc; } +// H3 fixed-D dot product. Cooperative, contiguous global loads convert each +// BF16 input once per CTA. Padding/transposition keeps shared-memory reads +// conflict-free for the (16 key, 16 query) thread layout. Each score retains +// the original ascending-d, separate-multiply/add FP32 reduction. +template +__global__ void h3_shared_qk_kernel( + const at::BFloat16* __restrict__ Q, + const at::BFloat16* __restrict__ K, + float* __restrict__ scores, + int64_t Sq, int64_t Skv, float scale) { + __shared__ float sq[16][129]; + __shared__ float sk[128][17]; + const int tid = threadIdx.y * 16 + threadIdx.x; + const int64_t q_start = (int64_t)blockIdx.y * 16; + const int64_t k_start = (int64_t)blockIdx.x * 16; + const int64_t bh = blockIdx.z; + for (int i = tid; i < 16 * 128; i += 256) { + const int row = i / 128; + const int d = i % 128; + sq[row][d] = q_start + row < Sq + ? (float)Q[(bh * Sq + q_start + row) * 128 + d] : 0.0f; + sk[d][row] = k_start + row < Skv + ? (float)K[(bh * Skv + k_start + row) * 128 + d] : 0.0f; + } + __syncthreads(); + const int64_t qi = q_start + threadIdx.y; + const int64_t ki = k_start + threadIdx.x; + if (qi >= Sq || ki >= Skv) return; + float acc = 0.0f; + #pragma unroll 8 + for (int d = 0; d < 128; ++d) { + acc += sq[threadIdx.y][d] * sk[d][threadIdx.x]; + } + scores[(bh * Sq + qi) * Skv + ki] = ApplyScale ? scale * acc : acc; +} + +// H3 PV/dQ and fused dK/dV reuse staged operands in shared memory. +// Each output retains the native scalar ascending reduction order. +// All threads take every barrier, including tail rows/channels. +template +__global__ void h3_staged_row_product( + const float* __restrict__ A, const at::BFloat16* __restrict__ X, + at::BFloat16* __restrict__ Y, int64_t Rows, int64_t Inner, float scale) { + __shared__ float a[8][33]; + __shared__ float x[32][33]; + const int tid = threadIdx.y * 32 + threadIdx.x; + const int64_t bh = blockIdx.z; + const int64_t row = (int64_t)blockIdx.y * 8 + threadIdx.y; + const int d = blockIdx.x * 32 + threadIdx.x; + float acc = 0.0f; + for (int64_t start = 0; start < Inner; start += 32) { + const int ar = tid / 32, ak = tid % 32; + const int64_t global_row = (int64_t)blockIdx.y * 8 + ar; + a[ar][ak] = global_row < Rows && start + ak < Inner + ? A[(bh * Rows + global_row) * Inner + start + ak] : 0.0f; + for (int i = tid; i < 32 * 32; i += 256) { + const int k = i / 32, c = i % 32; + x[k][c] = start + k < Inner + ? (float)X[(bh * Inner + start + k) * 128 + blockIdx.x * 32 + c] : 0.0f; + } + __syncthreads(); + // No padded contributions: even signed zero keeps the original operation count. + const int n = (int)min((int64_t)32, Inner - start); + #pragma unroll 8 + for (int j = 0; j < n; ++j) { + acc += a[threadIdx.y][j] * x[j][threadIdx.x]; + } + __syncthreads(); + } + if (row < Rows) Y[(bh * Rows + row) * 128 + d] = + (at::BFloat16)(Scale ? scale * acc : acc); +} + +__global__ void h3_staged_dkv( + const float* __restrict__ DS, const float* __restrict__ P, + const at::BFloat16* __restrict__ Q, const at::BFloat16* __restrict__ DO, + float* __restrict__ DK, float* __restrict__ DV, int64_t Sq, int64_t Skv) { + __shared__ float ds[32][9], p[32][9]; + __shared__ float q[32][33], dout[32][33]; + const int tid = threadIdx.y * 32 + threadIdx.x; + const int64_t bh = blockIdx.z; + const int64_t ki = (int64_t)blockIdx.y * 8 + threadIdx.y; + const int d = blockIdx.x * 32 + threadIdx.x; + const int64_t index = (bh * Skv + ki) * 128 + d; + float dk = ki < Skv ? DK[index] : 0.0f; + float dv = ki < Skv ? DV[index] : 0.0f; + for (int64_t start = 0; start < Sq; start += 32) { + const int r = tid / 8, c = tid % 8; + const int64_t k = (int64_t)blockIdx.y * 8 + c; + const bool valid = start + r < Sq && k < Skv; + ds[r][c] = valid ? DS[(bh * Sq + start + r) * Skv + k] : 0.0f; + p[r][c] = valid ? P[(bh * Sq + start + r) * Skv + k] : 0.0f; + for (int i = tid; i < 32 * 32; i += 256) { + const int row = i / 32, dim = i % 32; + const bool valid_q = start + row < Sq; + const int64_t offset = (bh * Sq + start + row) * 128 + blockIdx.x * 32 + dim; + q[row][dim] = valid_q ? (float)Q[offset] : 0.0f; + dout[row][dim] = valid_q ? (float)DO[offset] : 0.0f; + } + __syncthreads(); + const int n = (int)min((int64_t)32, Sq - start); + #pragma unroll 8 + for (int j = 0; j < n; ++j) { + dk += ds[j][threadIdx.y] * q[j][threadIdx.x]; + dv += p[j][threadIdx.y] * dout[j][threadIdx.x]; + } + __syncthreads(); + } + if (ki < Skv) { DK[index] = dk; DV[index] = dv; } +} + // --------------------------------------------------------------------------- // Masked Softmax + LSE Kernel // One CTA per (b, hq, q) row. Fixed 256 threads. @@ -263,7 +378,9 @@ std::vector deterministic_attention_forward_impl( bool causal, double scale, torch::optional key_padding_mask, - bool output_fp32) { + bool output_fp32, + bool compute_output = true, + bool h3_shared_qk = false) { check_deterministic_attention_inputs(q, k, v, key_padding_mask); const at::cuda::OptionalCUDAGuard device_guard(at::device_of(q)); @@ -295,17 +412,24 @@ std::vector deterministic_attention_forward_impl( (Skv + kQKTileK - 1) / kQKTileK, (Sq + kQKTileQ - 1) / kQKTileQ, B * Hq); - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, - q_contig.scalar_type(), "qk_kernel", [&] { - qk_kernel<<>>( - q_contig.data_ptr(), - k_contig.data_ptr(), - scores.data_ptr(), - B, Hq, Hkv, Sq, Skv, D, - (float)scale); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); + if (h3_shared_qk) { + h3_shared_qk_kernel<<>>( + q_contig.data_ptr(), k_contig.data_ptr(), + scores.data_ptr(), Sq, Skv, (float)scale); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_contig.scalar_type(), "qk_kernel", [&] { + qk_kernel<<>>( + q_contig.data_ptr(), + k_contig.data_ptr(), + scores.data_ptr(), + B, Hq, Hkv, Sq, Skv, D, + (float)scale); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } } // --- Launch Masked Softmax + LSE kernel --- @@ -326,6 +450,10 @@ std::vector deterministic_attention_forward_impl( C10_CUDA_KERNEL_LAUNCH_CHECK(); } + if (!compute_output) { + return {torch::Tensor(), lse, scores}; + } + // --- Launch PV kernel --- auto out = output_fp32 ? torch::empty(q_contig.sizes(), q_contig.options().dtype(at::kFloat)) @@ -336,24 +464,33 @@ std::vector deterministic_attention_forward_impl( (D + kPVTileD - 1) / kPVTileD, (Sq + kPVTileQ - 1) / kPVTileQ, B * Hq); - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, - q_contig.scalar_type(), "pv_kernel", [&] { - if (output_fp32) { - pv_kernel<<>>( - scores.data_ptr(), - v_contig.data_ptr(), - out.data_ptr(), - B, Hq, Hkv, Sq, Skv, D); - } else { - pv_kernel<<>>( - scores.data_ptr(), - v_contig.data_ptr(), - out.data_ptr(), - B, Hq, Hkv, Sq, Skv, D); - } - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); + if (h3_shared_qk && !output_fp32) { + dim3 h3block(32, 8); + dim3 h3grid(4, (Sq + 7) / 8, B * Hq); + h3_staged_row_product<<>>( + scores.data_ptr(), v_contig.data_ptr(), + out.data_ptr(), Sq, Skv, (float)scale); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_contig.scalar_type(), "pv_kernel", [&] { + if (output_fp32) { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } else { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } } return {out, lse, scores}; @@ -373,6 +510,63 @@ std::vector deterministic_attention_forward_fp32( q, k, v, causal, scale, key_padding_mask, true); } +namespace { +void check_h3_build() { +#if !defined(RL_KERNEL_H3_STRICT) || !RL_KERNEL_H3_STRICT || defined(USE_ROCM) + TORCH_CHECK(false, + "H3 requires an audited strict CUDA build; rebuild with RL_KERNEL_H3_STRICT=1"); +#endif +} + +void check_h3_inputs(const torch::Tensor& q, const torch::Tensor& k, + const torch::Tensor& v) { + check_h3_build(); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, + "H3 requires 4-D head-major Q/K/V"); + TORCH_CHECK(q.size(1) == 56 && k.size(1) == 56 && v.size(1) == 56, + "H3 requires 56 MHA heads"); + TORCH_CHECK(q.scalar_type() == at::kBFloat16, "H3 requires BF16"); + TORCH_CHECK(q.size(0) > 0 && q.size(0) <= 65535 / 56, + "H3 batch exceeds CUDA grid bounds"); + // H3's staged PV/dQ/dK/dV kernels use eight rows per grid block. Keep + // their y-dimension within CUDA's 16-bit grid limit as well as the QK + // kernel's sixteen-row limit. + TORCH_CHECK(q.size(2) <= 65535 * 8 && k.size(2) <= 65535 * 8, + "H3 sequence exceeds CUDA grid bounds"); + TORCH_CHECK(q.size(0) * 56 * q.size(2) <= 2147483647LL, + "H3 query chunk exceeds softmax row index bounds"); + check_deterministic_attention_inputs(q, k, v, torch::nullopt); + TORCH_CHECK(at::cuda::getDeviceProperties(q.get_device())->major >= 8, + "H3 BF16 CUDA profile requires SM80 or newer"); +} +} // namespace + +std::map h3_attention_build_info() { +#if defined(RL_KERNEL_H3_STRICT) && RL_KERNEL_H3_STRICT && !defined(USE_ROCM) + return {{"profile", "h3-cuda-strict-v1"}, {"fmad", "false"}, {"ftz", "false"}, + {"prec_div", "true"}, {"prec_sqrt", "true"}, {"fast_math", "false"}, + {"cuda_compiler", std::to_string(__CUDACC_VER_MAJOR__) + "." + + std::to_string(__CUDACC_VER_MINOR__) + "." + + std::to_string(__CUDACC_VER_BUILD__)}}; +#else + return {{"profile", "unqualified"}}; +#endif +} + +std::vector h3_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v) { + check_h3_inputs(q, k, v); + return deterministic_attention_forward_impl(q, k, v, false, 1.0 / std::sqrt(128.0), + torch::nullopt, false, true, true); +} + +// Recompute only probabilities during H3 backward; do not execute an unused PV. +torch::Tensor h3_attention_probabilities(torch::Tensor q, torch::Tensor k, torch::Tensor v) { + check_h3_inputs(q, k, v); + return deterministic_attention_forward_impl(q, k, v, false, 1.0 / std::sqrt(128.0), + torch::nullopt, false, false, true)[2]; +} + // =========================================================================== // BACKWARD // =========================================================================== @@ -686,6 +880,93 @@ std::vector deterministic_attention_backward( return {dQ, dK, dV}; } +// H3 backward continues FP32 accumulators across query chunks. +std::vector h3_attention_backward_accumulate( + torch::Tensor grad_output, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor P, // recomputed [B, 56, Sq, Skv] FP32 + torch::Tensor dK, torch::Tensor dV) { + check_h3_inputs(q, k, v); + TORCH_CHECK(grad_output.sizes() == q.sizes() && grad_output.device() == q.device() + && grad_output.scalar_type() == q.scalar_type(), "invalid H3 grad_output"); + TORCH_CHECK(P.dim() == 4 && P.size(0) == q.size(0) && P.size(1) == 56 + && P.size(2) == q.size(2) && P.size(3) == k.size(2) + && P.device() == q.device() && P.scalar_type() == at::kFloat, "invalid H3 P"); + for (const auto& acc : {dK, dV}) { + TORCH_CHECK(acc.device() == q.device() && acc.scalar_type() == at::kFloat + && acc.sizes() == k.sizes() && acc.is_contiguous(), "invalid H3 accumulator"); + } + + at::assert_no_overlap(dK, dV); + for (const auto& input : {q, k, v, P, grad_output}) { + at::assert_no_overlap(dK, input); + at::assert_no_overlap(dV, input); + } + + const at::cuda::OptionalCUDAGuard device_guard(at::device_of(q)); + + auto dO = grad_output.contiguous(); + auto q_c = q.contiguous(); + auto k_c = k.contiguous(); + auto v_c = v.contiguous(); + auto P_c = P.contiguous(); + + const int64_t B = q_c.size(0); + const int64_t Hq = q_c.size(1); + const int64_t Sq = q_c.size(2); + const int64_t Skv = k_c.size(2); + + auto stream = at::cuda::getCurrentCUDAStream(); + + // dP = dO @ V^T [B, Hq, Sq, Skv] + auto dP = torch::empty({B, Hq, Sq, Skv}, q_c.options().dtype(at::kFloat)); + { + dim3 block(kQKTileK, kQKTileQ); + dim3 grid( + (Skv + kQKTileK - 1) / kQKTileK, + (Sq + kQKTileQ - 1) / kQKTileQ, + B * Hq); + h3_shared_qk_kernel<<>>( + dO.data_ptr(), v_c.data_ptr(), + dP.data_ptr(), Sq, Skv, 1.0f); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + // Softmax backward: dS = P * (dP - delta), writes in-place over dP + { + const int64_t num_rows = B * Hq * Sq; + softmax_backward_kernel<<>>( + dP.data_ptr(), + P_c.data_ptr(), + Skv); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // dP buffer now contains dS + + auto dQ = torch::empty_like(q_c); + { + dim3 block(32, 8); + dim3 grid(4, (Sq + 7) / 8, B * Hq); + h3_staged_row_product<<>>( + dP.data_ptr(), k_c.data_ptr(), + dQ.data_ptr(), Sq, Skv, (float)(1.0 / std::sqrt(128.0))); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + { + dim3 block(32, 8); + dim3 grid(4, (Skv + 7) / 8, B * Hq); + h3_staged_dkv<<>>( + dP.data_ptr(), P_c.data_ptr(), + q_c.data_ptr(), dO.data_ptr(), + dK.data_ptr(), dV.data_ptr(), Sq, Skv); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + return {dQ, dK, dV}; +} + #if defined(USE_ROCM) namespace { diff --git a/docs/operators/README.md b/docs/operators/README.md index 065d754b6..c99451553 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -41,4 +41,5 @@ Every operator page should include: - [MiniMax-H3 Final AdaLN Output](h3-final-adaln-out.md) - [MiniMax-H3 Tensor-Parallel AdaLN Projection](h3-tp-adaln-3mod.md) - [MiniMax-H3 Sequence-Parallel Norm and AdaLN](h3-sp-norm-adaln.md) +- [MiniMax-H3 One Block, Forward and Backward](h3-one-block.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 52f8517ea..886e91935 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -271,3 +271,60 @@ for measured peak memory at representative shapes. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). - No FP8, no multi-GPU / sequence-parallel. + +## MiniMax-H3 full attention + +This implements the CUDA `h3_full_attention` work item in [RFC #420](https://github.com/RL-Align/RL-Kernel/issues/420), with an independent PyTorch reference. It covers the attention core after Q/K normalization and RoPE; it does not close the complete model roadmap. + +### Interface and scope + +```python +from rl_engine.backends.cuda.attention.h3_full_attention import H3FullAttentionOp + +out = H3FullAttentionOp(chunk_rows=128)(q, k, v) +``` + +Q/K/V must be matching contiguous BF16 `[B,S,56,128]` tensors. They must already have the checkpoint's QK RMSNorm and partial three-axis RoPE applied. Output has the same shape and dtype. `forward_with_lse` additionally returns diagnostic FP32 `[B,S,56]` LSE without gradients. + +Only non-causal MHA, single rank, TP1/CP1 and first derivatives are supported. No masks, GQA, KV cache, CP or projections are provided. CUDA requires NVIDIA SM80 or newer and an audited strict extension; unsupported inputs/builds fail closed. H20 is the measured device. ROCm is outside this PR. + +`kernel_registry.get_op("h3_full_attention", device="cuda")` selects CUDA. The semantic registry explicitly declares dtype/layout/topology and rejects unsupported requests. + +### Arithmetic and workspace + +The CUDA path retains #240's ascending-D scalar QK reduction, 256-thread softmax tree and ascending-key PV/dQ reductions. Shared-memory staging reuses converted Q/K and dO/V operands and rows of probabilities/dS. Fused dK/dV loading shares operands while maintaining two separate ascending-query FP32 accumulation chains. Each gradient element has one writer; there are no cross-CTA atomics or split-K partial-sum merges. + +Queries are processed in fixed-size chunks, while each query consumes all keys. Forward saves Q/K/V; backward recomputes probabilities using the same QK/softmax arithmetic, skipping unused PV. dK/dV retain FP32 accumulators across chunks, dK is scaled once, and both gradients are cast once at the final boundary. Fixed-chunk probability workspace is `[B,56,min(chunk_rows,S),S]` FP32; backward additionally needs dP/dS and two `[B,56,S,128]` FP32 accumulators. Workspace grows linearly with S at a fixed chunk size; compute remains quadratic. + +Strict compilation uses `--fmad=false --ftz=false --prec-div=true --prec-sqrt=true`. Fast math and external NVCC prepend/append flags are rejected. An isolated `_H3_C` extension enforces this profile without changing ordinary `_C` source/compiler/linker settings. Native H3 entry points and the wrapper verify build identity. Raw-bit promises apply within a qualified device/build, not across platforms or compilers. Independent mathematical accuracy is a separate tolerance check. + +### Build and validation + +Use an isolated checkout and environment. The ordinary build intentionally omits the strict extension. + +```bash +RL_KERNEL_H3_STRICT=1 RL_KERNEL_REQUIRE_EXT=1 TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=2 \ + python setup.py build_ext --inplace +python -m pytest -q tests/backends/cuda/test_h3_full_attention.py +python tools/validation/operators/check_operator.py --op h3_full_attention \ + --candidate cuda --device cuda --dtype bf16 --batch 2 --seq 129 --check-grad --json +``` + +Tests cover independent FP32/FP64 output/all-gradient accuracy, strict-native raw bits, +repeat/batch/chunk invariance, LSE, invalid inputs, checkpointing and streams. CUDA tests require the optional strict `_H3_C`; +CPU-only success or skipped GPU tests do not establish GPU qualification. + +### Performance reproduction + +```bash +python benchmarks/operators/attention/benchmark_deterministic_attention.py h3 \ + --output /path/to/local-evidence/matrix +``` + +The matrix compares strict native and CUDA chunks 64/128/256 on six shapes. +Output/all-gradient raw-bit checks precede timing; two warmups and six rotating samples +include public-wrapper copies, allocations and forward plus all QKV gradients. JSON records +raw samples, extra peak allocation, environment, backend readback and source/binary hashes. +At fixed chunk size, probability workspace grows linearly with sequence length while compute +remains quadratic. These are operator measurements, without +full-model, multi-GPU or ROCm hardware qualification. diff --git a/docs/operators/h3-one-block.md b/docs/operators/h3-one-block.md new file mode 100644 index 000000000..b2f1d4ce0 --- /dev/null +++ b/docs/operators/h3-one-block.md @@ -0,0 +1,183 @@ +# MiniMax-H3 One Block, Forward and Backward + +## Summary + +`ws1_one_h3_block` (RFC #420, WS1 closeout) runs **block 0 of the pinned MiniMax-H3 checkpoint** +node by node, forward and backward, and reports where it first departs from diffusers. The block +is a fixed graph of 17 nodes, from the AdaLN projection to the second gated residual. It runs on +packed FL2VA layouts that are **bitwise equal to the pinned diffusers pipeline's** `position_ids` +and modality tags. + +On one B200, at S = 264, 925 and 3160: + +- **Determinism.** Every node, every gradient reaching a node, and all 14 leaf gradients are + bitwise repeatable. +- **Batch invariance.** Batch row 0 of a two-row batch has the single-row bytes at every node, + forward and backward. +- **Agreement with diffusers.** Every node that promises it is bitwise equal to diffusers on + diffusers' inputs. The first drift from diffusers is `adaln_projection`, the first reduction + whose tree differs from cuBLAS. +- **Accuracy.** The block output is exactly as close to the FP32 golden as diffusers' (relative + error 5.2e-3 / 7.9e-3 / 9.0e-3). + +Only three of the 17 nodes are bound to this stack's row operators. The others use **interim** +deterministic operators until their own rows land, or the provider replay where no kernel exists +yet (see Bindings). The block-level promises above hold with whatever a node is bound to. Speed +and the exact rounding of interim nodes belong to those operators, not to this row. + +## Entry Point + +```python +from rl_engine.validation.models.h3_block import CandidateOps, run_block_backward, run_block_case +from rl_engine.validation.models.h3_weights import load_h3_block_weights + +weights = load_h3_block_weights("cuda") # conditioning + block-0 tensors, sha256-checked +ops = CandidateOps() +forward = run_block_case(weights, layout="small", ops=ops) # per node, first_*_drift +backward = run_block_backward(weights, layout="small", ops=ops) +``` + +The CLI writes the JSON evidence and the plot script renders it: + +```bash +python tools/weights/prepare_h3_weights.py --out --block # shard 1 of 14, ~4.8 GB +export RL_KERNEL_H3_WEIGHTS= +python tools/validation/models/h3_block_replay.py --layouts tiny,small,medium --out report.json +python tools/validation/models/plot_h3_block.py report.json +``` + +## Graph and Bindings + +`rl_engine/validation/models/h3_block.py` `NODES` (graph order): + +| Node | Owning RFC row | Binding today | +|---|---|---| +| `adaln_projection` | `adaln_projection_3mod` | **row**: `H3AdaLNProjectionCudaOp` (registry) | +| `norm1`, `norm2` | `h3_rmsnorm` | **row**: `H3RMSNormCudaOp.forward_modulated` (registry) | +| `q_proj`, `k_proj`, `v_proj` | `h3_qkv_gemm` | interim: Triton FP32-accumulate GEMM, one BF16 rounding | +| `q_norm`, `k_norm` | `h3_qk_rmsnorm_d128` | interim: `H3RMSNormCudaOp.forward` on the 128-wide heads | +| `rope_q`, `rope_k` | `h3_mm_rope_3axis_partial` | reference: the provider replay (no kernel exists yet) | +| `attention` | `h3_full_attention` | interim: `DeterministicAttentionOp(causal=False)` | +| `o_proj` | `h3_attention_o_gemm` | interim: Triton FP32-accumulate GEMM | +| `residual_attn`, `residual_mlp` | `adaln_gate_residual` | **row**: `H3GateResidualCudaOp` (registry) | +| `ffn_gate_up` | `h3_ffn_gate_up_gemm` | interim: Triton FP32-accumulate GEMM | +| `swiglu` | `h3_swiglu` | interim: `SwiGLUCudaOp` (FP32 math, one rounding) | +| `ffn_down` | `h3_ffn_down_gemm` | interim: Triton FP32-accumulate GEMM | + +When a row lands, replacing its node's `candidate` callable is the whole integration. The node's +report entries then become that operator's block-level evidence. + +The interim GEMMs use `_triton_gemm_fp32`, one fixed-order FP32 K loop with no split-K, and round +to BF16 once. They do not use `TritonDetGemmOp.__call__`, whose K tree has BF16 nodes. Measured on +B200: + +| Shape | Tree path err | FP32 K loop → BF16 | cuBLAS | Tree | FP32 K loop | +|---|---|---|---|---|---| +| 264 × 5376 → 7168 | 7.0e-3 | 2.4e-3 | 2.4e-3 | 10.4 ms | 1.2 ms | +| 925 × 14336 → 5376 | 7.5e-3 | 2.6e-3 | 2.6e-3 | 47 ms | 2.7 ms | + +Both paths are bitwise repeatable and row-invariant. Error is max|err| / max|ref| against FP64. + +## What Each Node Reports + +Every node runs three ways: the provider replay (`h3_provider`, diffusers@f53d552 op for op), the +candidate (chained, and *isolated* on the provider's inputs), and the FP32 golden +(`rl_engine/reference/minimax_h3/block.py`: FP32 math, TF32 off, no BF16 storage between nodes). + +- **Forward:** `repeat_bitwise_equal` and `batch_row_bitwise_equal` (row 0 of a B = 2 batch whose + row 1 has a different input). Also chained and isolated agreement with the provider (bitwise, + max abs, first differing packed row), the candidate's and provider's `max|err| / max|golden|`, + and the sha256 of the node's output. +- **Backward:** for the block input, `temb` and the 12 block parameters: repeat bitwise, and the + candidate's and provider's error against the golden. For the gradient reaching each node's + output, in reverse graph order: repeat bitwise and batch-row bitwise. Parameter gradients sum + over the batch, so only activation gradients take part in the batch check. +- **First drift:** `first_drift`, `first_isolated_drift`, `first_repeat_drift`, + `first_batch_drift`, `first_grad_repeat_drift` and `first_grad_batch_drift` each name the first + node where that comparison stops being bitwise. +- **Permutation probe (control H-C3, reported, not promised):** permute the packed rows together + with positions and AdaLN indices, and compare the output to the permuted original. + +Inputs: `temb` comes from the provider's conditioning path (the conditioning rows are checked by +`h3_chain`). The block input is a seeded BF16 stand-in, because the packed input projections are +other rows. Text rows use AdaLN table row 0 and media rows use row 1. + +## Evidence + +[`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr539/reports/experiments/h3-one-block-b200/report.json) was generated from a clean tree +at `41e4293` on an otherwise idle B200 (torch 2.13.0+cu130): + +![one-block evidence](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr539/reports/experiments/h3-one-block-b200/figure.png) + +| Layout | S | First drift / first isolated drift | Repeat / batch drift (fwd, grad) | Output err: RL-Kernel / diffusers | Fwd ms: RL-Kernel / diffusers | Fwd+bwd ms | +|---|---|---|---|---|---|---| +| tiny | 264 | `adaln_projection` / `adaln_projection` | none | 5.21e-3 / 5.21e-3 | 7.4 / 0.61 | 22.7 / 2.8 | +| small | 925 | `adaln_projection` / `adaln_projection` | none | 7.87e-3 / 7.87e-3 | 30.1 / 1.12 | 81.7 / 5.0 | +| medium | 3160 | `adaln_projection` / `adaln_projection` | none | 8.97e-3 / 8.97e-3 | 225 / 3.8 | 577 / 11.5 | + +- **Forward accuracy.** Per node, the candidate's error is within 0.93–1.20× diffusers'. The + largest is `attention` at S = 3160, which is the interim attention operator's. +- **Gradient accuracy.** The candidate is better than diffusers on `adaln_proj.weight`, + `adaln_proj.bias` and `temb`, by up to 7×: diffusers' `index_select` backward accumulates the + table gradient in BF16 atomics, while the row operators use FP32 segment sums. The candidate is + within 1.16× on the other leaves, except `attn.norm_q.weight` at S = 3160 (9.5e-3 vs 5.5e-3). + At S = 925 the same leaf is 2× *more* accurate than diffusers (1.3e-2 vs 2.4e-2), so this + 128-element gradient's error moves with BF16 rounding rather than favouring one side. +- **Permutation probe (H-C3).** The output changes in 79 / 450 / 2199 of the rows, by at most 8. + Attention sums over keys in packed order, so a consistent permutation is not byte-preserving with + the interim attention operator. This is for `h3_full_attention` to settle. +- **Speed.** The candidate is 12–60× slower than diffusers forward. Almost all of it is the + interim GEMMs and attention (FP32 scores materialised). The three row operators are each + faster than or equal to diffusers in their own evidence. + +## Existing implementations (RFC #420 reuse rule) + +- **Is there an existing block-level harness?** diffusers' `MiniMaxH3TransformerBlock` is the + executable reference. It is replayed op for op as the provider, so diffusers is not imported and + the replay pins the reference's dtypes and ordering. The Qwen3-8B chain gate from #315 + (`rl_engine/validation/models/chain_gate.py`) has per-node token digests and first-drift + localisation, but its graph, shapes and acceptance are Qwen3-specific (causal attention, no + AdaLN). This row reuses its method (per-node digests, graph-ordered first drift) on the H3 + `h3_chain` stage pattern. +- **Interim operators.** These are reused from RL-Kernel and are not new code. The only choice + made here is the FP32-accumulate entry of the Triton GEMM instead of its BF16 tree, for the + accuracy and speed measured above. + +## Tests + +`tests/models/minimax_h3/test_h3_one_block.py`: + +- **CPU:** + - the three layouts hash to the pinned diffusers `build_packed_sequence` output; + - the graph is well formed; + - the partial MM-RoPE leaves channels 96–127 unchanged. +- **GPU, `tiny` and `small` forward:** + - every node repeats and is batch-row bitwise; + - promised nodes are isolated-bitwise to diffusers; + - node and output errors are at most 1.25 × diffusers' + 1e-3. +- **GPU, `tiny` backward:** + - every leaf repeats; + - dX is batch-row bitwise; + - no gradient drift; + - leaf errors are at most 1.5 × diffusers' + 5e-3, or, for a BF16 leaf, at most 6 ulps + (of its largest golden element) above diffusers'. A stored gradient's worst element moves + in whole ulps, so the relative rule alone is seed- and torch-version-sensitive; + `tools/validation/models/h3_block_ulp_sweep.py` measured at most +4.15 ulps over + tiny/small × 10 seeds on torch 2.13 and 2.9.1, and +5.15 ulps over 34 single-node + candidate/provider swaps (reports: `reports/experiments/h3-one-block-b200/ulp-sweep-torch*.json`). + +```bash +RL_KERNEL_H3_WEIGHTS= CUDA_VISIBLE_DEVICES=0,1 python -m pytest tests/models/minimax_h3 -q +``` + +## Known Limitations + +- **Single platform.** CUDA only. `DeterministicAttentionOp` and the Triton GEMM also build for + ROCm, but the block has not been run there. +- **Interim nodes.** Eleven of the 17 nodes are interim or reference bindings. Their speed and + rounding are not the owning rows' final ones, and the MM-RoPE nodes are the provider replay + itself, so they carry no RL-Kernel evidence until `h3_mm_rope_3axis_partial` lands. +- **Stand-in inputs.** The block input is a seeded stand-in and `temb` is the provider's. + Real packed projections arrive with the input-projection rows. +- **Not covered here.** Block 0 only (`ws1_full_h3_transformer` covers 50 blocks and the refiner), + and one sequence per batch row (H3's batch is a replication axis). diff --git a/examples/models/minimax_h3/diagnose-block.sbatch b/examples/models/minimax_h3/diagnose-block.sbatch new file mode 100644 index 000000000..128ea1955 --- /dev/null +++ b/examples/models/minimax_h3/diagnose-block.sbatch @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +failed_result=$3 +result_root=$4 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +result_dir="$result_root/job-$SLURM_JOB_ID-attempt-${SLURM_RESTART_COUNT:-0}" +mkdir "$result_dir" +export PYTHONPATH="$source_dir" RL_KERNEL_H3_WEIGHTS="$failed_result/preflight/weights" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +scontrol show job "$SLURM_JOB_ID" > "$result_dir/slurm.txt" +nvidia-smi -q > "$result_dir/device.txt" +git rev-parse HEAD > "$result_dir/commit.txt" +"$python_bin" -m pip freeze > "$result_dir/dependencies.txt" +# Preserve the exact native binary from the failed attempt; record and recheck +# its identity instead of changing both the kernel and diagnostic at once. +"$python_bin" - "$failed_result" "$result_dir" <<'PY' +import json +import shutil +import sys +from pathlib import Path +from tools.validation.models.h3_full_model import digest +failed, output = map(Path, sys.argv[1:]) +identity = json.loads((failed / "preflight/identity.json").read_text()) +source = Path(identity["extension"]) +assert digest(source) == identity["extension_sha256"] +shutil.copy2(source, Path("rl_engine") / source.name) +(output / "native-origin.json").write_text(json.dumps(identity, indent=2) + "\n") +(output / "retry-of.txt").write_text(str(failed) + "\n") +PY +"$python_bin" -m tools.validation.models.h3_block_diagnostic --failed-job "$failed_result" --causal-sweep \ + --output "$result_dir/diagnostic" > "$result_dir/diagnostic.log" 2>&1 diff --git a/rl_engine/backends/cuda/attention/h3_full_attention.py b/rl_engine/backends/cuda/attention/h3_full_attention.py new file mode 100644 index 000000000..0a30fa04b --- /dev/null +++ b/rl_engine/backends/cuda/attention/h3_full_attention.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""H3 non-causal MHA with query-chunk recomputation, retaining #240 arithmetic.""" +from __future__ import annotations + +import importlib +import math +from types import ModuleType + +import torch +from torch.autograd import Function +from torch.autograd.function import once_differentiable + +_C: ModuleType | None +try: + _C = importlib.import_module("rl_engine._H3_C") +except ImportError: + _C = None + +_SCALE = 1.0 / math.sqrt(128) + + +def _extension(): + required = ( + "h3_attention_forward", + "h3_attention_build_info", + "h3_attention_backward_accumulate", + "h3_attention_probabilities", + ) + if _C is None or any(not hasattr(_C, name) for name in required): + raise RuntimeError("H3 CUDA extension unavailable; rebuild RL-Kernel (no fallback)") + if _C.h3_attention_build_info().get("profile") != "h3-cuda-strict-v1": + raise RuntimeError("H3 requires a strict CUDA build; rebuild with RL_KERNEL_H3_STRICT=1") + return _C + + +class _H3AttentionFn(Function): + @staticmethod + def forward(ctx, q, k, v, chunk_rows): + ext = _extension() + q, k, v = (x.contiguous() for x in (q, k, v)) + out = torch.empty_like(q) + lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32) + for start in range(0, q.shape[2], chunk_rows): + end = min(start + chunk_rows, q.shape[2]) + block_out, block_lse, p = ext.h3_attention_forward( + q[:, :, start:end].contiguous(), k, v + ) + out[:, :, start:end] = block_out + lse[:, :, start:end] = block_lse + del block_out, block_lse, p + ctx.save_for_backward(q, k, v) + ctx.chunk_rows = chunk_rows + ctx.mark_non_differentiable(lse) + return out, lse + + @staticmethod + @once_differentiable + def backward(ctx, grad_out, grad_lse): + ext = _extension() + q, k, v = ctx.saved_tensors + dq = torch.empty_like(q) + dk = torch.zeros_like(k, dtype=torch.float32) + dv = torch.zeros_like(v, dtype=torch.float32) + for start in range(0, q.shape[2], ctx.chunk_rows): + end = min(start + ctx.chunk_rows, q.shape[2]) + qt = q[:, :, start:end].contiguous() + # Deliberately retain the original softmax (not an LSE-based alternative). + p = ext.h3_attention_probabilities(qt, k, v) + dqt, _, _ = ext.h3_attention_backward_accumulate( + grad_out[:, :, start:end].contiguous(), + qt, + k, + v, + p, + dk, + dv, + ) + dq[:, :, start:end] = dqt + del qt, p, dqt + # No per-chunk rounding or partial-sum merge: scale/cast exactly once. + dk.mul_(_SCALE) + return dq, dk.to(k.dtype), dv.to(v.dtype), None + + +class H3FullAttentionOp: + """BF16 [B,S,56,128] -> [B,S,56,128], Q/K already normalized/rotated. + + Physical sequence order defines the fixed reduction order. + """ + + _function: type[Function] = _H3AttentionFn + + backend_id = "rlkernel.cuda.h3_full_attention.query_chunk.v1" + + def __init__(self, chunk_rows=128): + if isinstance(chunk_rows, bool) or not isinstance(chunk_rows, int) or chunk_rows <= 0: + raise ValueError("chunk_rows must be a positive integer") + self.chunk_rows = chunk_rows + + @staticmethod + def _validate(q, k, v): + if any(x.ndim != 4 for x in (q, k, v)): + raise ValueError("H3 Q/K/V must have shape [B,S,56,128]") + if q.shape != k.shape or q.shape != v.shape or q.shape[2:] != (56, 128): + raise ValueError("H3 requires equal Q/K/V shapes [B,S,56,128]") + if not 0 < q.shape[0] <= 65535 // 56 or q.shape[1] <= 0: + raise ValueError("H3 requires nonempty tensors and batch <= 1170") + # Native H3 staging launches eight query/key rows per grid block. + if q.shape[1] > 65535 * 8: + raise ValueError("H3 sequence exceeds CUDA grid bounds") + if any(x.dtype != torch.bfloat16 for x in (q, k, v)): + raise ValueError("H3 strict checkpoint profile requires BF16 Q/K/V") + if any(not x.is_contiguous() for x in (q, k, v)): + raise ValueError("H3 requires contiguous BSHD Q/K/V; call contiguous() explicitly") + if any(x.device != q.device for x in (k, v)): + raise ValueError("H3 Q/K/V must share a device") + if not q.is_cuda or torch.version.hip is not None: + raise ValueError("H3 CUDA implementation requires NVIDIA CUDA") + + def forward_with_lse(self, q, k, v): + self._validate(q, k, v) + out, lse = self._function.apply( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), self.chunk_rows + ) + return out.transpose(1, 2), lse.transpose(1, 2) + + def forward(self, q, k, v): + return self.forward_with_lse(q, k, v)[0] + + __call__ = forward + + def provenance(self): + return { + "backend": self.backend_id, + "build": dict(_extension().h3_attention_build_info()), + "fallback": False, + "layout": "BSHD", + "heads": 56, + "head_dim": 128, + "dtype": "bf16", + "causal": False, + "qk_preprocessing": "caller-owned RMSNorm + partial 3-axis RoPE", + "chunk_rows": self.chunk_rows, + "accumulator": "fp32", + "reduction": "#240: ascending D/K/Q, softmax 256-thread tree", + "backward": "recompute P; continuous FP32 dK/dV accumulators; final cast", + "supports_cp": False, + "supports_double_backward": False, + } diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py index 763dba546..17a3a3884 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py @@ -16,8 +16,6 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( _segment_tiles, adaln_row_gather_available, @@ -29,6 +27,8 @@ det_linear_forward, silu_backward_fp32, ) +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM from rl_engine.reference.minimax_h3.adaln_projection import validate_h3_adaln_projection from rl_engine.reference.minimax_h3.adaln_row_gather import ( diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py index 6c9f5bf49..f02431903 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py @@ -8,7 +8,6 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( ACT_NONE, CONTRACT, @@ -18,6 +17,7 @@ det_linear_forward, silu_backward_fp32, ) +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.adaln_projection import ( split_adaln_table, validate_h3_adaln_projection, diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py index 92835f0cc..ee681454d 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py @@ -7,8 +7,8 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM from rl_engine.reference.minimax_h3.adaln_row_gather import ( h3_adaln_indices, diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py index 2ff1fd16a..741a0a85e 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py @@ -8,8 +8,6 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( det_linear_available, @@ -19,6 +17,8 @@ silu_backward_fp32, ) from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import h3_rmsnorm_available +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.final_adaln_out import validate_h3_final_adaln_out from rl_engine.reference.minimax_h3.rmsnorm import H3_NORM_EPS diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index 2c6c6be55..3d040e9e8 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -7,9 +7,9 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C, _EXT_AVAILABLE from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.gate_residual import validate_h3_gate_residual KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py index a6ed3a28b..d9266e759 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py @@ -7,9 +7,9 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C, _EXT_AVAILABLE from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.rmsnorm import ( H3_NORM_EPS, validate_h3_modulation, diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py b/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py index adf01094e..37cea620b 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py @@ -26,13 +26,13 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C, _EXT_AVAILABLE from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( BACKWARD_TILE, _segment_tiles, ) from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.gate_residual import validate_h3_gate_residual from rl_engine.reference.minimax_h3.rmsnorm import ( H3_NORM_EPS, @@ -223,6 +223,8 @@ def sp_plan( families = [_dweight_tiles(layout.batch * layout.seq_len, device)] if index_full is None: return SPPlan(layout, _Plan(layout, families, device), None, None, None, None) + if num_rows is None: + raise ValueError("indexed SP plans require num_rows") *seg_tiles, seg_first_tile = _segment_tiles(index_full.repeat(layout.batch), num_rows) plan = _Plan(layout, [*families, tuple(seg_tiles)], device) local_index = index_full[layout.lo : layout.hi].clone(memory_format=torch.contiguous_format) @@ -370,6 +372,8 @@ def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | No f"index must be the full ({self.layout.seq_len},) row index, " f"got {tuple(index.shape)}" ) + if num_rows is None: + raise ValueError("indexed SP validation requires num_rows") if index.numel() and (int(index.min()) < 0 or int(index.max()) >= num_rows): raise IndexError(f"index must be in [0, {num_rows})") diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py index aff81eaa8..4d58f6d9f 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py @@ -7,7 +7,6 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( ACT_NONE, ACT_SILU, @@ -18,6 +17,7 @@ det_linear_forward, silu_backward_fp32, ) +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.timestep_mlp import validate_h3_timestep_mlp KERNEL_ID = "rl_engine._C.h3_det_linear_forward[silu]+rl_engine._C.h3_det_linear_forward" diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py index 177ce97f1..c4ceed334 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py @@ -31,7 +31,6 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( CONTRACT, DINPUT_CHUNK, @@ -43,6 +42,7 @@ silu_backward_fp32, ) from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM from rl_engine.reference.minimax_h3.adaln_projection import split_adaln_table diff --git a/rl_engine/reference/attention/standard_attn.py b/rl_engine/reference/attention/standard_attn.py index 804efb6ed..b0a84c225 100644 --- a/rl_engine/reference/attention/standard_attn.py +++ b/rl_engine/reference/attention/standard_attn.py @@ -189,3 +189,38 @@ def _strict_fp32_math(device_type: str): yield finally: torch.backends.cuda.matmul.allow_tf32 = prev_tf32 + + +class H3FullAttentionReference: + def forward_fp32(self, q, k, v): + with NativeAttentionOp._strict_fp32_math(q.device.type): + q, k, v = (x.float().transpose(1, 2) for x in (q, k, v)) + return ((q @ k.transpose(-1, -2) / 128**0.5).softmax(-1) @ v).transpose(1, 2) + + __call__ = forward_fp32 + + @torch.no_grad() + def value_and_grad(self, q, k, v, grad_output, chunk_rows=128): + """FP32 explicit VJP: dS=P*(dO V^T-sum(P*dO V^T)). + + Query chunks bound temporary memory; dK/dV merge GEMM partial sums. + This intentionally differs from the candidate's scalar reduction order. + """ + if isinstance(chunk_rows, bool) or not isinstance(chunk_rows, int) or chunk_rows < 1: + raise ValueError("chunk_rows must be a positive integer") + with NativeAttentionOp._strict_fp32_math(q.device.type): + q, k, v, dy = (x.float().transpose(1, 2) for x in (q, k, v, grad_output)) + out, dq = torch.empty_like(q), torch.empty_like(q) + dk, dv = torch.zeros_like(k), torch.zeros_like(v) + for start in range(0, q.shape[2], chunk_rows): + stop = min(start + chunk_rows, q.shape[2]) + qt, dyt = q[:, :, start:stop], dy[:, :, start:stop] + p = (qt @ k.transpose(-1, -2) / 128**0.5).softmax(-1) + out[:, :, start:stop] = p @ v + dv.add_(p.transpose(-1, -2) @ dyt) + dp = dyt @ v.transpose(-1, -2) + ds = p * (dp - (p * dp).sum(-1, keepdim=True)) + dq[:, :, start:stop] = (ds @ k) / 128**0.5 + dk.add_((ds.transpose(-1, -2) @ qt) / 128**0.5) + del p, dp, ds + return [x.transpose(1, 2) for x in (out, dq, dk, dv)] diff --git a/rl_engine/reference/minimax_h3/block.py b/rl_engine/reference/minimax_h3/block.py new file mode 100644 index 000000000..427d7de2e --- /dev/null +++ b/rl_engine/reference/minimax_h3/block.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""FP32 golden of one MiniMax-H3 transformer block (RFC #420 ``ws1_one_h3_block``). + +Every node takes FP32 inputs (BF16 weights and activations upcast exactly), +computes in FP32 with TF32 off, and returns FP32 without rounding back to +BF16, so the golden is the block's arithmetic without the BF16 storage the +provider and the candidate both have. It is the yardstick their errors are +measured against, not a bitwise target. The one declared cast kept from the +checkpoint contract is the AdaLN input: ``silu(temb)`` is rounded to BF16 +before the projection, as every implementation must do it. + +The partial MM-RoPE uses exact FP32 ``cos``/``sin`` tables on the leading 96 +of 128 head channels (the provider rounds the tables to BF16 first). +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +H3_HEADS = 56 +H3_HEAD_DIM = 128 +H3_ROTARY_DIM = 96 +H3_NORM_EPS = 1e-5 + + +def _f32(t: torch.Tensor) -> torch.Tensor: + return t.to(torch.float32) + + +def golden_adaln_tables(temb, weight, bias, hidden_size: int = 5376) -> tuple[torch.Tensor, ...]: + """Six ``(3T, H)`` FP32 tables; ``silu(temb)`` rounded to the weight dtype first.""" + + act = F.silu(_f32(temb)).to(weight.dtype) + table = F.linear(_f32(act), _f32(weight), _f32(bias)) + return table.view(-1, 6 * hidden_size).chunk(6, dim=-1) + + +def golden_rmsnorm(x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + x = _f32(x) + rstd = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) + return x * rstd * _f32(weight) + + +def golden_norm_modulate(x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + n = golden_rmsnorm(x, weight, eps) + return n * (1.0 + _f32(scale).index_select(0, index)) + _f32(shift).index_select(0, index) + + +def golden_linear(x, weight) -> torch.Tensor: + return F.linear(_f32(x), _f32(weight)) + + +def golden_rope_tables(position_ids, rope_freq_dim: int = 16, rope_theta: float = 10000.0): + """FP32 ``(S, 96)`` cos/sin, the same angles as the provider's rotary module.""" + + exponent = torch.arange( + 0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device + ) + inv_freq = 1.0 / (rope_theta ** (exponent / (2 * rope_freq_dim))) + freqs = position_ids.to(torch.float32).unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs = torch.cat(freqs.unbind(dim=1), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def golden_apply_rotary(x, cos, sin) -> torch.Tensor: + """Rotate-half on the leading ``cos.shape[-1]`` channels of ``(B, S, H, D)``; FP32 tables.""" + + x = _f32(x) + rot = cos.shape[-1] + head, tail = x[..., :rot], x[..., rot:] + x1, x2 = head.chunk(2, dim=-1) + rotated = torch.cat((-x2, x1), dim=-1) + head = head * cos[None, :, None, :] + rotated * sin[None, :, None, :] + return torch.cat((head, tail), dim=-1) + + +def golden_attention(q, k, v) -> torch.Tensor: + """Non-causal full attention over one packed document, ``(B, S, H, D)`` -> ``(B, S, H*D)``.""" + + q, k, v = (_f32(t).permute(0, 2, 1, 3) for t in (q, k, v)) + scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(q.shape[-1]) + out = torch.matmul(torch.softmax(scores, dim=-1), v) + return out.permute(0, 2, 1, 3).flatten(2, 3) + + +def golden_swiglu(projected) -> torch.Tensor: + hidden, gate = _f32(projected).chunk(2, dim=-1) + return hidden * F.silu(gate) + + +def golden_gate_residual(residual, gate, index, y) -> torch.Tensor: + return _f32(residual) + _f32(gate).index_select(0, index) * _f32(y) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 292ba453c..a733e95ce 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -76,6 +76,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): MUSA_FUSED_LOGP_GENERIC = "rl_engine.backends.musa.logprob.logp.MusaFusedLogpOp" CUDA_DETERMINISTIC_LOGP = "rl_engine.kernels.ops.cuda.loss.logp.DeterministicLogpCUDAOp" # Deterministic standard-softmax attention (issue #147); not FlashAttention. + CUDA_H3_FULL_ATTENTION = "rl_engine.backends.cuda.attention.h3_full_attention.H3FullAttentionOp" CUDA_DETERMINISTIC_ATTENTION = ( "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp" ) @@ -243,6 +244,34 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: return ( + OperatorBackendDescriptor( + semantic_op="h3_full_attention", + backend_id="rlkernel.cuda.h3_full_attention.query_chunk.v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cuda"}), + supported_dtypes=frozenset({"bfloat16"}), + supported_topologies={ + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + determinism_or_alignment_properties={ + "deterministic": True, + "batch_invariant": True, + "strict_observable": True, + "arithmetic_profile": "h3-cuda-strict-v1", + "layout": "BSHD", + "heads": 56, + "head_dim": 128, + "causal": False, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=( + "rl_engine.backends.cuda.attention.h3_full_attention.H3FullAttentionOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="h3-query-chunk-v1; binary profile checked at invocation", + ), OperatorBackendDescriptor( semantic_op="video_patch_pack_unpack", backend_id="pytorch-video-patch-pack-unpack-v1", @@ -637,6 +666,7 @@ def __init__(self): OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN, ], + "h3_full_attention": [OpBackend.CUDA_H3_FULL_ATTENTION], "attention": [ OpBackend.CUDA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION, @@ -1159,6 +1189,8 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: """Select the best legacy operator for the requested device.""" 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") candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: diff --git a/rl_engine/validation/models/h3_block.py b/rl_engine/validation/models/h3_block.py new file mode 100644 index 000000000..b20e885a1 --- /dev/null +++ b/rl_engine/validation/models/h3_block.py @@ -0,0 +1,718 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One real MiniMax-H3 transformer block, node by node (RFC #420 ``ws1_one_h3_block``). + +The block is a fixed graph of 17 nodes (``NODES``), from the AdaLN projection +through attention and the FFN to the second gated residual, on block 0's +pinned weights and a packed FL2VA layout. Every node runs three ways: + +* provider: the diffusers op sequence (``h3_provider``), fed provider inputs; +* candidate: the RL-Kernel binding of the node, fed candidate inputs + (chained) and, separately, provider inputs (isolated); +* golden: the FP32 reference (``rl_engine.reference.minimax_h3.block``). + +Each node also names the RFC row that owns its arithmetic and how it is bound +today (``Binding.status``): + +* ``row`` -- that row's operator, in this PR stack; +* ``interim`` -- an existing deterministic RL-Kernel operator standing in + until the owning row lands (its arithmetic, speed and rounding are that + operator's, not the row's); +* ``reference`` -- no kernel exists yet; the node runs the provider replay + itself, so it is trivially bitwise and contributes no RL-Kernel evidence. + +The report gives, per node: repeat and batch-row bitwise checks, chained and +isolated agreement with the provider, error against the golden for both the +candidate and the provider, and the backend that ran. ``first_*_drift`` name +the first node (graph order) where a comparison stops being bitwise; the +backward report does the same for the gradients reaching each node's output, +in reverse graph order. +""" + +from __future__ import annotations + +import hashlib +import time +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable + +import torch + +from rl_engine.reference.minimax_h3 import block as gold +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models import h3_provider as prov +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS, h3_fl2va_layout, h3_timesteps + +HIDDEN = 5376 +HEADS = 56 +HEAD_DIM = 128 +BLOCK = "transformer_blocks.0" +NUM_TIMESTEPS = 2 + + +@dataclass(frozen=True) +class Binding: + rfc_row: str + status: str # "row" | "interim" | "reference" + backend: str + + +@dataclass(frozen=True) +class Node: + name: str + inputs: tuple[str, ...] + binding: Binding + # (ops, ctx, *inputs) -> output + candidate: Callable[..., Any] + # (ctx, *inputs) -> output + provider: Callable[..., Any] + golden: Callable[..., Any] + # Promise checked by the tests: the isolated candidate equals the provider bitwise. + provider_bitwise_isolated: bool + + +class CandidateOps: + """The operators bound to the block's nodes, resolved once per run. + + The three rows of this stack are resolved through the registry, so the + report records what dispatch actually returned; the interim operators are + named directly because their registry entries (``det_gemm``, + ``attention``, ``swiglu``) prefer backends that are not available or not + deterministic on every GPU this replay runs on. + """ + + def __init__(self, registry: KernelRegistry | None = None) -> None: + from rl_engine.backends.cuda.activation.swiglu import SwiGLUCudaOp + from rl_engine.backends.cuda.attention.deterministic_attn import DeterministicAttentionOp + from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp + + registry = registry or KernelRegistry() + self.projection = registry.get_op("adaln_projection_3mod", device="cuda") + self.norm = registry.get_op("h3_rmsnorm", device="cuda") + self.gate = registry.get_op("adaln_gate_residual", device="cuda") + self.gemm = TritonDetGemmOp() + self.attention = DeterministicAttentionOp() + self.swiglu = SwiGLUCudaOp() + + def backends(self) -> dict[str, str]: + return { + name: f"{type(op).__module__}.{type(op).__name__}" for name, op in vars(self).items() + } + + +def _w(ctx: dict[str, Any], name: str) -> torch.Tensor: + return ctx["weights"][f"{BLOCK}.{name}"] + + +def _heads(x: torch.Tensor) -> torch.Tensor: + return x.unflatten(-1, (HEADS, HEAD_DIM)) + + +class _InterimLinear(torch.autograd.Function): + """``y = x @ W^T`` on the Triton FP32-accumulate kernel, rounded to BF16 once. + + Interim binding for the block's GEMM rows. ``TritonDetGemmOp.__call__`` + reduces K as a tree with BF16 nodes (about 3x cuBLAS's error at these + shapes); ``_triton_gemm_fp32`` runs one fixed-order FP32 K loop (no + split-K), so each output row depends only on its own input row. The + backward uses the same kernel for ``dX`` (reduction over the output + features, row-local) and ``dW`` (reduction over tokens). + """ + + @staticmethod + def forward(ctx, x2, weight): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + ctx.save_for_backward(x2, weight) + return _triton_gemm_fp32(x2, weight.t()).to(x2.dtype) + + @staticmethod + def backward(ctx, grad): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + x2, weight = ctx.saved_tensors + grad = grad.contiguous() + dx = _triton_gemm_fp32(grad, weight).to(x2.dtype) + dw = _triton_gemm_fp32(grad.t(), x2).to(weight.dtype) + return dx, dw + + +def _cand_linear(ops: CandidateOps, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + flat = x.reshape(-1, x.shape[-1]) + return _InterimLinear.apply(flat, weight).view(*x.shape[:-1], weight.shape[0]) + + +def _cand_attention(ops: CandidateOps, q, k, v) -> torch.Tensor: + q, k, v = (t.transpose(1, 2).contiguous() for t in (q, k, _heads(v))) + out = ops.attention(q, k, v, causal=False) + return out.transpose(1, 2).flatten(2, 3) + + +def _cand_swiglu(ops: CandidateOps, projected: torch.Tensor) -> torch.Tensor: + hidden, gate = projected.chunk(2, dim=-1) + return ops.swiglu(gate.contiguous(), hidden.contiguous()) + + +def _linear_node(name: str, source: str, weight: str, row: str, status: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding( + row, status, "triton _det_gemm_fp32_kernel (FP32 K loop, one BF16 rounding)" + ), + candidate=lambda ops, ctx, x: _cand_linear(ops, x, _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_linear(x, _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_linear(x, _w(ctx, weight)), + provider_bitwise_isolated=False, # tree reduction vs cuBLAS + ) + + +def _norm_modulate_node(name: str, source: str, weight: str, shift: int, scale: int) -> Node: + return Node( + name=name, + inputs=(source, "adaln_projection"), + binding=Binding("h3_rmsnorm", "row", "H3RMSNormCudaOp.forward_modulated"), + candidate=lambda ops, ctx, x, t: ops.norm.forward_modulated( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + provider=lambda ctx, x, t: prov.provider_norm_modulate( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + golden=lambda ctx, x, t: gold.golden_norm_modulate( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + provider_bitwise_isolated=True, + ) + + +def _head_norm_node(name: str, source: str, weight: str) -> Node: + return Node( + name=name, + inputs=(source,), + # The row is someone else's; until it lands, this stack's plain H3 RMSNorm + # (bitwise equal to nn.RMSNorm at any width) runs on the 128-wide heads. + binding=Binding("h3_qk_rmsnorm_d128", "interim", "H3RMSNormCudaOp.forward"), + candidate=lambda ops, ctx, x: ops.norm(_heads(x), _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_head_rmsnorm(_heads(x), _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_rmsnorm(_heads(x), _w(ctx, weight)), + provider_bitwise_isolated=True, + ) + + +def _rope_node(name: str, source: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding("h3_mm_rope_3axis_partial", "reference", "provider_apply_rotary"), + candidate=lambda ops, ctx, x: prov.provider_apply_rotary(x, *ctx["rope"]), + provider=lambda ctx, x: prov.provider_apply_rotary(x, *ctx["rope"]), + golden=lambda ctx, x: gold.golden_apply_rotary(x, *ctx["rope_fp32"]), + provider_bitwise_isolated=True, + ) + + +def _gate_node(name: str, residual: str, sublayer: str, gate: int) -> Node: + return Node( + name=name, + inputs=(residual, sublayer, "adaln_projection"), + binding=Binding("adaln_gate_residual", "row", "H3GateResidualCudaOp"), + candidate=lambda ops, ctx, r, y, t: ops.gate(r, y, t[gate], ctx["index"]), + provider=lambda ctx, r, y, t: prov.provider_gate_residual(r, t[gate], ctx["index"], y), + golden=lambda ctx, r, y, t: gold.golden_gate_residual(r, t[gate], ctx["index"], y), + provider_bitwise_isolated=True, + ) + + +def _adaln_params(ctx): + return _w(ctx, "adaln_proj.linear.weight"), _w(ctx, "adaln_proj.linear.bias") + + +NODES: list[Node] = [ + Node( + name="adaln_projection", + inputs=("temb",), + binding=Binding("adaln_projection_3mod", "row", "H3AdaLNProjectionCudaOp"), + candidate=lambda ops, ctx, temb: ops.projection(temb, *_adaln_params(ctx)), + provider=lambda ctx, temb: prov.provider_adaln_modulation(temb, *_adaln_params(ctx)), + golden=lambda ctx, temb: gold.golden_adaln_tables(temb, *_adaln_params(ctx)), + provider_bitwise_isolated=False, # tensor-core tree vs cuBLAS + ), + _norm_modulate_node("norm1", "hidden", "norm1.weight", shift=0, scale=1), + _linear_node("q_proj", "norm1", "attn.to_q.weight", "h3_qkv_gemm", "interim"), + _linear_node("k_proj", "norm1", "attn.to_k.weight", "h3_qkv_gemm", "interim"), + _linear_node("v_proj", "norm1", "attn.to_v.weight", "h3_qkv_gemm", "interim"), + _head_norm_node("q_norm", "q_proj", "attn.norm_q.weight"), + _head_norm_node("k_norm", "k_proj", "attn.norm_k.weight"), + _rope_node("rope_q", "q_norm"), + _rope_node("rope_k", "k_norm"), + Node( + name="attention", + inputs=("rope_q", "rope_k", "v_proj"), + binding=Binding("h3_full_attention", "interim", "DeterministicAttentionOp(causal=False)"), + candidate=lambda ops, ctx, q, k, v: _cand_attention(ops, q, k, v), + provider=lambda ctx, q, k, v: prov.provider_attention(q, k, _heads(v)), + golden=lambda ctx, q, k, v: gold.golden_attention(q, k, _heads(v)), + provider_bitwise_isolated=False, + ), + _linear_node("o_proj", "attention", "attn.to_out.0.weight", "h3_attention_o_gemm", "interim"), + _gate_node("residual_attn", "hidden", "o_proj", gate=2), + _norm_modulate_node("norm2", "residual_attn", "norm2.weight", shift=3, scale=4), + _linear_node("ffn_gate_up", "norm2", "ff.net.0.proj.weight", "h3_ffn_gate_up_gemm", "interim"), + Node( + name="swiglu", + inputs=("ffn_gate_up",), + binding=Binding("h3_swiglu", "interim", "SwiGLUCudaOp"), + candidate=lambda ops, ctx, x: _cand_swiglu(ops, x), + provider=lambda ctx, x: prov.provider_swiglu(x), + golden=lambda ctx, x: gold.golden_swiglu(x), + provider_bitwise_isolated=False, # one FP32 rounding vs the eager two + ), + _linear_node("ffn_down", "swiglu", "ff.net.2.weight", "h3_ffn_down_gemm", "interim"), + _gate_node("residual_mlp", "residual_attn", "ffn_down", gate=5), +] +NODE_NAMES = tuple(node.name for node in NODES) +OUTPUT = NODES[-1].name + +# Gradient leaves of the backward replay: the block input, temb and every block parameter. +PARAM_LEAVES = ( + "adaln_proj.linear.weight", + "adaln_proj.linear.bias", + "norm1.weight", + "attn.to_q.weight", + "attn.to_k.weight", + "attn.to_v.weight", + "attn.norm_q.weight", + "attn.norm_k.weight", + "attn.to_out.0.weight", + "norm2.weight", + "ff.net.0.proj.weight", + "ff.net.2.weight", +) + + +# --- Inputs --------------------------------------------------------------------------------- + + +def make_block_context( + weights: dict[str, torch.Tensor], *, layout: str, seed: int = 0, batch: int = 1 +) -> dict[str, Any]: + """Block-0 weights, a packed FL2VA layout, the timestep embedding and a seeded input. + + ``temb`` comes from the provider's conditioning path (``h3_chain`` covers + those rows); the block input is a seeded BF16 stand-in for the packed + projections, which are other rows. Batch rows ``b > 0`` get their own + seeded inputs; row 0 is the same tensor at every ``batch``. + """ + + spec = H3_BLOCK_LAYOUTS[layout] + packed = h3_fl2va_layout( + num_text_tokens=spec["num_text_tokens"], + num_latent_frames=spec["num_latent_frames"], + latent_height=spec["latent_height"], + latent_width=spec["latent_width"], + num_audio_latents=spec["num_audio_latents"], + ) + seq_len = packed["position_ids"].shape[0] + timestep = h3_timesteps(NUM_TIMESTEPS, seed=seed) + w = weights + with torch.no_grad(): + features = prov.provider_time_proj(timestep) + temb = prov.provider_time_embedder( + features, + w["time_embedder.linear_1.weight"], + w["time_embedder.linear_1.bias"], + w["time_embedder.linear_2.weight"], + w["time_embedder.linear_2.bias"], + ) + rows = [ + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 2 + b)) + for b in range(batch) + ] + hidden = torch.cat(rows).to(torch.bfloat16).cuda() + return { + "layout": layout, + "seq_len": seq_len, + "batch": batch, + "weights": weights, + "temb": temb, + "hidden": hidden, + "index": packed["timestep_indices"] * 3 + packed["token_tags"], + "position_ids": packed["position_ids"], + "rope": prov.provider_rope(packed["position_ids"]), + "rope_fp32": gold.golden_rope_tables(packed["position_ids"]), + } + + +@contextmanager +def strict_fp32_matmul(): + """Golden matmuls in full FP32: no TF32, no reduced-precision reductions.""" + + saved = (torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = saved + + +# --- Graph execution ------------------------------------------------------------------------ + + +def run_graph(mode: str, ctx: dict[str, Any], ops: CandidateOps | None = None) -> dict[str, Any]: + """Every node's output for one mode ("candidate", "provider" or "golden").""" + + values: dict[str, Any] = {"hidden": ctx["hidden"], "temb": ctx["temb"]} + # Diagnostic hybrids only: these candidate nodes run their provider replay instead. + provider_nodes = ctx.get("provider_nodes", ()) + with strict_fp32_matmul(): + for node in NODES: + args = [values[name] for name in node.inputs] + if mode == "candidate" and node.name in provider_nodes: + values[node.name] = node.provider(ctx, *args) + elif mode == "candidate": + values[node.name] = node.candidate(ops, ctx, *args) + elif mode == "provider": + values[node.name] = node.provider(ctx, *args) + else: + values[node.name] = node.golden(ctx, *args) + observer = ctx.get("node_observer") + if observer is not None: + observer(node.name) + return values + + +def run_isolated( + ctx: dict[str, Any], ops: CandidateOps, provider: dict[str, Any] +) -> dict[str, Any]: + """Each candidate node fed the provider's inputs to that node.""" + + return { + node.name: node.candidate(ops, ctx, *[provider[name] for name in node.inputs]) + for node in NODES + } + + +def _flat(value: Any) -> list[torch.Tensor]: + if isinstance(value, torch.Tensor): + return [value] + return [t for item in value for t in _flat(item)] + + +def _batch_row(value: Any, row: int, batch: int) -> Any: + """Row ``row`` of a batched activation; AdaLN tables have no batch axis.""" + + if isinstance(value, torch.Tensor): + return value[row : row + 1] if value.dim() >= 3 and value.shape[0] == batch else value + return tuple(_batch_row(v, row, batch) for v in value) + + +def bitwise(lhs: Any, rhs: Any) -> bool: + return all(torch.equal(a, b) for a, b in zip(_flat(lhs), _flat(rhs), strict=True)) + + +def max_abs(lhs: Any, rhs: Any) -> float: + return max( + float((a.float() - b.float()).abs().max()) if a.numel() else 0.0 + for a, b in zip(_flat(lhs), _flat(rhs), strict=True) + ) + + +def first_differing_row(lhs: Any, rhs: Any) -> int | None: + """First packed row (sequence position) where two node outputs differ, or None.""" + + for a, b in zip(_flat(lhs), _flat(rhs), strict=True): + if a.dim() < 3: + if not torch.equal(a, b): + return -1 # a table row, not a sequence position + continue + diff = (a != b).reshape(a.shape[0], a.shape[1], -1).any(dim=-1).any(dim=0) + rows = diff.nonzero() + if rows.numel(): + return int(rows[0]) + return None + + +def rel_err(value: Any, golden: Any) -> float: + """``max|value - golden| / max|golden|`` over a node's tensors.""" + + num = den = 0.0 + for a, g in zip(_flat(value), _flat(golden), strict=True): + g32 = g.float() + num = max(num, float((a.float() - g32).abs().max())) + den = max(den, float(g32.abs().max())) + return num / max(den, 1e-30) + + +def ulp_err(value: torch.Tensor, golden: torch.Tensor) -> float: + """``max|value - golden|`` in units of ``value.dtype``'s ulp at ``max|golden|``. + + ``rel_err`` of a stored gradient moves in steps of this ulp: its worst + element is a single rounded number, so two equally valid reduction orders + can differ by whole ulps there while every other element agrees. + """ + + g32 = golden.float() + top = float(g32.abs().max()) + if top == 0.0: + return 0.0 + ulp = torch.finfo(value.dtype).eps * 2.0 ** torch.floor(torch.log2(torch.tensor(top))).item() + return float((value.float() - g32).abs().max()) / ulp + + +def digest(value: Any) -> str: + """sha256 over a node's tensors' bytes, for comparing runs across processes or hosts.""" + + h = hashlib.sha256() + for t in _flat(value): + h.update(t.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return h.hexdigest() + + +def _time_ms(fn: Callable[[], Any], repeats: int) -> float: + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize() + samples.append((time.perf_counter() - start) * 1e3) + return sorted(samples)[len(samples) // 2] + + +# --- Forward report ------------------------------------------------------------------------- + + +def run_block_case( + weights: dict[str, torch.Tensor], + *, + layout: str, + seed: int = 0, + ops: CandidateOps | None = None, + timing_repeats: int = 0, +) -> dict[str, Any]: + """Forward replay of block 0: per-node agreement, invariance, accuracy and first drift.""" + + ops = ops or CandidateOps() + ctx = make_block_context(weights, layout=layout, seed=seed, batch=1) + ctx2 = make_block_context(weights, layout=layout, seed=seed, batch=2) + with torch.no_grad(): + provider = run_graph("provider", ctx) + golden = run_graph("golden", ctx) + chained = run_graph("candidate", ctx, ops) + repeat = run_graph("candidate", ctx, ops) + batched = run_graph("candidate", ctx2, ops) + isolated = run_isolated(ctx, ops, provider) + + report: dict[str, Any] = { + "layout": layout, + "seq_len": ctx["seq_len"], + "seed": seed, + "backends": ops.backends(), + "nodes": [], + } + firsts: dict[str, str | None] = {"first_drift": None, "first_isolated_drift": None} + firsts.update({"first_repeat_drift": None, "first_batch_drift": None}) + for node in NODES: + name = node.name + row0 = _batch_row(batched[name], 0, 2) + entry: dict[str, Any] = { + "node": name, + "rfc_row": node.binding.rfc_row, + "status": node.binding.status, + "backend": node.binding.backend, + "provider_bitwise_isolated_promised": node.provider_bitwise_isolated, + "repeat_bitwise_equal": bitwise(repeat[name], chained[name]), + "batch_row_bitwise_equal": bitwise(row0, chained[name]), + "chained_vs_provider": { + "bitwise_equal": bitwise(chained[name], provider[name]), + "max_abs": max_abs(chained[name], provider[name]), + "first_differing_row": first_differing_row(chained[name], provider[name]), + }, + "isolated_vs_provider": { + "bitwise_equal": bitwise(isolated[name], provider[name]), + "max_abs": max_abs(isolated[name], provider[name]), + }, + "candidate_rel_err_vs_golden": rel_err(chained[name], golden[name]), + "provider_rel_err_vs_golden": rel_err(provider[name], golden[name]), + "candidate_sha256": digest(chained[name]), + } + checks = { + "first_drift": entry["chained_vs_provider"]["bitwise_equal"], + "first_isolated_drift": entry["isolated_vs_provider"]["bitwise_equal"], + "first_repeat_drift": entry["repeat_bitwise_equal"], + "first_batch_drift": entry["batch_row_bitwise_equal"], + } + for key, ok in checks.items(): + if firsts[key] is None and not ok: + firsts[key] = name + report["nodes"].append(entry) + report.update(firsts) + report["output"] = { + "candidate_rel_err_vs_golden": rel_err(chained[OUTPUT], golden[OUTPUT]), + "provider_rel_err_vs_golden": rel_err(provider[OUTPUT], golden[OUTPUT]), + "candidate_vs_provider_max_abs": max_abs(chained[OUTPUT], provider[OUTPUT]), + } + report["permutation_probe"] = permutation_probe(ctx, ops, chained[OUTPUT], seed) + if timing_repeats: + with torch.no_grad(): + report["timing_ms"] = { + "candidate_forward": _time_ms( + lambda: run_graph("candidate", ctx, ops), timing_repeats + ), + "provider_forward": _time_ms(lambda: run_graph("provider", ctx), timing_repeats), + } + return report + + +def permutation_probe( + ctx: dict[str, Any], ops: CandidateOps, base_output: torch.Tensor, seed: int +) -> dict[str, Any]: + """RFC #420 control H-C3, reported and not promised. + + Permute the packed rows together with every per-row tensor (positions, + AdaLN indices); the block output should be the same permutation of the + original. Row-local nodes keep each row's bytes, but attention sums over + keys in packed order, so a different key order may round differently. + """ + + perm = torch.randperm(ctx["seq_len"], generator=torch.Generator().manual_seed(seed + 7)).cuda() + permuted = { + **ctx, + "hidden": ctx["hidden"][:, perm], + "index": ctx["index"][perm], + "position_ids": ctx["position_ids"][perm], + "rope": tuple(t[perm] for t in ctx["rope"]), + "rope_fp32": tuple(t[perm] for t in ctx["rope_fp32"]), + } + with torch.no_grad(): + out = run_graph("candidate", permuted, ops)[OUTPUT] + expected = base_output[:, perm] + return { + "bitwise_equal_after_permutation": bool(torch.equal(out, expected)), + "max_abs": max_abs(out, expected), + "differing_rows": int((out != expected).reshape(out.shape[1], -1).any(dim=-1).sum()), + } + + +# --- Backward report ------------------------------------------------------------------------ + + +def _backward( + mode: str, ctx: dict[str, Any], ops: CandidateOps | None, upstream: torch.Tensor +) -> tuple[dict[str, torch.Tensor], dict[str, list[torch.Tensor]]]: + """Gradients of every leaf, and the gradient reaching each node's output, for one mode. + + The golden runs on FP32 copies of the leaves, so its gradients are not + rounded to the BF16 parameter dtype. + """ + + def leaf(t: torch.Tensor) -> torch.Tensor: + t = t.detach().float() if mode == "golden" else t.detach() + return t.clone().requires_grad_(True) + + leaves = {name: leaf(ctx["weights"][f"{BLOCK}.{name}"]) for name in PARAM_LEAVES} + leaves["hidden"] = leaf(ctx["hidden"]) + leaves["temb"] = leaf(ctx["temb"]) + weights = {**ctx["weights"], **{f"{BLOCK}.{k}": leaves[k] for k in PARAM_LEAVES}} + run_ctx = {**ctx, "weights": weights, "hidden": leaves["hidden"], "temb": leaves["temb"]} + values = run_graph(mode, run_ctx, ops) + for node in NODES: + for t in _flat(values[node.name]): + if t.requires_grad: + t.retain_grad() + out = values[OUTPUT] + torch.autograd.backward(out, upstream.to(out.dtype)) + node_grads = { + node.name: [ + t.grad if t.grad is not None else torch.zeros_like(t) for t in _flat(values[node.name]) + ] + for node in NODES + } + return {name: t.grad for name, t in leaves.items()}, node_grads + + +def run_block_backward( + weights: dict[str, torch.Tensor], + *, + layout: str, + seed: int = 0, + ops: CandidateOps | None = None, + timing_repeats: int = 0, +) -> dict[str, Any]: + """Backward replay of block 0. + + Per leaf (block input, ``temb`` and the 12 block parameters): repeat + bitwise, and the candidate's and provider's error against the FP32 + golden. Per node, in reverse graph order: whether the gradient reaching + the node's output repeats bitwise, and whether batch row 0 of a two-row + batch gets the same bytes as the batch of one (rows are separate packed + documents, so every activation gradient is row-local). Parameter + gradients sum over the batch and are not part of the batch check. + """ + + ops = ops or CandidateOps() + ctx = make_block_context(weights, layout=layout, seed=seed, batch=1) + ctx2 = make_block_context(weights, layout=layout, seed=seed, batch=2) + generator = torch.Generator().manual_seed(seed + 11) + up_row0 = torch.randn((1, ctx["seq_len"], HIDDEN), generator=generator) + up_row1 = torch.randn((1, ctx["seq_len"], HIDDEN), generator=generator) + upstream = up_row0.bfloat16().cuda() + upstream2 = torch.cat([up_row0, up_row1]).bfloat16().cuda() + + golden, _ = _backward("golden", ctx, None, upstream) + provider, _ = _backward("provider", ctx, None, upstream) + first, first_nodes = _backward("candidate", ctx, ops, upstream) + second, second_nodes = _backward("candidate", ctx, ops, upstream) + batched, batched_nodes = _backward("candidate", ctx2, ops, upstream2) + + report: dict[str, Any] = { + "layout": layout, + "seq_len": ctx["seq_len"], + "seed": seed, + "leaves": {}, + "nodes": [], + } + for name in (*PARAM_LEAVES, "hidden", "temb"): + report["leaves"][name] = { + "repeat_bitwise_equal": bool(torch.equal(first[name], second[name])), + "candidate_rel_err_vs_golden": rel_err(first[name], golden[name]), + "provider_rel_err_vs_golden": rel_err(provider[name], golden[name]), + "storage_dtype": str(first[name].dtype).removeprefix("torch."), + "candidate_ulps_vs_golden": ulp_err(first[name], golden[name]), + "provider_ulps_vs_golden": ulp_err(provider[name], golden[name]), + } + report["leaves"]["hidden"]["batch_row_bitwise_equal"] = bool( + torch.equal(batched["hidden"][:1], first["hidden"]) + ) + first_repeat = first_batch = None + for node in reversed(NODES): + name = node.name + row0 = [_batch_row(g, 0, 2) for g in batched_nodes[name]] + entry: dict[str, Any] = { + "node": name, + "grad_repeat_bitwise_equal": bitwise(first_nodes[name], second_nodes[name]), + "grad_batch_row_bitwise_equal": ( + bitwise(row0, first_nodes[name]) if node.name != "adaln_projection" else None + ), + } + if first_repeat is None and not entry["grad_repeat_bitwise_equal"]: + first_repeat = name + if first_batch is None and entry["grad_batch_row_bitwise_equal"] is False: + first_batch = name + report["nodes"].append(entry) + report["first_grad_repeat_drift"] = first_repeat + report["first_grad_batch_drift"] = first_batch + if timing_repeats: + report["timing_ms"] = { + "candidate_forward_backward": _time_ms( + lambda: _backward("candidate", ctx, ops, upstream), timing_repeats + ), + "provider_forward_backward": _time_ms( + lambda: _backward("provider", ctx, None, upstream), timing_repeats + ), + } + return report diff --git a/rl_engine/validation/models/h3_cases.py b/rl_engine/validation/models/h3_cases.py index 62184cc75..92455bae0 100644 --- a/rl_engine/validation/models/h3_cases.py +++ b/rl_engine/validation/models/h3_cases.py @@ -62,3 +62,127 @@ def h3_block_layout( token_tags[lo : lo + text] = 1 token_tags[hi - audio : hi] = 2 return timestep_indices.to(device), token_tags.to(device) + + +# diffusers@f53d552 modular_pipelines/minimax_h3/before_denoise.py constants. +_ROPE_FRAME_RESCALE = 5.0 / 3.0 +_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) +_ROPE_SPATIAL_SCALE = 32 +VIDEO_TAG, TEXT_TAG, AUDIO_TAG = 0, 1, 2 + + +def _spatial_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: + """One aspect-normalised spatial axis, as numpy's ``linspace(endpoint=False)`` builds it.""" + + import numpy as np + + ratio = dim / sqrt_area + left = (1.0 - ratio) / 2.0 + grid = np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE + return torch.from_numpy(grid).to(torch.float64) + + +def h3_fl2va_layout( + *, + num_text_tokens: int, + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + audio_channels: int = 2, + patch: tuple[int, int] = (2, 2), + device: str = "cuda", +) -> dict[str, torch.Tensor]: + """Packed text/audio/video rows of a text-only FL2VA request, laid out as diffusers packs them. + + Replays ``position_ids``/``token_tags`` of the pinned pipeline's FL2VA + layout step (no keyframe anchors): text rows on the time axis at their row + index, channel-major audio rows pinned to the width extremes, then the video + frames on the non-uniform ``5/3 * (1, 4, 4, 4, 4)`` clock. ``position_ids`` + is ``(S, 3)`` float64 as the pipeline makes it; the rotary module casts it. + ``timestep_indices`` puts text on table row 0 and the generated media on + row 1, a two-timestep stand-in for the conditioning rows. + """ + + import numpy as np + + patch_h, patch_w = patch + rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w) + num_audio_rows = num_audio_latents * audio_channels + num_video_rows = num_latent_frames * rows_per_frame + seq_len = num_text_tokens + num_audio_rows + num_video_rows + audio_start = num_text_tokens + video_start = audio_start + num_audio_rows + + position_ids = torch.zeros(seq_len, 3, dtype=torch.float64) + position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) + sqrt_area = np.sqrt(latent_height * latent_width) + height_grid = _spatial_grid(latent_height, patch_h, sqrt_area) + width_grid = _spatial_grid(latent_width, patch_w, sqrt_area) + grids = torch.meshgrid(height_grid, width_grid, indexing="ij") + frame_grid = torch.stack([grid.reshape(-1) for grid in grids], dim=-1) + + audio_time = float(num_text_tokens) + torch.arange(num_audio_latents, dtype=torch.float64) + position_ids[audio_start:video_start, 0] = audio_time.repeat(audio_channels) + position_ids[audio_start:video_start, 2] = torch.cat( + [ + torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64), + torch.full( + (num_audio_rows - num_audio_latents,), float(width_grid[-1]), dtype=torch.float64 + ), + ] + ) + spans = torch.tensor( + [ + _ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[i % len(_ROPE_FRAMES_PER_LATENT)] + for i in range(num_latent_frames) + ], + dtype=torch.float64, + ) + frame_time = float(num_text_tokens) + torch.cat( + [torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)] + ) + video = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64) + video[:, :, 0] = frame_time[:, None] + video[:, :, 1:] = frame_grid[None] + position_ids[video_start:] = video.reshape(-1, 3) + + token_tags = torch.full((seq_len,), VIDEO_TAG, dtype=torch.long) + token_tags[:num_text_tokens] = TEXT_TAG + token_tags[audio_start:video_start] = AUDIO_TAG + timestep_indices = torch.ones(seq_len, dtype=torch.long) + timestep_indices[:num_text_tokens] = 0 + return { + "position_ids": position_ids.to(device), + "token_tags": token_tags.to(device), + "timestep_indices": timestep_indices.to(device), + } + + +# Named FL2VA layouts for the one-block replay: (text, latent frames, latent H, latent W, audio). +H3_BLOCK_LAYOUTS = { + # 32 + 2*20 + 2*(8*12) = 264 rows: the CI-sized case. + "tiny": dict( + num_text_tokens=32, + num_latent_frames=2, + latent_height=16, + latent_width=24, + num_audio_latents=20, + ), + # 77 + 2*40 + 3*(16*16) = 925 rows: a non-power-of-two, three-frame case. + "small": dict( + num_text_tokens=77, + num_latent_frames=3, + latent_height=32, + latent_width=32, + num_audio_latents=40, + ), + # 120 + 2*80 + 5*(18*32) = 3160 rows: a 16:9 five-frame case. + "medium": dict( + num_text_tokens=120, + num_latent_frames=5, + latent_height=36, + latent_width=64, + num_audio_latents=80, + ), +} diff --git a/rl_engine/validation/models/h3_chain.py b/rl_engine/validation/models/h3_chain.py index 72cbc760a..3aa1aa191 100644 --- a/rl_engine/validation/models/h3_chain.py +++ b/rl_engine/validation/models/h3_chain.py @@ -307,10 +307,12 @@ def run_case( "seed": seed, "stages": [], } - chained = provider = golden = None + chained: Any = None + provider: Any = None + golden: Any = None first_drift = first_isolated = None # Each mode sees its own earlier outputs (a stage may read any earlier stage). - history = {"candidate": {}, "provider": {}, "golden": {}} + history: dict[str, dict[str, Any]] = {"candidate": {}, "provider": {}, "golden": {}} ctx_c = {**ctx, "history": history["candidate"]} ctx_p = {**ctx, "history": history["provider"]} ctx_g = {**ctx, "history": history["golden"]} @@ -379,6 +381,8 @@ def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torc value = stage.candidate(op, run_ctx, value) else: value = stage.golden(golden_op(registry, stage.op_type), run_ctx, value) + if value is None: + raise ValueError("conditioning chain produced no output") outputs = list(value) torch.autograd.backward( outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs, strict=True)] diff --git a/rl_engine/validation/models/h3_manifest.json b/rl_engine/validation/models/h3_manifest.json index 55f7e5220..e590a8473 100644 --- a/rl_engine/validation/models/h3_manifest.json +++ b/rl_engine/validation/models/h3_manifest.json @@ -79,6 +79,40 @@ "sha256": "3b12cdc7bbda772e7e31539ffa9a32f529dba7f4148273d2128bdc621fd27438" } }, + "block_tensors": { + "transformer_blocks.0.attn.to_q.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "196f9120ad018784bcefd4224f0e08fae88ba1adf079ecdca18410a414a6a2a4" + }, + "transformer_blocks.0.attn.to_k.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "8e9171f3b8cf894904d1ddbd96510739d7291d2992310de6a52e1d733ec895cd" + }, + "transformer_blocks.0.attn.to_v.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "275bbc993b1ba738831e04d86f290c9988108d96fe8c97c3b576a043c7ff3fe1" + }, + "transformer_blocks.0.attn.norm_q.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "34a3f13557a2c0a9af260e48d431041e381779a0f61616cfe9800ec839693519" + }, + "transformer_blocks.0.attn.norm_k.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "8991949ddea3bc2136a0fcf0b980fa63bd0df760207e3116f7ce96df4e6ca058" + }, + "transformer_blocks.0.attn.to_out.0.weight": { + "dtype": "bfloat16", "shape": [5376, 7168], + "sha256": "822cb426aaebe57d3d219b3c4f6855ca5c13821cb84ebe0a7474fe3a87869fd6" + }, + "transformer_blocks.0.ff.net.0.proj.weight": { + "dtype": "bfloat16", "shape": [28672, 5376], + "sha256": "49587044f325aae73f046c4cda9637052670227ec3c2c289359d9d1ebf29f95d" + }, + "transformer_blocks.0.ff.net.2.weight": { + "dtype": "bfloat16", "shape": [5376, 14336], + "sha256": "e19d4e1649e4b55b1931842e5cc4bb46892c2e635b9315b2996630a86284e88f" + } + }, "reference_implementation": { "repo": "huggingface/diffusers", "commit": "f53d552036a0d1bd5570782a39cd40cfabf112bc", diff --git a/rl_engine/validation/models/h3_provider.py b/rl_engine/validation/models/h3_provider.py index 19c2b5034..3eec1818f 100644 --- a/rl_engine/validation/models/h3_provider.py +++ b/rl_engine/validation/models/h3_provider.py @@ -118,3 +118,75 @@ def provider_final_adaln_out( return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( 0, timestep_indices ) + + +# --- One transformer block (RFC #420 ``ws1_one_h3_block``) ------------------------------------ + + +def provider_rope( + position_ids: torch.Tensor, rope_freq_dim: int = 16, rope_theta: float = 10000.0 +) -> tuple[torch.Tensor, torch.Tensor]: + """``MiniMaxH3RotaryPosEmbed``: FP32 ``(S, 96)`` cos/sin over the ``(t, h, w)`` axes.""" + + inv_freq = 1.0 / ( + rope_theta + ** ( + torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device) + / (2 * rope_freq_dim) + ) + ) + position_ids = position_ids.to(torch.float32) + freqs = position_ids.unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs_t, freqs_h, freqs_w = freqs.unbind(dim=1) + freqs = torch.cat((freqs_t, freqs_h, freqs_w), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def provider_apply_rotary( + hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + """``_apply_rotary_emb``: rotate-half on the leading 96 of 128 channels of ``(B, S, H, D)``. + + The FP32 tables are cast to the activation dtype (BF16) before the rotation. + """ + + rotary_dim = cos.shape[-1] + hidden_states_rotary = hidden_states[..., :rotary_dim] + hidden_states_pass = hidden_states[..., rotary_dim:] + cos = cos.to(hidden_states.dtype)[None, :, None, :] + sin = sin.to(hidden_states.dtype)[None, :, None, :] + x1, x2 = hidden_states_rotary.chunk(2, dim=-1) + hidden_states_rotated = torch.cat((-x2, x1), dim=-1) + hidden_states_rotary = hidden_states_rotary * cos + hidden_states_rotated * sin + return torch.cat((hidden_states_rotary, hidden_states_pass), dim=-1).contiguous() + + +def provider_linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """A bias-free ``nn.Linear`` of the block (Q/K/V, O and both FFN projections).""" + + return F.linear(x, weight) + + +def provider_head_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-5) -> torch.Tensor: + """``norm_q``/``norm_k``: ``nn.RMSNorm(128)`` on ``(B, S, H, D)``.""" + + return F.rms_norm(x, (x.shape[-1],), weight, eps) + + +def provider_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + """The native diffusers backend: SDPA on ``(B, H, S, D)`` views, non-causal, no mask. + + Returns ``(B, S, H*D)`` after ``flatten(2, 3).type_as(query)``. + """ + + q, k, v = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + return out.permute(0, 2, 1, 3).flatten(2, 3).type_as(query) + + +def provider_swiglu(projected: torch.Tensor) -> torch.Tensor: + """diffusers ``SwiGLU`` after its projection: ``hidden * silu(gate)`` on the two halves.""" + + hidden_states, gate = projected.chunk(2, dim=-1) + return hidden_states * F.silu(gate) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 9acee0c32..16b212e6d 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -100,7 +100,9 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. """ - row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} + row: dict[str, Any] = { + key: case[key] for key in ("op", "case", "backend", "bytes") if key in case + } keys = [key for key in TIMED_KEYS if key in case] if not keys: return row @@ -116,7 +118,7 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st for key in orders[iteration % 2]: _prepared_call(case[key], case.get(f"{key}_setup"))() torch.cuda.synchronize() - samples = {key: [] for key in keys} + samples: dict[str, list[float]] = {key: [] for key in keys} for iteration in range(iters): for key in orders[iteration % 2]: samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) @@ -163,6 +165,8 @@ def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: golden = registry._get_or_create_backend( registry._priority_map["cpu"]["timestep_sinusoid_h3"][-1] ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") rows = [] for num in (1, 2, 3, 7, 64, 1000, 4097): t = h3_timesteps(num, seed=num) @@ -227,6 +231,8 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: op = registry.get_op("timestep_mlp_fp32", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["timestep_mlp_fp32"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") sinusoid = registry.get_op("timestep_sinusoid_h3", device="cuda") params = h3_params(MLP_NAMES, MLP_SHAPES) cuda_err, provider_err = [], [] @@ -301,6 +307,8 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, golden = registry._get_or_create_backend( registry._priority_map["cpu"]["adaln_projection_3mod"][-1] ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") weight, bias = _adaln_params() cuda_frac, provider_frac, early_frac = [], [], [] for seed in range(draws): @@ -389,6 +397,8 @@ def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("adaln_row_gather", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") # Seeded so the op-backward fractions below are reproducible from this commit. gen = torch.Generator(device="cuda").manual_seed(0) rows = torch.randn(9, 6 * 5376, device="cuda", generator=gen).bfloat16() @@ -552,7 +562,7 @@ def golden(x_, w_, sh_, sc_): backward = {} for name, fn in ( ("cuda", lambda *t: op.forward_modulated(*t, index)), - ("provider", lambda *t: provider_norm_modulate(*t, index)), + ("provider", lambda *t: provider_norm_modulate(t[0], t[1], t[2], t[3], index)), ): first, second = grads(fn), grads(fn) backward[name] = { @@ -726,14 +736,16 @@ def backward(state): # read x + norm_out.linear, write out "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), - "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), + "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out( + t[0], t[1], t[2], t[3], t[4], i + ), "candidate_backward": backward, "provider_backward": backward, "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( lambda *t: op(*t, i) ), "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( - lambda *t: provider_final_adaln_out(*t, i) + lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], i) ), } ) @@ -745,6 +757,8 @@ def _final_accuracy(registry: KernelRegistry) -> dict[str, Any]: golden_op = registry._get_or_create_backend( registry._priority_map["cpu"]["final_adaln_out"][-1] ) + if golden_op is None: + raise RuntimeError("required H3 golden backend is unavailable") x, nw, temb, w, b, ti = _final_inputs(4097, seed=1) ours = op(x, nw, temb, w, b, ti) theirs = provider_final_adaln_out(x, nw, temb, w, b, ti) @@ -778,7 +792,7 @@ def golden_backward(x_, nw_, t_, w_, b_): backward = {} for name, fn in ( ("cuda", lambda *t: op(*t, ti)), - ("provider", lambda *t: provider_final_adaln_out(*t, ti)), + ("provider", lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], ti)), ): first, second = grads(fn), grads(fn) backward[name] = { diff --git a/rl_engine/validation/models/h3_weights.py b/rl_engine/validation/models/h3_weights.py index c6e5942c2..f587f2726 100644 --- a/rl_engine/validation/models/h3_weights.py +++ b/rl_engine/validation/models/h3_weights.py @@ -6,7 +6,9 @@ Only the tensors listed in ``h3_manifest.json`` are used. They all live in shard 1 of 14, so ``tools/weights/prepare_h3_weights.py`` downloads that shard, checks it against the manifest and extracts the tensors into -``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. +``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. With ``--block`` it also +extracts block 0's attention and FFN weights (manifest ``block_tensors``, +about 1.1 GB) into ``h3_block0.safetensors`` for ``ws1_one_h3_block``. """ from __future__ import annotations @@ -22,6 +24,7 @@ MANIFEST_PATH = Path(__file__).with_name("h3_manifest.json") WEIGHTS_ENV = "RL_KERNEL_H3_WEIGHTS" EXTRACTED_FILE = "h3_conditioning.safetensors" +BLOCK_FILE = "h3_block0.safetensors" def load_h3_manifest() -> dict[str, Any]: @@ -54,23 +57,24 @@ def sha256_tensor(tensor: torch.Tensor) -> str: return hashlib.sha256(memoryview(tensor_bytes)).hexdigest() -def load_h3_conditioning_weights( - device: torch.device | str = "cpu", names: list[str] | None = None +def _load_pinned( + filename: str, section: str, device: torch.device | str, names: list[str] | None ) -> dict[str, torch.Tensor]: - """Load the extracted pinned tensors; raise if they are missing or off-manifest.""" + """Load one extracted file and check every tensor of a manifest section against it.""" from safetensors.torch import load_file root = h3_weights_dir() if root is None: raise FileNotFoundError(f"set {WEIGHTS_ENV} to the prepare_h3_weights.py output dir") - path = root / EXTRACTED_FILE + path = root / filename if not path.is_file(): - raise FileNotFoundError(f"{path} missing; run tools/weights/prepare_h3_weights.py") - manifest = load_h3_manifest() + flag = " --block" if section == "block_tensors" else "" + raise FileNotFoundError(f"{path} missing; run tools/weights/prepare_h3_weights.py{flag}") + specs = load_h3_manifest()[section] tensors = load_file(str(path), device="cpu") - wanted = names if names is not None else list(manifest["tensors"]) - for name, spec in manifest["tensors"].items(): + wanted = names if names is not None else list(specs) + for name, spec in specs.items(): tensor = tensors[name] if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: raise ValueError(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") @@ -81,7 +85,24 @@ def load_h3_conditioning_weights( raise ValueError(f"{name}: sha256 {actual} != manifest {spec['sha256']}") out: dict[str, torch.Tensor] = {} for name in wanted: - if name not in manifest["tensors"]: + if name not in specs: raise KeyError(name) out[name] = tensors[name].to(device=device) return out + + +def load_h3_conditioning_weights( + device: torch.device | str = "cpu", names: list[str] | None = None +) -> dict[str, torch.Tensor]: + """Load the extracted pinned tensors; raise if they are missing or off-manifest.""" + + return _load_pinned(EXTRACTED_FILE, "tensors", device, names) + + +def load_h3_block_weights(device: torch.device | str = "cpu") -> dict[str, torch.Tensor]: + """Conditioning tensors plus block 0's attention/FFN weights, every one sha-checked.""" + + return { + **load_h3_conditioning_weights(device), + **_load_pinned(BLOCK_FILE, "block_tensors", device, None), + } diff --git a/rl_engine/validation/models/h3_ws2.py b/rl_engine/validation/models/h3_ws2.py index 7342774a0..9ec7b22c3 100644 --- a/rl_engine/validation/models/h3_ws2.py +++ b/rl_engine/validation/models/h3_ws2.py @@ -369,7 +369,7 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa ] started = [] deadline = time.monotonic() + timeout - status = {} + status: dict[int, object] = {} try: for proc in procs: proc.start() diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index 16e282514..15d7adc91 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -37,6 +37,7 @@ def make_operator_inputs( "matmul": _make_matmul_inputs, "det_gemm": _make_det_gemm_inputs, "attention": _make_attention_inputs, + "h3_full_attention": _make_h3_attention_inputs, "prefix_shared_attention": _make_prefix_shared_attention_inputs, "cp_attention": _make_cp_attention_inputs, "logp": _make_logp_inputs, @@ -77,6 +78,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "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}", + "h3_full_attention": f"{batch}x{seq}x56x128", "prefix_shared_attention": f"{batch}x{_arg_int(args, 'n_heads', DEFAULT_N_HEADS)}" f"x{seq}x{DEFAULT_HEAD_DIM}", "cp_attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}xcp2", @@ -600,3 +602,11 @@ def _arg_int(args: argparse.Namespace, name: str, default: int) -> int: def _arg_str(args: argparse.Namespace, name: str, default: str) -> str: return str(getattr(args, name, default)) + + +def _make_h3_attention_inputs(args, dtype, device): + batch, seq = _batch_seq(args) + return { + name: _floating_tensor((batch, seq, 56, 128), args, dtype, device, offset) + for offset, name in enumerate(("q", "k", "v")) + } diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index c0f174982..d5e6a7c85 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -32,6 +32,16 @@ def _load_object(path: str) -> Any: OP_SPECS = { + "h3_full_attention": OperatorSpec( + name="h3_full_attention", + op_class="attention", + gold_path=("rl_engine.reference.attention.standard_attn.H3FullAttentionReference"), + gold_method="forward_fp32", + candidate_paths={ + "cuda": "rl_engine.backends.cuda.attention.h3_full_attention.H3FullAttentionOp", + }, + grad_input_names=("q", "k", "v"), + ), "rms_norm": OperatorSpec( name="rms_norm", op_class="reduction", diff --git a/tests/backends/cuda/test_h3_full_attention.py b/tests/backends/cuda/test_h3_full_attention.py new file mode 100644 index 000000000..5ee28e8f0 --- /dev/null +++ b/tests/backends/cuda/test_h3_full_attention.py @@ -0,0 +1,344 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +import pytest +import torch + +from benchmarks.operators.attention.benchmark_deterministic_attention import H3StrictNative +from benchmarks.operators.attention.benchmark_deterministic_attention import ( + h3_bits_equal as bits_equal, +) +from benchmarks.operators.attention.benchmark_deterministic_attention import ( + h3_value_grad as value_grad, +) +from rl_engine.backends.cuda.attention.h3_full_attention import _C, H3FullAttentionOp +from rl_engine.contracts.numerical import load_contract, resolve_tolerance +from rl_engine.reference.attention.standard_attn import H3FullAttentionReference + + +def case_data(b, s, seed): + torch.manual_seed(seed) + return [ + torch.randn(b, s, 56, 128, device="cuda", dtype=torch.bfloat16).requires_grad_() + for _ in range(3) + ] + + +def assert_accuracy(actual, reference): + contract = load_contract() + for i, (a, g) in enumerate(zip(actual, reference)): + tol = resolve_tolerance( + contract, + op_class="attention", + dtype="bfloat16", + judgment="forward_accuracy" if i == 0 else "gradient_accuracy", + backend_profile="cuda_bf16", + ) + torch.testing.assert_close(a.to(g.dtype), g, atol=tol.atol, rtol=tol.rtol) + + +@pytest.fixture +def gpu_op(): + if not torch.cuda.is_available(): + pytest.skip("requires a GPU") + if torch.version.hip is not None or _C is None: + pytest.skip("requires the optional strict NVIDIA CUDA extension") + assert _C.h3_attention_build_info()["profile"] == "h3-cuda-strict-v1" + return H3FullAttentionOp + + +@pytest.mark.parametrize("value", [0, -1, True, 1.5]) +def test_invalid_chunk(value): + with pytest.raises(ValueError, match="positive integer"): + H3FullAttentionOp(value) + + +@pytest.mark.parametrize( + "shape,dtype", + [ + ((1, 2, 55, 128), torch.bfloat16), + ((1, 2, 56, 128), torch.float32), + ((1, 0, 56, 128), torch.bfloat16), + ], +) +def test_invalid_profile(shape, dtype): + q = torch.empty(shape, dtype=dtype) + with pytest.raises(ValueError): + H3FullAttentionOp()(q, q, q) + + +def test_cpu_fails_closed(): + q = torch.empty(1, 2, 56, 128, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="NVIDIA CUDA"): + H3FullAttentionOp()(q, q, q) + + +@pytest.mark.parametrize("platform", ["cpu", "rocm", "musa", "ascend"]) +def test_registry_rejects_unsupported_platform(monkeypatch, platform): + from rl_engine.runtime.registry import kernel_registry + + monkeypatch.setattr(kernel_registry, "_platform_for_device", lambda _: platform) + with pytest.raises(RuntimeError, match="no fallback"): + kernel_registry.get_op("h3_full_attention", device="cpu") + + +@pytest.mark.parametrize( + "seq,chunk,amplitude", + [ + (1, 1, 1), + (17, 16, 1), + (65, 64, 1), + (129, 128, 1), + (257, 128, 1), + (1024, 128, 1), + (1025, 128, 4), + (4096, 128, 1), + ], +) +def test_accuracy_and_chunk_invariance(gpu_op, seq, chunk, amplitude): + xs = [(x.detach() * amplitude).requires_grad_() for x in case_data(1, seq, 100 + seq)] + dy = torch.randn_like(xs[0]) + op = gpu_op(chunk) + actual = value_grad(op, xs, dy) + golden = H3FullAttentionReference().value_and_grad(*xs, dy) + references = [golden] + if seq <= 257: + + def fp64(q, k, v): + q, k, v = (x.double().transpose(1, 2) for x in (q, k, v)) + return ((q @ k.transpose(-1, -2) / 128**0.5).softmax(-1) @ v).transpose(1, 2) + + references.append( + value_grad(fp64, [x.detach().double().requires_grad_() for x in xs], dy.double()) + ) + for reference in references: + assert_accuracy(actual, reference) + other = value_grad(gpu_op(257), xs, dy) + assert all(bits_equal(a, b) for a, b in zip(actual, other)) + if _C is not None and torch.version.hip is None: + native = value_grad(lambda *x: H3StrictNative.apply(*x)[0], xs, dy) + assert all(bits_equal(a, b) for a, b in zip(actual, native)) + _, lse = op.forward_with_lse(*xs) + _, native_lse = H3StrictNative.apply(*xs) + assert bits_equal(lse, native_lse) + + +def test_batch_checkpoint_and_stream_invariance(gpu_op): + from torch.utils.checkpoint import checkpoint + + op = gpu_op(64) + xs = case_data(3, 129, 781) + dy = torch.randn_like(xs[0]) + expected = value_grad(op, xs, dy) + for inputs, grad, golden in ( + (xs, dy, expected), + ([x[1:2].detach().requires_grad_() for x in xs], dy[1:2], [x[1:2] for x in expected]), + ): + actual = value_grad(op, inputs, grad) + assert all(bits_equal(a, b) for a, b in zip(golden, actual)) + perm = torch.tensor([2, 0, 1], device="cuda") + actual = value_grad(op, [x[perm].detach().requires_grad_() for x in xs], dy[perm]) + assert all(bits_equal(a[perm], b) for a, b in zip(expected, actual)) + unrelated = [x.detach().clone().requires_grad_() for x in xs] + with torch.no_grad(): + for x in unrelated: + x[0].fill_(7) + x[2].fill_(-3) + changed = value_grad(op, unrelated, dy) + assert all(bits_equal(a[1], b[1]) for a, b in zip(expected, changed)) + # Real strided views, not merely a logically permuted contiguous copy. + strided = [] + for x in xs: + backing = torch.empty(3, 258, 56, 128, device="cuda", dtype=torch.bfloat16) + view = backing[:, ::2] + view.copy_(x.detach()) + strided.append(view.requires_grad_()) + with pytest.raises(ValueError, match="contiguous"): + op(*strided) + actual = value_grad(lambda *args: checkpoint(op, *args, use_reentrant=False), xs, dy) + assert all(bits_equal(a, b) for a, b in zip(expected, actual)) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + actual = value_grad(op, xs, dy) + torch.cuda.current_stream().wait_stream(stream) + assert all(bits_equal(a, b) for a, b in zip(expected, actual)) + + +@pytest.mark.parametrize("seq,chunk", [(1, 1), (7, 3), (17, 16)]) +def test_bounded_reference_vjp(seq, chunk): + torch.manual_seed(241) + xs = [torch.randn(1, seq, 56, 128, requires_grad=True) for _ in range(3)] + dy = torch.randn_like(xs[0]) + ref = H3FullAttentionReference() + out = ref(*xs) + expected = [out, *torch.autograd.grad(out, xs, dy)] + for actual, golden in zip(ref.value_and_grad(*xs, dy, chunk), expected): + torch.testing.assert_close(actual, golden, atol=2e-6, rtol=2e-5) + + +def test_gtest_registration(): + from argparse import Namespace + + from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case + + args = Namespace(op="h3_full_attention", candidate="cuda", arch_key="sm90", batch=1, seq=7) + case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) + assert case.inputs["q"].shape == (1, 7, 56, 128) + assert isinstance(make_candidate(args).fn, H3FullAttentionOp) + assert case.gold_fn(**case.inputs).dtype == torch.float32 + + +def test_strict_build_flags(monkeypatch): + from build_tools import envs + + for name in ( + "RL_KERNEL_H3_STRICT", + "KERNEL_ALIGN_USE_FAST_MATH", + "NVCC_PREPEND_FLAGS", + "NVCC_APPEND_FLAGS", + ): + monkeypatch.delenv(name, raising=False) + assert envs.h3_strict_cuda_flags() == [] + monkeypatch.setenv("RL_KERNEL_H3_STRICT", "1") + assert set(envs.h3_strict_cuda_flags()) == { + "-DRL_KERNEL_H3_STRICT=1", + "--fmad=false", + "--ftz=false", + "--prec-div=true", + "--prec-sqrt=true", + } + for name, value in ( + ("KERNEL_ALIGN_USE_FAST_MATH", "1"), + ("NVCC_PREPEND_FLAGS", "--use_fast_math"), + ("NVCC_APPEND_FLAGS", "--fmad=true"), + ): + monkeypatch.setenv(name, value) + with pytest.raises(ValueError): + envs.h3_strict_cuda_flags() + monkeypatch.delenv(name) + + +@pytest.fixture +def extension_config(monkeypatch): + from build_tools import extensions + + for name in ( + "RL_KERNEL_H3_STRICT", + "KERNEL_ALIGN_USE_FAST_MATH", + "NVCC_PREPEND_FLAGS", + "NVCC_APPEND_FLAGS", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(extensions, "_load_torch_extension_tools", lambda: (torch, None, dict)) + monkeypatch.setattr(extensions, "_musa_build_available", lambda _: False) + monkeypatch.setattr(extensions, "_ascend_extensions", lambda: []) + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (8, 0)) + return extensions.get_extensions + + +def test_strict_extension_isolation(monkeypatch, extension_config): + baseline = extension_config() + monkeypatch.setenv("RL_KERNEL_H3_STRICT", "1") + ordinary, h3 = extension_config() + assert baseline == [ordinary] + assert ordinary["name"] == "rl_engine._C" + assert h3["name"] == "rl_engine._H3_C" + assert "-DRL_KERNEL_H3_STRICT=1" not in ordinary["extra_compile_args"]["nvcc"] + flags = h3["extra_compile_args"]["nvcc"] + assert "-DRL_KERNEL_H3_STRICT=1" in flags + assert "--fmad=false" in flags and "--ftz=false" in flags + assert "--use_fast_math" not in flags + + +def test_strict_extension_rejects_rocm(monkeypatch, extension_config): + monkeypatch.setenv("RL_KERNEL_H3_STRICT", "1") + monkeypatch.setattr(torch.version, "hip", "test") + with pytest.raises(RuntimeError, match="NVIDIA CUDA only"): + extension_config() + + +def test_unqualified_extension_rejected(monkeypatch): + from types import SimpleNamespace + + from rl_engine.backends.cuda.attention import h3_full_attention as h3 + + fake = SimpleNamespace( + h3_attention_forward=lambda: None, + h3_attention_backward_accumulate=lambda: None, + h3_attention_probabilities=lambda: None, + h3_attention_build_info=lambda: {"profile": "unqualified"}, + ) + monkeypatch.setattr(h3, "_C", fake) + with pytest.raises(RuntimeError, match="RL_KERNEL_H3_STRICT"): + h3._extension() + + +@pytest.mark.parametrize( + "device,dtype,world,cp", + [ + ("cuda", "bfloat16", 1, 1), + ("cpu", "bfloat16", 1, 1), + ("rocm", "bfloat16", 1, 1), + ("cuda", "float16", 1, 1), + ("cuda", "bfloat16", 2, 1), + ("cuda", "bfloat16", 1, 2), + ], +) +def test_semantic_dispatch_contract(device, dtype, world, cp): + from rl_engine.runtime.registry import kernel_registry + from rl_engine.runtime.semantic_registry import OperatorRequirements, OperatorResolutionError + + session = kernel_registry.semantic.session() + kwargs = dict( + semantic_op="h3_full_attention", + requested_backend=H3FullAttentionOp.backend_id, + target="training", + requirements=OperatorRequirements( + device=device, + dtype=dtype, + topology={"world_size": world, "tensor_parallel_size": 1, "context_parallel_size": cp}, + alignment_properties={"deterministic": True}, + ), + ) + if (device, dtype, world, cp) == ("cuda", "bfloat16", 1, 1): + resolution = session.resolve(**kwargs) + instance = session.instantiate(resolution) + assert isinstance(instance, H3FullAttentionOp) + assert ( + session.instance_provenance(resolution, instance).backend_id + == H3FullAttentionOp.backend_id + ) + else: + with pytest.raises(OperatorResolutionError): + session.resolve(**kwargs) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None, reason="requires NVIDIA CUDA" +) +def test_cuda_build_and_invalid_inputs(): + if _C is None: + with pytest.raises(RuntimeError, match="extension unavailable"): + H3FullAttentionOp()(*case_data(1, 17, 9)) + return + assert hasattr(_C, "h3_attention_build_info") + xs = case_data(1, 17, 9) + op = H3FullAttentionOp() + if _C.h3_attention_build_info()["profile"] != "h3-cuda-strict-v1": + with pytest.raises(RuntimeError, match="RL_KERNEL_H3_STRICT"): + op(*xs) + with pytest.raises(RuntimeError, match="RL_KERNEL_H3_STRICT"): + _C.h3_attention_forward(*(x.transpose(1, 2) for x in xs)) + return + for bad in (xs[1].cpu(), xs[1].float(), xs[1][..., :64], xs[1][:, :0]): + with pytest.raises(ValueError): + op(xs[0], bad, xs[2]) + q, k, v = [x.transpose(1, 2).contiguous() for x in xs] + with pytest.raises(RuntimeError): + _C.h3_attention_probabilities(q, k[:, :0], v[:, :0]) + p = _C.h3_attention_probabilities(q, k, v) + acc = torch.zeros_like(k, dtype=torch.float32) + with pytest.raises(RuntimeError): + _C.h3_attention_backward_accumulate(q, q, k, v, p, acc, acc) diff --git a/tests/models/minimax_h3/conftest.py b/tests/models/minimax_h3/conftest.py index 8bc18d0df..f9c2b000f 100644 --- a/tests/models/minimax_h3/conftest.py +++ b/tests/models/minimax_h3/conftest.py @@ -16,6 +16,7 @@ from rl_engine.validation.models.h3_weights import ( WEIGHTS_ENV, h3_weights_dir, + load_h3_block_weights, load_h3_conditioning_weights, ) @@ -32,6 +33,20 @@ def h3_weights_cpu() -> dict[str, torch.Tensor]: pytest.skip(str(exc)) +@pytest.fixture(scope="session") +def h3_block_weights_cuda() -> dict[str, torch.Tensor]: + """Conditioning plus block-0 tensors on CUDA, skipping when they were not extracted.""" + + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py --block") + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + try: + return load_h3_block_weights("cuda") + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + def h3_kernels_built() -> bool: """False for KERNEL_ALIGN_USE_FAST_MATH=1 builds, which omit every H3 kernel.""" return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) diff --git a/tests/models/minimax_h3/test_h3_adaln_row_gather.py b/tests/models/minimax_h3/test_h3_adaln_row_gather.py index 2b7df3803..0f9fec6c3 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -149,7 +149,8 @@ def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): """Reject semantic index violations without poisoning pytest's CUDA context.""" _cuda_op() - probe = textwrap.dedent(""" + probe = textwrap.dedent( + """ import sys import torch from rl_engine.backends.extension import _C @@ -167,7 +168,8 @@ def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): else: H3AdaLNRowGatherCudaOp().forward(rows, ti, tags, check_range=False) torch.cuda.synchronize() - """) + """ + ) completed = subprocess.run( [sys.executable, "-c", probe, entrypoint, str(timestep), str(tag), index_dtype], cwd=Path(__file__).resolve().parents[3], diff --git a/tests/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py index f9bc79829..ca11c4e53 100644 --- a/tests/models/minimax_h3/test_h3_det_linear.py +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -8,8 +8,8 @@ import pytest import torch -from rl_engine.backends.extension import _C from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available +from rl_engine.backends.extension import _C pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), diff --git a/tests/models/minimax_h3/test_h3_native_validation.py b/tests/models/minimax_h3/test_h3_native_validation.py index bcb06f1f2..03a70f1a3 100644 --- a/tests/models/minimax_h3/test_h3_native_validation.py +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -286,7 +286,8 @@ def test_native_table_bounds_fail_in_isolated_process(native_extension, operatio import sys import textwrap - script = textwrap.dedent(""" + script = textwrap.dedent( + """ import sys import torch from rl_engine import _C @@ -318,7 +319,8 @@ def test_native_table_bounds_fail_in_isolated_process(native_extension, operatio sys.exit(0) raise raise AssertionError("invalid index reached native table lookup without an assertion") - """) + """ + ) result = subprocess.run( [sys.executable, "-c", script, operation, str(bad_index)], capture_output=True, @@ -352,7 +354,8 @@ def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metada import sys import textwrap - script = textwrap.dedent(""" + script = textwrap.dedent( + """ import sys import torch from rl_engine import _C @@ -372,7 +375,8 @@ def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metada sys.exit(0) raise raise AssertionError("invalid metadata was not rejected") - """) + """ + ) result = subprocess.run( [sys.executable, "-c", script, metadata, bound], capture_output=True, text=True, timeout=60 ) diff --git a/tests/models/minimax_h3/test_h3_one_block.py b/tests/models/minimax_h3/test_h3_one_block.py new file mode 100644 index 000000000..8f64a4508 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_one_block.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One real MiniMax-H3 block, forward and backward (RFC #420 ``ws1_one_h3_block``). + +GPU tests run block 0 of the pinned checkpoint on packed FL2VA layouts and +check the block-level promises of ``rl_engine.validation.models.h3_block``: + +* every node is bitwise repeatable, and batch row 0 of a two-row batch has the + bytes of the single-row run (forward, and the gradient reaching every node); +* nodes that promise it are bitwise equal to diffusers on diffusers' inputs; +* every node, the block output and every gradient are about as close to the + FP32 golden as diffusers itself. + +CPU tests pin the packed layouts to the pinned diffusers pipeline and check +the graph's structure. +""" + +from __future__ import annotations + +import hashlib + +import pytest +import torch + +from rl_engine.reference.minimax_h3 import block as gold +from rl_engine.validation.models import h3_block +from rl_engine.validation.models import h3_provider as prov +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS, h3_fl2va_layout + +needs_gpu = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) + +# sha256(position_ids float64 bytes + token_tags int64 bytes) of diffusers@f53d552's +# ``MiniMaxH3PrepareLayoutStep.build_packed_sequence`` for each layout (text tags 1, patch +# (1, 2, 2), two audio channels, audio tag 2, video tag 0, no keyframe anchors). +PINNED_LAYOUTS = { + "tiny": "5329a0a13a9996c4fea7a0e676b65aa24302d65bea2e2248ba68d12184cc7c37", + "small": "64e4e79d8fb91d5a68ae7d15bb29049e08add38da2639030e089c6ef80968931", + "medium": "16e448800695eedfb85e0aacbca9bd5415ed53fbbe25f701bcb60df5f4754a31", +} + +# Accuracy: the candidate's error against the FP32 golden may exceed diffusers' by this +# factor plus an absolute floor (relative to the golden's max magnitude). +FORWARD_FACTOR, FORWARD_FLOOR = 1.25, 1e-3 +BACKWARD_FACTOR, BACKWARD_FLOOR = 1.5, 5e-3 +# A BF16 leaf gradient may instead exceed diffusers' error by this many ulps of its largest +# golden element. Its max|err| / max|golden| is set by one rounded element, so it moves in +# whole ulps (0.004 relative for attn.norm_q.weight), and the relative rule alone failed on +# seed 2 with torch 2.13 and on seed 0 with torch 2.9.1. Sized from +# tools/validation/models/h3_block_ulp_sweep.py over tiny/small x seeds 0-9 on both torch +# versions (max excess +4.15 ulp) and 34 single-node candidate/provider swaps (max +5.15 ulp); +# reports in reports/experiments/h3-one-block-b200/ulp-sweep-torch*.json. +BACKWARD_BF16_EXCESS_ULPS = 6.0 + + +@pytest.mark.parametrize("layout", sorted(PINNED_LAYOUTS)) +def test_layout_matches_pinned_pipeline(layout): + packed = h3_fl2va_layout(**H3_BLOCK_LAYOUTS[layout], device="cpu") + data = packed["position_ids"].numpy().tobytes() + packed["token_tags"].numpy().tobytes() + assert hashlib.sha256(data).hexdigest() == PINNED_LAYOUTS[layout] + + +def test_graph_is_well_formed(): + seen = {"hidden", "temb"} + for node in h3_block.NODES: + assert set(node.inputs) <= seen, node.name + assert node.name not in seen, node.name + assert node.binding.status in ("row", "interim", "reference"), node.name + seen.add(node.name) + assert h3_block.OUTPUT == "residual_mlp" + rows = {n.binding.rfc_row for n in h3_block.NODES if n.binding.status == "row"} + assert rows == {"adaln_projection_3mod", "h3_rmsnorm", "adaln_gate_residual"} + + +def test_partial_rope_rotates_96_of_128_channels(): + packed = h3_fl2va_layout(**H3_BLOCK_LAYOUTS["tiny"], device="cpu") + cos, sin = prov.provider_rope(packed["position_ids"]) + assert cos.shape == (packed["position_ids"].shape[0], 96) and cos.dtype == torch.float32 + gold_cos, gold_sin = gold.golden_rope_tables(packed["position_ids"]) + assert torch.equal(cos, gold_cos) and torch.equal(sin, gold_sin) + x = torch.randn(1, cos.shape[0], 2, 128).bfloat16() + for out in (prov.provider_apply_rotary(x, cos, sin), gold.golden_apply_rotary(x, cos, sin)): + assert torch.equal(out[..., 96:].float(), x[..., 96:].float()) + assert not torch.equal(out[..., :96].float(), x[..., :96].float()) + + +@pytest.fixture(scope="module") +def ops(): + return h3_block.CandidateOps() + + +@needs_gpu +def test_stack_rows_dispatch_to_cuda(require_h3_kernels, ops): + backends = ops.backends() + for name in ("projection", "norm", "gate"): + assert backends[name].endswith("CudaOp"), backends + + +@needs_gpu +@pytest.mark.parametrize("layout", ["tiny", "small"]) +def test_block_forward(h3_block_weights_cuda, ops, layout): + report = h3_block.run_block_case(h3_block_weights_cuda, layout=layout, ops=ops) + nodes = {node.name: node for node in h3_block.NODES} + for entry in report["nodes"]: + name = entry["node"] + assert entry["repeat_bitwise_equal"], name + assert entry["batch_row_bitwise_equal"], name + if nodes[name].provider_bitwise_isolated: + assert entry["isolated_vs_provider"]["bitwise_equal"], name + bound = FORWARD_FACTOR * entry["provider_rel_err_vs_golden"] + FORWARD_FLOOR + assert entry["candidate_rel_err_vs_golden"] <= bound, entry + assert report["first_repeat_drift"] is None + assert report["first_batch_drift"] is None + drift = report["first_isolated_drift"] + assert drift is None or not nodes[drift].provider_bitwise_isolated + out = report["output"] + assert out["candidate_rel_err_vs_golden"] <= ( + FORWARD_FACTOR * out["provider_rel_err_vs_golden"] + FORWARD_FLOOR + ) + + +@needs_gpu +def test_block_backward(h3_block_weights_cuda, ops): + report = h3_block.run_block_backward(h3_block_weights_cuda, layout="tiny", ops=ops) + for name, entry in report["leaves"].items(): + assert entry["repeat_bitwise_equal"], name + bound = BACKWARD_FACTOR * entry["provider_rel_err_vs_golden"] + BACKWARD_FLOOR + within_rel = entry["candidate_rel_err_vs_golden"] <= bound + within_ulps = entry["storage_dtype"] == "bfloat16" and ( + entry["candidate_ulps_vs_golden"] + <= entry["provider_ulps_vs_golden"] + BACKWARD_BF16_EXCESS_ULPS + ) + assert within_rel or within_ulps, (name, entry) + assert report["leaves"]["hidden"]["batch_row_bitwise_equal"] + assert report["first_grad_repeat_drift"] is None + assert report["first_grad_batch_drift"] is None diff --git a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py index 4866e68f3..a60489a04 100644 --- a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py +++ b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py @@ -284,7 +284,9 @@ def _worker_never_reports(rank, world, init, queue, path, target, kwargs): @pytest.mark.parametrize("mode", ["clean_exit", "crash", "reported_error", "timeout"]) def test_run_world_cleans_up_failed_workers(monkeypatch, mode): import time + import torch.multiprocessing as mp + from rl_engine.validation.models import h3_ws2 workers = { diff --git a/tools/migration/layout.json b/tools/migration/layout.json index 8d469480c..6c7fc9898 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -477,6 +477,8 @@ "docs/usage/evidence/h3-final-adaln-out-b200/chain_replay.json": "reports/experiments/h3-final-adaln-out-b200/chain_replay.json", "docs/usage/evidence/h3-final-adaln-out-b200/figure.png": "reports/experiments/h3-final-adaln-out-b200/figure.png", "docs/usage/evidence/h3-final-adaln-out-b200/report.json": "reports/experiments/h3-final-adaln-out-b200/report.json", + "docs/usage/evidence/h3-one-block-b200/figure.png": "reports/experiments/h3-one-block-b200/figure.png", + "docs/usage/evidence/h3-one-block-b200/report.json": "reports/experiments/h3-one-block-b200/report.json", "docs/usage/evidence/h3-prior-art-b200/adaln_projection.json": "reports/experiments/h3-prior-art-b200/adaln_projection.json", "docs/usage/evidence/h3-prior-art-b200/adaln_projection.png": "reports/experiments/h3-prior-art-b200/adaln_projection.png", "docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json": "reports/experiments/h3-prior-art-b200/adaln_row_gather.json", diff --git a/tools/validation/models/h3_block_diagnostic.py b/tools/validation/models/h3_block_diagnostic.py new file mode 100644 index 000000000..19e5b111a --- /dev/null +++ b/tools/validation/models/h3_block_diagnostic.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Localize a failed block-0 gradient gate without changing its numerical contract.""" +import argparse +import json +import os +import subprocess +from dataclasses import replace +from pathlib import Path + +import torch + +from rl_engine.backends.extension import _C +from rl_engine.validation.models import h3_block as block +from rl_engine.validation.models.h3_weights import load_h3_block_weights +from tools.validation.models.h3_full_model import digest + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--failed-job", required=True) + parser.add_argument("--causal-sweep", action="store_true") + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=False) + if subprocess.check_output(["git", "status", "--porcelain"], text=True).strip(): + raise ValueError("diagnostic requires clean source") + ops = block.CandidateOps() + weights = load_h3_block_weights("cuda") + ctx = block.make_block_context(weights, layout="tiny", seed=0, batch=1) + upstream = ( + torch.randn((1, ctx["seq_len"], block.HIDDEN), generator=torch.Generator().manual_seed(11)) + .bfloat16() + .cuda() + ) + report = dict( + status="diagnostic_only_not_acceptance", + failed_job=args.failed_job, + job_id=os.environ.get("SLURM_JOB_ID"), + commit=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + extension_sha256=digest(Path(_C.__file__)), + torch=torch.__version__, + device=torch.cuda.get_device_name(), + backends=ops.backends(), + default_matmul_tf32=torch.backends.cuda.matmul.allow_tf32, + bf16_reduced_precision_reduction=( + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction + ), + leaves={}, + nodes=[], + ) + # Keep FP32 arithmetic strict through backward as well as forward. + with block.strict_fp32_matmul(): + golden, golden_nodes = block._backward("golden", ctx, None, upstream) + provider, provider_nodes = block._backward("provider", ctx, None, upstream) + candidate, candidate_nodes = block._backward("candidate", ctx, ops, upstream) + for name in golden: + report["leaves"][name] = dict( + candidate=block.rel_err(candidate[name], golden[name]), + provider=block.rel_err(provider[name], golden[name]), + ) + for node in reversed(block.NODES): + report["nodes"].append( + dict( + node=node.name, + errors=[ + dict(candidate=block.rel_err(c, g), provider=block.rel_err(p, g)) + for g, p, c in zip( + golden_nodes[node.name], + provider_nodes[node.name], + candidate_nodes[node.name], + strict=True, + ) + ], + ) + ) + # Same input and upstream gradient isolate the RMSNorm implementation from + # error accumulated through other candidate nodes. + with torch.no_grad(): + values = block.run_graph("provider", ctx) + isolated = {} + for name, source, weight_name in ( + ("q_norm", "q_proj", "attn.norm_q.weight"), + ("k_norm", "k_proj", "attn.norm_k.weight"), + ): + grads = {} + node = next(n for n in block.NODES if n.name == name) + for mode in ("golden", "provider", "candidate"): + x = values[source].detach().clone() + w = weights[block.BLOCK + "." + weight_name].detach().clone() + if mode == "golden": + x, w = x.float(), w.float() + x.requires_grad_() + w.requires_grad_() + local = {**ctx, "weights": {**weights, block.BLOCK + "." + weight_name: w}} + y = ( + node.candidate(ops, local, x) + if mode == "candidate" + else getattr(node, mode)(local, x) + ) + torch.autograd.backward(y, provider_nodes[name][0].to(y.dtype)) + grads[mode] = w.grad + isolated[name] = dict( + candidate=block.rel_err(grads["candidate"], grads["golden"]), + provider=block.rel_err(grads["provider"], grads["golden"]), + ) + report["isolated_norm_weight_gradients"] = isolated + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + if args.causal_sweep: + # These hybrids intentionally use provider nodes for localization only. + # They cannot be submitted to the candidate acceptance comparator. + report["causal_sweep"] = [] + original_nodes = block.NODES + try: + for node in original_nodes: + for baseline in ("candidate", "provider"): + mixed = [] + for item in original_nodes: + use_provider = (item.name == node.name) == (baseline == "candidate") + mixed.append( + replace( + item, + candidate=lambda ops, ctx, *xs, item=item: item.provider(ctx, *xs), + ) + if use_provider + else item + ) + block.NODES = tuple(mixed) + with block.strict_fp32_matmul(): + leaves, _ = block._backward("candidate", ctx, ops, upstream) + report["causal_sweep"].append( + dict( + baseline=baseline, + replaced_node=node.name, + errors={ + name: block.rel_err(leaves[name], golden[name]) for name in golden + }, + ) + ) + del leaves + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + finally: + block.NODES = original_nodes + report["causal_sweep_complete"] = True + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_replay.py b/tools/validation/models/h3_block_replay.py new file mode 100644 index 000000000..7d8230294 --- /dev/null +++ b/tools/validation/models/h3_block_replay.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay one real MiniMax-H3 block forward and backward and write JSON evidence (RFC #420). + +See ``rl_engine/validation/models/h3_block.py`` for what is compared. + + python tools/weights/prepare_h3_weights.py --out --block + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_replay.py --out block_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_block_weights, + load_h3_manifest, +) + + +def main() -> None: + """Run the forward (and optionally backward) block replay per layout and save the report.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--layouts", + default="tiny,small,medium", + help=f"comma list of packed FL2VA layouts: {', '.join(H3_BLOCK_LAYOUTS)}", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--no-backward", action="store_true", help="forward replay only") + parser.add_argument("--timing-repeats", type=int, default=5) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + layouts = args.layouts.split(",") + unknown = sorted(set(layouts) - set(H3_BLOCK_LAYOUTS)) + if unknown: + parser.error(f"unknown layouts: {unknown}") + if not torch.cuda.is_available(): + raise SystemExit("the block replay needs a CUDA device") + + weights = load_h3_block_weights("cuda") + manifest = load_h3_manifest() + ops = h3_block.CandidateOps() + + forward, backward = [], [] + for layout in layouts: + case = h3_block.run_block_case( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + forward.append(case) + print( + f"{layout} S={case['seq_len']}: first_drift={case['first_drift']} " + f"first_isolated_drift={case['first_isolated_drift']} " + f"repeat_drift={case['first_repeat_drift']} batch_drift={case['first_batch_drift']} " + f"out err cand {case['output']['candidate_rel_err_vs_golden']:.2e} " + f"prov {case['output']['provider_rel_err_vs_golden']:.2e}" + ) + if not args.no_backward: + case = h3_block.run_block_backward( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + backward.append(case) + det = all(v["repeat_bitwise_equal"] for v in case["leaves"].values()) + print( + f"{layout} backward: leaves {'det' if det else 'NONDET'}, " + f"grad_repeat_drift={case['first_grad_repeat_drift']} " + f"grad_batch_drift={case['first_grad_batch_drift']} " + f"dX batch-row {case['leaves']['hidden']['batch_row_bitwise_equal']}" + ) + torch.cuda.empty_cache() + + evidence = { + "kind": "h3_one_block_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "nodes": [ + { + "node": n.name, + "inputs": list(n.inputs), + "rfc_row": n.binding.rfc_row, + "status": n.binding.status, + "backend": n.binding.backend, + } + for n in h3_block.NODES + ], + "forward_cases": forward, + "backward_cases": backward, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_ulp_sweep.py b/tools/validation/models/h3_block_ulp_sweep.py new file mode 100644 index 000000000..7fe66646b --- /dev/null +++ b/tools/validation/models/h3_block_ulp_sweep.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sweep the one-block backward over layouts and seeds and record leaf errors in ulps. + +Backs the ulp allowance of ``tests/models/minimax_h3/test_h3_one_block.py``: +for every leaf gradient it records the candidate's and diffusers' error against +the FP32 golden both as ``max|err| / max|golden|`` and in ulps of the leaf's +storage dtype at ``max|golden|``, and the per-leaf excess +``candidate_ulps - provider_ulps``. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_ulp_sweep.py --layouts tiny,small --seeds 0,1,2,3,4 \\ + --out ulp_sweep.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_weights import load_h3_block_weights # noqa: E402 + + +def main() -> None: + """Run the backward replay for every (layout, seed) and write per-leaf ulp errors.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--layouts", default="tiny,small") + parser.add_argument("--seeds", default="0,1,2,3,4") + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("the sweep needs a CUDA device") + + weights = load_h3_block_weights("cuda") + ops = h3_block.CandidateOps() + cases = [] + for layout in args.layouts.split(","): + for seed in (int(s) for s in args.seeds.split(",")): + report = h3_block.run_block_backward(weights, layout=layout, seed=seed, ops=ops) + leaves = { + name: { + key: entry[key] + for key in ( + "storage_dtype", + "candidate_rel_err_vs_golden", + "provider_rel_err_vs_golden", + "candidate_ulps_vs_golden", + "provider_ulps_vs_golden", + ) + } + | { + "excess_ulps": entry["candidate_ulps_vs_golden"] + - entry["provider_ulps_vs_golden"] + } + for name, entry in report["leaves"].items() + } + worst = max(leaves, key=lambda n: leaves[n]["excess_ulps"]) + print( + f"{layout} S={report['seq_len']} seed={seed}: worst excess " + f"{leaves[worst]['excess_ulps']:+.2f} ulp ({worst})" + ) + cases.append( + {"layout": layout, "seq_len": report["seq_len"], "seed": seed, "leaves": leaves} + ) + torch.cuda.empty_cache() + + bf16 = [ + leaf["excess_ulps"] + for case in cases + for leaf in case["leaves"].values() + if leaf["storage_dtype"] == "bfloat16" + ] + evidence = { + "kind": "h3_one_block_backward_ulp_sweep", + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "cases": cases, + "max_bf16_excess_ulps": max(bf16), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"max BF16 excess: {max(bf16):+.2f} ulp; wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_evidence_integrity.py b/tools/validation/models/h3_evidence_integrity.py new file mode 100644 index 000000000..42d9b383d --- /dev/null +++ b/tools/validation/models/h3_evidence_integrity.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Validate immutable reference inventories before consuming any tensor evidence.""" +from __future__ import annotations + +import json +from pathlib import Path + +from tools.validation.models.h3_full_model import digest + + +def verify_inventory(root: Path, manifest: Path): + records = json.loads(manifest.read_text())["files"] + for name, expected in records.items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()): + raise ValueError("evidence path escapes root") + if digest(path) != expected: + raise ValueError(f"reference evidence changed: {name}") + + +def verify_generation(root): + report = json.loads((root / "report.json").read_text()) + if report["status"] != "generation_complete_comparison_pending": + raise ValueError("generation did not complete") + required = {"initial.pt", "video.npy", "audio.npy"} | { + f"step-{i:03}.pt" for i in range(report["steps"]) + } + if not required.issubset(report["artifacts"]): + raise ValueError("incomplete generation artifact inventory") + for name, expected in report["artifacts"].items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()) or digest(path) != expected: + raise ValueError(f"generation artifact changed: {name}") + return report diff --git a/tools/validation/models/h3_first_step_diagnostic.py b/tools/validation/models/h3_first_step_diagnostic.py new file mode 100644 index 000000000..e433b8481 --- /dev/null +++ b/tools/validation/models/h3_first_step_diagnostic.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay the first denoising step of a generation under one execution configuration. + +The first transformer call of a generation sees the saved packed inputs (``initial.pt``) +and nothing has fed back yet, so the difference to the provider's saved ``step-000.pt`` +is the model's own error on a real generation input. A configuration is the provider +(optionally with one forced SDPA backend or cuBLASLt), the candidate, or the candidate +with chosen block nodes run as their provider replay (``--provider-nodes``). Replaying +the provider default must reproduce ``step-000.pt`` bitwise; ``baseline_bitwise`` records +whether it did, which validates the replay itself. Diagnostic only: nothing here accepts +or rejects a candidate. + + python -m tools.validation.models.h3_first_step_diagnostic --checkpoint \\ + --generation --seed 43 --mode candidate \\ + --provider-nodes attention --output out.json +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch + +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution +from tools.validation.models.h3_full_generation import QueryChunkAttention +from tools.validation.models.h3_full_model import PROVIDER_VARIANTS, provider_variant_context + + +def errors(value: torch.Tensor, reference: torch.Tensor) -> dict[str, float]: + """max_abs and relative_l2 of ``value`` against ``reference`` (float64 norms).""" + + delta = value.double() - reference.double() + return dict( + max_abs=float(delta.abs().max()), + relative_l2=float(delta.norm() / reference.double().norm().clamp_min(1e-12)), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--generation", type=Path, required=True) + parser.add_argument("--contract", type=Path, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--mode", choices=["provider", "candidate"], required=True) + parser.add_argument("--provider-variant", choices=PROVIDER_VARIANTS, default="default") + parser.add_argument("--provider-nodes", default="") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.mode == "candidate" and args.provider_variant != "default": + parser.error("--provider-variant applies to --mode provider only") + if args.mode == "provider" and args.provider_nodes: + parser.error("--provider-nodes applies to --mode candidate only") + nodes = tuple(n for n in args.provider_nodes.split(",") if n) + + from diffusers import MiniMaxH3Transformer3DModel + + directory = args.generation / str(args.seed) + inputs = torch.load(directory / "initial.pt", map_location="cpu", weights_only=True) + saved = torch.load(directory / "step-000.pt", map_location="cpu", weights_only=True) + contract = json.loads(args.contract.read_text()) + + started = time.perf_counter() + model = MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), torch_dtype=torch.bfloat16, local_files_only=True + ) + model = model.to("cuda").eval() + loaded = time.perf_counter() - started + kwargs = {k: v.to("cuda") for k, v in inputs.items()} + ops = None + if args.mode == "candidate": + ops = CandidateOps() + ops.attention = QueryChunkAttention(ops.attention, contract["attention_query_chunk"]) + + trace = FullModelTrace() + started = time.perf_counter() + with ( + torch.inference_mode(), + full_model_execution( + model, + candidate=args.mode == "candidate", + trace=trace, + ops=ops, + provider_nodes=nodes, + ), + provider_variant_context(args.provider_variant), + ): + video, audio = model(**kwargs, attention_kwargs=None, return_dict=False) + torch.cuda.synchronize() + forward_seconds = time.perf_counter() - started + video, audio = video.float().cpu(), audio.float().cpu() + result = dict( + seed=args.seed, + mode=args.mode, + provider_variant=args.provider_variant, + provider_nodes=list(nodes), + video=errors(video, saved[0]), + audio=errors(audio, saved[1]), + baseline_bitwise=bool(torch.equal(video, saved[0]) and torch.equal(audio, saved[1])), + load_seconds=loaded, + forward_seconds=forward_seconds, + torch=torch.__version__, + device=torch.cuda.get_device_name(), + ) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_block.py b/tools/validation/models/plot_h3_block.py new file mode 100644 index 000000000..228f3a289 --- /dev/null +++ b/tools/validation/models/plot_h3_block.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render the RFC #420 one-block replay (``tools/validation/models/h3_block_replay.py``) as a PNG. + +Needs only the report JSON and matplotlib: + + python tools/validation/models/plot_h3_block.py \ + reports/experiments/h3-one-block-b200/report.json + +writes ``figure.png`` next to the report. Panels: the per-node invariance and +provider checks over every layout; the per-node error against the FP32 golden +for the candidate and diffusers; the same for every backward leaf. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.colors import ListedColormap # noqa: E402 + +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +PASS = "#1baf7a" +FAIL = "#d64545" +NA = "#e4e3dd" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" +STATUS_MARK = {"row": "", "interim": "*", "reference": "†"} + + +def _style(ax, title: str) -> None: + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=10.5, color=INK, pad=8) + ax.tick_params(colors=MUTED, labelsize=7.5, length=0) + for side in ("top", "right", "left", "bottom"): + ax.spines[side].set_visible(False) + + +def _checks_panel(ax, report: dict[str, Any]) -> None: + """Rows: nodes. Columns: forward/backward checks, ANDed over every layout.""" + + nodes = [n["node"] for n in report["nodes"]] + status = {n["node"]: n["status"] for n in report["nodes"]} + cols = [ + ("repeat", "forward", "repeat_bitwise_equal"), + ("batch row", "forward", "batch_row_bitwise_equal"), + ("= diffusers\n(isolated)", "forward", None), + ("grad\nrepeat", "backward", "grad_repeat_bitwise_equal"), + ("grad\nbatch row", "backward", "grad_batch_row_bitwise_equal"), + ] + grid = [] + for node in nodes: + row = [] + for _label, kind, key in cols: + cases = report["forward_cases" if kind == "forward" else "backward_cases"] + entries = [e for case in cases for e in case["nodes"] if e["node"] == node] + if key is None: + promised = entries[0]["provider_bitwise_isolated_promised"] + ok = all(e["isolated_vs_provider"]["bitwise_equal"] for e in entries) + row.append(2 if not promised and not ok else (1 if ok else 0)) + else: + values = [e[key] for e in entries] + row.append(2 if not values or None in values else (1 if all(values) else 0)) + grid.append(row) + ax.imshow(grid, cmap=ListedColormap([FAIL, PASS, NA]), vmin=0, vmax=2, aspect="auto") + ax.set_xticks(range(len(cols)), [c[0] for c in cols], fontsize=7.5, color=INK) + ax.set_yticks( + range(len(nodes)), [f"{n}{STATUS_MARK[status[n]]}" for n in nodes], fontsize=7.5, color=INK + ) + ax.xaxis.tick_top() + _style(ax, "Bitwise checks, all layouts") + ax.set_xlabel( + "green = holds grey = not promised (reduction differs from diffusers)\n" + "or not applicable (AdaLN tables have no batch axis)\n" + "* interim operator † provider replay (no kernel yet)", + fontsize=7, + color=MUTED, + ) + + +def _err_panel(ax, labels, cand, prov, title: str) -> None: + xs = range(len(labels)) + ax.plot( + xs, [math.log10(v) for v in cand], "o-", color=RL_KERNEL, lw=1.8, ms=5, label="RL-Kernel" + ) + ax.plot( + xs, [math.log10(v) for v in prov], "s--", color=PROVIDER, lw=1.4, ms=4, label="diffusers" + ) + ax.set_xticks(list(xs), labels, rotation=55, ha="right", fontsize=7, color=INK) + lo = math.floor(min(math.log10(v) for v in cand + prov)) + hi = math.ceil(max(math.log10(v) for v in cand + prov)) + ax.set_yticks(range(lo, hi + 1), [f"1e{v}" for v in range(lo, hi + 1)]) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_ylabel("max|err| / max|golden|", color=MUTED, fontsize=8) + ax.legend(frameon=False, fontsize=7.5, labelcolor=INK, loc="upper left") + _style(ax, title) + + +def plot(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(16, 5.6), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.0, 1.6, 1.3]} + ) + case = max(report["forward_cases"], key=lambda c: c["seq_len"]) + back = ( + max(report["backward_cases"], key=lambda c: c["seq_len"]) + if report["backward_cases"] + else None + ) + _checks_panel(axes[0], report) + _err_panel( + axes[1], + [e["node"] for e in case["nodes"]], + [e["candidate_rel_err_vs_golden"] for e in case["nodes"]], + [e["provider_rel_err_vs_golden"] for e in case["nodes"]], + f"Forward error vs FP32 golden, S = {case['seq_len']}", + ) + if back is not None: + leaves = list(back["leaves"]) + _err_panel( + axes[2], + [f"d {name}" for name in leaves], + [back["leaves"][n]["candidate_rel_err_vs_golden"] for n in leaves], + [back["leaves"][n]["provider_rel_err_vs_golden"] for n in leaves], + f"Gradient error vs FP32 golden, S = {back['seq_len']}", + ) + env = report["environment"] + sizes = ", ".join(str(c["seq_len"]) for c in report["forward_cases"]) + fig.suptitle( + "MiniMax-H3 block 0, forward and backward (RFC #420 ws1_one_h3_block)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.905, + f"{env['gpu']} | torch {env['torch']} | layouts S = {sizes} | " + f"first drift vs diffusers: {case['first_drift']} | " + f"rl-kernel {report['rl_kernel_commit'][:7]}", + fontsize=8, + color=MUTED, + ) + fig.tight_layout(rect=(0, 0, 1, 0.89)) + return fig + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + out = args.report.with_name("figure.png") + plot(report).savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py index 5c6f163f6..f6e32bcb5 100644 --- a/tools/weights/prepare_h3_weights.py +++ b/tools/weights/prepare_h3_weights.py @@ -8,9 +8,11 @@ manifest tensors (shard 1 of 14, ~4.8 GB) at the pinned revision, verifies every sha256 against ``rl_engine/validation/models/h3_manifest.json`` and writes the tensors to ``/h3_conditioning.safetensors``. Point the tests at it with -``export RL_KERNEL_H3_WEIGHTS=``. +``export RL_KERNEL_H3_WEIGHTS=``. ``--block`` also writes block 0's +attention/FFN weights (manifest ``block_tensors``, same shard) to +``/h3_block0.safetensors`` for the ``ws1_one_h3_block`` replay. - python tools/weights/prepare_h3_weights.py --out ~/.cache/rl-kernel/h3 + python tools/weights/prepare_h3_weights.py --out ~/.cache/rl-kernel/h3 [--block] """ from __future__ import annotations @@ -24,6 +26,7 @@ sys.path.insert(0, str(REPO_ROOT)) from rl_engine.validation.models.h3_weights import ( # noqa: E402 + BLOCK_FILE, EXTRACTED_FILE, load_h3_manifest, sha256_file, @@ -51,6 +54,9 @@ def main() -> None: default=None, help="where the raw shard goes (default: /hf)", ) + parser.add_argument( + "--block", action="store_true", help="also extract block 0 (manifest block_tensors)" + ) args = parser.parse_args() from huggingface_hub import hf_hub_download @@ -80,30 +86,38 @@ def fetch(filename: str) -> Path: _check(index_path, identity["index_sha256"]) weight_map = json.loads(index_path.read_text())["weight_map"] - tensors = {} - for shard in sorted({weight_map[name] for name in manifest["tensors"]}): - if shard not in manifest["weight_shards"]: - raise SystemExit(f"shard {shard} is not pinned in the manifest") - shard_path = fetch(shard) - _check(shard_path, manifest["weight_shards"][shard]["sha256"]) - with safe_open(str(shard_path), framework="pt") as handle: - for name in manifest["tensors"]: - if weight_map[name] == shard: - tensors[name] = handle.get_tensor(name).contiguous() - - for name, spec in manifest["tensors"].items(): - tensor = tensors[name] - if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: - raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") - if list(tensor.shape) != spec["shape"]: - raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") - actual = sha256_tensor(tensor) - if actual != spec["sha256"]: - raise SystemExit(f"{name}: sha256 {actual} != manifest {spec['sha256']}") - - target = out_dir / EXTRACTED_FILE - save_file(tensors, str(target), metadata={"revision": identity["revision"]}) - print(f"wrote {len(tensors)} tensors to {target}") + def extract(section: str, filename: str) -> None: + """Pull one manifest section's tensors from their pinned shards, verify, and save.""" + + specs = manifest[section] + tensors = {} + for shard in sorted({weight_map[name] for name in specs}): + if shard not in manifest["weight_shards"]: + raise SystemExit(f"shard {shard} is not pinned in the manifest") + shard_path = fetch(shard) + _check(shard_path, manifest["weight_shards"][shard]["sha256"]) + with safe_open(str(shard_path), framework="pt") as handle: + for name in specs: + if weight_map[name] == shard: + tensors[name] = handle.get_tensor(name).contiguous() + + for name, spec in specs.items(): + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise SystemExit(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + + target = out_dir / filename + save_file(tensors, str(target), metadata={"revision": identity["revision"]}) + print(f"wrote {len(tensors)} tensors to {target}") + + extract("tensors", EXTRACTED_FILE) + if args.block: + extract("block_tensors", BLOCK_FILE) print(f"export RL_KERNEL_H3_WEIGHTS={out_dir}")