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 001/139] [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 002/139] [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 003/139] [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 004/139] [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 005/139] [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 006/139] [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 007/139] [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 008/139] [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 009/139] [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 010/139] [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 011/139] [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 012/139] [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 013/139] [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 014/139] [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 015/139] 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 016/139] 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 017/139] 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 018/139] 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 019/139] 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 020/139] 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 021/139] 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 022/139] 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 023/139] 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 024/139] 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 025/139] 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 026/139] 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 027/139] 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 028/139] 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 029/139] 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 030/139] 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 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 031/139] 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 032/139] 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 033/139] 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 034/139] 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 035/139] 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 036/139] 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 037/139] 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 038/139] 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 039/139] 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 040/139] 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 041/139] 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 042/139] 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 043/139] 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 044/139] 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 045/139] 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 046/139] 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 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 047/139] 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 048/139] 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 049/139] 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 050/139] 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 051/139] 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 052/139] 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 053/139] 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 054/139] 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 055/139] 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 056/139] 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 057/139] 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 058/139] 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 059/139] 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 060/139] 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 061/139] 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 062/139] 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 063/139] 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 064/139] 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 065/139] 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 066/139] 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 067/139] 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 068/139] 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 069/139] 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 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 070/139] 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 071/139] 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 072/139] 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 073/139] 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 074/139] 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 075/139] 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 076/139] 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 077/139] 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 9ba6ce68ca2cb2eeac4f2f36eda905d3f7de3a1b Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:08 +0000 Subject: [PATCH 078/139] 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-timestep-sinusoid.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) 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 From 8132d061c7e21ccfceedfb386e97475612c347f6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:10 +0000 Subject: [PATCH 079/139] 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-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) 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 From 06bf572c136b2529b26f2ca871e2aaa0bf069814 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:12 +0000 Subject: [PATCH 080/139] 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-projection.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a76cc0417..a9c8832e2 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) | 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 | |---|---|---|---|---| @@ -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 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 @@ -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_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-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 From e26e369bda5a5ba8d9c2366fb6237fea3d14f9a0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:14 +0000 Subject: [PATCH 081/139] 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-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 4 files changed, 25 insertions(+), 25 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a76cc0417..a9c8832e2 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) | 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 | |---|---|---|---|---| @@ -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 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 @@ -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_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 9fa8a690e..a6cd10d0d 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -120,13 +120,13 @@ six output gradients and the per-tile FP32 partials are materialised. ## 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/pr486/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/pr486/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 rewritten from a clean tree at commit `40f4400c`, after the op-backward check @@ -140,7 +140,7 @@ is unchanged; only timings and the diffusers backward errors, which vary from ru ## 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 | |---|---|---|---| @@ -151,7 +151,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 @@ -168,10 +168,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-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 From 2985fd20e3ce9d9d0bd567716350bbfb38f2e3bd Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:16 +0000 Subject: [PATCH 082/139] 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-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 5 files changed, 31 insertions(+), 31 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a76cc0417..a9c8832e2 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) | 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 | |---|---|---|---|---| @@ -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 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 @@ -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_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-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 7a16f4a57..60f111647 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -90,9 +90,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; @@ -101,7 +101,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 | |---|---|---|---| @@ -117,7 +117,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 @@ -134,10 +134,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-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 From a8e5901b5d194bf846c5ad670ea9780850087c04 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:18 +0000 Subject: [PATCH 083/139] 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-rmsnorm.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 6 files changed, 37 insertions(+), 37 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index e003fd5b8..8847d6161 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/pr488/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/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/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](../../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 a76cc0417..a9c8832e2 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) | 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 | |---|---|---|---|---| @@ -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 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 @@ -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_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-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-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 From 39fb03ddf1447c5ac8a23eaab3df1298476b6a1a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:19 +0000 Subject: [PATCH 084/139] 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-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 7 files changed, 44 insertions(+), 44 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index e003fd5b8..8847d6161 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/pr488/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/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/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](../../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 a76cc0417..a9c8832e2 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) | 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 | |---|---|---|---|---| @@ -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 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 @@ -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_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 44e08e507..bb1194db8 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -71,17 +71,17 @@ projection backward. ## 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/pr489/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/pr489/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 `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 +- [`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`. Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` @@ -90,7 +90,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 | |---|---|---|---| @@ -101,7 +101,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 @@ -118,10 +118,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-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 From e5040298e531379b970258ca212882b5d6b62f09 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:52 +0000 Subject: [PATCH 085/139] 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-timestep-sinusoid.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 \ From 3b74ce356e99170e561126f337bdb895162e24b2 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:53 +0000 Subject: [PATCH 086/139] 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-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) 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 \ From 441e00004bc856365e8734ca607e2839a8811c4c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:54 +0000 Subject: [PATCH 087/139] 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-projection.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a9c8832e2..3f79732cd 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 @@ -172,7 +172,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-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 \ From 1a8e7b26e67861866ce4cfacf35b043e3dbef760 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:56 +0000 Subject: [PATCH 088/139] 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-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 6 +++--- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a9c8832e2..3f79732cd 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 @@ -172,7 +172,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 a6cd10d0d..a6bec87f8 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -129,10 +129,10 @@ 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 rewritten from a clean tree at commit `40f4400c`, after the op-backward check +`report.json` was rewritten from a clean tree at commit `a6c85451`, 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`. +the previous report). `chain_replay.json` was written from a clean tree at commit `61d23ade`. 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). Apart from the seeded op-backward fraction, every CUDA field and every pass/fail flag @@ -171,7 +171,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-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 \ From 26a7cab3cad361ef47b50059cd10c4abe5f77289 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:57 +0000 Subject: [PATCH 089/139] 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-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a9c8832e2..3f79732cd 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 @@ -172,7 +172,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-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 60f111647..b64c5e44e 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -93,7 +93,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; @@ -137,7 +137,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-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 \ From 38018081392ba4c1222a1bb32e0261a6122cf804 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:58 +0000 Subject: [PATCH 090/139] 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 | 4 ++-- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 6 files changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 8847d6161..b35bc41c1 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -78,7 +78,7 @@ tile setup. ![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `e7bc53a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) @@ -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 a9c8832e2..3f79732cd 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 @@ -172,7 +172,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-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-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 \ From 428baa1139daa781927a800be85ee5e705ab2ce7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:06:00 +0000 Subject: [PATCH 091/139] 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 | 4 ++-- 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-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 7 files changed, 16 insertions(+), 16 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 8847d6161..b35bc41c1 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -78,7 +78,7 @@ tile setup. ![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `e7bc53a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) @@ -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 a9c8832e2..3f79732cd 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 @@ -172,7 +172,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 bb1194db8..73eac42d6 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -77,14 +77,14 @@ There are two data files: - [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/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 `50e913fa` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + tree at `a4bc99f9` 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`. Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` + tree at `a4bc99f9`. 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). @@ -121,7 +121,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-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 \ From cb9828f8b2477728e707473d570d33fb9e56445a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Sun, 25 Oct 2026 23:25:43 +0000 Subject: [PATCH 092/139] fix(h3): time competing callables in a shuffled order and fix the sinusoid latency column Review findings on the timing and documentation of the H3 evidence: - measure() timed every sample of the candidate before the provider's, so clock boost, thermal state and L2 contents could bias the gap. Each iteration now runs all callables once in a random order (seeded) and the executed order is stored next to the medians. - The sinusoid table called its latency column fwd+bwd although the op has no backward and every cell is forward-only; the header and the sentence below it now say forward latency. - Move the noqa: E402 of prepare_h3_weights.py to the first line of the multi-line import, where the linters report it. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-sinusoid.md | 4 +- rl_engine/validation/models/h3_report.py | 57 +++++++++++++++--- .../minimax_h3/test_h3_report_timing.py | 59 +++++++++++++++++++ tools/weights/prepare_h3_weights.py | 4 +- 4 files changed, 113 insertions(+), 11 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_report_timing.py diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index ff6579cd8..5696f5df7 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -98,14 +98,14 @@ records that commit and the environment. ![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 | +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd latency | size 256: fwd err / worst grad err / fwd latency | size 2048: fwd err / worst grad err / fwd latency | |---|---|---|---|---| | 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 +forward time (the op has no backward) 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`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index f7bd01520..5bd5c507b 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -10,6 +10,7 @@ from __future__ import annotations +import random import statistics from typing import Any, Callable @@ -57,14 +58,56 @@ def peak_mib(fn: Callable[[], Any]) -> float: return (torch.cuda.max_memory_allocated() - base) / 2**20 +def time_interleaved( + fns: dict[str, Callable[[], Any]], warmup: int = 20, iters: int = 200, seed: int = 0 +) -> tuple[dict[str, float], list[str]]: + """Median CUDA-event time of every callable, each sample in a freshly shuffled order. + + Clock boost, thermal state and L2 contents drift while a benchmark runs, so timing all + samples of one callable before the other's would bias the comparison. Every iteration runs + each callable once, in a random order from ``seed``. Returns the medians in microseconds and + the executed order of every iteration (the digits are indices into ``list(fns)``). + """ + + names = list(fns) + rng = random.Random(seed) + for fn in fns.values(): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples: dict[str, list[float]] = {name: [] for name in names} + executed: list[str] = [] + for _ in range(iters): + order = names[:] + rng.shuffle(order) + executed.append("".join(str(names.index(name)) for name in order)) + for name in order: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fns[name]() + end.record() + end.synchronize() + samples[name].append(start.elapsed_time(end) * 1e3) + return {name: statistics.median(values) for name, values in samples.items()}, executed + + 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]) + 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] + medians, executed = time_interleaved({key: case[key] for key in keys}, warmup, iters) + for key in keys: + row[f"{key}_us"] = medians[key] + row[f"{key}_gbps"] = case["bytes"] / (medians[key] * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key]) + row["timing"] = { + "order": "shuffled per sample, all callables timed in every iteration", + "seed": 0, + "callables": keys, + "executed_order": executed, + } return row diff --git a/tests/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py new file mode 100644 index 000000000..6efba02ef --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Timing harness of the H3 evidence: competing callables are timed in a shuffled order. + +Run on a real CUDA device with real kernels (torch matmuls of two sizes); the check is about +the executed order and the medians, not about absolute times. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.validation.models.h3_report import measure, time_interleaved + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _callables(): + a = torch.randn(256, 256, device="cuda") + b = torch.randn(1024, 1024, device="cuda") + return {"small": lambda: a @ a, "large": lambda: b @ b} + + +def test_every_callable_runs_once_per_iteration_in_both_orders(): + fns = _callables() + medians, executed = time_interleaved(fns, warmup=2, iters=40, seed=0) + assert set(medians) == {"small", "large"} + assert all(value > 0 for value in medians.values()) + assert len(executed) == 40 + # digits index list(fns): each iteration runs both callables exactly once + assert {"".join(sorted(order)) for order in executed} == {"01"} + # the order is shuffled, so neither callable always goes first + assert {"01", "10"} <= set(executed) + + +def test_the_executed_order_is_reproducible_from_the_seed(): + fns = _callables() + _, first = time_interleaved(fns, warmup=1, iters=25, seed=3) + _, second = time_interleaved(fns, warmup=1, iters=25, seed=3) + _, other = time_interleaved(fns, warmup=1, iters=25, seed=4) + assert first == second + assert first != other + + +def test_measure_records_the_order_next_to_the_timings(): + case = { + "op": "matmul", + "case": "T=1", + "backend": "torch", + "bytes": 1 << 20, + **{key: fn for key, fn in zip(("candidate", "provider"), _callables().values())}, + } + row = measure(case, warmup=2, iters=20) + assert row["candidate_us"] > 0 and row["provider_us"] > 0 + assert row["timing"]["callables"] == ["candidate", "provider"] + assert len(row["timing"]["executed_order"]) == 20 + assert {"01", "10"} <= set(row["timing"]["executed_order"]) diff --git a/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py index e5470ef18..c21f56c2d 100644 --- a/tools/weights/prepare_h3_weights.py +++ b/tools/weights/prepare_h3_weights.py @@ -23,11 +23,11 @@ REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT)) -from rl_engine.validation.models.h3_weights import ( +from rl_engine.validation.models.h3_weights import ( # noqa: E402 EXTRACTED_FILE, load_h3_manifest, sha256_file, -) # noqa: E402 +) def _check(path: Path, expected: str) -> None: From a74b1a5af85ecbebeedcedb8c6b79e1f4c138f61 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Sun, 25 Oct 2026 23:27:38 +0000 Subject: [PATCH 093/139] fix(h3): record which weights the timestep MLP evidence used h3_params fell back to seeded random tensors when RL_KERNEL_H3_WEIGHTS was unset and nothing recorded it, while the figure title always claimed pinned weights. h3_params now returns the source (pinned or random); the accuracy report and every perf row store it, the title is derived from it, and h3_evidence.py refuses to write a timestep_mlp_fp32 report from random weights unless --allow-random-weights is given (smoke runs). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 20 +++++--- .../minimax_h3/test_h3_weight_source.py | 51 +++++++++++++++++++ tools/validation/models/h3_evidence.py | 22 ++++++++ tools/validation/models/plot_h3_evidence.py | 13 ++++- 4 files changed, 98 insertions(+), 8 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_weight_source.py diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 5bd5c507b..a157688aa 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -94,7 +94,7 @@ def time_interleaved( def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: row: dict[str, Any] = { - key: case[key] for key in ("op", "case", "backend", "bytes") if key in case + key: case[key] for key in ("op", "case", "backend", "bytes", "weights") if key in case } keys = [key for key in TIMED_KEYS if key in case] medians, executed = time_interleaved({key: case[key] for key in keys}, warmup, iters) @@ -155,17 +155,21 @@ 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.""" +def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> tuple[list[torch.Tensor], str]: + """Pinned checkpoint tensors when available, otherwise same-shape random ones. + + Returns the tensors and their source, ``"pinned"`` or ``"random"``; the evidence stores it so a + report made without the pinned checkpoint cannot pass for one that used it. + """ if h3_weights_dir() is not None: weights = load_h3_conditioning_weights("cuda", names) - return [weights[name] for name in names] + return [weights[name] for name in names], "pinned" generator = torch.Generator(device="cuda").manual_seed(0) return [ torch.randn(shape, device="cuda", generator=generator) * shape[-1] ** -0.5 for shape in shapes - ] + ], "random" # --------------------------------------------------------------------------- # @@ -178,7 +182,7 @@ def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Ten 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) + params, weights_source = 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): @@ -190,6 +194,7 @@ def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: "case": f"T={num}", "backend": type(op).__name__, "bytes": weight_bytes, + "weights": weights_source, "candidate": lambda x=x: op(x, *params), "provider": lambda x=x: provider_time_embedder(x, *params), } @@ -201,7 +206,7 @@ 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) + params, weights_source = h3_params(MLP_NAMES, MLP_SHAPES) cuda_err, provider_err = [], [] for seed in range(draws): x = sinusoid(h3_timesteps(4, seed=seed)) @@ -214,6 +219,7 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: invariant = all(torch.equal(op(x[i : i + 1], *params)[0], full[i]) for i in range(9)) return { "contract_atol": 1e-4, + "weights_source": weights_source, "draws": draws, "num_timesteps": 4, "cuda_max_abs_vs_fp64": cuda_err, diff --git a/tests/models/minimax_h3/test_h3_weight_source.py b/tests/models/minimax_h3/test_h3_weight_source.py new file mode 100644 index 000000000..251a728a2 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_weight_source.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The H3 MLP evidence says which weights it was measured with and refuses unpinned runs.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.validation.models.h3_report import MLP_NAMES, MLP_SHAPES, h3_params +from rl_engine.validation.models.h3_weights import WEIGHTS_ENV, h3_weights_dir +from tools.validation.models.h3_evidence import check_weight_source + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_without_the_pinned_checkpoint_the_params_are_marked_random(monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + params, source = h3_params(MLP_NAMES, MLP_SHAPES) + assert source == "random" + assert [tuple(p.shape) for p in params] == MLP_SHAPES + assert all(p.is_cuda for p in params) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_with_the_pinned_checkpoint_the_params_are_marked_pinned(): + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py") + params, source = h3_params(MLP_NAMES, MLP_SHAPES) + assert source == "pinned" + assert [tuple(p.shape) for p in params] == MLP_SHAPES + + +def test_mlp_evidence_is_refused_without_pinned_weights(): + random_run = {"accuracy": {"weights_source": "random"}} + pinned_run = {"accuracy": {"weights_source": "pinned"}} + with pytest.raises(SystemExit, match="pinned time_embedder weights"): + check_weight_source("timestep_mlp_fp32", random_run, allow_random=False) + check_weight_source("timestep_mlp_fp32", random_run, allow_random=True) + check_weight_source("timestep_mlp_fp32", pinned_run, allow_random=False) + # the sinusoid op has no weights, so it records none + check_weight_source("timestep_sinusoid_h3", {"accuracy": {}}, allow_random=False) + + +def test_the_figure_title_names_the_recorded_weight_source(): + pytest.importorskip("matplotlib") + from tools.validation.models.plot_h3_evidence import mlp_title + + assert "pinned time_embedder weights" in mlp_title({"accuracy": {"weights_source": "pinned"}}) + assert "random" in mlp_title({"accuracy": {"weights_source": "random"}}) + assert "not recorded" in mlp_title({"accuracy": {}}) diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index b4852d61f..41157682a 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -29,10 +29,31 @@ from rl_engine.validation.models.h3_weights import load_h3_manifest # noqa: E402 +def check_weight_source(op: str, report: dict, allow_random: bool) -> None: + """Refuse a timestep_mlp_fp32 report that was measured without the pinned checkpoint.""" + + if ( + op == "timestep_mlp_fp32" + and report["accuracy"]["weights_source"] != "pinned" + and not allow_random + ): + raise SystemExit( + "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " + "RL_KERNEL_H3_WEIGHTS (tools/weights/prepare_h3_weights.py) " + "or pass --allow-random-weights" + ) + + 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) + parser.add_argument( + "--allow-random-weights", + action="store_true", + help="write a timestep_mlp_fp32 report even without the pinned checkpoint " + "(smoke runs only)", + ) args = parser.parse_args() if not torch.cuda.is_available(): raise SystemExit("needs a CUDA device") @@ -52,6 +73,7 @@ def main() -> None: "accuracy": ACCURACY[args.op](registry), "perf": [measure(case) for case in PERF_CASES[args.op](registry)], } + check_weight_source(args.op, report, args.allow_random_weights) args.out.parent.mkdir(parents=True, exist_ok=True) args.out.write_text(json.dumps(report, indent=2) + "\n") print( diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py index ddf1ee947..bb8aee771 100644 --- a/tools/validation/models/plot_h3_evidence.py +++ b/tools/validation/models/plot_h3_evidence.py @@ -209,8 +209,19 @@ def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], con ax.set_xlim(-0.6, len(series) - 0.1) +def mlp_title(report: dict[str, Any]) -> str: + """Figure title that names the weights the report was measured with.""" + + source = report["accuracy"].get("weights_source") + label = { + "pinned": "pinned time_embedder weights", + "random": "random same-shape weights, not the checkpoint", + }.get(source, "weight source not recorded") + return f"timestep_mlp_fp32 ({label})" + + def plot_mlp(report: dict[str, Any]): - fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") + fig, (left, right) = _figure(report, mlp_title(report)) _latency_panel(left, report["perf"], "diffusers (cuBLAS)") acc = report["accuracy"] _error_boxes( From 95df2a9246dbd536a5d5007bf2a12a8f217b6dc0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Sun, 25 Oct 2026 23:25:43 +0000 Subject: [PATCH 094/139] fix(h3): time competing callables in a shuffled order and fix the sinusoid latency column Review findings on the timing and documentation of the H3 evidence: - measure() timed every sample of the candidate before the provider's, so clock boost, thermal state and L2 contents could bias the gap. Each iteration now runs all callables once in a random order (seeded) and the executed order is stored next to the medians. - The sinusoid table called its latency column fwd+bwd although the op has no backward and every cell is forward-only; the header and the sentence below it now say forward latency. - Move the noqa: E402 of prepare_h3_weights.py to the first line of the multi-line import, where the linters report it. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-sinusoid.md | 4 +- rl_engine/validation/models/h3_report.py | 57 +++++++++++++++--- .../minimax_h3/test_h3_report_timing.py | 59 +++++++++++++++++++ tools/weights/prepare_h3_weights.py | 4 +- 4 files changed, 113 insertions(+), 11 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_report_timing.py diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index ff6579cd8..5696f5df7 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -98,14 +98,14 @@ records that commit and the environment. ![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 | +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd latency | size 256: fwd err / worst grad err / fwd latency | size 2048: fwd err / worst grad err / fwd latency | |---|---|---|---|---| | 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 +forward time (the op has no backward) 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`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 960e88ee3..f84a94347 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -10,6 +10,7 @@ from __future__ import annotations +import random import statistics from typing import Any, Callable @@ -56,14 +57,56 @@ def peak_mib(fn: Callable[[], Any]) -> float: return (torch.cuda.max_memory_allocated() - base) / 2**20 +def time_interleaved( + fns: dict[str, Callable[[], Any]], warmup: int = 20, iters: int = 200, seed: int = 0 +) -> tuple[dict[str, float], list[str]]: + """Median CUDA-event time of every callable, each sample in a freshly shuffled order. + + Clock boost, thermal state and L2 contents drift while a benchmark runs, so timing all + samples of one callable before the other's would bias the comparison. Every iteration runs + each callable once, in a random order from ``seed``. Returns the medians in microseconds and + the executed order of every iteration (the digits are indices into ``list(fns)``). + """ + + names = list(fns) + rng = random.Random(seed) + for fn in fns.values(): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples: dict[str, list[float]] = {name: [] for name in names} + executed: list[str] = [] + for _ in range(iters): + order = names[:] + rng.shuffle(order) + executed.append("".join(str(names.index(name)) for name in order)) + for name in order: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fns[name]() + end.record() + end.synchronize() + samples[name].append(start.elapsed_time(end) * 1e3) + return {name: statistics.median(values) for name, values in samples.items()}, executed + + 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]) + 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] + medians, executed = time_interleaved({key: case[key] for key in keys}, warmup, iters) + for key in keys: + row[f"{key}_us"] = medians[key] + row[f"{key}_gbps"] = case["bytes"] / (medians[key] * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key]) + row["timing"] = { + "order": "shuffled per sample, all callables timed in every iteration", + "seed": 0, + "callables": keys, + "executed_order": executed, + } return row diff --git a/tests/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py new file mode 100644 index 000000000..6efba02ef --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Timing harness of the H3 evidence: competing callables are timed in a shuffled order. + +Run on a real CUDA device with real kernels (torch matmuls of two sizes); the check is about +the executed order and the medians, not about absolute times. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.validation.models.h3_report import measure, time_interleaved + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _callables(): + a = torch.randn(256, 256, device="cuda") + b = torch.randn(1024, 1024, device="cuda") + return {"small": lambda: a @ a, "large": lambda: b @ b} + + +def test_every_callable_runs_once_per_iteration_in_both_orders(): + fns = _callables() + medians, executed = time_interleaved(fns, warmup=2, iters=40, seed=0) + assert set(medians) == {"small", "large"} + assert all(value > 0 for value in medians.values()) + assert len(executed) == 40 + # digits index list(fns): each iteration runs both callables exactly once + assert {"".join(sorted(order)) for order in executed} == {"01"} + # the order is shuffled, so neither callable always goes first + assert {"01", "10"} <= set(executed) + + +def test_the_executed_order_is_reproducible_from_the_seed(): + fns = _callables() + _, first = time_interleaved(fns, warmup=1, iters=25, seed=3) + _, second = time_interleaved(fns, warmup=1, iters=25, seed=3) + _, other = time_interleaved(fns, warmup=1, iters=25, seed=4) + assert first == second + assert first != other + + +def test_measure_records_the_order_next_to_the_timings(): + case = { + "op": "matmul", + "case": "T=1", + "backend": "torch", + "bytes": 1 << 20, + **{key: fn for key, fn in zip(("candidate", "provider"), _callables().values())}, + } + row = measure(case, warmup=2, iters=20) + assert row["candidate_us"] > 0 and row["provider_us"] > 0 + assert row["timing"]["callables"] == ["candidate", "provider"] + assert len(row["timing"]["executed_order"]) == 20 + assert {"01", "10"} <= set(row["timing"]["executed_order"]) diff --git a/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py index e5470ef18..c21f56c2d 100644 --- a/tools/weights/prepare_h3_weights.py +++ b/tools/weights/prepare_h3_weights.py @@ -23,11 +23,11 @@ REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT)) -from rl_engine.validation.models.h3_weights import ( +from rl_engine.validation.models.h3_weights import ( # noqa: E402 EXTRACTED_FILE, load_h3_manifest, sha256_file, -) # noqa: E402 +) def _check(path: Path, expected: str) -> None: From 9a748eb25977069d596452132df86b06efc8a586 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:22:38 +0000 Subject: [PATCH 095/139] fix(h3): time candidate and provider with the rotating harness of the later PRs The first fix for the timing-order review finding interleaved the callables in a seeded random order. #485 and #486 already carry a different harness that rotates the first callable each iteration and stores the raw CUDA-event samples and the executed orders, which keeps more of the evidence. Use that same measure() here so the stacked PRs agree and later merges do not conflict: drop time_interleaved, take measure() verbatim from #485, add its test (test_h3_report.py) and retarget the real-GPU test to the rotating harness. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 66 ++++++--------- tests/models/minimax_h3/test_h3_report.py | 82 +++++++++++++++++++ .../minimax_h3/test_h3_report_timing.py | 73 ++++++++--------- 3 files changed, 142 insertions(+), 79 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_report.py diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index a157688aa..ddc58a958 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -10,7 +10,6 @@ from __future__ import annotations -import random import statistics from typing import Any, Callable @@ -58,56 +57,43 @@ def peak_mib(fn: Callable[[], Any]) -> float: return (torch.cuda.max_memory_allocated() - base) / 2**20 -def time_interleaved( - fns: dict[str, Callable[[], Any]], warmup: int = 20, iters: int = 200, seed: int = 0 -) -> tuple[dict[str, float], list[str]]: - """Median CUDA-event time of every callable, each sample in a freshly shuffled order. +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. - Clock boost, thermal state and L2 contents drift while a benchmark runs, so timing all - samples of one callable before the other's would bias the comparison. Every iteration runs - each callable once, in a random order from ``seed``. Returns the medians in microseconds and - the executed order of every iteration (the digits are indices into ``list(fns)``). + Record raw CUDA-event samples, executed orders, median microseconds, + decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. """ - names = list(fns) - rng = random.Random(seed) - for fn in fns.values(): - for _ in range(warmup): - fn() + 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: + return row + for _ in range(warmup): + for key in keys: + case[key]() torch.cuda.synchronize() - samples: dict[str, list[float]] = {name: [] for name in names} - executed: list[str] = [] - for _ in range(iters): - order = names[:] - rng.shuffle(order) - executed.append("".join(str(names.index(name)) for name in order)) - for name in order: + 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() - fns[name]() + case[key]() end.record() end.synchronize() - samples[name].append(start.elapsed_time(end) * 1e3) - return {name: statistics.median(values) for name, values in samples.items()}, executed - - -def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: - row: dict[str, Any] = { - key: case[key] for key in ("op", "case", "backend", "bytes", "weights") if key in case - } - keys = [key for key in TIMED_KEYS if key in case] - medians, executed = time_interleaved({key: case[key] for key in keys}, warmup, iters) + samples[key].append(start.elapsed_time(end) * 1e3) + row["timing_samples_us"] = samples + row["timing_order"] = orders for key in keys: - row[f"{key}_us"] = medians[key] - row[f"{key}_gbps"] = case["bytes"] / (medians[key] * 1e-6) / 1e9 + 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]) - row["timing"] = { - "order": "shuffled per sample, all callables timed in every iteration", - "seed": 0, - "callables": keys, - "executed_order": executed, - } return row diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py new file mode 100644 index 000000000..98e6ab714 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report.py @@ -0,0 +1,82 @@ +# 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.validation.models import h3_report + + +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 + + 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/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py index 6efba02ef..853680a67 100644 --- a/tests/models/minimax_h3/test_h3_report_timing.py +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -1,59 +1,54 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Timing harness of the H3 evidence: competing callables are timed in a shuffled order. +"""``measure`` on a real CUDA device with real kernels. -Run on a real CUDA device with real kernels (torch matmuls of two sizes); the check is about -the executed order and the medians, not about absolute times. +``test_h3_report.py`` checks the bookkeeping with a stubbed clock; this one runs two real matmuls +of different size and checks that every callable is timed in every iteration, that the first +callable rotates, and that the stored samples reproduce the reported medians. """ from __future__ import annotations +import statistics + import pytest import torch -from rl_engine.validation.models.h3_report import measure, time_interleaved +from rl_engine.validation.models.h3_report import measure pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def _callables(): - a = torch.randn(256, 256, device="cuda") - b = torch.randn(1024, 1024, device="cuda") - return {"small": lambda: a @ a, "large": lambda: b @ b} - - -def test_every_callable_runs_once_per_iteration_in_both_orders(): - fns = _callables() - medians, executed = time_interleaved(fns, warmup=2, iters=40, seed=0) - assert set(medians) == {"small", "large"} - assert all(value > 0 for value in medians.values()) - assert len(executed) == 40 - # digits index list(fns): each iteration runs both callables exactly once - assert {"".join(sorted(order)) for order in executed} == {"01"} - # the order is shuffled, so neither callable always goes first - assert {"01", "10"} <= set(executed) - - -def test_the_executed_order_is_reproducible_from_the_seed(): - fns = _callables() - _, first = time_interleaved(fns, warmup=1, iters=25, seed=3) - _, second = time_interleaved(fns, warmup=1, iters=25, seed=3) - _, other = time_interleaved(fns, warmup=1, iters=25, seed=4) - assert first == second - assert first != other - - -def test_measure_records_the_order_next_to_the_timings(): - case = { +def _case(): + small = torch.randn(256, 256, device="cuda") + large = torch.randn(1024, 1024, device="cuda") + return { "op": "matmul", "case": "T=1", "backend": "torch", "bytes": 1 << 20, - **{key: fn for key, fn in zip(("candidate", "provider"), _callables().values())}, + "candidate": lambda: small @ small, + "provider": lambda: large @ large, } - row = measure(case, warmup=2, iters=20) - assert row["candidate_us"] > 0 and row["provider_us"] > 0 - assert row["timing"]["callables"] == ["candidate", "provider"] - assert len(row["timing"]["executed_order"]) == 20 - assert {"01", "10"} <= set(row["timing"]["executed_order"]) + + +def test_every_callable_is_timed_in_every_iteration_with_a_rotating_first_callable(): + row = measure(_case(), warmup=2, iters=20) + assert len(row["timing_samples_us"]["candidate"]) == 20 + assert len(row["timing_samples_us"]["provider"]) == 20 + assert row["timing_order"][0] == ["candidate", "provider"] + assert row["timing_order"][1] == ["provider", "candidate"] + assert all(sorted(order) == ["candidate", "provider"] for order in row["timing_order"]) + assert {order[0] for order in row["timing_order"]} == {"candidate", "provider"} + + +def test_the_reported_medians_come_from_the_stored_samples(): + row = measure(_case(), warmup=2, iters=20) + for key in ("candidate", "provider"): + assert row[f"{key}_us"] == statistics.median(row["timing_samples_us"][key]) + assert row[f"{key}_us"] > 0 + assert row[f"{key}_gbps"] == pytest.approx(row["bytes"] / (row[f"{key}_us"] * 1e-6) / 1e9) + assert row[f"{key}_peak_mib"] >= 0 + # the larger matmul is slower than the smaller one in the same run + assert row["provider_us"] > row["candidate_us"] From 9beb1407eb5ef104ab809984f55cb984f0f2fc93 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:23:01 +0000 Subject: [PATCH 096/139] fix(h3): keep the weight source on every timestep MLP perf row measure() is the shared rotating harness and does not copy the case's weights field, so the evidence writer stores it with the row it gets back. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/test_h3_weight_source.py | 19 ++++++++++++++++++- tools/validation/models/h3_evidence.py | 11 ++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/tests/models/minimax_h3/test_h3_weight_source.py b/tests/models/minimax_h3/test_h3_weight_source.py index 251a728a2..0a44d84ef 100644 --- a/tests/models/minimax_h3/test_h3_weight_source.py +++ b/tests/models/minimax_h3/test_h3_weight_source.py @@ -10,7 +10,7 @@ from rl_engine.validation.models.h3_report import MLP_NAMES, MLP_SHAPES, h3_params from rl_engine.validation.models.h3_weights import WEIGHTS_ENV, h3_weights_dir -from tools.validation.models.h3_evidence import check_weight_source +from tools.validation.models.h3_evidence import check_weight_source, measure_case @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") @@ -49,3 +49,20 @@ def test_the_figure_title_names_the_recorded_weight_source(): assert "pinned time_embedder weights" in mlp_title({"accuracy": {"weights_source": "pinned"}}) assert "random" in mlp_title({"accuracy": {"weights_source": "random"}}) assert "not recorded" in mlp_title({"accuracy": {}}) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_a_perf_row_keeps_the_weight_source_of_its_case(): + a = torch.randn(128, 128, device="cuda") + case = { + "op": "matmul", + "case": "T=1", + "backend": "torch", + "bytes": 1 << 16, + "weights": "random", + "candidate": lambda: a @ a, + } + row = measure_case(case) + assert row["weights"] == "random" + assert row["candidate_us"] > 0 + assert "weights" not in measure_case({k: v for k, v in case.items() if k != "weights"}) diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index 41157682a..e67333b8b 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -29,6 +29,15 @@ from rl_engine.validation.models.h3_weights import load_h3_manifest # noqa: E402 +def measure_case(case: dict) -> dict: + """Time one perf case and keep the weight source the case was built with next to the timings.""" + + row = measure(case) + if "weights" in case: + row["weights"] = case["weights"] + return row + + def check_weight_source(op: str, report: dict, allow_random: bool) -> None: """Refuse a timestep_mlp_fp32 report that was measured without the pinned checkpoint.""" @@ -71,7 +80,7 @@ def main() -> None: **git_state(), "environment": environment(), "accuracy": ACCURACY[args.op](registry), - "perf": [measure(case) for case in PERF_CASES[args.op](registry)], + "perf": [measure_case(case) for case in PERF_CASES[args.op](registry)], } check_weight_source(args.op, report, args.allow_random_weights) args.out.parent.mkdir(parents=True, exist_ok=True) From c19a6e50f0df9903cc7a044b8e1d206f0cefdfdd Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:23:51 +0000 Subject: [PATCH 097/139] fix(h3): time candidate and provider with the rotating harness of the later PRs The first fix for the timing-order review finding interleaved the callables in a seeded random order. #485 and #486 already carry a different harness that rotates the first callable each iteration and stores the raw CUDA-event samples and the executed orders, which keeps more of the evidence. Use that same measure() here so the stacked PRs agree and later merges do not conflict: drop time_interleaved, take measure() verbatim from #485, add its test (test_h3_report.py) and retarget the real-GPU test to the rotating harness. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 66 ++++++--------- tests/models/minimax_h3/test_h3_report.py | 82 +++++++++++++++++++ .../minimax_h3/test_h3_report_timing.py | 73 ++++++++--------- 3 files changed, 142 insertions(+), 79 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_report.py diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index f84a94347..b62a2e81b 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -10,7 +10,6 @@ from __future__ import annotations -import random import statistics from typing import Any, Callable @@ -57,56 +56,43 @@ def peak_mib(fn: Callable[[], Any]) -> float: return (torch.cuda.max_memory_allocated() - base) / 2**20 -def time_interleaved( - fns: dict[str, Callable[[], Any]], warmup: int = 20, iters: int = 200, seed: int = 0 -) -> tuple[dict[str, float], list[str]]: - """Median CUDA-event time of every callable, each sample in a freshly shuffled order. +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. - Clock boost, thermal state and L2 contents drift while a benchmark runs, so timing all - samples of one callable before the other's would bias the comparison. Every iteration runs - each callable once, in a random order from ``seed``. Returns the medians in microseconds and - the executed order of every iteration (the digits are indices into ``list(fns)``). + Record raw CUDA-event samples, executed orders, median microseconds, + decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. """ - names = list(fns) - rng = random.Random(seed) - for fn in fns.values(): - for _ in range(warmup): - fn() + 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: + return row + for _ in range(warmup): + for key in keys: + case[key]() torch.cuda.synchronize() - samples: dict[str, list[float]] = {name: [] for name in names} - executed: list[str] = [] - for _ in range(iters): - order = names[:] - rng.shuffle(order) - executed.append("".join(str(names.index(name)) for name in order)) - for name in order: + 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() - fns[name]() + case[key]() end.record() end.synchronize() - samples[name].append(start.elapsed_time(end) * 1e3) - return {name: statistics.median(values) for name, values in samples.items()}, executed - - -def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: - 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] - medians, executed = time_interleaved({key: case[key] for key in keys}, warmup, iters) + samples[key].append(start.elapsed_time(end) * 1e3) + row["timing_samples_us"] = samples + row["timing_order"] = orders for key in keys: - row[f"{key}_us"] = medians[key] - row[f"{key}_gbps"] = case["bytes"] / (medians[key] * 1e-6) / 1e9 + 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]) - row["timing"] = { - "order": "shuffled per sample, all callables timed in every iteration", - "seed": 0, - "callables": keys, - "executed_order": executed, - } return row diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py new file mode 100644 index 000000000..98e6ab714 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report.py @@ -0,0 +1,82 @@ +# 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.validation.models import h3_report + + +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 + + 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/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py index 6efba02ef..853680a67 100644 --- a/tests/models/minimax_h3/test_h3_report_timing.py +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -1,59 +1,54 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Timing harness of the H3 evidence: competing callables are timed in a shuffled order. +"""``measure`` on a real CUDA device with real kernels. -Run on a real CUDA device with real kernels (torch matmuls of two sizes); the check is about -the executed order and the medians, not about absolute times. +``test_h3_report.py`` checks the bookkeeping with a stubbed clock; this one runs two real matmuls +of different size and checks that every callable is timed in every iteration, that the first +callable rotates, and that the stored samples reproduce the reported medians. """ from __future__ import annotations +import statistics + import pytest import torch -from rl_engine.validation.models.h3_report import measure, time_interleaved +from rl_engine.validation.models.h3_report import measure pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def _callables(): - a = torch.randn(256, 256, device="cuda") - b = torch.randn(1024, 1024, device="cuda") - return {"small": lambda: a @ a, "large": lambda: b @ b} - - -def test_every_callable_runs_once_per_iteration_in_both_orders(): - fns = _callables() - medians, executed = time_interleaved(fns, warmup=2, iters=40, seed=0) - assert set(medians) == {"small", "large"} - assert all(value > 0 for value in medians.values()) - assert len(executed) == 40 - # digits index list(fns): each iteration runs both callables exactly once - assert {"".join(sorted(order)) for order in executed} == {"01"} - # the order is shuffled, so neither callable always goes first - assert {"01", "10"} <= set(executed) - - -def test_the_executed_order_is_reproducible_from_the_seed(): - fns = _callables() - _, first = time_interleaved(fns, warmup=1, iters=25, seed=3) - _, second = time_interleaved(fns, warmup=1, iters=25, seed=3) - _, other = time_interleaved(fns, warmup=1, iters=25, seed=4) - assert first == second - assert first != other - - -def test_measure_records_the_order_next_to_the_timings(): - case = { +def _case(): + small = torch.randn(256, 256, device="cuda") + large = torch.randn(1024, 1024, device="cuda") + return { "op": "matmul", "case": "T=1", "backend": "torch", "bytes": 1 << 20, - **{key: fn for key, fn in zip(("candidate", "provider"), _callables().values())}, + "candidate": lambda: small @ small, + "provider": lambda: large @ large, } - row = measure(case, warmup=2, iters=20) - assert row["candidate_us"] > 0 and row["provider_us"] > 0 - assert row["timing"]["callables"] == ["candidate", "provider"] - assert len(row["timing"]["executed_order"]) == 20 - assert {"01", "10"} <= set(row["timing"]["executed_order"]) + + +def test_every_callable_is_timed_in_every_iteration_with_a_rotating_first_callable(): + row = measure(_case(), warmup=2, iters=20) + assert len(row["timing_samples_us"]["candidate"]) == 20 + assert len(row["timing_samples_us"]["provider"]) == 20 + assert row["timing_order"][0] == ["candidate", "provider"] + assert row["timing_order"][1] == ["provider", "candidate"] + assert all(sorted(order) == ["candidate", "provider"] for order in row["timing_order"]) + assert {order[0] for order in row["timing_order"]} == {"candidate", "provider"} + + +def test_the_reported_medians_come_from_the_stored_samples(): + row = measure(_case(), warmup=2, iters=20) + for key in ("candidate", "provider"): + assert row[f"{key}_us"] == statistics.median(row["timing_samples_us"][key]) + assert row[f"{key}_us"] > 0 + assert row[f"{key}_gbps"] == pytest.approx(row["bytes"] / (row[f"{key}_us"] * 1e-6) / 1e9) + assert row[f"{key}_peak_mib"] >= 0 + # the larger matmul is slower than the smaller one in the same run + assert row["provider_us"] > row["candidate_us"] From a5d155784564a902fe6e3afe164da5d432afcb88 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:43:48 +0000 Subject: [PATCH 098/139] fix(h3): use one timing harness across the stacked PRs The stacked PRs carried three versions of measure(). #487-#494 have the most recent one (alternating order, setup outside the timed region, backward-only scope); take that block and its test here so the whole stack agrees and merging the PRs in order does not conflict. Also move the noqa: E402 of prepare_h3_weights.py to the first line of the import. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 91 ++++++++++++++--------- tests/models/minimax_h3/test_h3_report.py | 2 +- 2 files changed, 56 insertions(+), 37 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 2dff72a8a..6069de83a 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -34,37 +34,59 @@ ) -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) + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) -def peak_mib(fn: Callable[[], Any]) -> float: - """Synchronize one CUDA call and report peak allocated memory above baseline in MiB.""" +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]: - """Time case callables in rotating order after warmup and return their evidence. + """Interleave timed keys, reversing their execution order on alternate iterations. Record raw CUDA-event samples, executed orders, median microseconds, decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. @@ -74,32 +96,29 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st 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]() + 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} - 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) + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) row["timing_samples_us"] = samples - row["timing_order"] = orders + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] 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]) + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) return row diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py index 98e6ab714..b79973eb1 100644 --- a/tests/models/minimax_h3/test_h3_report.py +++ b/tests/models/minimax_h3/test_h3_report.py @@ -53,7 +53,7 @@ def provider(): 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) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) report = h3_report.measure( {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 ) From 343fa40bd64386ad0bdc11ead4891377666ab4a5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:43:51 +0000 Subject: [PATCH 099/139] fix(h3): use one timing harness across the stacked PRs The stacked PRs carried three versions of measure(). #487-#494 have the most recent one (alternating order, setup outside the timed region, backward-only scope); take that block and its test here so the whole stack agrees and merging the PRs in order does not conflict. Also move the noqa: E402 of prepare_h3_weights.py to the first line of the import. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 91 ++++++++++++++--------- tests/models/minimax_h3/test_h3_report.py | 2 +- 2 files changed, 56 insertions(+), 37 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 347f3f856..e15f837a8 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -35,37 +35,59 @@ ) -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) + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) -def peak_mib(fn: Callable[[], Any]) -> float: - """Synchronize one CUDA call and report peak allocated memory above baseline in MiB.""" +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]: - """Time case callables in rotating order after warmup and return their evidence. + """Interleave timed keys, reversing their execution order on alternate iterations. Record raw CUDA-event samples, executed orders, median microseconds, decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. @@ -75,32 +97,29 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st 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]() + 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} - 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) + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) row["timing_samples_us"] = samples - row["timing_order"] = orders + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] 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]) + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) return row diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py index 98e6ab714..b79973eb1 100644 --- a/tests/models/minimax_h3/test_h3_report.py +++ b/tests/models/minimax_h3/test_h3_report.py @@ -53,7 +53,7 @@ def provider(): 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) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) report = h3_report.measure( {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 ) From c80704ffceff2f7265fca95a37517aaa893553b3 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:49:08 +0000 Subject: [PATCH 100/139] fix(h3): use one timing harness across the stacked PRs The stacked PRs carried three versions of measure(). #487-#494 have the most recent one (alternating order, setup outside the timed region, backward-only scope); take that block and its test here so the whole stack agrees and merging the PRs in order does not conflict. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 95 ++++++++++++++--------- tests/models/minimax_h3/test_h3_report.py | 2 +- 2 files changed, 59 insertions(+), 38 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index b62a2e81b..3d640f9e2 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -29,35 +29,59 @@ ) -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]: - """Time case callables in rotating order after warmup and return their evidence. + """Interleave timed keys, reversing their execution order on alternate iterations. Record raw CUDA-event samples, executed orders, median microseconds, decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. @@ -67,32 +91,29 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st 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]() + 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} - 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) + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) row["timing_samples_us"] = samples - row["timing_order"] = orders + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] 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]) + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) return row diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py index 98e6ab714..b79973eb1 100644 --- a/tests/models/minimax_h3/test_h3_report.py +++ b/tests/models/minimax_h3/test_h3_report.py @@ -53,7 +53,7 @@ def provider(): 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) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) report = h3_report.measure( {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 ) From 31f342d76ed8eb6892420178757c9705c93b1620 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:49:09 +0000 Subject: [PATCH 101/139] fix(h3): use one timing harness and record the MLP weight source additively The stacked PRs carried three versions of measure(). #487-#494 have the most recent one (alternating order, setup outside the timed region, backward-only scope); take that block and its test here so the whole stack agrees and merging the PRs in order does not conflict. The later PRs reword the docstrings around h3_params and plot_mlp, so the weight-source change of the previous commits conflicts there. Keep h3_params as it was, name the source with a new h3_weights_source(), keep it on the accuracy report and the perf rows, and set the figure title from it after _figure() instead of editing the line that calls it. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 121 +++++++++++------- tests/models/minimax_h3/test_h3_report.py | 2 +- .../minimax_h3/test_h3_weight_source.py | 17 +-- tools/validation/models/plot_h3_evidence.py | 4 +- 4 files changed, 81 insertions(+), 63 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index ddc58a958..5d9437a01 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -30,35 +30,59 @@ ) -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]: - """Time case callables in rotating order after warmup and return their evidence. + """Interleave timed keys, reversing their execution order on alternate iterations. Record raw CUDA-event samples, executed orders, median microseconds, decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. @@ -68,32 +92,29 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st 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]() + 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} - 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) + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) row["timing_samples_us"] = samples - row["timing_order"] = orders + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] 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]) + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) return row @@ -141,21 +162,23 @@ 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, ...]]) -> tuple[list[torch.Tensor], str]: - """Pinned checkpoint tensors when available, otherwise same-shape random ones. - - Returns the tensors and their source, ``"pinned"`` or ``"random"``; the evidence stores it so a - report made without the pinned checkpoint cannot pass for one that used it. - """ +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], "pinned" + 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 - ], "random" + ] + + +def h3_weights_source() -> str: + """Source ``h3_params`` draws from: ``pinned`` checkpoint tensors or ``random`` ones.""" + + return "pinned" if h3_weights_dir() is not None else "random" # --------------------------------------------------------------------------- # @@ -168,7 +191,7 @@ def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> tuple[list[tor def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: op = registry.get_op("timestep_mlp_fp32", device="cuda") - params, weights_source = h3_params(MLP_NAMES, MLP_SHAPES) + 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): @@ -180,7 +203,7 @@ def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: "case": f"T={num}", "backend": type(op).__name__, "bytes": weight_bytes, - "weights": weights_source, + "weights": h3_weights_source(), "candidate": lambda x=x: op(x, *params), "provider": lambda x=x: provider_time_embedder(x, *params), } @@ -192,7 +215,7 @@ 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, weights_source = h3_params(MLP_NAMES, MLP_SHAPES) + params = h3_params(MLP_NAMES, MLP_SHAPES) cuda_err, provider_err = [], [] for seed in range(draws): x = sinusoid(h3_timesteps(4, seed=seed)) @@ -205,7 +228,7 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: invariant = all(torch.equal(op(x[i : i + 1], *params)[0], full[i]) for i in range(9)) return { "contract_atol": 1e-4, - "weights_source": weights_source, + "weights_source": h3_weights_source(), "draws": draws, "num_timesteps": 4, "cuda_max_abs_vs_fp64": cuda_err, diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py index 98e6ab714..b79973eb1 100644 --- a/tests/models/minimax_h3/test_h3_report.py +++ b/tests/models/minimax_h3/test_h3_report.py @@ -53,7 +53,7 @@ def provider(): 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) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) report = h3_report.measure( {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 ) diff --git a/tests/models/minimax_h3/test_h3_weight_source.py b/tests/models/minimax_h3/test_h3_weight_source.py index 0a44d84ef..a8cb005eb 100644 --- a/tests/models/minimax_h3/test_h3_weight_source.py +++ b/tests/models/minimax_h3/test_h3_weight_source.py @@ -8,27 +8,20 @@ import pytest import torch -from rl_engine.validation.models.h3_report import MLP_NAMES, MLP_SHAPES, h3_params +from rl_engine.validation.models.h3_report import h3_weights_source from rl_engine.validation.models.h3_weights import WEIGHTS_ENV, h3_weights_dir from tools.validation.models.h3_evidence import check_weight_source, measure_case -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def test_without_the_pinned_checkpoint_the_params_are_marked_random(monkeypatch): +def test_without_the_pinned_checkpoint_the_source_is_random(monkeypatch): monkeypatch.delenv(WEIGHTS_ENV, raising=False) - params, source = h3_params(MLP_NAMES, MLP_SHAPES) - assert source == "random" - assert [tuple(p.shape) for p in params] == MLP_SHAPES - assert all(p.is_cuda for p in params) + assert h3_weights_source() == "random" -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def test_with_the_pinned_checkpoint_the_params_are_marked_pinned(): +def test_with_the_pinned_checkpoint_the_source_is_pinned(): if h3_weights_dir() is None: pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py") - params, source = h3_params(MLP_NAMES, MLP_SHAPES) - assert source == "pinned" - assert [tuple(p.shape) for p in params] == MLP_SHAPES + assert h3_weights_source() == "pinned" def test_mlp_evidence_is_refused_without_pinned_weights(): diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py index bb8aee771..ecea9ab1e 100644 --- a/tools/validation/models/plot_h3_evidence.py +++ b/tools/validation/models/plot_h3_evidence.py @@ -221,7 +221,9 @@ def mlp_title(report: dict[str, Any]) -> str: def plot_mlp(report: dict[str, Any]): - fig, (left, right) = _figure(report, mlp_title(report)) + fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") + # name the weights the report was measured with instead of always claiming the pinned ones + fig.suptitle(mlp_title(report), x=0.01, ha="left", fontsize=12, color=INK, fontweight="bold") _latency_panel(left, report["perf"], "diffusers (cuBLAS)") acc = report["accuracy"] _error_boxes( From 07d319996358ac2c1a89f6c6bdd82c314f62eecf Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:28:37 +0000 Subject: [PATCH 102/139] fix(h3): count every output in the bitwise heatmap and keep results when a helper fails - plot_h3_bitwise_reuse: a cell now adds up differ/numel over all forward outputs (adaln_row_gather has six) and shows a shape or dtype mismatch as a full difference instead of reading output 0 only. - h3_prior_art: a mistyped --megatron-src no longer raises after the modes have finished and their part files were deleted; the Megatron commit is left empty. - h3_bitwise_reuse: a failing fast-math sinusoid build is recorded as unavailable instead of ending the child before the other candidates are written. - tools/migration/layout.json: map scripts/h3_bitwise_reuse.py and scripts/plot_h3_bitwise_reuse.py. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/test_h3_bitwise_reuse_plot.py | 36 +++++++++++++++++++ tools/migration/layout.json | 2 ++ tools/validation/models/h3_bitwise_reuse.py | 7 +++- tools/validation/models/h3_prior_art.py | 21 ++++++----- .../models/plot_h3_bitwise_reuse.py | 9 +++-- 5 files changed, 62 insertions(+), 13 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py diff --git a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py new file mode 100644 index 000000000..8fa02de7e --- /dev/null +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""A heatmap cell counts every forward output of the operator, not only the first.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("matplotlib") +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "tools/validation/models")) + +from plot_h3_bitwise_reuse import _cell # noqa: E402 + + +def _out(differ: int, numel: int = 100) -> dict: + return {"equal": differ == 0, "differ": differ, "numel": numel} + + +def test_cell_sums_over_all_outputs(): + assert _cell({"out": [_out(0), _out(0), _out(0)]}) == 0.0 + # a difference that is only in the fifth of six outputs is not hidden by output 0 + assert _cell({"out": [_out(0)] * 4 + [_out(30)] + [_out(0)]}) == pytest.approx(30 / 600) + + +def test_cell_marks_shape_or_dtype_mismatch_as_full_difference(): + assert _cell({"out": [_out(0), {"note": "shape mismatch"}]}) == 1.0 + + +def test_cell_without_data_is_not_applicable(): + assert _cell(None) is None + assert _cell({"out": []}) is None + assert _cell({"out": [None]}) is None diff --git a/tools/migration/layout.json b/tools/migration/layout.json index 206cbda6e..0b04ce765 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -819,9 +819,11 @@ "scripts/check_rocm_env.py": "tools/env/check_rocm_env.py", "scripts/check_stateful_kv.py": "tools/validation/operators/check_stateful_kv.py", "scripts/ci_smoke.py": "tools/checks/ci_smoke.py", + "scripts/h3_bitwise_reuse.py": "tools/validation/models/h3_bitwise_reuse.py", "scripts/h3_chain_replay.py": "tools/validation/models/h3_chain_replay.py", "scripts/h3_evidence.py": "tools/validation/models/h3_evidence.py", "scripts/h3_prior_art.py": "tools/validation/models/h3_prior_art.py", + "scripts/plot_h3_bitwise_reuse.py": "tools/validation/models/plot_h3_bitwise_reuse.py", "scripts/plot_h3_evidence.py": "tools/validation/models/plot_h3_evidence.py", "scripts/plot_h3_prior_art.py": "tools/validation/models/plot_h3_prior_art.py", "scripts/prepare_h3_weights.py": "tools/weights/prepare_h3_weights.py", diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py index d43e5c70e..345aa18cc 100644 --- a/tools/validation/models/h3_bitwise_reuse.py +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -118,7 +118,12 @@ def dump(op: str, mode: str, parts: Path, megatron_src: str | None, fast_math: b 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")}) + try: + cands.append({"name": "rl_kernel_fast_math", **_fast_math_sinusoid(parts / "build")}) + except Exception as exc: # noqa: BLE001 + (parts / "rl_kernel_fast_math.unavailable").write_text( + f"{type(exc).__name__}: {exc}"[:300] + ) for cand in cands: name = cand.pop("name") if "unavailable" in cand: diff --git a/tools/validation/models/h3_prior_art.py b/tools/validation/models/h3_prior_art.py index 9012542e2..07555ad26 100644 --- a/tools/validation/models/h3_prior_art.py +++ b/tools/validation/models/h3_prior_art.py @@ -911,15 +911,18 @@ def main() -> None: part.unlink(missing_ok=True) 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 - ) + try: + megatron_commit = ( + subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + cwd=args.megatron_src, + capture_output=True, + text=True, + ).stdout.strip() + or None + ) + except OSError: # a bad path must not discard the finished modes + megatron_commit = None report = { "kind": "h3_prior_art", "rfc": "RL-Align/RL-Kernel#420", diff --git a/tools/validation/models/plot_h3_bitwise_reuse.py b/tools/validation/models/plot_h3_bitwise_reuse.py index 85983fc0a..6da21441f 100644 --- a/tools/validation/models/plot_h3_bitwise_reuse.py +++ b/tools/validation/models/plot_h3_bitwise_reuse.py @@ -29,10 +29,13 @@ 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: + outs = [o for o in entry["out"] if o is not None] + if not outs: return None - return 0.0 if out["equal"] else out["differ"] / out["numel"] + if any("differ" not in o for o in outs): # shape or dtype mismatch + return 1.0 + numel = sum(o["numel"] for o in outs) + return sum(o["differ"] for o in outs) / numel if numel else None def _text(frac: float | None) -> str: From 92175dbe23351b8a2f4190e1a38bc6c69048cdcd Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:31:25 +0000 Subject: [PATCH 103/139] fix(h3): refuse a random-weight MLP evidence run before measuring h3_evidence.py checked the weight source only after the accuracy and every performance case had run, so a run without RL_KERNEL_H3_WEIGHTS and without --allow-random-weights spent the whole GPU time and then wrote nothing. require_pinned_weights() makes the same decision from h3_weights_source() right after the arguments are parsed; the check on the finished report stays. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/test_h3_weight_source.py | 14 +++++++++- tools/validation/models/h3_evidence.py | 27 ++++++++++++++----- 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/tests/models/minimax_h3/test_h3_weight_source.py b/tests/models/minimax_h3/test_h3_weight_source.py index a8cb005eb..ad1cf09ae 100644 --- a/tests/models/minimax_h3/test_h3_weight_source.py +++ b/tests/models/minimax_h3/test_h3_weight_source.py @@ -10,7 +10,11 @@ from rl_engine.validation.models.h3_report import h3_weights_source from rl_engine.validation.models.h3_weights import WEIGHTS_ENV, h3_weights_dir -from tools.validation.models.h3_evidence import check_weight_source, measure_case +from tools.validation.models.h3_evidence import ( + check_weight_source, + measure_case, + require_pinned_weights, +) def test_without_the_pinned_checkpoint_the_source_is_random(monkeypatch): @@ -35,6 +39,14 @@ def test_mlp_evidence_is_refused_without_pinned_weights(): check_weight_source("timestep_sinusoid_h3", {"accuracy": {}}, allow_random=False) +def test_mlp_evidence_is_refused_before_measuring_without_pinned_weights(monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + with pytest.raises(SystemExit, match="pinned time_embedder weights"): + require_pinned_weights("timestep_mlp_fp32", allow_random=False) + require_pinned_weights("timestep_mlp_fp32", allow_random=True) + require_pinned_weights("timestep_sinusoid_h3", allow_random=False) + + def test_the_figure_title_names_the_recorded_weight_source(): pytest.importorskip("matplotlib") from tools.validation.models.plot_h3_evidence import mlp_title diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index e67333b8b..5aa5e4f73 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -25,9 +25,20 @@ from rl_engine.runtime.registry import KernelRegistry # noqa: E402 from rl_engine.validation.models.h3_chain import environment, git_state # noqa: E402 -from rl_engine.validation.models.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 +from rl_engine.validation.models.h3_report import ( # noqa: E402 + ACCURACY, + PERF_CASES, + h3_weights_source, + measure, +) from rl_engine.validation.models.h3_weights import load_h3_manifest # noqa: E402 +_NEEDS_PINNED_MLP = ( + "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " + "RL_KERNEL_H3_WEIGHTS (tools/weights/prepare_h3_weights.py) " + "or pass --allow-random-weights" +) + def measure_case(case: dict) -> dict: """Time one perf case and keep the weight source the case was built with next to the timings.""" @@ -46,11 +57,14 @@ def check_weight_source(op: str, report: dict, allow_random: bool) -> None: and report["accuracy"]["weights_source"] != "pinned" and not allow_random ): - raise SystemExit( - "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " - "RL_KERNEL_H3_WEIGHTS (tools/weights/prepare_h3_weights.py) " - "or pass --allow-random-weights" - ) + raise SystemExit(_NEEDS_PINNED_MLP) + + +def require_pinned_weights(op: str, allow_random: bool) -> None: + """Stop before any measurement when timestep_mlp_fp32 would run on random weights.""" + + if op == "timestep_mlp_fp32" and h3_weights_source() != "pinned" and not allow_random: + raise SystemExit(_NEEDS_PINNED_MLP) def main() -> None: @@ -64,6 +78,7 @@ def main() -> None: "(smoke runs only)", ) args = parser.parse_args() + require_pinned_weights(args.op, args.allow_random_weights) if not torch.cuda.is_available(): raise SystemExit("needs a CUDA device") torch.backends.cuda.matmul.allow_tf32 = False From 7bb309240f831a5720c31c3baec98ad77210d12f Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:32:33 +0000 Subject: [PATCH 104/139] docs(h3): parameterize the sinusoid numerics and say which commits wrote the reports - Name the current-history commits of the two report files and note that the JSON files record the pre-rewrite hashes (7a82917, 0eb5ddd5). - State the frequency count and the gradient reduction width as num_channels / 2 and num_channels (128 and 256 for H3). - The latency column is forward-only because the prior-art candidates have no backward, not because the operator lacks one. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-sinusoid.md | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 5696f5df7..03b9b489d 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -53,8 +53,8 @@ the host synchronization. Pass `check_range=False` to `forward` or ## 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`. + runs on CUDA: `(float)(-ln 10000) * k`, multiplied by the FP32 reciprocal of `num_channels / 2` + (128 for H3; 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. @@ -64,7 +64,7 @@ the host synchronization. Pass `check_range=False` to `forward` or 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 + the `num_channels` channels (256 for H3) 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): @@ -105,11 +105,11 @@ records that commit and the environment. | 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 time (the op has no backward) on an otherwise idle B200. Batch invariance is bitwise and covers +forward time (the prior-art candidates were benchmarked forward-only) 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`](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 +was written from a clean tree at `0d3794de` (`7a82917` in the JSON, see the note before Tests) by ```bash python tools/validation/models/h3_prior_art.py --op timestep_sinusoid --out reports/experiments/h3-prior-art-b200/timestep_sinusoid.json @@ -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 `be6a94ae` by +was written from a clean tree at `be6a94ae` (`0eb5ddd5` in the JSON, see the note before Tests) by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ @@ -141,6 +141,9 @@ python tools/validation/models/plot_h3_bitwise_reuse.py \ --out reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png ``` +The JSON files record the commits as they were before the branch history was rewritten on 2026-10-23 +(`7a82917`, `0eb5ddd5`); `0d3794de` and `be6a94ae` are the same commits in the current history. + ## Tests ```bash From 90990842c2046bc65d6c6e276d60a9776389f1d4 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:35:00 +0000 Subject: [PATCH 105/139] fix(h3): refuse the projection op below SM80 and state the activation cast of the weight dtype - H3AdaLNProjectionCudaOp checks the device in __init__ (det_linear_device_supported), so on a pre-SM80 GPU the registry marks the backend failed and falls back to the PyTorch reference instead of failing at every BF16 call. - The formula in the operator doc casts the activation to the weight dtype like the wrapper does (BF16 for the checkpoint weights, FP32 is allowed by the contract). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 2 +- .../cuda/model_specific/minimax_h3/adaln_projection.py | 5 ++++- .../cuda/model_specific/minimax_h3/det_linear.py | 6 ++++++ tests/models/minimax_h3/test_h3_adaln_projection.py | 9 +++++++++ 4 files changed, 20 insertions(+), 2 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 3f79732cd..5de612b3f 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -7,7 +7,7 @@ modulation tensors of one transformer block, for all three modalities: ```text -act = silu(temb).to(bf16) SiLU in FP32, one declared cast +act = silu(temb).to(W.dtype) SiLU in FP32, one declared cast (BF16 for the checkpoint weights) 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 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..1363d8a44 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,16 +8,17 @@ 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, det_linear_available, det_linear_backward_input, det_linear_backward_weight, + det_linear_device_supported, 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, @@ -98,6 +99,8 @@ def __init__(self) -> None: raise RuntimeError( "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" ) + if not det_linear_device_supported(): + raise RuntimeError("the BF16 h3 det_linear forward needs SM80 or newer") def __call__(self, temb, weight, bias): """Return six weight-dtype ``(3T, H)`` modulation views for CUDA inputs.""" diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py index 34d523464..d288d2fa9 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py @@ -29,6 +29,12 @@ def det_linear_available() -> bool: return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) +def det_linear_device_supported() -> bool: + """Return whether the current CUDA device can run the BF16 forward (``mma.sync``, SM80+).""" + + return not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] >= 8 + + def _require() -> None: """Raise when a deterministic linear entry point is unavailable.""" diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 5fb10a73f..f35b13a79 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -41,6 +41,15 @@ def _cuda_op(): return H3AdaLNProjectionCudaOp() +def test_pre_sm80_device_is_refused_at_construction_so_the_registry_can_fall_back(monkeypatch): + from rl_engine.backends.cuda.model_specific.minimax_h3 import adaln_projection as module + + monkeypatch.setattr(module, "det_linear_available", lambda: True) + monkeypatch.setattr(module, "det_linear_device_supported", lambda: False) + with pytest.raises(RuntimeError, match="SM80"): + module.H3AdaLNProjectionCudaOp() + + def _temb(num, dim=2688, seed=0, device="cuda"): """Generate seeded FP32 timestep embeddings on CPU and move them to the test device.""" From 9ed5d9c1296828b76b99914f9f0580f280d121c7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:41:03 +0000 Subject: [PATCH 106/139] fix(h3): reject indices on another device than temb, pre-SM80 devices, and an elided gather example - H3AdaLNModulationCudaOp validated the indices against a stand-in allocated on the indices' own device, so the device check was always true; the stand-in now sits on temb's device and a mismatch raises ValueError before the native call (also with check_range=False). - H3AdaLNModulationCudaOp refuses a pre-SM80 device in __init__ like the projection op, so the registry falls back to the PyTorch reference. - The row-gather doc passes all six tensors to gather_chunks instead of an Ellipsis. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 4 +++- .../minimax_h3/adaln_modulation.py | 18 ++++++++++-------- .../minimax_h3/test_h3_adaln_modulation.py | 17 +++++++++++++++++ 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index a6bec87f8..8d0eb0584 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -22,7 +22,9 @@ from rl_engine.runtime.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 +outs = op.gather_chunks( + (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp), timestep_indices, token_tags +) # drop-in, six (T*3, 5376) tensors ``` ## Backends 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..9e7cbaf70 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, @@ -26,9 +24,12 @@ det_linear_available, det_linear_backward_input, det_linear_backward_weight, + det_linear_device_supported, 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 ( @@ -108,12 +109,14 @@ class H3AdaLNModulationCudaOp: 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") + if not det_linear_device_supported(): + raise RuntimeError("the BF16 h3 det_linear forward needs SM80 or newer") 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) + _validate_indices(timestep_indices, token_tags, temb.shape[0], temb.device, check_range) packed = _H3AdaLNModulationCuda.apply( temb, weight, @@ -125,9 +128,8 @@ def __call__(self, temb, weight, bias, timestep_indices, token_tags, *, check_ra 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 - ) +def _validate_indices(timestep_indices, token_tags, num_timesteps, device, check_range): + # Reuse the gather's checks against a stand-in of the right row count on temb's device, so + # indices on another device are rejected here. + stand_in = torch.empty((num_timesteps * H3_MODALITY_NUM, H3_ADALN_CHUNKS), device=device) validate_h3_adaln_row_gather(stand_in, timestep_indices, token_tags, check_range=check_range) diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py index 066445fc0..dc27820dd 100644 --- a/tests/models/minimax_h3/test_h3_adaln_modulation.py +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -56,8 +56,25 @@ def _leaf_grads(fn, temb, weight, bias, grads): return [leaf.grad for leaf in leaves] +def test_pre_sm80_device_is_refused_at_construction_so_the_registry_can_fall_back(monkeypatch): + from rl_engine.backends.cuda.model_specific.minimax_h3 import adaln_modulation as module + + monkeypatch.setattr(module, "det_linear_available", lambda: True) + monkeypatch.setattr(module, "adaln_row_gather_available", lambda: True) + monkeypatch.setattr(module, "det_linear_device_supported", lambda: False) + with pytest.raises(RuntimeError, match="SM80"): + module.H3AdaLNModulationCudaOp() + + @requires_cuda class TestFusedModulation: + def test_indices_on_another_device_than_temb_are_rejected(self): + fused, _, _ = _ops() + temb, weight, bias = _params(64, 2688) + ti, tags = h3_packed_layout(1000, 3) + with pytest.raises(ValueError, match="rows are on"): + fused(temb, weight, bias, ti.cpu(), tags.cpu(), check_range=False) + def test_forward_bitwise_equal_to_separate_ops(self): fused, proj, gather = _ops() temb, weight, bias = _params(64, 2688) From 0f67234e5f3dffa6820c8b31801bbe096cce31c0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:06:55 +0000 Subject: [PATCH 107/139] fix(h3): bound-check the RMSNorm modulation row and the row-gather backward metadata - rmsnorm_modulate.cu: Modulation carries the table row count; the forward, dx and dweight kernels assert the selected row is in [0, R) before computing a shift/scale offset, as the row-gather kernel does for check_range=False callers. - adaln_row_gather.cu: h3_adaln_row_gather_backward requires the tile metadata on grad's device and a non-empty seg_first_tile (an empty one made num_rows -1). - docs/operators/h3-rmsnorm.md names the commit that wrote report.json in the current history (d2a6546, recorded as 80e4609 in the report and the figure footer) and the prior-art commit. - Tests: device-side assert for rows -1 and R in forward and backward (subprocess), metadata checks for the gather backward. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 7 +-- csrc/cuda/h3/rmsnorm_modulate.cu | 18 +++++-- docs/operators/h3-rmsnorm.md | 5 +- .../minimax_h3/test_h3_adaln_row_gather.py | 47 ++++++++++++++++- .../test_h3_rmsnorm_native_validation.py | 51 +++++++++++++++++++ 5 files changed, 117 insertions(+), 11 deletions(-) diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..838142b1e 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -207,10 +207,11 @@ 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() && - t->scalar_type() == at::kLong, - "tile metadata must be contiguous 1-D int64 CUDA tensors"); + 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 tensors on grad's device"); } + TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); 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); diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 4eef07be1..0165a1bb1 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -105,8 +105,17 @@ 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 rows; // R: table rows; an index outside [0, R) is a caller error }; +// Table row of x row `r`. check_range=False callers skip the Python range check, so the native side +// asserts the index before it is used as an offset (as the row-gather kernel does). +__device__ __forceinline__ int64_t modulation_row(const Modulation& mod, int64_t r) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.rows); + return i; +} + template __global__ void __launch_bounds__(32 * kWarps) rmsnorm_modulate_fwd_kernel(const T* __restrict__ x, const T* __restrict__ w, T* __restrict__ y, @@ -118,7 +127,7 @@ __global__ void __launch_bounds__(32 * kWarps) const T* shift = nullptr; const T* scale = nullptr; if (mod.index != nullptr) { - const int64_t i = mod.index[r % mod.seq]; + const int64_t i = modulation_row(mod, r); shift = static_cast(mod.shift) + i * mod.row_stride; scale = static_cast(mod.scale) + i * mod.row_stride; } @@ -156,7 +165,7 @@ __global__ void __launch_bounds__(32 * kWarps) 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 + ? static_cast(mod.scale) + modulation_row(mod, r) * mod.row_stride : nullptr; const int numx = blockDim.x * blockDim.y; const int thrx = threadIdx.x + threadIdx.y * blockDim.x; @@ -193,7 +202,7 @@ __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* 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 + modulation_row(mod, r) * mod.row_stride : nullptr; acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); } @@ -265,7 +274,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; @@ -291,6 +300,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.rows = sh.size(0); return mod; } diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index b64c5e44e..50fd4bf17 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -93,7 +93,8 @@ 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 `d2a6546`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546` (`80e4609` in the report and the figure footer, which predate +the history rewrite of 2026-10-23). It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -118,7 +119,7 @@ forward + backward time on an otherwise idle B200. Batch invariance is bitwise a 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`](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 +was written from a clean tree at `f6ad7a23` (`ee83dec` in the JSON, same rewrite) by ```bash python tools/validation/models/h3_prior_art.py --op norm_modulate --out reports/experiments/h3-prior-art-b200/norm_modulate.json 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..1a5d548ab 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], @@ -296,3 +298,44 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("adaln_row_gather", device="cuda") assert type(op).__name__ == "H3AdaLNRowGatherCudaOp" + + +def _backward_args(device="cuda"): + kw = {"device": device, "dtype": torch.int64} + return { + "grad": torch.ones((6, 4, 8), device="cuda"), + "sorted_pos": torch.tensor([0, 2, 1, 3], **kw), + "tile_begin": torch.tensor([0, 2], **kw), + "tile_end": torch.tensor([2, 4], **kw), + "seg_first_tile": torch.tensor([0, 1, 2], **kw), + "out_dtype": torch.float32, + } + + +@requires_cuda +class TestNativeBackwardMetadata: + def test_a_valid_call_runs(self): + _cuda_op() + from rl_engine.backends.extension import _C + + assert _C.h3_adaln_row_gather_backward(**_backward_args()).shape == (2, 48) + + def test_empty_segment_table_is_rejected(self): + _cuda_op() + from rl_engine.backends.extension import _C + + args = _backward_args() + args["seg_first_tile"] = args["seg_first_tile"][:0] + with pytest.raises(RuntimeError, match="seg_first_tile"): + _C.h3_adaln_row_gather_backward(**args) + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") + @pytest.mark.parametrize("name", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) + def test_metadata_on_another_device_than_grad_is_rejected(self, name): + _cuda_op() + from rl_engine.backends.extension import _C + + args = _backward_args() + args[name] = args[name].to("cuda:1") + with pytest.raises(RuntimeError, match="grad's device"): + _C.h3_adaln_row_gather_backward(**args) diff --git a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py index ff94c7924..a2b032344 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py @@ -5,6 +5,12 @@ from __future__ import annotations +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest import torch @@ -156,3 +162,48 @@ 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("entrypoint", ["forward", "backward"]) +@pytest.mark.parametrize("bad_row", [-1, 2], ids=["negative", "past-table"]) +def test_modulation_row_outside_the_table_asserts_in_the_kernel(native, entrypoint, bad_row): + """check_range=False callers get a device-side assert, not a read past the shift/scale table.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + entrypoint, bad_row = sys.argv[1], int(sys.argv[2]) + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device="cuda") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) # a table of two rows + index = torch.tensor([0, bad_row], dtype=torch.int64, device="cuda") + if entrypoint == "forward": + _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + else: + _C.h3_rmsnorm_backward( + grad=torch.ones_like(x), x=x, weight=weight, rstd=rstd, shift=shift, scale=scale, + index=index, + sorted_pos=torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cuda"), + tile_begin=torch.tensor([0, 2], dtype=torch.int64, device="cuda"), + tile_end=torch.tensor([2, 4], dtype=torch.int64, device="cuda"), + seg_first_tile=torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"), + ) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(bad_row)], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + assert completed.returncode != 0, "the kernel accepted a modulation row outside the table" + assert ( + "device-side assert triggered" in completed.stderr + ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" From 669513d223d33fbc8d0af14b1b00733a50c858fd Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:07:41 +0000 Subject: [PATCH 108/139] fix(h3): keep the gate-residual forward out of the timed backward _gate_perf had no candidate_backward_setup/provider_backward_setup, so each timed backward call also created the leaves, ran the forward and built the graph while measure() reported backward_timing_scope = backward_only. The forward moves into the setup like in _gather_perf and _norm_perf. A test checks that every perf case with a timed backward has its setup. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 13 +++++++++---- .../models/minimax_h3/test_h3_report_timing.py | 18 ++++++++++++++++++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 492bfc1b6..aa1b53498 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -606,9 +606,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( { @@ -622,10 +625,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_) ), } diff --git a/tests/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py index 853680a67..9b77b16cf 100644 --- a/tests/models/minimax_h3/test_h3_report_timing.py +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -52,3 +52,21 @@ def test_the_reported_medians_come_from_the_stored_samples(): assert row[f"{key}_peak_mib"] >= 0 # the larger matmul is slower than the smaller one in the same run assert row["provider_us"] > row["candidate_us"] + + +def test_backward_timings_of_every_perf_case_exclude_the_forward_pass(monkeypatch): + """A case with a timed backward must prepare the graph outside the CUDA-event window.""" + + from rl_engine.runtime.registry import KernelRegistry + from rl_engine.validation.models import h3_report + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (257,)) + registry = KernelRegistry() + checked = 0 + for op, build in h3_report.PERF_CASES.items(): + for case in build(registry): + for side in ("candidate", "provider"): + if f"{side}_backward" in case: + checked += 1 + assert f"{side}_backward_setup" in case, f"{op} {case['case']} {side}" + assert checked From e5318ad2ab1a75100336762b08d057b65d5a853d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:11:52 +0000 Subject: [PATCH 109/139] fix(h3): do not ask for the checkpoint when the gate-residual evidence never reads it h3_evidence.py treated every op but timestep_sinusoid_h3 as weighted. adaln_gate_residual runs on seeded inputs only, so the CLI refused to run without the 4.8 GB shard and a run with it would have been labelled pinned_checkpoint. The weight-free ops are now an explicit set. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_cli.py | 1 + tools/validation/models/h3_evidence.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_cli.py b/tests/models/minimax_h3/test_h3_cli.py index 8e97fc6a2..3f9b726ed 100644 --- a/tests/models/minimax_h3/test_h3_cli.py +++ b/tests/models/minimax_h3/test_h3_cli.py @@ -145,6 +145,7 @@ def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, we "op, weight_source", [ ("timestep_sinusoid_h3", "not_applicable"), + ("adaln_gate_residual", "not_applicable"), ("timestep_mlp_fp32", "pinned_checkpoint"), ("adaln_projection_3mod", "pinned_checkpoint"), ], diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index ed8267a94..892366786 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -37,6 +37,9 @@ load_h3_manifest, ) +# Operators whose evidence runs on seeded inputs only and never reads the checkpoint. +WEIGHT_FREE_OPS = frozenset({"timestep_sinusoid_h3", "adaln_gate_residual"}) + _NEEDS_PINNED_MLP = ( "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " "RL_KERNEL_H3_WEIGHTS (tools/weights/prepare_h3_weights.py) " @@ -85,7 +88,7 @@ def main() -> None: ) args = parser.parse_args() require_pinned_weights(args.op, args.allow_random_weights) - uses_weights = args.op != "timestep_sinusoid_h3" + uses_weights = args.op not in WEIGHT_FREE_OPS smoke_run = args.allow_random_weights and args.op == "timestep_mlp_fp32" if uses_weights and h3_weights_dir() is None and not smoke_run: raise SystemExit(f"set {WEIGHTS_ENV} to the tools/weights/prepare_h3_weights.py output dir") From e6ee9327d26129f8a1fd4dfe32af7041fbc83ccc Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:12:13 +0000 Subject: [PATCH 110/139] test(h3): give the stubbed MLP accuracy a weight source check_weight_source reads accuracy["weights_source"] for timestep_mlp_fp32 since the early check of #484 was merged in; the stub returned an empty dict and the evidence CLI test failed with a KeyError. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_cli.py b/tests/models/minimax_h3/test_h3_cli.py index 8e97fc6a2..0732760b6 100644 --- a/tests/models/minimax_h3/test_h3_cli.py +++ b/tests/models/minimax_h3/test_h3_cli.py @@ -166,7 +166,7 @@ def test_evidence_records_weight_source(monkeypatch, tmp_path, op, weight_source 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.ACCURACY, op, lambda registry: {"weights_source": "pinned"}) 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)]) From 0f66d75515b48375a9dcc9f0ec91498ee5d29050 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:52:11 +0000 Subject: [PATCH 111/139] fix(h3): do not consume the weight registers of columns past n_out The forward kernel loads weights only for n0 + c < n_out, but the loop that unpacks them and feeds the FMAs ran for every c < kCols. On the last column tile of an N that is not a multiple of kCols (any odd N for FP32, kCols = 2; N % 4 for BF16) it read raw[u][c] uninitialized. Those columns are never stored, but the read is undefined behaviour. The loop is guarded with n0 + c < n_out like the load. test_h3_det_linear_odd_width.py covers N in {1, 2, 3, 5, 7, 9, 33, 65} for FP32 and BF16: forward, input gradient, weight gradient and bias gradient against an FP64 golden, and bitwise against the same product with zero columns padded on. Run on a fresh build; the other kernels of det_linear.cu index their columns by n < n_out already. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 10 ++- .../test_h3_det_linear_odd_width.py | 86 +++++++++++++++++++ 2 files changed, 93 insertions(+), 3 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_det_linear_odd_width.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index ddaf199ee..7af6592b4 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -149,10 +149,14 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) 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); + // raw[u][c] is only loaded for columns inside the output; the last column tile of + // an N that is not a multiple of kCols must not read the other registers. + if (n0 + c < n_out) { + 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]); + for (int j = 0; j < kVec; ++j) acc[c][r] = fmaf(xv[j], wv[j], acc[c][r]); + } } } } diff --git a/tests/models/minimax_h3/test_h3_det_linear_odd_width.py b/tests/models/minimax_h3/test_h3_det_linear_odd_width.py new file mode 100644 index 000000000..19f95784c --- /dev/null +++ b/tests/models/minimax_h3/test_h3_det_linear_odd_width.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Output widths that are not a multiple of the column tile of the deterministic linear. + +Each warp of the forward kernel owns kCols adjacent output columns (2 for FP32, 4 for BF16), so the +last tile of an odd or short N has columns that are not loaded and must not be consumed. The +results for such widths must match an FP64 golden, and must be bitwise the leading columns of the +same product with zero columns padded on (the tail does not influence the real columns). +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, +) + +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_*"), +] + +DTYPES = [torch.float32, torch.bfloat16] +WIDTHS = [1, 2, 3, 5, 7, 9, 33, 65] # odd widths, and tails past a 2- or 4-column tile +K_IN = 64 +ROWS = [1, 3, 5] + + +def _inputs(dtype, rows, n_out, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed + n_out) + x = torch.randn(rows, K_IN, generator=g).to(dtype).cuda() + weight = (torch.randn(n_out, K_IN, generator=g) / K_IN**0.5).to(dtype).cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + return x, weight, bias + + +def _pad(weight, bias, extra): + zeros = torch.zeros(extra, weight.shape[1], dtype=weight.dtype, device=weight.device) + return torch.cat([weight, zeros]), torch.cat([bias, bias.new_zeros(extra)]) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("n_out", WIDTHS) +@pytest.mark.parametrize("rows", ROWS) +def test_forward_matches_golden_and_ignores_the_tail(dtype, n_out, rows): + x, weight, bias = _inputs(dtype, rows, n_out) + out = det_linear_forward(x, weight, bias)[0] + assert out.shape == (rows, n_out) and torch.isfinite(out.float()).all() + golden = x.double() @ weight.double().T + bias.double() + tol = 1e-5 if dtype == torch.float32 else 1e-2 + torch.testing.assert_close(out.double(), golden, atol=tol, rtol=tol) + padded = det_linear_forward(x, *_pad(weight, bias, 3))[0] + assert torch.equal(out, padded[:, :n_out]) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("n_out", WIDTHS) +def test_backward_matches_golden_and_ignores_the_tail(dtype, n_out): + rows = 3 + x, weight, _ = _inputs(dtype, rows, n_out) + g = torch.Generator(device="cpu").manual_seed(1) + grad = torch.randn(rows, n_out, generator=g).cuda() + tol = 1e-5 if dtype == torch.float32 else 1e-2 + + d_input = det_linear_backward_input(grad, weight, torch.float32) + torch.testing.assert_close( + d_input.double(), grad.double() @ weight.double(), atol=tol, rtol=tol + ) + d_weight, d_bias = det_linear_backward_weight(grad, x, dtype, with_bias=True) + torch.testing.assert_close(d_weight.double(), grad.double().T @ x.double(), atol=tol, rtol=tol) + torch.testing.assert_close(d_bias.double(), grad.double().sum(0), atol=tol, rtol=tol) + + zero = torch.zeros(rows, 3, device="cuda") + padded_weight, _ = _pad(weight, weight[:, 0], 3) + padded_grad = torch.cat([grad, zero], dim=1) + assert torch.equal( + d_input, det_linear_backward_input(padded_grad, padded_weight, torch.float32) + ) + d_weight_p, d_bias_p = det_linear_backward_weight(padded_grad, x, dtype, with_bias=True) + assert torch.equal(d_weight, d_weight_p[:n_out]) and torch.equal(d_bias, d_bias_p[:n_out]) From c0b4cd83d786453da94062a578c8b0ea687aeb58 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:09:22 +0000 Subject: [PATCH 112/139] fix(build): decide the H3 sources and bindings from one predicate, and reject relaxed math from the environment - build_tools/extensions.py: h3_enabled is computed once and gates both the csrc/cuda/h3 sources and -DKERNEL_ALIGN_WITH_H3_KERNELS, so the bindings cannot outlive the kernels. - NVCC_PREPEND_FLAGS / NVCC_APPEND_FLAGS / TORCH_NVCC_FLAGS are appended by PyTorch's BuildExtension after the flags built here; when they carry --use_fast_math, --ftz=true or --prec-div/--prec-sqrt=false the H3 kernels are left out (with a warning) like in a KERNEL_ALIGN_USE_FAST_MATH=1 build, and dispatch falls back to the PyTorch reference. - H3TimestepSinusoid backward checks the native output shape before slicing it as [cos | sin]. - tests/build: environment flags exclude the kernels; unrelated flags keep them. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- build_tools/extensions.py | 39 ++++++++++++++++++- .../minimax_h3/timestep_sinusoid.py | 6 +++ .../build/test_build_platform_collectives.py | 28 +++++++++++++ 3 files changed, 71 insertions(+), 2 deletions(-) diff --git a/build_tools/extensions.py b/build_tools/extensions.py index 9ae188071..6362a1c58 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -3,6 +3,7 @@ import importlib.util import os +import shlex import sysconfig import warnings from distutils.errors import CompileError @@ -21,6 +22,29 @@ # 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" +# nvcc options that relax FP32 math. They can also arrive through the environment, where PyTorch's +# BuildExtension appends them after the flags built here, so the H3 kernels check for them too. +_RELAXED_MATH_NVCC_FLAGS = frozenset( + { + "--use_fast_math", + "-use_fast_math", + "--ftz=true", + "-ftz=true", + "--prec-div=false", + "-prec-div=false", + "--prec-sqrt=false", + "-prec-sqrt=false", + } +) + + +def _external_relaxed_math() -> bool: + """Return whether NVCC_*_FLAGS / TORCH_NVCC_FLAGS turn on relaxed FP32 math.""" + + for name in ("NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS", "TORCH_NVCC_FLAGS"): + if _RELAXED_MATH_NVCC_FLAGS & set(shlex.split(os.environ.get(name, ""))): + return True + return False def _resolve_package_data_source( @@ -228,6 +252,17 @@ def get_extensions(): if is_rocm or torch.cuda.is_available(): use_fast_math = envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH) + # One predicate decides both the H3 sources and the define that enables their bindings, so + # the bindings can never outlive the kernels. + relaxed_math_env = _external_relaxed_math() + if relaxed_math_env and not is_rocm: + warnings.warn( + "NVCC_*_FLAGS / TORCH_NVCC_FLAGS relax FP32 math; the bitwise MiniMax-H3 kernels " + "are left out of this build and dispatch falls back to the PyTorch reference.", + RuntimeWarning, + stacklevel=2, + ) + h3_enabled = not is_rocm and not use_fast_math and not relaxed_math_env cuda_sources = [ "csrc/bindings/ops.cpp", "csrc/cuda/logprob/fused_logp_kernel.cu", @@ -250,7 +285,7 @@ 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") - if not use_fast_math: + if h3_enabled: 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. @@ -344,7 +379,7 @@ 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 not is_rocm and not use_fast_math: + if h3_enabled: cxx_flags.append(H3_KERNELS_DEFINE) nvcc_flags.append(H3_KERNELS_DEFINE) 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 b792d3f73..5f3ef468a 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 @@ -40,6 +40,12 @@ def forward(ctx, timestep: torch.Tensor, num_channels: int, check_range: bool): def backward(ctx, grad_out: torch.Tensor): (out,) = ctx.saved_tensors half = ctx.num_channels // 2 + # The VJP slices the native output as [cos | sin] over frequencies_fp32 order. + if out.dim() != 2 or out.shape[-1] != ctx.num_channels: + raise RuntimeError( + f"native timestep features have shape {tuple(out.shape)}, expected " + f"(T, {ctx.num_channels}) in [cos | sin] order" + ) freq = NativeH3TimestepSinusoidOp.frequencies_fp32(ctx.num_channels, out.device) cos_part, sin_part = out[:, :half], out[:, half:] g = grad_out.float() diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index 820347673..9d4ebcfde 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -6,6 +6,7 @@ import runpy from typing import Any +import pytest import setuptools import torch from torch.utils import cpp_extension @@ -113,3 +114,30 @@ def test_fast_math_build_excludes_bitwise_h3_kernels(monkeypatch) -> None: 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"] + + +@pytest.mark.parametrize( + "variable", ["NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS", "TORCH_NVCC_FLAGS"] +) +def test_relaxed_math_nvcc_flags_from_the_environment_exclude_bitwise_h3_kernels( + monkeypatch, variable +) -> None: + # The flags are appended after the ones built here, so a build with + # NVCC_APPEND_FLAGS='--use_fast_math' would otherwise ship H3 kernels with approximated math. + monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) + monkeypatch.setenv(variable, "-O2 --use_fast_math --ftz=true") + with pytest.warns(RuntimeWarning, match="bitwise MiniMax-H3 kernels"): + extension = _load_extension_config(monkeypatch, hip=None) + + 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"] + + +def test_unrelated_nvcc_flags_keep_the_h3_kernels(monkeypatch) -> None: + monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) + monkeypatch.setenv("NVCC_APPEND_FLAGS", "-lineinfo --fmad=true") + extension = _load_extension_config(monkeypatch, hip=None) + + assert set(H3_BITWISE_SOURCES) <= set(extension["sources"]) + assert H3_KERNELS_DEFINE in extension["extra_compile_args"]["nvcc"] From 0e911bed680d8932407106ef9650b25f63af7316 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:11:36 +0000 Subject: [PATCH 113/139] fix(h3): fail at build time or with a clear error when det_linear launch shapes overflow - det_linear_forward_kernel: static_assert(kCols * kRows <= kWarp). The epilogue stores through lane == c * kRows + r, so a tuned config above a warp would silently drop stores. - h3_det_linear_backward_input: TORCH_CHECK on the chunk count (blockIdx.y <= 65535) and on the fold launch size instead of 'invalid configuration argument'. - test_h3_det_linear_odd_width.py: K that is a vector multiple but not a multiple of the 32-lane stride (BF16 8, 24, 40, 72, 264, 520; FP32 4, 12, 20, 36, 132, 1044) against the FP64 golden. The same shapes, with the backward kernels, ran under compute-sanitizer memcheck with 0 errors. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 9 +++++++++ .../test_h3_det_linear_odd_width.py | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 7af6592b4..13798a0e8 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -35,6 +35,7 @@ #include #include +#include #include namespace { @@ -110,6 +111,8 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) int64_t k_in, int64_t activation) { using WV = Vec16; using XV = Vec16; + static_assert(kCols * kRows <= kWarp, + "the epilogue stores through lane == c * kRows + r, so kCols * kRows must fit a warp"); constexpr int kVec = WV::kN; static_assert(XV::kN == kVec, "x and weight share a dtype"); constexpr int64_t kStride = static_cast(kWarp) * kVec; @@ -405,6 +408,12 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei 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; + // chunks is blockIdx.y and rows * k_in / threads is blockIdx.x of the fold: fail with a clear message + // instead of "invalid configuration argument" when a launch dimension would overflow. + TORCH_CHECK(chunks <= 65535, "N=", n_out, " needs ", chunks, + " input-gradient chunks; at most 65535 are supported"); + TORCH_CHECK((rows * k_in + 255) / 256 <= std::numeric_limits::max(), "rows * K = ", + rows * k_in, " is too large for one input-gradient fold launch"); 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(); diff --git a/tests/models/minimax_h3/test_h3_det_linear_odd_width.py b/tests/models/minimax_h3/test_h3_det_linear_odd_width.py index 19f95784c..2b0ee1aac 100644 --- a/tests/models/minimax_h3/test_h3_det_linear_odd_width.py +++ b/tests/models/minimax_h3/test_h3_det_linear_odd_width.py @@ -84,3 +84,22 @@ def test_backward_matches_golden_and_ignores_the_tail(dtype, n_out): ) d_weight_p, d_bias_p = det_linear_backward_weight(padded_grad, x, dtype, with_bias=True) assert torch.equal(d_weight, d_weight_p[:n_out]) and torch.equal(d_bias, d_bias_p[:n_out]) + + +@pytest.mark.parametrize( + "dtype, k_in", + [(torch.bfloat16, k) for k in (8, 24, 40, 72, 264, 520)] + + [(torch.float32, k) for k in (4, 12, 20, 36, 132, 1044)], +) +@pytest.mark.parametrize("n_out", [1, 3, 33]) +def test_k_that_is_not_a_multiple_of_the_lane_stride_matches_the_golden(dtype, k_in, n_out): + """Row lengths that are a vector multiple but not a multiple of the 32-lane stride.""" + + g = torch.Generator(device="cpu").manual_seed(k_in + n_out) + x = torch.randn(3, k_in, generator=g).to(dtype).cuda() + weight = (torch.randn(n_out, k_in, generator=g) / k_in**0.5).to(dtype).cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + out = det_linear_forward(x, weight, bias)[0] + tol = 1e-5 if dtype == torch.float32 else 1e-2 + golden = x.double() @ weight.double().T + bias.double() + torch.testing.assert_close(out.double(), golden, atol=tol, rtol=tol) From f2f51f3b28719097e5e00d1db6fd6e8d6c32dd46 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:14:55 +0000 Subject: [PATCH 114/139] fix(h3): state the position contract of the gather backward and the no-activation projection - _repeat_positions() builds the table row of every flattened position and raises when the row count is not a whole number of repeats of the index (the native forward already requires it); the backwards of the stacked ops use it instead of index.repeat(rows // len(index)), which truncates silently. - The fused modulation passes ACT_NONE to det_linear_forward explicitly, like the projection op. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/adaln_modulation.py | 3 ++- .../minimax_h3/adaln_row_gather.py | 17 ++++++++++++++++- .../minimax_h3/test_h3_adaln_row_gather.py | 17 +++++++++++++++-- 3 files changed, 33 insertions(+), 4 deletions(-) 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 9e7cbaf70..a856ce616 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 @@ -21,6 +21,7 @@ adaln_row_gather_available, ) from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + ACT_NONE, det_linear_available, det_linear_backward_input, det_linear_backward_weight, @@ -55,7 +56,7 @@ class _H3AdaLNModulationCuda(torch.autograd.Function): 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) + (table,) = det_linear_forward(act, weight, bias, activation=ACT_NONE) 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 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..2298ad294 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, @@ -46,6 +46,21 @@ def _segment_tiles(index: torch.Tensor, num_rows: int): return sorted_pos, tile_begin.contiguous(), tile_end.contiguous(), seg_first_tile +def _repeat_positions(index: torch.Tensor, rows: int) -> torch.Tensor: + """Table row of every flattened position: ``index`` repeated ``rows // len(index)`` times. + + The native forward only accepts a row count that is a multiple of ``len(index)``; the backward + states the same contract instead of letting floor division build a shorter position list. + """ + + if index.numel() == 0 or rows % index.shape[0] != 0: + raise ValueError( + f"{rows} flattened rows are not a whole number of repeats of an index of " + f"length {index.shape[0]}" + ) + return index.repeat(rows // index.shape[0]) + + class _H3AdaLNRowGatherCuda(torch.autograd.Function): @staticmethod def forward(ctx, rows, timestep_indices, token_tags): 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..a8de11866 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -83,6 +83,17 @@ def test_rejects_bad_rows(self): NativeH3AdaLNRowGatherOp().forward(torch.zeros(3, 25), ti, tags) # not 6 * H +def test_backward_positions_are_a_whole_number_of_index_repeats(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _repeat_positions + + index = torch.tensor([2, 0, 1]) + assert _repeat_positions(index, 6).tolist() == [2, 0, 1, 2, 0, 1] + with pytest.raises(ValueError, match="whole number of repeats"): + _repeat_positions(index, 7) + with pytest.raises(ValueError, match="whole number of repeats"): + _repeat_positions(index[:0], 4) + + @requires_cuda class TestCudaForward: @pytest.mark.parametrize("seq", [1, 2, 3, 257, 4097, 32768]) @@ -149,7 +160,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 +179,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], From fb72001cdded4d7f17056d4b582c5c456e65ac22 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:19:52 +0000 Subject: [PATCH 115/139] fix(h3): check the position repeat count in the modulated RMSNorm backward Uses _repeat_positions() instead of index.repeat(rows // len(index)), so a row count that is not a whole number of index repeats raises instead of tiling a shorter position list. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../backends/cuda/model_specific/minimax_h3/rmsnorm.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) 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..5188f3f80 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,12 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _repeat_positions, + _segment_tiles, +) 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.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.rmsnorm import ( H3_NORM_EPS, validate_h3_modulation, @@ -46,7 +49,7 @@ def backward(ctx, grad): 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]) + positions = _repeat_positions(index, x2.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 From 239ef808e279b6f62dec4592c4311c7cc1b0c32f Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:27:36 +0000 Subject: [PATCH 116/139] fix(h3): assert the gate row in the native kernels and check the backward position repeat - gate_residual.cu: Gate carries the table row count and the row kernel asserts 0 <= index < rows before using it as an offset, like the RMSNorm and row-gather kernels (check_range=False callers and direct _C calls get a device-side assert instead of an out-of-range read). - The backward uses _repeat_positions() instead of index.repeat(rows // len(index)). - Test: rows -1 and R in forward and backward expect the device-side assert (subprocess). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 7 ++- .../minimax_h3/gate_residual.py | 9 ++-- .../minimax_h3/test_h3_adaln_gate_residual.py | 52 +++++++++++++++++++ 3 files changed, 63 insertions(+), 5 deletions(-) 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/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..ad9bb6d8f 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,12 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _repeat_positions, + _segment_tiles, +) 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.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" @@ -40,7 +43,7 @@ def forward(ctx, residual, y, gate, index): 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]) + positions = _repeat_positions(index, y2.shape[0]) tiles = _segment_tiles(positions, gate.shape[0]) dy, dgate = _C.h3_gate_residual_backward(g2, y2, gate, index, *tiles) record_backward( diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index e4c303e10..4dd9521f8 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -11,6 +11,12 @@ from __future__ import annotations +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest import torch @@ -113,3 +119,49 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("adaln_gate_residual", device="cuda") assert type(op).__name__ == "H3GateResidualCudaOp" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +@pytest.mark.parametrize("bad_row", [-1, 2], ids=["negative", "past-table"]) +def test_gate_row_outside_the_table_asserts_in_the_kernel(entrypoint, bad_row): + """check_range=False callers get a device-side assert, not a read past the gate table.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + if not hasattr(_C, "h3_gate_residual_forward"): + sys.exit(77) + entrypoint, bad_row = sys.argv[1], int(sys.argv[2]) + x = torch.randn(4, 8, device="cuda") + gate = torch.ones(2, 8, device="cuda") # a table of two rows + index = torch.tensor([0, bad_row], dtype=torch.int64, device="cuda") + if entrypoint == "forward": + _C.h3_gate_residual_forward(x, x, gate, index) + else: + i64 = dict(dtype=torch.int64, device="cuda") + _C.h3_gate_residual_backward( + torch.ones_like(x), x, gate, index, + torch.tensor([0, 2, 1, 3], **i64), torch.tensor([0, 2], **i64), + torch.tensor([2, 4], **i64), torch.tensor([0, 1, 2], **i64), + ) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(bad_row)], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + if completed.returncode == 77: + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + assert completed.returncode != 0, "the kernel accepted a gate row outside the table" + assert ( + "device-side assert triggered" in completed.stderr + ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" From 1d26287f4550ce64a2f52ecc615fbdbde3d4ab3a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:36:20 +0000 Subject: [PATCH 117/139] fix(h3): check the position repeat in final_adaln_out backward and keep its forward out of the timed backward - final_adaln_out backward uses _repeat_positions() instead of floor division. - _final_perf had no backward setup, so the timed backward also ran the forward and built the graph (the guard test added with the gate-residual fix finds exactly this); the forward moves into the setup like in the other perf cases. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../model_specific/minimax_h3/final_adaln_out.py | 11 +++++++---- rl_engine/validation/models/h3_report.py | 15 +++++++++++---- 2 files changed, 18 insertions(+), 8 deletions(-) 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..55993caad 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,9 +8,10 @@ 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.adaln_row_gather import ( + _repeat_positions, + _segment_tiles, +) from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( det_linear_available, det_linear_backward_input, @@ -19,6 +20,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 @@ -54,7 +57,7 @@ def forward(ctx, x, norm_weight, temb, weight, bias, timestep_indices, eps): 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]) + positions = _repeat_positions(timestep_indices, x2.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), diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 4e1033945..36eb1c713 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -717,9 +717,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( { @@ -730,8 +733,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 b=prepare_backward, i=ti: b( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda b=prepare_backward, i=ti: b( lambda *t: provider_final_adaln_out(*t, i) ), } From a2d0674bd8b8e975c3493abff6351f6429b7b8e2 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:16:56 +0000 Subject: [PATCH 118/139] docs(h3): gate-residual latencies from the backward-only rerun The table and the report links now come from the B200 report rerun at 239ef808, where the timed backward no longer includes the forward pass: backward 1.02 / 1.08 / 2.21 ms for S = 4097 / 32768 / 131072 against 0.55 / 2.24 / 8.92 ms for diffusers (the speedup at S = 131072 grows from 3.2x to 4.0x). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index b35bc41c1..e869f4363 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -65,9 +65,9 @@ B200, B = 1, H = 5376, BF16: | 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.06 ms | 0.07 ms | 1.02 ms | 0.55 ms | +| 32768 | 0.24 ms | 0.43 ms | 1.08 ms | 2.24 ms | +| 131072 | 0.88 ms | 1.61 ms | 2.21 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 @@ -75,10 +75,11 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/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/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `e7bc53a`. It also records +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `239ef808` on an otherwise idle B200 (backward timings exclude the leaf +creation, the forward pass and the graph construction; the execution order alternates and is recorded). It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) From 337ee6041a96d5f8691a58d2b4eeea27420c8fc0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:17:08 +0000 Subject: [PATCH 119/139] docs(h3): final_adaln_out latencies from the backward-only rerun The table and the report links now come from the B200 report rerun at 1d26287f (pinned norm_out weights), where the timed backward no longer includes the forward pass: backward 1.38 / 2.56 / 7.04 ms for S = 4097 / 32768 / 131072 against 0.86 / 4.47 / 17.7 ms for diffusers. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 73eac42d6..90fbbfd17 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -62,22 +62,23 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 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 | +| 4097 | 0.18 ms | 0.18 ms | 1.38 ms | 0.86 ms | +| 32768 | 0.42 ms | 0.82 ms | 2.56 ms | 4.47 ms | +| 131072 | 1.28 ms | 2.96 ms | 7.04 ms | 17.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](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/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/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr489/reports/experiments/h3-final-adaln-out-b200/figure.png) There are two data files: -- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr489/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 `a4bc99f9` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + tree at `1d26287f` on an otherwise idle B200 (backward timings exclude the leaf creation, the + forward pass and the graph construction; the previous revision came from `a4bc99f9`), 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`. From 2caf0c829fb6f3f7ae23e08bf6370428e41564d9 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:00:44 +0000 Subject: [PATCH 120/139] fix(h3): read every spelling of relaxed nvcc math, and keep sinusoid backward/forward_fp32 aligned with the reference build_tools/extensions.py: _external_relaxed_math() now splits NVCC_PREPEND_FLAGS, NVCC_APPEND_FLAGS and TORCH_NVCC_FLAGS like a shell, then on whitespace and ';' (a quoted "-O2 --use_fast_math" or a ';'-separated TORCH_NVCC_FLAGS is several options), expands @response files, strips one or two leading dashes, and reads the value of --ftz / --prec-div / --prec-sqrt from "=value" or from the next word. An option that nvcc would reject for lack of a value, an unreadable response file and an unparsable value count as relaxed: the H3 kernels are left out rather than built with flags nobody checked. Tests: the build tests clear all three variables so a runner that sets one cannot change the outcome, the exclusion test is paired with a control build without the variable, and the predicate is tested directly for every spelling (positive and negative), for response files and for the fail-closed cases. timestep_sinusoid: forward_fp32 of the CUDA op takes check_range like the reference, so validation can call either with the same arguments. Tests pin that the gradient has the timestep's dtype and the reference's value for fp32/bf16/fp16 timesteps, that forward_fp32 accepts check_range, and that an odd num_channels is rejected even with check_range=False (the structural checks backward's [cos | sin] split relies on always run). The reference documents that bf16/fp16 timesteps are upcast first and that check_range costs one device-to-host sync. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- build_tools/extensions.py | 81 +++++++++++--- .../minimax_h3/timestep_sinusoid.py | 19 +++- .../reference/minimax_h3/timestep_sinusoid.py | 7 +- .../build/test_build_platform_collectives.py | 104 +++++++++++++++++- .../minimax_h3/test_h3_timestep_sinusoid.py | 26 +++++ 5 files changed, 212 insertions(+), 25 deletions(-) diff --git a/build_tools/extensions.py b/build_tools/extensions.py index 6362a1c58..5caf96f50 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -24,25 +24,78 @@ H3_KERNELS_DEFINE = "-DKERNEL_ALIGN_WITH_H3_KERNELS" # nvcc options that relax FP32 math. They can also arrive through the environment, where PyTorch's # BuildExtension appends them after the flags built here, so the H3 kernels check for them too. -_RELAXED_MATH_NVCC_FLAGS = frozenset( - { - "--use_fast_math", - "-use_fast_math", - "--ftz=true", - "-ftz=true", - "--prec-div=false", - "-prec-div=false", - "--prec-sqrt=false", - "-prec-sqrt=false", - } -) +# nvcc takes ``--ftz``, ``--prec-div`` and ``--prec-sqrt`` with a value, as ``--opt=value`` or +# ``--opt value``, and accepts one or two leading dashes; ``--use_fast_math`` takes none. +_RELAXED_MATH_FLAG = "use_fast_math" +_RELAXED_MATH_VALUE_FLAGS = { + "ftz": {"true", "1"}, + "prec-div": {"false", "0"}, + "prec-sqrt": {"false", "0"}, +} +_RESPONSE_FILE_DEPTH = 4 + + +def _nvcc_flag_tokens(text: str, *, depth: int = 0) -> list[str] | None: + """Split an NVCC_*_FLAGS value into tokens; None when a response file cannot be read. + + Tokens are split like a shell would, then on whitespace and ``;`` (TORCH_NVCC_FLAGS is + ';'-separated), so ``"-O2 --use_fast_math"`` in one quoted word is still two options. ``@file`` + response files are expanded. + """ + + tokens: list[str] = [] + for word in shlex.split(text): + for token in word.replace(";", " ").split(): + if not token.startswith("@"): + tokens.append(token) + continue + if depth >= _RESPONSE_FILE_DEPTH: + return None + try: + inner = _nvcc_flag_tokens(Path(token[1:]).read_text(), depth=depth + 1) + except (OSError, ValueError): + return None + if inner is None: + return None + tokens.extend(inner) + return tokens + + +def _tokens_relax_math(tokens: list[str]) -> bool: + for i, token in enumerate(tokens): + if not token.startswith("-"): + continue + name, has_value, value = token.lstrip("-").partition("=") + if name == _RELAXED_MATH_FLAG: + return True + relaxed_values = _RELAXED_MATH_VALUE_FLAGS.get(name) + if relaxed_values is None: + continue + if not has_value: + # "--ftz true": the value is the next word. Without one, nvcc rejects the option; treat + # it as relaxed rather than guess. + nxt = tokens[i + 1] if i + 1 < len(tokens) else "" + if nxt.startswith("-") or not nxt: + return True + value = nxt + if value.lower() in relaxed_values: + return True + return False def _external_relaxed_math() -> bool: - """Return whether NVCC_*_FLAGS / TORCH_NVCC_FLAGS turn on relaxed FP32 math.""" + """Return whether NVCC_*_FLAGS / TORCH_NVCC_FLAGS turn on relaxed FP32 math. + + Fails closed: a value that cannot be parsed or a response file that cannot be read counts as + relaxed, because the H3 kernels would otherwise be built with flags nobody checked. + """ for name in ("NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS", "TORCH_NVCC_FLAGS"): - if _RELAXED_MATH_NVCC_FLAGS & set(shlex.split(os.environ.get(name, ""))): + try: + tokens = _nvcc_flag_tokens(os.environ.get(name, "")) + except ValueError: + return True + if tokens is None or _tokens_relax_math(tokens): return True return False 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 5f3ef468a..86b838403 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 @@ -54,6 +54,8 @@ def backward(ctx, grad_out: torch.Tensor): contrib = torch.cat( [-(g[:, :half] * sin_part) * freq, (g[:, half:] * cos_part) * freq], dim=-1 ) + # Autograd hands back a gradient in the input's dtype, as it does through the reference's + # ``.float()`` cast (a bf16 timestep gets a bf16 gradient from either path). grad_t = tree_sum_lastdim_fp32(contrib).to(ctx.timestep_dtype) record_backward( "timestep_sinusoid_h3", @@ -93,12 +95,21 @@ def forward( num_channels: int = H3_FREQ_DIM, check_range: bool = True, ) -> torch.Tensor: - # Native code owns value validation, avoiding a second host sync here. + # Dtype, rank, emptiness and even num_channels are checked here on every call (they are + # what backward's [cos | sin] split relies on); only the value check is left to the native + # entry point, which enforces check_range itself and avoids 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, 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. - return self.forward(timestep, num_channels=num_channels) + def forward_fp32( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + # The op is FP32 end to end; the FP32 path is the op itself. The signature matches the + # reference's, so validation can call either one with the same arguments. + return self.forward(timestep, num_channels=num_channels, check_range=check_range) diff --git a/rl_engine/reference/minimax_h3/timestep_sinusoid.py b/rl_engine/reference/minimax_h3/timestep_sinusoid.py index 2e468985e..f20aa77e9 100644 --- a/rl_engine/reference/minimax_h3/timestep_sinusoid.py +++ b/rl_engine/reference/minimax_h3/timestep_sinusoid.py @@ -34,7 +34,8 @@ def validate_h3_timesteps( ``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. + to the host (one device-to-host sync per call), which is cheap for the handful of distinct + timesteps; ``check_range=False`` skips it and keeps the structural checks. """ if not isinstance(timestep, torch.Tensor): @@ -63,6 +64,10 @@ class NativeH3TimestepSinusoidOp: ``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. + + bf16/fp16 timesteps are accepted and upcast to FP32 first (``timesteps[:, None].float()``), as + diffusers does; the features are then computed from the rounded low-precision value, so a caller + who needs a different input rounds it before the call. The CUDA op does the same. """ op_class = "elementwise" diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index 9d4ebcfde..dbabf9404 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -14,10 +14,20 @@ from build_tools.extensions import ( H3_BITWISE_SOURCES, H3_KERNELS_DEFINE, + _external_relaxed_math, _resolve_package_data_source, get_cmdclass, ) +_NVCC_ENV = ("NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS", "TORCH_NVCC_FLAGS") + + +def _clear_nvcc_env(monkeypatch) -> None: + # The build reads these variables; a runner that sets one must not change what a test sees. + monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) + for name in _NVCC_ENV: + monkeypatch.delenv(name, raising=False) + def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: captured: dict[str, Any] = {} @@ -92,7 +102,7 @@ def test_package_data_real_json_is_not_reinterpreted(tmp_path) -> None: def test_cuda_build_compiles_h3_kernels_without_fast_math(monkeypatch) -> None: - monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) + _clear_nvcc_env(monkeypatch) extension = _load_extension_config(monkeypatch, hip=None) assert set(H3_BITWISE_SOURCES) <= set(extension["sources"]) @@ -106,6 +116,7 @@ def test_fast_math_build_excludes_bitwise_h3_kernels(monkeypatch) -> None: # 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. + _clear_nvcc_env(monkeypatch) monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", "1") extension = _load_extension_config(monkeypatch, hip=None) @@ -116,16 +127,29 @@ def test_fast_math_build_excludes_bitwise_h3_kernels(monkeypatch) -> None: assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["nvcc"] +@pytest.mark.parametrize("variable", _NVCC_ENV) @pytest.mark.parametrize( - "variable", ["NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS", "TORCH_NVCC_FLAGS"] + "flags", + [ + "-O2 --use_fast_math --ftz=true", + "--ftz true", + "-ftz=1", + "--prec-div false", + "-prec-sqrt=false", + "--use_fast_math;--prec-div=false", + "'-O2 --use_fast_math'", + ], ) def test_relaxed_math_nvcc_flags_from_the_environment_exclude_bitwise_h3_kernels( - monkeypatch, variable + monkeypatch, variable, flags ) -> None: # The flags are appended after the ones built here, so a build with # NVCC_APPEND_FLAGS='--use_fast_math' would otherwise ship H3 kernels with approximated math. - monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) - monkeypatch.setenv(variable, "-O2 --use_fast_math --ftz=true") + _clear_nvcc_env(monkeypatch) + control = _load_extension_config(monkeypatch, hip=None) + assert set(H3_BITWISE_SOURCES) <= set(control["sources"]) # the flags are the only difference + + monkeypatch.setenv(variable, flags) with pytest.warns(RuntimeWarning, match="bitwise MiniMax-H3 kernels"): extension = _load_extension_config(monkeypatch, hip=None) @@ -134,8 +158,76 @@ def test_relaxed_math_nvcc_flags_from_the_environment_exclude_bitwise_h3_kernels assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["nvcc"] +@pytest.mark.parametrize( + "flags", + [ + "--use_fast_math", + "-use_fast_math", + "--use_fast_math=true", + "--ftz=true", + "-ftz=true", + "--ftz true", + "-ftz 1", + "--ftz", + "--prec-div=false", + "--prec-div false", + "-prec-div 0", + "--prec-sqrt=false", + "-prec-sqrt false", + "-O2;--use_fast_math", + "-O2 ; --prec-div=false", + "'-O2 --use_fast_math'", + "-lineinfo --ftz=true", + ], +) +def test_predicate_sees_every_spelling_of_relaxed_math(monkeypatch, flags) -> None: + _clear_nvcc_env(monkeypatch) + monkeypatch.setenv("NVCC_APPEND_FLAGS", flags) + assert _external_relaxed_math() + + +@pytest.mark.parametrize( + "flags", + [ + "", + "-lineinfo --fmad=true", + "--fmad=false", + "--ftz=false", + "--ftz false", + "-ftz=0", + "--prec-div=true", + "--prec-div true -O2", + "-prec-sqrt 1", + "-O2 -gencode arch=compute_90,code=sm_90", + "--use_fast_mathematically", + ], +) +def test_predicate_keeps_the_h3_kernels_for_unrelated_flags(monkeypatch, flags) -> None: + _clear_nvcc_env(monkeypatch) + monkeypatch.setenv("NVCC_APPEND_FLAGS", flags) + assert not _external_relaxed_math() + + +def test_predicate_reads_response_files_and_fails_closed(monkeypatch, tmp_path) -> None: + _clear_nvcc_env(monkeypatch) + relaxed = tmp_path / "relaxed.opts" + relaxed.write_text("-O2\n--use_fast_math\n") + plain = tmp_path / "plain.opts" + plain.write_text("-lineinfo\n") + + monkeypatch.setenv("NVCC_APPEND_FLAGS", f"@{relaxed}") + assert _external_relaxed_math() + monkeypatch.setenv("NVCC_APPEND_FLAGS", f"@{plain}") + assert not _external_relaxed_math() + # Nothing is known about a file that cannot be read or a value that cannot be split. + monkeypatch.setenv("NVCC_APPEND_FLAGS", f"@{tmp_path / 'missing.opts'}") + assert _external_relaxed_math() + monkeypatch.setenv("NVCC_APPEND_FLAGS", "'unterminated") + assert _external_relaxed_math() + + def test_unrelated_nvcc_flags_keep_the_h3_kernels(monkeypatch) -> None: - monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) + _clear_nvcc_env(monkeypatch) monkeypatch.setenv("NVCC_APPEND_FLAGS", "-lineinfo --fmad=true") extension = _load_extension_config(monkeypatch, hip=None) diff --git a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py index 5f1ade3ee..caad9ffcb 100644 --- a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py +++ b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py @@ -202,6 +202,32 @@ def test_backward_matches_autograd_reference(self): NativeH3TimestepSinusoidOp().forward_fp32(t_ref).backward(grad) torch.testing.assert_close(t.grad, t_ref.grad, atol=1e-4, rtol=1e-4) + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) + def test_gradient_dtype_and_value_follow_the_reference_for_every_timestep_dtype(self, dtype): + base = h3_timesteps(9).to(dtype) + t, t_ref = base.clone().requires_grad_(True), base.clone().requires_grad_(True) + grad = torch.randn(9, 256, device="cuda", generator=torch.Generator("cuda").manual_seed(3)) + _cuda_op()(t).backward(grad) + NativeH3TimestepSinusoidOp().forward(t_ref).backward(grad) + assert t.grad.dtype == t_ref.grad.dtype == dtype + torch.testing.assert_close(t.grad.float(), t_ref.grad.float(), atol=2e-2, rtol=2e-2) + + def test_forward_fp32_takes_check_range_like_the_reference(self): + t = torch.tensor([0.5, 2.0], device="cuda") + op, ref = _cuda_op(), NativeH3TimestepSinusoidOp() + for fn in (op.forward_fp32, ref.forward_fp32): + with pytest.raises((ValueError, RuntimeError)): + fn(t) + assert op.forward_fp32(t[:1], check_range=False).shape == (1, 256) + assert ref.forward_fp32(t[:1], check_range=False).shape == (1, 256) + + def test_odd_num_channels_is_rejected_before_the_native_call(self): + # check_range=False only skips the value check: the structural checks that backward's + # [cos | sin] split depends on still run on every call. + t = torch.tensor([0.5], device="cuda", requires_grad=True) + with pytest.raises(ValueError, match="even"): + _cuda_op().forward(t, num_channels=255, check_range=False) + def test_backward_is_row_invariant(self): op = _cuda_op() grad_row = torch.randn(256, device="cuda", generator=torch.Generator("cuda").manual_seed(2)) From e1274b0e61b486ffcdb7a3624a27e0973fe6e2e6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:01:25 +0000 Subject: [PATCH 121/139] docs(h3): state the K and N tail contract of det_linear next to the loads The 16-byte loads of the forward kernels start at multiples of VEC below K and the host requires K % VEC == 0, so k < K implies k + VEC <= K; say so in the contract header and at the load, together with the way N and the row count tails are masked. Also note that db is the dW chain with the multiplier 1 (a plain FP32 add) and that check_matrix admits float32 and bfloat16 only before the dW dispatch. No code change. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 13798a0e8..d9a7d5653 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -26,7 +26,14 @@ // 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. +// cross-row reduction; its order is the logical row order. db is the same +// chain with the multiplier 1, which is a plain FP32 add (fmaf(g, 1, a) == a + g). +// +// Tails: the host requires K % VEC == 0 (VEC = 16 / sizeof(w)), and every 16-byte +// load in the forward kernels starts at a multiple of VEC, so `k < K` implies +// `k + VEC <= K`: no load straddles the end of a row, and every row start stays +// 16-byte aligned. N and the row count have no such requirement; their tails are +// masked per column / per row. #include #include @@ -135,6 +142,8 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) uint4 raw[kLoadAhead][kCols]; #pragma unroll for (int u = 0; u < kLoadAhead; ++u) { + // K % kVec == 0 (host check) and k0 is a multiple of kVec, so k0 < k_in means the whole + // 16-byte chunk lies inside the row. const int64_t k0 = g0 + u * kStride; #pragma unroll for (int c = 0; c < kCols; ++c) { @@ -464,6 +473,7 @@ std::vector h3_det_linear_backward_weight(torch::Tensor grad, tor const int64_t total = n_out * k_in; const unsigned blocks = static_cast(grid_for(total, threads)); const float* g = grad.data_ptr(); + // check_matrix(x) above admits float32 and bfloat16 only, so the last branch is bfloat16. #define H3_DW_LAUNCH(XT, XC, WT, WC) \ det_linear_dweight_kernel<<>>( \ g, cptr(x), mptr(dw), rows, n_out, k_in) From ca14fadec7502be0b2d626c0279007102676b381 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:16:35 +0000 Subject: [PATCH 122/139] fix(h3): check the SM80 requirement of the BF16 forward for the device of the tensors det_linear_device_supported() took the current CUDA device and returned True without CUDA. It now takes the device in question (default: the current one) and is False when CUDA is unavailable or the device is not a CUDA device, so a constructor that calls it before any tensor exists fails and the registry falls back to the PyTorch reference. det_linear_forward() repeats the check for the device of the BF16 inputs and raises a RuntimeError naming SM80 instead of letting the launch fail inside the native call; FP32 is not affected. Tests pin that the capability is read for the requested device, that no CUDA means no support, and that a BF16 forward on a pre-SM80 device is refused while FP32 runs. The odd-width tests skip their BF16 forward cases on such a device, and their docstring now describes the 16-column MMA tile of the BF16 forward. Also drop the duplicated static_assert(kCols * kRows <= kWarp) of the FP32 forward kernel. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 1 - .../model_specific/minimax_h3/det_linear.py | 19 +++++++-- .../minimax_h3/test_h3_adaln_projection.py | 2 +- tests/models/minimax_h3/test_h3_det_linear.py | 40 ++++++++++++++++++- .../test_h3_det_linear_odd_width.py | 19 +++++++-- 5 files changed, 71 insertions(+), 10 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index d0ed82851..fb514abcc 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -134,7 +134,6 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) static_assert(kCols * kRows <= kWarp, "the epilogue stores through lane == c * kRows + r, so kCols * kRows must fit a warp"); 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; diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py index d288d2fa9..646685dc2 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py @@ -29,10 +29,19 @@ def det_linear_available() -> bool: return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) -def det_linear_device_supported() -> bool: - """Return whether the current CUDA device can run the BF16 forward (``mma.sync``, SM80+).""" +def det_linear_device_supported(device: torch.device | int | str | None = None) -> bool: + """Return whether ``device`` can run the BF16 forward (``mma.sync``, SM80+). - return not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] >= 8 + ``device`` defaults to the current CUDA device. Without CUDA, or for a non-CUDA device, the + answer is False: constructors use this as an early check so the registry falls back to the + PyTorch reference, and ``det_linear_forward`` repeats it for the device of the actual tensors. + """ + + if not torch.cuda.is_available(): + return False + if device is not None and torch.device(device).type != "cuda": + return False + return torch.cuda.get_device_capability(device)[0] >= 8 def _require() -> None: @@ -62,6 +71,10 @@ def det_linear_forward( """ _require() + if x.dtype == torch.bfloat16 and not det_linear_device_supported(x.device): + raise RuntimeError( + f"the BF16 h3 det_linear forward needs SM80 or newer (tensors are on {x.device})" + ) return _C.h3_det_linear_forward( x.contiguous(), weight.contiguous(), diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index f35b13a79..b22c5dea1 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -45,7 +45,7 @@ def test_pre_sm80_device_is_refused_at_construction_so_the_registry_can_fall_bac from rl_engine.backends.cuda.model_specific.minimax_h3 import adaln_projection as module monkeypatch.setattr(module, "det_linear_available", lambda: True) - monkeypatch.setattr(module, "det_linear_device_supported", lambda: False) + monkeypatch.setattr(module, "det_linear_device_supported", lambda *args: False) with pytest.raises(RuntimeError, match="SM80"): module.H3AdaLNProjectionCudaOp() diff --git a/tests/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py index f9bc79829..f1a2e3175 100644 --- a/tests/models/minimax_h3/test_h3_det_linear.py +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -8,8 +8,12 @@ import pytest import torch +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_device_supported, + det_linear_forward, +) from rl_engine.backends.extension import _C -from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), @@ -66,3 +70,37 @@ def test_backward_weight_handles_empty_columns(dtype, w_dtype, k_in, n_out, with 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) + + +def test_device_support_follows_the_device_that_is_asked_about(monkeypatch): + """The capability is read for the given device, not for the current one.""" + + seen = [] + + def capability(device=None): + seen.append(device) + return (7, 5) if device == torch.device("cuda:0") else (9, 0) + + monkeypatch.setattr(torch.cuda, "get_device_capability", capability) + assert not det_linear_device_supported(torch.device("cuda:0")) + assert det_linear_device_supported(torch.device("cuda:1")) + assert not det_linear_device_supported(torch.device("cpu")) + assert seen == [torch.device("cuda:0"), torch.device("cuda:1")] + + +def test_no_cuda_means_no_support(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + assert not det_linear_device_supported() + + +def test_bf16_forward_is_refused_on_a_pre_sm80_device_with_a_clear_error(monkeypatch): + """The wrapper checks the tensors' device, so no opaque launch error reaches the caller.""" + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (7, 5)) + x = torch.zeros(1, 64, device="cuda", dtype=torch.bfloat16) + w = torch.zeros(8, 64, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="SM80"): + det_linear_forward(x, w, None) + # FP32 does not use mma.sync and is not refused. + out = det_linear_forward(x.float(), w.float(), None)[0] + assert out.shape == (1, 8) diff --git a/tests/models/minimax_h3/test_h3_det_linear_odd_width.py b/tests/models/minimax_h3/test_h3_det_linear_odd_width.py index 2b0ee1aac..08036db8f 100644 --- a/tests/models/minimax_h3/test_h3_det_linear_odd_width.py +++ b/tests/models/minimax_h3/test_h3_det_linear_odd_width.py @@ -3,10 +3,12 @@ """Output widths that are not a multiple of the column tile of the deterministic linear. -Each warp of the forward kernel owns kCols adjacent output columns (2 for FP32, 4 for BF16), so the -last tile of an odd or short N has columns that are not loaded and must not be consumed. The -results for such widths must match an FP64 golden, and must be bitwise the leading columns of the -same product with zero columns padded on (the tail does not influence the real columns). +Each warp of the FP32 forward kernel owns kCols = 2 adjacent output columns, and each warp of the +BF16 forward (an ``mma.sync`` kernel) owns a 16-column tile (kMmaCols), so the last tile of an odd +or short N has columns that are not loaded and must not be consumed. The results for such widths +must match an FP64 golden, and must be bitwise the leading columns of the same product with zero +columns padded on (the tail does not influence the real columns). The BF16 forward needs SM80, so +its cases are skipped on older devices. """ from __future__ import annotations @@ -18,6 +20,7 @@ det_linear_available, det_linear_backward_input, det_linear_backward_weight, + det_linear_device_supported, det_linear_forward, ) @@ -32,6 +35,11 @@ ROWS = [1, 3, 5] +def _skip_unsupported(dtype): + if dtype == torch.bfloat16 and not det_linear_device_supported(): + pytest.skip("the BF16 forward needs SM80 or newer") + + def _inputs(dtype, rows, n_out, seed=0): g = torch.Generator(device="cpu").manual_seed(seed + n_out) x = torch.randn(rows, K_IN, generator=g).to(dtype).cuda() @@ -49,6 +57,7 @@ def _pad(weight, bias, extra): @pytest.mark.parametrize("n_out", WIDTHS) @pytest.mark.parametrize("rows", ROWS) def test_forward_matches_golden_and_ignores_the_tail(dtype, n_out, rows): + _skip_unsupported(dtype) x, weight, bias = _inputs(dtype, rows, n_out) out = det_linear_forward(x, weight, bias)[0] assert out.shape == (rows, n_out) and torch.isfinite(out.float()).all() @@ -95,6 +104,8 @@ def test_backward_matches_golden_and_ignores_the_tail(dtype, n_out): def test_k_that_is_not_a_multiple_of_the_lane_stride_matches_the_golden(dtype, k_in, n_out): """Row lengths that are a vector multiple but not a multiple of the 32-lane stride.""" + _skip_unsupported(dtype) + g = torch.Generator(device="cpu").manual_seed(k_in + n_out) x = torch.randn(3, k_in, generator=g).to(dtype).cuda() weight = (torch.randn(n_out, k_in, generator=g) / k_in**0.5).to(dtype).cuda() From 1adec0431149018537449833c1fd9ac9195ead71 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:22:26 +0000 Subject: [PATCH 123/139] fix(h3): report a w1 that is not a matrix as a dimension error, and say which MLP layer has the SiLU validate_h3_timestep_mlp unpacked (None, None) for a w1 that is not 2-D, so the next checks reported "w1 must be (H, K)" and "b1 must be (None,)". It now rejects the rank first with a message that says so. The operator doc stated that SiLU follows the bias addition of every output; it applies to linear_1 only, linear_2 returns its biased linear output. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 5 +++-- rl_engine/reference/minimax_h3/timestep_mlp.py | 4 +++- tests/models/minimax_h3/test_h3_timestep_mlp.py | 7 +++++++ 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index b65b19fca..69e6683f5 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -56,8 +56,9 @@ Forward, for each output `y[t, n]`: - 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. +- The bias is added once. In `linear_1` SiLU then runs in FP32 as `v / (1 + expf(-v))` + (PyTorch's formula); `linear_2` returns its biased linear output without an activation. + 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 diff --git a/rl_engine/reference/minimax_h3/timestep_mlp.py b/rl_engine/reference/minimax_h3/timestep_mlp.py index ac54d295f..f8348d99c 100644 --- a/rl_engine/reference/minimax_h3/timestep_mlp.py +++ b/rl_engine/reference/minimax_h3/timestep_mlp.py @@ -40,7 +40,9 @@ def validate_h3_timestep_mlp( 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 w1.dim() != 2: + raise ValueError(f"w1 must be a 2-D (H, {x.shape[1]}) matrix, got {tuple(w1.shape)}") + hidden, k_in = w1.shape if k_in != x.shape[1]: raise ValueError(f"w1 must be (H, {x.shape[1]}), got {tuple(w1.shape)}") if b1.shape != (hidden,): diff --git a/tests/models/minimax_h3/test_h3_timestep_mlp.py b/tests/models/minimax_h3/test_h3_timestep_mlp.py index bc6059125..09d1b2dd4 100644 --- a/tests/models/minimax_h3/test_h3_timestep_mlp.py +++ b/tests/models/minimax_h3/test_h3_timestep_mlp.py @@ -90,6 +90,13 @@ def test_rejects_shape_mismatches(self): with pytest.raises(ValueError): op.forward(x[:0], w1, b1, w2, b2) + @pytest.mark.parametrize("bad_w1", [torch.zeros(256), torch.zeros(2, 5, 256)]) + def test_a_w1_that_is_not_a_matrix_gets_a_dimension_error(self, bad_w1): + x = _features(2, device="cpu") + _, b1, w2, b2 = _synthetic_params(256, 64, 32, device="cpu") + with pytest.raises(ValueError, match="w1 must be a 2-D"): + NativeH3TimestepMLPOp().forward(x, bad_w1, b1, w2, b2) + @requires_cuda class TestCudaRealWeights: From 9c9789fd97f6a99d286e2615c89198230b3955ca Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:24:18 +0000 Subject: [PATCH 124/139] fix(h3): keep the bitwise-reuse report when a reference lacks a case, mark missing cells, reject zero-width AdaLN input h3_bitwise_reuse.compare indexed runs[ref][case] for every case of every candidate. A reference child that died after saving only some of its cases (an out-of-memory at the largest size) raised a KeyError before the report was written, losing the comparisons that had completed and leaving the .parts directory behind. A case the reference never saved is now recorded as unavailable. The plot treats an entry without outputs as not applicable, writes "n/a" in the cell and draws it grey instead of the white of "bitwise equal". validate_h3_adaln_projection rejects a weight with zero input width: F.linear over an empty reduction returns the broadcast bias instead of failing. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../reference/minimax_h3/adaln_projection.py | 3 ++ .../minimax_h3/test_h3_adaln_projection.py | 5 +++ .../minimax_h3/test_h3_bitwise_reuse_plot.py | 39 ++++++++++++++++++- tools/validation/models/h3_bitwise_reuse.py | 26 +++++++++---- .../models/plot_h3_bitwise_reuse.py | 8 ++-- 5 files changed, 69 insertions(+), 12 deletions(-) diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py index be83239c5..a6b9fd714 100644 --- a/rl_engine/reference/minimax_h3/adaln_projection.py +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -58,6 +58,9 @@ def validate_h3_adaln_projection(temb: torch.Tensor, weight: torch.Tensor, bias: f"weight/bias must share a dtype in {_WEIGHT_DTYPES}, got {weight.dtype}/{bias.dtype}" ) hidden = adaln_hidden_size(weight) + if weight.shape[1] == 0: + # F.linear over an empty reduction returns the broadcast bias instead of failing. + raise ValueError(f"AdaLN weight must have a nonzero input width, got {tuple(weight.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)}" diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index b22c5dea1..23aff80f6 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -107,6 +107,11 @@ def test_rejects_bad_shapes(self): with pytest.raises(TypeError): op.forward(_temb(2, 16, device="cpu"), weight, bias.float()) + def test_rejects_a_zero_width_input_instead_of_returning_the_bias(self): + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(ValueError, match="nonzero input width"): + NativeH3AdaLNProjectionOp().forward(_temb(2, 0, device="cpu"), weight[:, :0], bias) + @requires_cuda class TestLayout: diff --git a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py index 8fa02de7e..ffb87bd65 100644 --- a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -13,7 +13,8 @@ pytest.importorskip("matplotlib") sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "tools/validation/models")) -from plot_h3_bitwise_reuse import _cell # noqa: E402 +from h3_bitwise_reuse import _compare_case # noqa: E402 +from plot_h3_bitwise_reuse import _cell, _panel # noqa: E402 def _out(differ: int, numel: int = 100) -> dict: @@ -34,3 +35,39 @@ def test_cell_without_data_is_not_applicable(): assert _cell(None) is None assert _cell({"out": []}) is None assert _cell({"out": [None]}) is None + + +def test_a_case_the_reference_never_saved_is_unavailable_and_does_not_abort_the_comparison(): + import torch + + run = {"out": [torch.ones(3)], "grads": {"x": torch.zeros(2)}} + ref = {"out": [torch.ones(3)], "grads": {"x": torch.zeros(2)}} + done = _compare_case(run, ref) + assert done["out"][0]["equal"] and done["grads"]["x"]["equal"] + missing = _compare_case(run, None) + assert "unavailable" in missing and _cell(missing) is None + + +def test_missing_comparisons_are_drawn_grey_not_as_bitwise_equal(): + import matplotlib.pyplot as plt + + def result(frac_differ): + return {"out": [_out(frac_differ)], "grads": {}} + + report = { + "op": "probe", + "rl_kernel_commit": "0123456789", + "cases": ["a", "b"], + "results": { + "rl_kernel": {"vs diffusers": {"a": result(0), "b": {"unavailable": "no run"}}}, + "diffusers": {}, + }, + } + fig, ax = plt.subplots() + _panel(ax, report) + grid = ax.images[0].get_array() + assert grid.mask[0, 1] and not grid.mask[0, 0] # the made comparison is data, the other masked + assert ax.images[0].cmap(ax.images[0].norm(0.0)) != ax.images[0].cmap._rgba_bad + texts = [t.get_text() for t in ax.texts] + assert "n/a" in texts and "0" in texts + plt.close(fig) diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py index 345aa18cc..81a80b75c 100644 --- a/tools/validation/models/h3_bitwise_reuse.py +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -161,6 +161,23 @@ def _same(a: torch.Tensor | None, b: torch.Tensor | None) -> dict[str, Any] | No } +def _compare_case(run: dict[str, Any], ref_run: dict[str, Any] | None) -> dict[str, Any]: + """Compare one case with the reference's run of it. + + A case the reference never saved is unavailable. + + A reference child can die after saving only some of its cases (an out-of-memory at the largest + size, say); the comparisons that did complete must still reach the report. + """ + + if ref_run is None: + return {"unavailable": "the reference run has no saved result for this case"} + return { + "out": [_same(a, b) for a, b in zip(run["out"], ref_run["out"])], + "grads": {g: _same(run["grads"].get(g), ref_run["grads"].get(g)) for g in run["grads"]}, + } + + def compare(parts: Path) -> dict[str, Any]: runs: dict[str, dict[str, Any]] = {} for f in sorted(parts.glob("*.pt")): @@ -174,14 +191,7 @@ def compare(parts: Path) -> dict[str, Any]: 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() + case: _compare_case(run, runs[ref].get(case)) for case, run in by_case.items() } out[name] = entry for f in parts.glob("*.unavailable"): diff --git a/tools/validation/models/plot_h3_bitwise_reuse.py b/tools/validation/models/plot_h3_bitwise_reuse.py index 6da21441f..777b1edbf 100644 --- a/tools/validation/models/plot_h3_bitwise_reuse.py +++ b/tools/validation/models/plot_h3_bitwise_reuse.py @@ -27,7 +27,7 @@ def _cell(entry: dict | None) -> float | None: - if not entry: + if not entry or "out" not in entry: # no entry, or a case the reference never saved return None outs = [o for o in entry["out"] if o is not None] if not outs: @@ -62,14 +62,16 @@ def _panel(ax, report: dict) -> None: ax.text( j, i, - "—" if row == ref else _text(frac), + "—" if row == ref else ("n/a" if frac is None 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") + # A comparison that was not made is grey, so it cannot be read as the white of "bitwise equal". + cmap = plt.get_cmap("Reds").with_extremes(bad="0.82") + ax.imshow(np.ma.masked_invalid(np.sqrt(grid)), cmap=cmap, 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() From c416a10168892bbd7dd61bacd0d7ede80bd77832 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:45:06 +0000 Subject: [PATCH 125/139] fix(h3): check the row-gather tile metadata in one pre-pass, time the gather backward without the forward, scan rows once adaln_row_gather.cu: a small pre-pass kernel (check_tile_metadata in the new csrc/cuda/h3/tile_metadata.cuh, which the other segment-tile kernels of the stack reuse) checks the tile metadata before anything reads through it: every sorted position is a position of the sequence, every tile lies inside sorted_pos, and the segment boundaries are non-decreasing and inside the tile list (including the case of no tiles with a non-trivial seg_first_tile, where the partial buffer holds one uninitialized row). A stale or short list ends in a device-side assert instead of an out-of-bounds read. The checks are not made inside the partial and fold kernels: an assert in or before their loops changes how the compiler schedules the loads and slowed the native backward by 6-12% (and, inside the loop, by several times); the pre-pass costs about 4 us per call (native backward, idle kernel timing, min of 400 on an idle B200: 105.4 -> 109.2 us at S = 4097, 662.9 -> 665.8 us at S = 32768, 2578.0 -> 2581.3 us at S = 131072). A seg_first_tile with no boundary at all is a TORCH_CHECK, a table with no rows returns an empty gradient instead of launching a zero-sized grid, and the forward and backward check the grid.y limit (65535) for chunks and for the column blocks with a message that names the shape. Tests run corrupted metadata in a child process and expect the device assert, plus the host checks. h3_report: the row-gather backward cases build the forward graph in candidate_backward_setup / provider_backward_setup, so the timed region and the peak memory hold only autograd.backward, which is what backward_timing_scope = backward_only says. A test pins that every perf case with a timed backward has its setup. The numbers in docs/operators/h3-adaln-row-gather.md and the B200/H100 reports predate this and still include the forward pass. The reference backward (_DeterministicRowGather) takes the positions of all 3T rows from one stable sort and one host read (positions_by_row) instead of one torch.nonzero scan and sync per row; the sums are over the same ascending positions, and a test compares the gradient bitwise with the loop it replaces. adaln_modulation: det_linear_backward_weight skips the bias fold when only the weight needs a gradient (a test freezes each input in turn and compares the others bitwise). chain_grads zips with strict=True. The gather_chunks output shape in the doc is (S, 5376), not (3T, 5376). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 17 ++- csrc/cuda/h3/tile_metadata.cuh | 61 ++++++++++ docs/operators/h3-adaln-row-gather.md | 2 +- .../minimax_h3/adaln_modulation.py | 4 +- .../reference/minimax_h3/adaln_row_gather.py | 26 +++- rl_engine/validation/models/h3_chain.py | 4 +- rl_engine/validation/models/h3_report.py | 19 ++- .../minimax_h3/test_h3_adaln_modulation.py | 19 +++ .../minimax_h3/test_h3_adaln_row_gather.py | 112 ++++++++++++++++++ .../minimax_h3/test_h3_report_timing.py | 21 ++++ 10 files changed, 269 insertions(+), 16 deletions(-) create mode 100644 csrc/cuda/h3/tile_metadata.cuh diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..f8c2d3e31 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -28,6 +28,8 @@ #include +#include "tile_metadata.cuh" + namespace { __device__ __forceinline__ float to_float(float v) { return v; } @@ -109,8 +111,10 @@ __global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, const int64_t c = j / hidden; const int64_t h = j - c * hidden; const g_t* g = grad + c * seq * hidden + h; + const int64_t begin = tile_begin[tile]; + const int64_t end = tile_end[tile]; float acc = 0.0f; - for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + for (int64_t p = begin; p < end; ++p) { acc += to_float(g[sorted_pos[p] * hidden]); } partial[tile * width + j] = acc; @@ -151,6 +155,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); + TORCH_CHECK(chunks <= 65535, "chunks=", chunks, " exceeds the grid.y limit of 65535"); 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"); @@ -216,16 +221,24 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor const int64_t seq = grad.size(1); const int64_t hidden = grad.size(2); const int64_t width = chunks * hidden; + TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); 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)); + if (num_rows == 0) return out; // a zero-sized grid is not a valid launch 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); + const int64_t col_blocks64 = (width + threads - 1) / threads; + TORCH_CHECK(col_blocks64 <= 65535, "width ", width, " needs ", col_blocks64, + " column blocks, more than the grid.y limit of 65535"); + const unsigned col_blocks = static_cast(col_blocks64); + h3_tiles::check_tile_metadata(stream, sorted_pos.data_ptr(), seq, + tile_begin.data_ptr(), tile_end.data_ptr(), tiles, + seg_first_tile.data_ptr(), num_rows); if (tiles > 0) { AT_DISPATCH_FLOATING_TYPES_AND2( at::kHalf, at::kBFloat16, grad.scalar_type(), "h3_adaln_row_gather_partial", [&] { diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh new file mode 100644 index 000000000..a959054e2 --- /dev/null +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Segment-tile metadata shared by the H3 backward kernels that sum the rows of a gradient into the +// rows of a small table (adaln_row_gather, rmsnorm modulation, gate residual). +// +// The caller sorts the S positions by table row (`sorted_pos`), cuts every row's run into tiles +// `[tile_begin, tile_end)` over `sorted_pos`, and gives the first tile of each row in +// `seg_first_tile` (R + 1 entries). The partial kernels read gradients through `sorted_pos`, and the +// fold kernels read the partial rows of a segment's tiles, so a stale or short list would read out +// of bounds. +// +// check_tile_metadata() verifies the list on the device before any of that happens. It is a separate +// pass on purpose: an assert in or before the loops of the partial and fold kernels changes how the +// compiler schedules their loads and slowed them by 6-12% (and by several times inside the loop), +// while this pass costs one small launch (about 4 us). +#pragma once + +#include +#include + +#include +#include + +namespace h3_tiles { + +static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ sorted_pos, + int64_t positions, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + int64_t tiles, + const int64_t* __restrict__ seg_first_tile, + int64_t segments) { + const int64_t first_thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = first_thread; i < positions; i += stride) { + CUDA_KERNEL_ASSERT(sorted_pos[i] >= 0 && sorted_pos[i] < positions); + } + for (int64_t i = first_thread; i < tiles; i += stride) { + CUDA_KERNEL_ASSERT(tile_begin[i] >= 0 && tile_begin[i] <= tile_end[i] && + tile_end[i] <= positions); + } + for (int64_t i = first_thread; i <= segments; i += stride) { + CUDA_KERNEL_ASSERT(seg_first_tile[i] >= 0 && seg_first_tile[i] <= tiles); + CUDA_KERNEL_ASSERT(i == segments || seg_first_tile[i] <= seg_first_tile[i + 1]); + } +} + +// `positions` entries in sorted_pos (each in [0, positions)), `tiles` tiles inside sorted_pos, and +// `segments` row segments (seg_first_tile has segments + 1 entries). +inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, int64_t positions, + const int64_t* tile_begin, const int64_t* tile_end, int64_t tiles, + const int64_t* seg_first_tile, int64_t segments) { + const int64_t items = std::max({positions, tiles, segments + 1}); + const unsigned blocks = static_cast(std::min((items + 255) / 256, 1024)); + check_tile_metadata_kernel<<>>(sorted_pos, positions, tile_begin, + tile_end, tiles, seg_first_tile, segments); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace h3_tiles diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 8d0eb0584..2ce2990c1 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -24,7 +24,7 @@ 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, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp), timestep_indices, token_tags -) # drop-in, six (T*3, 5376) tensors +) # drop-in: six (3T, 5376) tables in, six (S, 5376) tensors out (S = len(timestep_indices)) ``` ## Backends 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 a856ce616..4c42eb606 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 @@ -76,7 +76,9 @@ def backward(ctx, grad_packed): 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) + # The bias fold is skipped when only the weight needs a gradient. + dw, *rest = det_linear_backward_weight(d_table, act, weight.dtype, with_bias=need_bias) + db = rest[0] if need_bias else None if need_temb: d_temb = silu_backward_fp32( det_linear_backward_input(d_table, weight, torch.float32), temb diff --git a/rl_engine/reference/minimax_h3/adaln_row_gather.py b/rl_engine/reference/minimax_h3/adaln_row_gather.py index 2cd713d45..937177f1b 100644 --- a/rl_engine/reference/minimax_h3/adaln_row_gather.py +++ b/rl_engine/reference/minimax_h3/adaln_row_gather.py @@ -78,14 +78,31 @@ def h3_adaln_indices(timestep_indices: torch.Tensor, token_tags: torch.Tensor) - return timestep_indices.long() * H3_MODALITY_NUM + token_tags.long() +def positions_by_row(index: torch.Tensor, num_rows: int) -> list[torch.Tensor]: + """Ascending positions of every table row, with one host read for all rows. + + Equal to ``torch.nonzero(index == row).flatten()`` for each ``row``: a stable sort keeps the + positions of a row in ascending order, and the per-row counts split the sorted list. It + replaces one device scan and one host sync per row, which grew with the 3T table rows. + """ + + order = torch.argsort(index, stable=True) + counts = torch.bincount(index, minlength=num_rows).tolist() + positions, start = [], 0 + for count in counts: + positions.append(order[start : start + count]) + start += count + return positions + + 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. + only 3T table rows (T is the handful of distinct timesteps), 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. The positions of all rows come from one sort and one host read. """ @staticmethod @@ -102,8 +119,7 @@ def backward(ctx, grad): 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() + for row, positions in enumerate(positions_by_row(index, num_rows)): if positions.numel(): out[row] = grad32.index_select(1, positions).sum(dim=1) return out.reshape(num_rows, -1).to(dtype), None diff --git a/rl_engine/validation/models/h3_chain.py b/rl_engine/validation/models/h3_chain.py index 8927c18dc..7b53eb56f 100644 --- a/rl_engine/validation/models/h3_chain.py +++ b/rl_engine/validation/models/h3_chain.py @@ -258,7 +258,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/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 6d4d82a59..8f8ba5501 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -354,9 +354,12 @@ def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: torch.randn(seq, hidden, device="cuda", generator=gen).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( { @@ -367,11 +370,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 + ) ), } ) diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py index dc27820dd..9adc6b4aa 100644 --- a/tests/models/minimax_h3/test_h3_adaln_modulation.py +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -124,6 +124,25 @@ def rel(grad, golden): for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad)): assert (grad == golden.to(torch.bfloat16)).float().mean() > 0.9999 + @pytest.mark.parametrize("frozen", ["bias", "weight", "temb", "bias_and_temb"]) + def test_frozen_inputs_get_no_gradient_and_the_others_are_unchanged(self, frozen): + fused, _, _ = _ops() + temb, weight, bias = _params(64, 2688, seed=5) + ti, tags = h3_packed_layout(1500, 3, seed=5) + grads = _grads(1500, 64, seed=5) + full = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + + leaves = [t.detach().clone() for t in (temb, weight, bias)] + names = ("temb", "weight", "bias") + for leaf, name in zip(leaves, names): + leaf.requires_grad_(name not in frozen) + torch.autograd.backward(list(fused(*leaves, ti, tags)), grads) + for leaf, name, expected in zip(leaves, names, full): + if name in frozen: + assert leaf.grad is None + else: + assert torch.equal(leaf.grad, expected) + def test_backward_repeat_bitwise(self): fused, _, _ = _ops() temb, weight, bias = _params(64, 2688, seed=3) 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 a8de11866..ff9662f75 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -94,6 +94,43 @@ def test_backward_positions_are_a_whole_number_of_index_repeats(): _repeat_positions(index[:0], 4) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_positions_by_row_equal_the_per_row_scan(device): + from rl_engine.reference.minimax_h3.adaln_row_gather import positions_by_row + + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("needs a CUDA device") + g = torch.Generator(device="cpu").manual_seed(0) + index = torch.randint(0, 7, (500,), generator=g).to(device) # rows 7..11 are never referenced + positions = positions_by_row(index, 12) + assert len(positions) == 12 + for row, found in enumerate(positions): + assert torch.equal(found, torch.nonzero(index == row).flatten()) + + +@requires_cuda +def test_reference_backward_is_bitwise_the_per_row_scan_it_replaced(): + from rl_engine.reference.minimax_h3.adaln_row_gather import _DeterministicRowGather + + rows = _rows(3, hidden=64, seed=7) + ti, tags = h3_packed_layout(777, 3, seed=7) + index = ti.long() * 3 + tags.long() + gen = torch.Generator(device="cuda").manual_seed(7) + grad = torch.randn(6, 777, 64, device="cuda", generator=gen).bfloat16() + + leaf = rows.detach().clone().requires_grad_(True) + outputs = _DeterministicRowGather.apply(leaf, index) + outputs.backward(grad) + + grad32 = grad.float() + expected = grad32.new_zeros((rows.shape[0], 6, 64)) + for row in range(rows.shape[0]): # the loop of the earlier implementation + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + expected[row] = grad32.index_select(1, positions).sum(dim=1) + assert torch.equal(leaf.grad, expected.reshape(rows.shape[0], -1).to(rows.dtype)) + + @requires_cuda class TestCudaForward: @pytest.mark.parametrize("seq", [1, 2, 3, 257, 4097, 32768]) @@ -303,6 +340,81 @@ def ref_grad(): first.float(), self._cuda_grad(rows, ti, tags, grads).float(), atol=1e-1, rtol=2e-2 ) + @staticmethod + def _backward_probe(mutation: str): + """Run the native backward on corrupted tile metadata in a child process.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + + mutation = sys.argv[1] + seq, hidden = 40, 8 + index = (torch.arange(seq, device="cuda") % 3) + grad = torch.ones(6, seq, hidden, device="cuda") + sorted_pos, begin, end, first = _segment_tiles(index, 3) + if mutation == "tile_end_past_seq": + end = end.clone(); end[-1] = seq + 5 + elif mutation == "position_past_seq": + sorted_pos = sorted_pos.clone(); sorted_pos[0] = seq + elif mutation == "segments_past_tiles": + first = first.clone(); first[-1] = begin.numel() + 3 + elif mutation == "no_tiles": + begin = begin[:0]; end = end[:0] + _C.h3_adaln_row_gather_backward(grad, sorted_pos, begin, end, first, torch.float32) + torch.cuda.synchronize() + """ + ) + return subprocess.run( + [sys.executable, "-c", probe, mutation], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + + def test_the_backward_stops_on_metadata_that_does_not_describe_the_tiles(self): + _cuda_op() + control = self._backward_probe("none") + assert control.returncode == 0, control.stderr + for mutation in ( + "tile_end_past_seq", + "position_past_seq", + "segments_past_tiles", + "no_tiles", + ): + done = self._backward_probe(mutation) + assert done.returncode != 0, mutation + assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) + + def test_backward_with_no_table_rows_returns_an_empty_table_and_rejects_no_boundaries(self): + from rl_engine.backends.extension import _C + + _cuda_op() + grad = torch.ones(6, 4, 8, device="cuda") + zero = torch.zeros(0, dtype=torch.int64, device="cuda") + pos = torch.arange(4, device="cuda") + one = torch.zeros(1, dtype=torch.int64, device="cuda") # R = 0 + empty = _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, one, torch.float32) + assert empty.shape == (0, 48) + with pytest.raises(RuntimeError, match="R \\+ 1 entries"): + _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, zero, torch.float32) + + def test_forward_rejects_more_chunks_than_one_grid_dimension_holds(self): + from rl_engine.backends.extension import _C + + _cuda_op() + rows = torch.zeros(3, 65536, device="cuda") + ti = torch.zeros(1, dtype=torch.int64, device="cuda") + with pytest.raises(RuntimeError, match="grid.y"): + _C.h3_adaln_row_gather_forward(rows, ti, ti, 65536, 3) + def test_registry_dispatches_cuda(self): from rl_engine.runtime.registry import KernelRegistry diff --git a/tests/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py index 853680a67..576a96f1e 100644 --- a/tests/models/minimax_h3/test_h3_report_timing.py +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -52,3 +52,24 @@ def test_the_reported_medians_come_from_the_stored_samples(): assert row[f"{key}_peak_mib"] >= 0 # the larger matmul is slower than the smaller one in the same run assert row["provider_us"] > row["candidate_us"] + + +def test_backward_timings_of_every_perf_case_exclude_the_forward_pass(monkeypatch): + """A case with a timed backward must prepare the graph outside the CUDA-event window.""" + + from rl_engine.runtime.registry import KernelRegistry + from rl_engine.validation.models import h3_report + + # The structure is what is checked; the cases need not be as large as the evidence runs. + for name in ("GATHER_SEQ_LENS", "NORM_SEQ_LENS"): + if hasattr(h3_report, name): + monkeypatch.setattr(h3_report, name, (257,)) + registry = KernelRegistry() + checked = 0 + for op, build in h3_report.PERF_CASES.items(): + for case in build(registry): + for side in ("candidate", "provider"): + if f"{side}_backward" in case: + checked += 1 + assert f"{side}_backward_setup" in case, f"{op} {case['case']} {side}" + assert checked From 3e4c84915a3042647b2f7fadc7636184ffba843b Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:53:08 +0000 Subject: [PATCH 126/139] fix(h3): check the rmsnorm table layout and tile metadata, scan rows once, and align the reference validators rmsnorm_modulate.cu: make_modulation rejects a shift/scale table whose row stride is smaller than N (an expanded or overlapping view is not a table of rows; a view of a wider table, such as table.chunk(2, dim=-1) with row stride 2N, is still read by its row pitch and a test pins that it is bitwise the copy). The modulated backward runs the shared tile-metadata pre-pass (tile_metadata.cuh, #486) before the segment kernels: sorted positions, tiles and segment boundaries that do not describe the rows end in a device-side assert instead of an out-of-bounds read, and the fold launch is skipped for zero segments. N is checked against the int and grid.y limits of the segment kernels. The column indices of the segment kernels were already int64 in the tile term, so nothing changes there. Tests run corrupted metadata in a child process. rmsnorm.py (CUDA op): a table that is not made of dense rows (stride(0) < H) is copied before the call, like mixed layouts; a wider-table view goes through as it is. Reference: the modulation backward takes the positions of all rows from positions_by_row (one sort and one host read instead of two torch.nonzero scans per row; a test compares the gradient bitwise with the loop it replaces). forward_modulated_fp32 takes check_range like forward_modulated. A device mismatch of shift/scale is a ValueError, as in the other H3 validators; a dtype mismatch stays a TypeError. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/rmsnorm_modulate.cu | 25 ++++++- .../cuda/model_specific/minimax_h3/rmsnorm.py | 11 ++- rl_engine/reference/minimax_h3/rmsnorm.py | 17 +++-- tests/models/minimax_h3/test_h3_rmsnorm.py | 50 +++++++++++++ .../test_h3_rmsnorm_native_validation.py | 73 +++++++++++++++++++ 5 files changed, 165 insertions(+), 11 deletions(-) diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 0165a1bb1..3f352f417 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -29,8 +29,11 @@ #include #include +#include #include +#include "tile_metadata.cuh" + namespace { constexpr int kVec = 4; // PyTorch's layer-norm vec_size @@ -291,6 +294,10 @@ 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"); + // Rows are read as dense N-element spans at `row_stride`: a view of a wider table is fine, an + // expanded or overlapping one (row_stride < N) is not a table of rows. + TORCH_CHECK(sh.size(0) <= 1 || sh.stride(0) >= n, "shift/scale row stride ", sh.stride(0), + " is smaller than N=", n); 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 (", @@ -416,9 +423,16 @@ std::vector h3_rmsnorm_backward( const int64_t segments = seg_first_tile->numel() - 1; const int64_t seg_tiles = tile_begin->numel(); const int64_t width = 2 * n; + TORCH_CHECK(width <= std::numeric_limits::max(), "N=", n, " is too large for the kernels"); + TORCH_CHECK((width + threads - 1) / threads <= 65535, "N=", n, + " needs more column blocks than the grid.y limit of 65535"); 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); + // The tile list is caller-built metadata: check it before the kernels read through it. + h3_tiles::check_tile_metadata(stream, sorted_pos->data_ptr(), rows, + tile_begin->data_ptr(), tile_end->data_ptr(), + seg_tiles, seg_first_tile->data_ptr(), segments); if (seg_tiles > 0) { H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_modulation_grad", [&] { modulation_grad_partial_kernel @@ -430,9 +444,12 @@ std::vector h3_rmsnorm_backward( }); 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(); + if (segments > 0) { // a zero-sized grid is not a valid launch + 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/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py index 5188f3f80..61e4a17b6 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py @@ -98,7 +98,16 @@ def forward_modulated( 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): + # A view of a wider table (``table.chunk(2, dim=-1)``: unit column stride, row stride 2H) + # goes to the kernel as it is, which indexes rows by the stride. Anything that is not a + # table of dense rows (mixed layouts, or rows that overlap because the stride is below + # H) is copied first. + if ( + shift.stride(1) != 1 + or scale.stride(1) != 1 + or shift.stride(0) != scale.stride(0) + or (shift.shape[0] > 1 and shift.stride(0) < shift.shape[1]) + ): shift, scale = shift.contiguous(), scale.contiguous() return _H3RMSNormCuda.apply(x, weight, shift, scale, index.contiguous(), eps) diff --git a/rl_engine/reference/minimax_h3/rmsnorm.py b/rl_engine/reference/minimax_h3/rmsnorm.py index 1727fd1d0..ddbdaf6bf 100644 --- a/rl_engine/reference/minimax_h3/rmsnorm.py +++ b/rl_engine/reference/minimax_h3/rmsnorm.py @@ -21,6 +21,8 @@ import torch import torch.nn.functional as F +from rl_engine.reference.minimax_h3.adaln_row_gather import positions_by_row + H3_NORM_EPS = 1e-5 _DTYPES = (torch.float32, torch.bfloat16, torch.float16) @@ -55,8 +57,10 @@ def validate_h3_modulation( ) -> 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.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if same_dtype and tensor.dtype != x.dtype: + raise TypeError(f"{name} must match x's dtype, got {tensor.dtype} for {x.dtype}") 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: @@ -93,8 +97,7 @@ def backward(ctx, grad): 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() + for row, positions in enumerate(positions_by_row(index, ctx.rows)): 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) @@ -134,9 +137,11 @@ def forward_modulated( 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): + def forward_modulated_fp32( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ): # 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_modulation(x, shift, scale, index, check_range=check_range, 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() diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py index 826b75700..6d91f60bb 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm.py @@ -74,6 +74,27 @@ def test_modulated_with_prevalidated_indices(self): expected, ) + def test_modulation_backward_is_bitwise_the_per_row_scan_it_replaced(self): + from rl_engine.reference.minimax_h3.rmsnorm import _DeterministicModulation + + g = torch.Generator(device="cpu").manual_seed(11) + n = torch.randn(300, 16, generator=g) + shift, scale = torch.randn(7, 16, generator=g), torch.randn(7, 16, generator=g) + index = torch.randint(0, 5, (300,), generator=g) # rows 5 and 6 are never referenced + grad = torch.randn(300, 16, generator=g) + leaves = [t.clone().requires_grad_(True) for t in (n, shift, scale)] + _DeterministicModulation.apply(*leaves, index).backward(grad) + + per_shift = grad.reshape(-1, 300, 16).sum(dim=0) + per_scale = (grad * n).reshape(-1, 300, 16).sum(dim=0) + d_shift, d_scale = torch.zeros(7, 16), torch.zeros(7, 16) + for row in range(7): # the loop of the earlier implementation + 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) + assert torch.equal(leaves[1].grad, d_shift) and torch.equal(leaves[2].grad, d_scale) + def test_rejects_bad_inputs(self): op = NativeH3RMSNormOp() x, w = torch.randn(2, 8), torch.ones(8) @@ -88,6 +109,13 @@ def test_rejects_bad_inputs(self): 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(IndexError): # the golden takes check_range like the provider path + op.forward_modulated_fp32(x, w, shift, scale, torch.tensor([0, 3])) + op.forward_modulated_fp32(x, w, shift, scale, torch.tensor([0, 2]), check_range=False) + with pytest.raises(ValueError, match="is on meta"): # a device mismatch is a ValueError + op.forward_modulated(x, w, shift.to("meta"), scale, torch.tensor([0, 1])) + with pytest.raises(TypeError, match="dtype"): # a dtype mismatch stays a TypeError + op.forward_modulated(x, w, shift.double(), scale, torch.tensor([0, 1])) with pytest.raises(ValueError): # one index per position S op.forward_modulated(x, w, shift, scale, torch.tensor([0])) @@ -139,6 +167,28 @@ def test_modulation_bitwise_in_every_dtype(self, dtype): ours = _cuda_op().forward_modulated(x, w, shift, scale, index) assert torch.equal(ours, provider_norm_modulate(x, w, shift, scale, index)) + def test_overlapping_table_rows_are_copied_and_wider_table_views_are_not(self): + w = (torch.rand(5376, device="cuda") + 0.5).bfloat16() + table = _x((9, 2 * 5376), seed=21, scale=0.5) + wide_shift, wide_scale = table.chunk(2, dim=-1) # row stride 2H, handed over as views + ti, tags = h3_packed_layout(100, 3, seed=21) + index = ti * 3 + tags + x = _x((1, 100, 5376), seed=21) + op = _cuda_op() + expected = op.forward_modulated( + x, w, wide_shift.contiguous(), wide_scale.contiguous(), index + ) + assert torch.equal(op.forward_modulated(x, w, wide_shift, wide_scale, index), expected) + # An expanded table (row stride 0) is a table of equal rows: copied, not rejected. + row = _x((1, 5376), seed=22, scale=0.5) + expanded = row.expand(9, 5376) + assert expanded.stride(0) == 0 + dense = row.repeat(9, 1) + assert torch.equal( + op.forward_modulated(x, w, expanded, expanded, index), + op.forward_modulated(x, w, dense, dense, 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) diff --git a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py index a2b032344..dda99680f 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py @@ -207,3 +207,76 @@ def test_modulation_row_outside_the_table_asserts_in_the_kernel(native, entrypoi assert ( "device-side assert triggered" in completed.stderr ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + + +def test_rejects_modulation_rows_that_overlap(native, inputs): + """A table of rows needs a row stride of at least N; an overlapping view is not one.""" + + table = torch.arange(32, dtype=torch.float32, device="cuda") + inputs["shift"] = inputs["scale"] = table.as_strided((2, 8), (4, 1)) + with pytest.raises(RuntimeError, match="row stride 4 is smaller than N=8"): + _call(native, inputs, "forward") + + +def test_a_view_of_a_wider_table_is_read_by_its_row_pitch(native, inputs): + """``table.chunk(2, dim=-1)`` has row stride 2N; the kernel must index rows by it.""" + + g = torch.Generator(device="cpu").manual_seed(0) + table = torch.randn(2, 16, generator=g).cuda() + views = table.chunk(2, dim=-1) + copies = tuple(v.contiguous() for v in views) + assert views[0].stride(0) == 16 and copies[0].stride(0) == 8 + out = [] + for shift, scale in (views, copies): + inputs["shift"], inputs["scale"] = shift, scale + out.append(_call(native, inputs, "forward")[0]) + assert torch.equal(out[0], out[1]) + + +def _backward_in_child(mutation: str) -> subprocess.CompletedProcess: + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + mutation = sys.argv[1] + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device="cuda") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) + args = dict( + grad=torch.ones_like(x), x=x, weight=weight, rstd=rstd, shift=shift, scale=scale, + index=torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda"), + sorted_pos=torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cuda"), + tile_begin=torch.tensor([0, 2], dtype=torch.int64, device="cuda"), + tile_end=torch.tensor([2, 4], dtype=torch.int64, device="cuda"), + seg_first_tile=torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"), + ) + if mutation == "tile_end_past_rows": + args["tile_end"] = torch.tensor([2, 9], dtype=torch.int64, device="cuda") + elif mutation == "position_past_rows": + args["sorted_pos"] = torch.tensor([0, 2, 1, 4], dtype=torch.int64, device="cuda") + elif mutation == "segments_past_tiles": + args["seg_first_tile"] = torch.tensor([0, 1, 5], dtype=torch.int64, device="cuda") + _C.h3_rmsnorm_backward(**args) + torch.cuda.synchronize() + """ + ) + return subprocess.run( + [sys.executable, "-c", probe, mutation], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + + +def test_the_backward_stops_on_tile_metadata_that_does_not_describe_the_rows(native): + control = _backward_in_child("none") + assert control.returncode == 0, control.stderr + for mutation in ("tile_end_past_rows", "position_past_rows", "segments_past_tiles"): + done = _backward_in_child(mutation) + assert done.returncode != 0, mutation + assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) From 1857fa69aed91b8333b9ba7cb16e985d0bfd1eba Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:58:21 +0000 Subject: [PATCH 127/139] fix(h3): check the gate backward's tile metadata and lengths, and scan the reference's rows once gate_residual.cu: the backward runs the shared tile-metadata pre-pass (tile_metadata.cuh, #486) before the segment kernels, so sorted positions, tiles and segment boundaries that do not describe the rows end in a device-side assert instead of an out-of-bounds read, including the case of no tiles with a non-trivial seg_first_tile (the partial buffer then holds one uninitialized row). Host checks: sorted_pos has M entries, tile_begin and tile_end have the same length, seg_first_tile has at least one entry, and the column blocks fit grid.y (65535); the fold launch is skipped for zero segments. Native gate backward, idle B200, min of 400: 86.6 -> 90.7 us at S = 4097, 362.7 -> 367.5 us at S = 32768, 1315.0 -> 1319.7 us at S = 131072 (one extra launch of about 4 us). Tests run corrupted metadata in a child process and the host checks. gate_residual.py: dy is viewed as the shape of y (validate_h3_gate_residual makes it the residual's shape, so nothing changes for valid calls) and the saved shapes are named for the tensor they belong to. Reference: the gate gradient takes the positions of all rows from positions_by_row (one sort and one host read instead of one torch.nonzero scan and sync per row); a test compares the gradient bitwise with the loop it replaces. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 22 ++++- .../minimax_h3/gate_residual.py | 6 +- .../reference/minimax_h3/gate_residual.py | 5 +- .../minimax_h3/test_h3_adaln_gate_residual.py | 87 +++++++++++++++++++ 4 files changed, 112 insertions(+), 8 deletions(-) diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 6e0b5a85b..123b434fa 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -27,6 +27,8 @@ #include #include +#include "tile_metadata.cuh" + namespace { __device__ __forceinline__ float to_f(float v) { return v; } @@ -216,8 +218,11 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: } const int64_t rows = grad.size(0); const int64_t n = grad.size(1); + TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); const int64_t segments = seg_first_tile.numel() - 1; TORCH_CHECK(segments == gate.size(0), "segments must equal the gate rows"); + TORCH_CHECK(sorted_pos.numel() == rows, "sorted_pos must have M entries"); + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); const int64_t tiles = tile_begin.numel(); const c10::cuda::CUDAGuard guard(grad.device()); auto stream = at::cuda::getCurrentCUDAStream(); @@ -225,7 +230,14 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: 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); + const int64_t col_blocks64 = (n + threads - 1) / threads; + TORCH_CHECK(col_blocks64 <= 65535, "N=", n, " needs ", col_blocks64, + " column blocks, more than the grid.y limit of 65535"); + const unsigned col_blocks = static_cast(col_blocks64); + // The tile list is caller-built metadata: check it before the kernels read through it. + h3_tiles::check_tile_metadata(stream, sorted_pos.data_ptr(), rows, + tile_begin.data_ptr(), tile_end.data_ptr(), tiles, + seg_first_tile.data_ptr(), segments); 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()); @@ -237,9 +249,11 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: 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); + if (segments > 0) { // a zero-sized grid is not a valid launch + 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/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index ad9bb6d8f..04f6bdf7d 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 @@ -36,7 +36,9 @@ def forward(ctx, residual, y, gate, index): 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 + # validate_h3_gate_residual makes y and residual one shape; each gradient is still viewed + # as the shape of its own input. + ctx.residual_shape, ctx.y_shape = residual.shape, y.shape return out.view(residual.shape) @staticmethod @@ -49,7 +51,7 @@ def backward(ctx, grad): record_backward( "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) - return grad, dy.view(ctx.shape), dgate, None + return grad, dy.view(ctx.y_shape), dgate, None class H3GateResidualCudaOp: diff --git a/rl_engine/reference/minimax_h3/gate_residual.py b/rl_engine/reference/minimax_h3/gate_residual.py index f68a0fdf4..0ad2c1b0e 100644 --- a/rl_engine/reference/minimax_h3/gate_residual.py +++ b/rl_engine/reference/minimax_h3/gate_residual.py @@ -18,6 +18,8 @@ import torch +from rl_engine.reference.minimax_h3.adaln_row_gather import positions_by_row + _DTYPES = (torch.float32, torch.bfloat16, torch.float16) @@ -69,8 +71,7 @@ def backward(ctx, grad): 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() + for row, positions in enumerate(positions_by_row(index, ctx.gate_rows)): 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 diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index 4dd9521f8..b26ab8a54 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -52,6 +52,25 @@ def _gate(table, chunk=2): class TestReference: + def test_gate_backward_is_bitwise_the_per_row_scan_it_replaced(self): + from rl_engine.reference.minimax_h3.gate_residual import _DeterministicGateResidual + + g = torch.Generator(device="cpu").manual_seed(5) + residual, y = torch.randn(2, 300, 16, generator=g), torch.randn(2, 300, 16, generator=g) + gate = torch.randn(7, 16, generator=g) + index = torch.randint(0, 5, (300,), generator=g) # rows 5 and 6 are never referenced + grad = torch.randn(2, 300, 16, generator=g) + leaf = gate.clone().requires_grad_(True) + _DeterministicGateResidual.apply(residual, y, leaf, index).backward(grad) + + per_position = (grad * y).reshape(-1, 300, 16).sum(dim=0) + expected = torch.zeros(7, 16) + for row in range(7): # the loop of the earlier implementation + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + expected[row] = per_position.index_select(0, positions).sum(dim=0) + assert torch.equal(leaf.grad, expected) + def test_rejects_bad_inputs(self): op = NativeH3GateResidualOp() res, y, gate = torch.zeros(4, 8), torch.zeros(4, 8), torch.zeros(3, 8) @@ -165,3 +184,71 @@ def test_gate_row_outside_the_table_asserts_in_the_kernel(entrypoint, bad_row): assert ( "device-side assert triggered" in completed.stderr ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + + +def _gate_backward_in_child(mutation: str) -> subprocess.CompletedProcess: + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + if not hasattr(_C, "h3_gate_residual_backward"): + sys.exit(77) + mutation = sys.argv[1] + x = torch.randn(4, 8, device="cuda") + gate = torch.ones(2, 8, device="cuda") + index = torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda") + i64 = dict(dtype=torch.int64, device="cuda") + sorted_pos = torch.tensor([0, 2, 1, 3], **i64) + begin, end = torch.tensor([0, 2], **i64), torch.tensor([2, 4], **i64) + first = torch.tensor([0, 1, 2], **i64) + if mutation == "tile_end_past_rows": + end = torch.tensor([2, 9], **i64) + elif mutation == "position_past_rows": + sorted_pos = torch.tensor([0, 2, 1, 4], **i64) + elif mutation == "segments_past_tiles": + first = torch.tensor([0, 1, 5], **i64) + elif mutation == "no_tiles": + begin, end = begin[:0], end[:0] + elif mutation == "short_sorted_pos": + sorted_pos = sorted_pos[:3] + elif mutation == "no_boundary": + first = first[:0] + grad = torch.ones_like(x) + _C.h3_gate_residual_backward(grad, x, gate, index, sorted_pos, begin, end, first) + torch.cuda.synchronize() + """ + ) + return subprocess.run( + [sys.executable, "-c", probe, mutation], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_the_gate_backward_stops_on_tile_metadata_that_does_not_describe_the_rows(): + control = _gate_backward_in_child("none") + if control.returncode == 77: + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + assert control.returncode == 0, control.stderr + for mutation in ( + "tile_end_past_rows", + "position_past_rows", + "segments_past_tiles", + "no_tiles", + ): + done = _gate_backward_in_child(mutation) + assert done.returncode != 0, mutation + assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) + # Lengths are checked on the host, before any launch. + for mutation, message in ( + ("short_sorted_pos", "sorted_pos must have M entries"), + ("no_boundary", "seg_first_tile must hold R + 1 entries"), + ): + done = _gate_backward_in_child(mutation) + assert done.returncode != 0 and message in done.stderr, (mutation, done.stderr) From fcf0a588b191d7b321757911e9e9331de26e4c7a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:04:07 +0000 Subject: [PATCH 128/139] fix(h3): run the final AdaLN backward only for the inputs that need a gradient, refuse pre-SM80 at construction, fix three doc statements final_adaln_out.py: backward reads ctx.needs_input_grad. The projection's weight and bias folds run only when one of them needs a gradient (the bias fold only when the bias does), d_temb (a SiLU VJP over a full input-gradient pass) only when temb does, and every input that does not need a gradient gets None instead of a tensor; a test freezes each input in turn and compares the others bitwise with the all-trainable gradients. H3FinalAdaLNOutCudaOp.__init__ calls det_linear_device_supported() like the projection and modulation ops, so the registry falls back to the PyTorch reference on a pre-SM80 device instead of every forward raising. docs/operators/h3-final-adaln-out.md: the projection is split with chunk(2, dim=-1); the gtest sentence says the gtest failed on d_norm_w from S = 257 and on d_scale/dW/d_temb from S = 1024 because the golden did not round where the model does; the gradient statement of the bitwise comparison names the exception recorded in the report (all five gradients differ at S = 257, 4097 and 32768; at S = 1 the bias gradient is equal). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 9 ++--- .../minimax_h3/final_adaln_out.py | 29 +++++++++++++--- .../minimax_h3/test_h3_final_adaln_out.py | 33 +++++++++++++++++++ 3 files changed, 63 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 90fbbfd17..067e8ec1e 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -6,7 +6,7 @@ blocks (RFC #420, WS1 step 7): ```text -shift, scale = norm_out.linear(silu(temb).to(bf16)).chunk(2) # (T, 5376) each, shift first +shift, scale = norm_out.linear(silu(temb).to(bf16)).chunk(2, dim=-1) # (T, 5376) each, shift first out = norm_out.norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] ``` @@ -49,8 +49,9 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste 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. + systematically with S, and the gtest failed on `d_norm_w` from S = 257 because the golden did not + round where the model does. `norm_out.norm`'s rounding enters `d_scale`, `dW` and `d_temb` + summed over S, and the gtest failed on those from S = 1024 for the same reason. ## Performance Notes @@ -121,7 +122,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. +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 the gradients differ as expected for the deterministic FP32 reductions: all five at S = 257, 4097 and 32768, and all but the bias gradient at S = 1, where one position leaves a single term to sum. [`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 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 55993caad..042e50f7d 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 @@ -16,6 +16,7 @@ det_linear_available, det_linear_backward_input, det_linear_backward_weight, + det_linear_device_supported, det_linear_forward, silu_backward_fp32, ) @@ -56,6 +57,7 @@ def forward(ctx, x, norm_weight, temb, weight, bias, timestep_indices, eps): @staticmethod def backward(ctx, grad): x2, norm_weight, rstd, temb, act, weight, table, timestep_indices = ctx.saved_tensors + need_x, need_norm_weight, need_temb, need_weight, need_bias = ctx.needs_input_grad[:5] shift, scale = table.chunk(2, dim=-1) positions = _repeat_positions(timestep_indices, x2.shape[0]) tiles = _segment_tiles(positions, table.shape[0]) @@ -69,14 +71,29 @@ def backward(ctx, grad): timestep_indices, *tiles, ) - # The table gradient stays FP32 into the projection backward. + # The table gradient stays FP32 into the projection backward; the projection's weight, + # bias and input folds run only for the inputs that need a gradient. 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) + d_temb = dw = db = None + if need_weight or need_bias: + dw, *rest = det_linear_backward_weight(d_table, act, weight.dtype, with_bias=need_bias) + db = rest[0] if need_bias else None + if need_temb: + 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 + return ( + dx.view(ctx.x_shape) if need_x else None, + d_norm_weight if need_norm_weight else None, + d_temb, + dw if need_weight else None, + db if need_bias else None, + None, + None, + ) class H3FinalAdaLNOutCudaOp: @@ -95,6 +112,10 @@ class H3FinalAdaLNOutCudaOp: 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") + if not det_linear_device_supported(): + # The BF16 projection is an mma.sync kernel; refuse here so the registry falls back to + # the PyTorch reference instead of failing in every forward. + raise RuntimeError("the BF16 h3 det_linear forward needs SM80 or newer") 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) 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 5fb964677..dc2bb8153 100644 --- a/tests/models/minimax_h3/test_h3_final_adaln_out.py +++ b/tests/models/minimax_h3/test_h3_final_adaln_out.py @@ -130,6 +130,39 @@ def grads(fn, dtype=None): rel = ((g.double() - r).abs().max() / r.abs().max()).item() assert rel < 1e-2, (name, rel) # BF16 output rounding and round(1 + scale) only + @pytest.mark.parametrize( + "frozen", ["x", "norm_weight", "temb", "weight", "bias", "weight_and_bias", "temb_and_bias"] + ) + def test_frozen_inputs_get_no_gradient_and_the_others_are_unchanged(self, norm_out, frozen): + x, temb, ti = _inputs(1500, seed=6) + grad = torch.randn_like(x) + names = ("x", "norm_weight", "temb", "weight", "bias") + tensors = (x, norm_out[0], temb, *norm_out[1:]) + + def run(frozen_names): + leaves = [ + t.detach().clone().requires_grad_(name not in frozen_names) + for name, t in zip(names, tensors) + ] + _cuda_op()(*leaves, ti).backward(grad) + return [leaf.grad for leaf in leaves] + + full = run(()) + partial = run(frozen.split("_and_") if "_and_" in frozen else [frozen]) + frozen_set = frozen.split("_and_") if "_and_" in frozen else [frozen] + for name, got, expected in zip(names, partial, full): + if name in frozen_set: + assert got is None, name + else: + assert torch.equal(got, expected), name + + def test_unsupported_device_is_refused_at_construction(self, monkeypatch): + from rl_engine.backends.cuda.model_specific.minimax_h3 import final_adaln_out as module + + monkeypatch.setattr(module, "det_linear_device_supported", lambda *args: False) + with pytest.raises(RuntimeError, match="SM80"): + module.H3FinalAdaLNOutCudaOp() + def test_report_backward_uses_fp64_golden(self, monkeypatch): from types import SimpleNamespace From 86d6f11c42af16532bd7644721133f0c3bf2595c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:58:48 +0000 Subject: [PATCH 129/139] docs(h3): row-gather backward timings measured without the forward pass The B200 report and figure are regenerated from a clean tree at c416a101 on an otherwise idle GPU (utilization 0% before and after) with the backward cases built by candidate_backward_setup / provider_backward_setup. The backward columns of the latency table are now backward-only: 1.01 vs 1.30 ms at S = 4097, 1.76 vs 8.00 ms at S = 32768, 5.59 vs 32.2 ms at S = 131072 (CUDA vs index_select); the peak-memory sentence describes the backward-only increment. The H100 report and chain_replay.json are not regenerated and still carry the earlier timing. Links point at the pr_assets commit that holds the new files. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 2ce2990c1..93c66a4a5 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -113,25 +113,33 @@ 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 | 1.01 ms | 1.30 ms | +| 32768 | 0.34 ms | 0.63 ms | 1.76 ms | 8.00 ms | +| 131072 | 1.24 ms (6.8 TB/s write) | 2.40 ms | 5.59 ms | 32.2 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. +The backward columns time `autograd.backward` only: the forward pass and the graph construction run +outside the timed region (`candidate_backward_setup` / `provider_backward_setup`), and the execution +order alternates and is recorded. + +Backward-only peak memory (increment over the forward graph): 255, 2034 and 8130 MiB for the CUDA op +at S = 4097, 32768 and 131072, against about 1 MiB for the provider. Autograd stacks the six +output gradients into one `(6, S, H)` tensor for the CUDA op and the per-tile FP32 partials are +materialised, while the provider's atomic scatter-add writes straight into the `(3T, 6H)` table +gradient. The peak is the price of the deterministic FP32 segment sum. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/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/999c0c75078605229ad74dc600dc963d89c7ed9e/pr486/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/999c0c75078605229ad74dc600dc963d89c7ed9e/pr486/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`](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 rewritten from a clean tree at commit `a6c85451`, after the op-backward check +`report.json` was rewritten from a clean tree at commit `c416a101` on an otherwise idle B200 (backward +timings exclude the forward pass; the previous revision came from `a6c85451`), 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 `61d23ade`. From ac225c70bd0c43c41933b7de1f59b90ba5198d96 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:01:01 +0000 Subject: [PATCH 130/139] fix(build): read nvcc options files, say why a build left the H3 kernels out, and keep the sinusoid sources self-contained * NVCC_*_FLAGS can switch fast math on through `--options-file` / `-optf` (a comma-separated list, as `--options-file=a,b` or `--options-file a,b`): nvcc expands those files itself, while it does not expand `@file`. The predicate read only the `@file` spelling, so `NVCC_APPEND_FLAGS="--options-file flags.txt"` with `--use_fast_math` inside built the H3 kernels with approximated math. Both spellings are read now; a file that cannot be opened, or an option nvcc rejects for lack of a file name, counts as relaxed (the build would fail anyway). * The build warning names the variable and the reason (`NVCC_APPEND_FLAGS contains --ftz true`, `... cannot read options file ...`, `... cannot be parsed`), so an unreadable value is not mistaken for a relaxed one. * A test uses nvcc as the oracle: `nvcc --dryrun` prints the cicc flags it would use, and the predicate must agree with them for every spelling in the list (plain, split, dashes, options files, nested files, PREPEND and APPEND). Another test checks the nvcc list handed to the compiler: a relaxed flag in it and the H3 define never appear together. * The dispatch fallback test removes the symbol from the extension (instead of replacing the probe) and checks that the CUDA backend was rejected. * `timestep_sinusoid.cu` includes for std::min; `frequencies_fp32` rejects an odd or non-positive channel count like the validator; the H3 block of `_C.pyi` says the symbols exist only in builds that compile the H3 sources. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- build_tools/extensions.py | 115 +++++++++----- csrc/cuda/h3/timestep_sinusoid.cu | 1 + rl_engine/_C.pyi | 7 + .../reference/minimax_h3/timestep_sinusoid.py | 2 + .../build/test_build_platform_collectives.py | 146 ++++++++++++++++++ .../minimax_h3/test_h3_timestep_sinusoid.py | 8 + tests/runtime/test_dispatch.py | 12 +- 7 files changed, 252 insertions(+), 39 deletions(-) diff --git a/build_tools/extensions.py b/build_tools/extensions.py index 5caf96f50..cbd64b53e 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -25,79 +25,118 @@ # nvcc options that relax FP32 math. They can also arrive through the environment, where PyTorch's # BuildExtension appends them after the flags built here, so the H3 kernels check for them too. # nvcc takes ``--ftz``, ``--prec-div`` and ``--prec-sqrt`` with a value, as ``--opt=value`` or -# ``--opt value``, and accepts one or two leading dashes; ``--use_fast_math`` takes none. +# ``--opt value``, and accepts one or two leading dashes; ``--use_fast_math`` takes none. nvcc only +# takes true / false for them and stops on ``--ftz 1``; 1 / 0 are listed too, so that a spelling +# nvcc rejects is never read as the safe one. _RELAXED_MATH_FLAG = "use_fast_math" _RELAXED_MATH_VALUE_FLAGS = { "ftz": {"true", "1"}, "prec-div": {"false", "0"}, "prec-sqrt": {"false", "0"}, } -_RESPONSE_FILE_DEPTH = 4 +# nvcc reads more options from the files named by ``--options-file`` / ``-optf`` (a comma-separated +# list, as ``--options-file=a,b`` or ``--options-file a,b``), also when the option comes from +# NVCC_*_FLAGS. nvcc 13.0 does not expand ``@file`` (it takes it for an input file and stops); that +# spelling is read as well, so a driver that does expand it is covered. +_OPTIONS_FILE_FLAGS = {"optf", "options-file"} +_OPTIONS_FILE_DEPTH = 4 -def _nvcc_flag_tokens(text: str, *, depth: int = 0) -> list[str] | None: - """Split an NVCC_*_FLAGS value into tokens; None when a response file cannot be read. +class _UncheckableFlags(Exception): + """The flags cannot be read well enough to say whether they relax math.""" + + +def _nvcc_flag_tokens(text: str, *, depth: int = 0) -> list[str]: + """Split an NVCC_*_FLAGS value into tokens, expanding options files. Tokens are split like a shell would, then on whitespace and ``;`` (TORCH_NVCC_FLAGS is - ';'-separated), so ``"-O2 --use_fast_math"`` in one quoted word is still two options. ``@file`` - response files are expanded. + ';'-separated), so ``"-O2 --use_fast_math"`` in one quoted word is still two options. Raises + _UncheckableFlags for an options file that cannot be read or an option nvcc rejects for lack + of its file name, and ValueError for a string that cannot be split. """ + words = [part for word in shlex.split(text) for part in word.replace(";", " ").split()] tokens: list[str] = [] - for word in shlex.split(text): - for token in word.replace(";", " ").split(): - if not token.startswith("@"): - tokens.append(token) - continue - if depth >= _RESPONSE_FILE_DEPTH: - return None + i = 0 + while i < len(words): + token = words[i] + i += 1 + files: list[str] + if token.startswith("@"): + files = [token[1:]] + elif token.startswith("-") and token.lstrip("-").partition("=")[0] in _OPTIONS_FILE_FLAGS: + _, has_value, value = token.lstrip("-").partition("=") + if not has_value: + if i >= len(words) or words[i].startswith("-"): + raise _UncheckableFlags(f"{token} has no file name") + value = words[i] + i += 1 + files = value.split(",") + else: + tokens.append(token) + continue + if depth >= _OPTIONS_FILE_DEPTH: + raise _UncheckableFlags(f"options files nest deeper than {_OPTIONS_FILE_DEPTH} levels") + for name in files: try: - inner = _nvcc_flag_tokens(Path(token[1:]).read_text(), depth=depth + 1) - except (OSError, ValueError): - return None - if inner is None: - return None - tokens.extend(inner) + content = Path(name).read_text() + except OSError as exc: + raise _UncheckableFlags( + f"cannot read options file {name!r}: {exc.strerror}" + ) from exc + tokens.extend(_nvcc_flag_tokens(content, depth=depth + 1)) return tokens -def _tokens_relax_math(tokens: list[str]) -> bool: +def _tokens_relax_math(tokens: list[str]) -> str | None: + """Return the first option that relaxes FP32 math, spelled as given, or None.""" + for i, token in enumerate(tokens): if not token.startswith("-"): continue name, has_value, value = token.lstrip("-").partition("=") if name == _RELAXED_MATH_FLAG: - return True + return token relaxed_values = _RELAXED_MATH_VALUE_FLAGS.get(name) if relaxed_values is None: continue + spelled = token if not has_value: # "--ftz true": the value is the next word. Without one, nvcc rejects the option; treat # it as relaxed rather than guess. nxt = tokens[i + 1] if i + 1 < len(tokens) else "" if nxt.startswith("-") or not nxt: - return True + return f"{token} without a value" value = nxt + spelled = f"{token} {nxt}" if value.lower() in relaxed_values: - return True - return False + return spelled + return None -def _external_relaxed_math() -> bool: - """Return whether NVCC_*_FLAGS / TORCH_NVCC_FLAGS turn on relaxed FP32 math. +def _relaxed_math_reason() -> str | None: + """Say why NVCC_*_FLAGS / TORCH_NVCC_FLAGS turn on relaxed FP32 math, or None if they do not. - Fails closed: a value that cannot be parsed or a response file that cannot be read counts as - relaxed, because the H3 kernels would otherwise be built with flags nobody checked. + Fails closed: a value that cannot be parsed or an options file that cannot be read gives a + reason too, because the H3 kernels would otherwise be built with flags nobody checked. """ for name in ("NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS", "TORCH_NVCC_FLAGS"): try: - tokens = _nvcc_flag_tokens(os.environ.get(name, "")) - except ValueError: - return True - if tokens is None or _tokens_relax_math(tokens): - return True - return False + relaxed = _tokens_relax_math(_nvcc_flag_tokens(os.environ.get(name, ""))) + except ValueError as exc: + return f"{name} cannot be parsed ({exc})" + except _UncheckableFlags as exc: + return f"{name}: {exc}" + if relaxed is not None: + return f"{name} contains {relaxed}" + return None + + +def _external_relaxed_math() -> bool: + """Return whether NVCC_*_FLAGS / TORCH_NVCC_FLAGS relax FP32 math or cannot be checked.""" + + return _relaxed_math_reason() is not None def _resolve_package_data_source( @@ -307,11 +346,13 @@ def get_extensions(): use_fast_math = envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH) # One predicate decides both the H3 sources and the define that enables their bindings, so # the bindings can never outlive the kernels. - relaxed_math_env = _external_relaxed_math() + relaxed_reason = _relaxed_math_reason() + relaxed_math_env = relaxed_reason is not None if relaxed_math_env and not is_rocm: warnings.warn( - "NVCC_*_FLAGS / TORCH_NVCC_FLAGS relax FP32 math; the bitwise MiniMax-H3 kernels " - "are left out of this build and dispatch falls back to the PyTorch reference.", + f"NVCC_*_FLAGS / TORCH_NVCC_FLAGS relax FP32 math or cannot be checked " + f"({relaxed_reason}); the bitwise MiniMax-H3 kernels are left out of this build " + "and dispatch falls back to the PyTorch reference.", RuntimeWarning, stacklevel=2, ) diff --git a/csrc/cuda/h3/timestep_sinusoid.cu b/csrc/cuda/h3/timestep_sinusoid.cu index 2acfd516c..c618a0ad3 100644 --- a/csrc/cuda/h3/timestep_sinusoid.cu +++ b/csrc/cuda/h3/timestep_sinusoid.cu @@ -22,6 +22,7 @@ #include #include +#include #include #include diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 6c435b850..3371735f8 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -325,6 +325,13 @@ def deterministic_collective_rocm_ipc_all_gather_input( input: torch.Tensor, output: torch.Tensor, ) -> None: ... + +# MiniMax-H3 kernels (RFC #420). These exist only in a build that compiles the H3 sources: not on +# ROCm, and not with KERNEL_ALIGN_USE_FAST_MATH=1 or relaxed NVCC_*_FLAGS (build_tools/extensions.py +# leaves the sources and the KERNEL_ALIGN_WITH_H3_KERNELS define out together). Like the ROCm and +# SM90 symbols above, they are declared here for type checking; the Python backends test for them +# (the ``*_available()`` helpers, called by every constructor) so dispatch falls back to the PyTorch +# reference instead of raising AttributeError. def h3_timestep_sinusoid_forward( timestep: torch.Tensor, num_channels: int = 256, diff --git a/rl_engine/reference/minimax_h3/timestep_sinusoid.py b/rl_engine/reference/minimax_h3/timestep_sinusoid.py index f20aa77e9..030159836 100644 --- a/rl_engine/reference/minimax_h3/timestep_sinusoid.py +++ b/rl_engine/reference/minimax_h3/timestep_sinusoid.py @@ -112,6 +112,8 @@ def forward_fp32( def frequencies_fp32(num_channels: int, device: torch.device | str) -> torch.Tensor: """The FP32 frequency table the provider path multiplies by.""" + if num_channels <= 0 or num_channels % 2 != 0: + raise ValueError(f"num_channels must be a positive even number, got {num_channels}") half = num_channels // 2 exponent = -math.log(H3_MAX_PERIOD) * torch.arange( start=0, end=half, dtype=torch.float32, device=device diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index dbabf9404..ce2abea84 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -3,7 +3,12 @@ from __future__ import annotations +import os +import re import runpy +import shutil +import subprocess +from pathlib import Path from typing import Any import pytest @@ -15,7 +20,9 @@ H3_BITWISE_SOURCES, H3_KERNELS_DEFINE, _external_relaxed_math, + _relaxed_math_reason, _resolve_package_data_source, + _tokens_relax_math, get_cmdclass, ) @@ -226,6 +233,145 @@ def test_predicate_reads_response_files_and_fails_closed(monkeypatch, tmp_path) assert _external_relaxed_math() +_OPTION_FILES = { + "fast.opts": "--use_fast_math\n", + "ftz.opts": "-O2\n--ftz=true\n--prec-div=false\n", + "sqrt.opts": "--prec-sqrt=false", + "plain.opts": "-lineinfo\n--fmad=false\n", + "nested.opts": "-optf fast.opts\n", +} + + +@pytest.mark.parametrize( + ("flags", "relaxed"), + [ + ("-optf fast.opts", True), + ("--options-file fast.opts", True), + ("--options-file=fast.opts", True), + ("-options-file=ftz.opts", True), + ("--options-file plain.opts,sqrt.opts", True), + ("-optf=plain.opts,fast.opts", True), + ("-O2 -optf nested.opts", True), + ("-optf plain.opts", False), + ("--options-file=plain.opts,plain.opts", False), + ("@fast.opts", True), + ("@plain.opts", False), + ], +) +def test_predicate_reads_the_options_files_nvcc_reads( + monkeypatch, tmp_path, flags, relaxed +) -> None: + # nvcc expands ``--options-file`` / ``-optf`` (a comma-separated list) wherever the option + # comes from, so ``NVCC_APPEND_FLAGS='--options-file flags.txt'`` can switch fast math on. + _clear_nvcc_env(monkeypatch) + for name, content in _OPTION_FILES.items(): + (tmp_path / name).write_text(content) + monkeypatch.chdir(tmp_path) + for variable in _NVCC_ENV: + monkeypatch.setenv(variable, flags) + assert _external_relaxed_math() is relaxed, (variable, flags) + monkeypatch.delenv(variable) + + +@pytest.mark.parametrize("flags", ["-optf", "--options-file -O2", "-optf missing.opts"]) +def test_options_files_nvcc_cannot_open_count_as_relaxed(monkeypatch, tmp_path, flags) -> None: + _clear_nvcc_env(monkeypatch) + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("NVCC_PREPEND_FLAGS", flags) + assert _external_relaxed_math() + + +def test_the_build_warning_names_the_variable_and_the_reason(monkeypatch, tmp_path) -> None: + _clear_nvcc_env(monkeypatch) + missing = str(tmp_path / "missing.opts") + for flags, reason in [ + ("-O2 --ftz true", "NVCC_APPEND_FLAGS contains --ftz true"), + ("--prec-div", "NVCC_APPEND_FLAGS contains --prec-div without a value"), + ("--use_fast_math", "NVCC_APPEND_FLAGS contains --use_fast_math"), + (f"-optf {missing}", f"NVCC_APPEND_FLAGS: cannot read options file '{missing}'"), + ("'unterminated", "NVCC_APPEND_FLAGS cannot be parsed"), + ]: + monkeypatch.setenv("NVCC_APPEND_FLAGS", flags) + assert _relaxed_math_reason().startswith(reason), (flags, _relaxed_math_reason()) + with pytest.warns(RuntimeWarning, match=re.escape(reason)): + _load_extension_config(monkeypatch, hip=None) + monkeypatch.delenv("NVCC_APPEND_FLAGS") + assert _relaxed_math_reason() is None + + +def _nvcc() -> str | None: + for root in (os.environ.get("CUDA_HOME"), "/usr/local/cuda"): + if root and (Path(root) / "bin" / "nvcc").exists(): + return str(Path(root) / "bin" / "nvcc") + return shutil.which("nvcc") + + +@pytest.mark.skipif(_nvcc() is None, reason="needs nvcc") +@pytest.mark.parametrize( + "flags", + [ + "", + "-lineinfo --fmad=false", + "--use_fast_math", + "-use_fast_math", + "--ftz=true", + "--ftz=false", + "-ftz true", + "--ftz false", + "--prec-div=false", + "--prec-div true", + "-prec-div false", + "--prec-sqrt=false", + "-prec-sqrt=true", + "-optf fast.opts", + "--options-file=ftz.opts", + "--options-file plain.opts,sqrt.opts", + "-optf plain.opts", + "-optf nested.opts", + ], +) +@pytest.mark.parametrize("variable", ["NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS"]) +def test_predicate_agrees_with_what_nvcc_compiles_with( + monkeypatch, tmp_path, flags, variable +) -> None: + # The oracle is nvcc itself: ``--dryrun`` prints the cicc command it would run, whose + # -ftz / -prec_div / -prec_sqrt say whether the flags turned FP32 math into the fast variants. + # (nvcc takes only true / false for these options and stops on ``--ftz 1``; the predicate reads + # 1 / 0 as relaxed anyway, which is the safe side of a build that fails.) + _clear_nvcc_env(monkeypatch) + for name, content in _OPTION_FILES.items(): + (tmp_path / name).write_text(content) + (tmp_path / "k.cu").write_text("__global__ void k(float* a) { a[0] = 1.0f / a[1]; }\n") + monkeypatch.chdir(tmp_path) + monkeypatch.setenv(variable, flags) + + done = subprocess.run( + [_nvcc(), "-arch=compute_80", "-ptx", "k.cu", "-o", "k.ptx", "--dryrun"], + capture_output=True, + text=True, + env=os.environ.copy(), + ) + assert done.returncode == 0, done.stderr + cicc = [line for line in done.stderr.splitlines() if "cicc" in line] + assert cicc, done.stderr + nvcc_relaxes = any( + f" {opt} " in f" {cicc[0]} " for opt in ("-ftz=1", "-prec_div=0", "-prec_sqrt=0") + ) + assert _external_relaxed_math() is nvcc_relaxes, (variable, flags, cicc[0]) + + +def test_no_flag_built_here_relaxes_math_while_the_h3_define_is_set(monkeypatch) -> None: + # The predicate looks at the environment; this looks at the nvcc list that is handed to the + # compiler, so a relaxed flag added to the list later cannot ship together with the H3 define. + _clear_nvcc_env(monkeypatch) + default = _load_extension_config(monkeypatch, hip=None)["extra_compile_args"]["nvcc"] + assert H3_KERNELS_DEFINE in default and _tokens_relax_math(default) is None + + monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", "1") + fast = _load_extension_config(monkeypatch, hip=None)["extra_compile_args"]["nvcc"] + assert H3_KERNELS_DEFINE not in fast and _tokens_relax_math(fast) == "--use_fast_math" + + def test_unrelated_nvcc_flags_keep_the_h3_kernels(monkeypatch) -> None: _clear_nvcc_env(monkeypatch) monkeypatch.setenv("NVCC_APPEND_FLAGS", "-lineinfo --fmad=true") diff --git a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py index caad9ffcb..5ccca25ef 100644 --- a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py +++ b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py @@ -84,6 +84,14 @@ def test_rejects_odd_channels(self): with pytest.raises(ValueError): NativeH3TimestepSinusoidOp().forward(torch.tensor([0.5]), num_channels=255) + @pytest.mark.parametrize("num_channels", [0, -2, 1, 255]) + def test_frequency_table_rejects_what_the_validator_rejects(self, num_channels): + # The table is public API (the backward builds it): it must not truncate an odd count to + # half = (n - 1) / 2 or return an empty table where forward raises. + with pytest.raises(ValueError, match="positive even"): + NativeH3TimestepSinusoidOp.frequencies_fp32(num_channels, "cpu") + assert NativeH3TimestepSinusoidOp.frequencies_fp32(256, "cpu").shape == (128,) + @requires_cuda class TestCuda: diff --git a/tests/runtime/test_dispatch.py b/tests/runtime/test_dispatch.py index 2d095aa14..6dfb5a6ce 100644 --- a/tests/runtime/test_dispatch.py +++ b/tests/runtime/test_dispatch.py @@ -252,12 +252,20 @@ def test_h3_timestep_falls_back_when_fast_math_build_omits_kernel(monkeypatch): 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) + # Simulate the omitted binding itself (an extension without the symbol), not the probe that + # looks for it, so the test follows the real trigger. + class ExtensionWithoutH3: + pass + + monkeypatch.setattr(timestep_sinusoid, "_C", ExtensionWithoutH3()) + monkeypatch.setattr(timestep_sinusoid, "_EXT_AVAILABLE", True) monkeypatch.setattr(torch.version, "hip", None, raising=False) - op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda:0") + registry = KernelRegistry() + op = registry.get_op("timestep_sinusoid_h3", device="cuda:0") assert isinstance(op, NativeH3TimestepSinusoidOp) + assert "CUDA_H3_TIMESTEP_SINUSOID" in registry._failed_backends def test_npu_available_handles_runtime_failure(monkeypatch): From 2ee5abc074f8f27ed889069b344e36eee1cf3f2e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:05:21 +0000 Subject: [PATCH 131/139] fix(h3): skip the MLP bias fold when no bias needs a gradient, and pin the device and K = 0 behaviour of det_linear * `timestep_mlp` backward passes `with_bias=` to `det_linear_backward_weight`, so the bias fold of a layer whose weight alone needs a gradient is not run and not allocated (`adaln_modulation` already did this). A test freezes each input in turn, compares the remaining gradients bitwise with the all-trainable run and records the `with_bias` argument of every call. * The det_linear entry points are registered as plain function pointers; the `CUDAGuard` inside each of them is what puts the launch on the tensors' device. A two-GPU test calls forward and both backwards with the tensors on cuda:1 while cuda:0 is current and compares the bytes with the same call on cuda:0, for FP32 and BF16. * `h3_det_linear_backward_weight` with K = 0 returns an empty dW and the real bias gradient (the column sums of grad); a test pins that, so an early return for K = 0 cannot hand back an uninitialised bias. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../model_specific/minimax_h3/timestep_mlp.py | 11 ++- .../models/minimax_h3/test_h3_timestep_mlp.py | 97 +++++++++++++++++++ 2 files changed, 105 insertions(+), 3 deletions(-) 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 3244d49d6..87763d308 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" @@ -45,13 +45,18 @@ def backward(ctx, grad_out): 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) + # The bias fold is skipped when only the weight needs a gradient. + dw2, *rest = det_linear_backward_weight( + grad_out, hidden, torch.float32, with_bias=need_b2 + ) + db2 = rest[0] if need_b2 else None 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) + dw1, *rest = det_linear_backward_weight(d_pre, x, torch.float32, with_bias=need_b1) + db1 = rest[0] if need_b1 else None if need_x: dx = det_linear_backward_input(d_pre, w1, torch.float32) record_backward( diff --git a/tests/models/minimax_h3/test_h3_timestep_mlp.py b/tests/models/minimax_h3/test_h3_timestep_mlp.py index 09d1b2dd4..28359f0c2 100644 --- a/tests/models/minimax_h3/test_h3_timestep_mlp.py +++ b/tests/models/minimax_h3/test_h3_timestep_mlp.py @@ -150,6 +150,25 @@ def test_other_sizes(self, k_in, hidden, out): rtol=1e-4, ) + @pytest.mark.parametrize("with_bias", [True, False]) + def test_weight_gradient_with_no_input_columns_keeps_the_bias_gradient(self, with_bias): + # K = 0: dW is an (N, 0) tensor, but the bias gradient is still the column sums of grad, so + # it must be computed (not skipped like the empty dW). + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_weight, + ) + + _cuda_op() + grad = torch.randn(5, 7, device="cuda") + x = torch.empty(5, 0, device="cuda") + result = det_linear_backward_weight(grad, x, torch.float32, with_bias=with_bias) + assert len(result) == (2 if with_bias else 1) + assert result[0].shape == (7, 0) + if with_bias: + torch.testing.assert_close( + result[1].double(), grad.double().sum(0), atol=1e-5, rtol=1e-5 + ) + def test_rejects_unaligned_k(self): params = _synthetic_params(6, 8, 8) with pytest.raises(RuntimeError, match="multiple of 4"): @@ -192,6 +211,51 @@ def dx(rows): for i in range(6): assert torch.equal(dx(slice(i, i + 1))[0], full[i]) + @pytest.mark.parametrize( + "frozen", + [("b1",), ("b2",), ("b1", "b2"), ("w1",), ("w2",), ("x",), ("w1", "b1", "w2", "b2")], + ) + def test_frozen_inputs_get_no_gradient_and_do_not_change_the_others(self, frozen, monkeypatch): + from rl_engine.backends.cuda.model_specific.minimax_h3 import timestep_mlp + + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=6) + x = _features(3, seed=6) + grad_out = torch.randn(3, 2688, device="cuda") + + def grads(frozen_names): + leaves = [t.detach().clone() for t in (x, *params)] + for name, leaf in zip(NAMES, leaves, strict=True): + leaf.requires_grad_(name not in frozen_names) + if all(not leaf.requires_grad for leaf in leaves): + return leaves + op(*leaves).backward(grad_out) + return leaves + + full = grads(()) + calls = [] + real = timestep_mlp.det_linear_backward_weight + + def recording(*args, **kwargs): + calls.append(kwargs.get("with_bias", True)) + return real(*args, **kwargs) + + monkeypatch.setattr(timestep_mlp, "det_linear_backward_weight", recording) + part = grads(frozen) + for name, a, b in zip(NAMES, full, part, strict=True): + if name in frozen: + assert b.grad is None, name + else: + assert torch.equal(a.grad, b.grad), name + # The bias fold runs only for a layer whose bias needs a gradient; a layer none of whose + # parameters need one is not folded at all. + expected = [] + if "w2" not in frozen or "b2" not in frozen: + expected.append("b2" not in frozen) + if "w1" not in frozen or "b1" not in frozen: + expected.append("b1" not in frozen) + assert calls == expected, (frozen, calls) + def test_registry_dispatches_cuda(self): from rl_engine.runtime.registry import KernelRegistry @@ -235,3 +299,36 @@ def test_backward_weight_rejects_input_on_another_device(self, dtype): 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) + + def test_ops_run_on_the_device_of_their_tensors_not_the_current_one(self, dtype): + # The entry points are registered as plain function pointers; the device guard inside + # each of them (not the binding) is what puts the launch on the tensors' device. + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + ) + + _cuda_op() + g = torch.Generator("cpu").manual_seed(0) + + def run(device): + x = torch.randn(3, 64, generator=g).to(device=device, dtype=dtype) + weight = torch.randn(40, 64, generator=g).to(device=device, dtype=dtype) + bias = torch.randn(40, generator=g).to(device=device, dtype=dtype) + grad = torch.randn(3, 40, generator=g).to(device) + (out,) = det_linear_forward(x, weight, bias) + dx = det_linear_backward_input(grad, weight, torch.float32) + dw, db = det_linear_backward_weight(grad, x, dtype) + return [t.cpu() for t in (out, dx, dw, db)] + + # Same bytes from the same generator state on both devices, called while device 0 is + # current: nothing may run on (or touch the memory of) the wrong device. + torch.cuda.set_device(0) + g.manual_seed(0) + on_zero = run("cuda:0") + g.manual_seed(0) + on_one = run("cuda:1") + assert torch.cuda.current_device() == 0 + for a, b in zip(on_zero, on_one, strict=True): + assert torch.equal(a, b) From 984f287df703c02abbd616a7a2cd0de59639b34c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:15:02 +0000 Subject: [PATCH 132/139] fix(h3): fail closed on a device the probe cannot query, skip the projection bias fold, reject a weight without rows * `det_linear_device_supported` returns False where torch raises (an index outside this host's devices, a device string it cannot parse, a CUDA runtime that fails to start), so a constructor that uses it falls back to the PyTorch reference instead of failing; integer indices keep working (`torch.device(0)` is `cuda:0`). Tests cover each case. * `adaln_projection` backward passes `with_bias=` to `det_linear_backward_weight`, like the fused modulation op, so the bias fold does not run when only the weight needs a gradient. A test freezes each input in turn, compares the remaining gradients bitwise with the all-trainable run and records the argument of every call. * `adaln_hidden_size` rejects a (0, D) weight: it divided by 18 and gave H = 0 and six empty outputs. * `det_linear.cu`: the two FP32 launch-shape knobs (`H3_FWD_F32_COLS`, `H3_FWD_F32_AHEAD`) must be overridden together; a build that sets one fails in the preprocessor with that message instead of with an undeclared identifier. The test runs the preprocessor on stand-ins for the torch headers. * `h3_bitwise_reuse._compare_case` records a different number of outputs as a difference (a bare zip stopped at the shorter list and could call the case bitwise equal) and zips strictly otherwise. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 6 ++ .../minimax_h3/adaln_projection.py | 6 +- .../model_specific/minimax_h3/det_linear.py | 16 +++-- .../reference/minimax_h3/adaln_projection.py | 2 + .../minimax_h3/test_h3_adaln_projection.py | 47 +++++++++++++++ .../minimax_h3/test_h3_bitwise_reuse_plot.py | 16 +++++ tests/models/minimax_h3/test_h3_det_linear.py | 60 +++++++++++++++++++ tools/validation/models/h3_bitwise_reuse.py | 10 +++- 8 files changed, 156 insertions(+), 7 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index fb514abcc..c8625d0bf 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -63,6 +63,12 @@ namespace { +// The FP32 forward launch shape is one pair of knobs, and the single compiled case below is built +// from the same pair, so the shape picked at run time is always one that was compiled. A build that +// overrides one knob has to override the other (and H3_FWD_CONFIGS for a multi-case sweep). +#if defined(H3_FWD_F32_COLS) != defined(H3_FWD_F32_AHEAD) +#error "define H3_FWD_F32_COLS and H3_FWD_F32_AHEAD together" +#endif #ifndef H3_FWD_F32_COLS #define H3_FWD_F32_COLS 2 #define H3_FWD_F32_AHEAD 6 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 1363d8a44..1946cf6a9 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 @@ -63,7 +63,11 @@ def backward(ctx, grad_table): 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) + # The bias fold is skipped when only the weight needs a gradient. + dw, *rest = det_linear_backward_weight( + grad_table, act, weight.dtype, with_bias=need_bias + ) + db = rest[0] if need_bias else None 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) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py index 646685dc2..ff35ee019 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py @@ -32,16 +32,22 @@ def det_linear_available() -> bool: def det_linear_device_supported(device: torch.device | int | str | None = None) -> bool: """Return whether ``device`` can run the BF16 forward (``mma.sync``, SM80+). - ``device`` defaults to the current CUDA device. Without CUDA, or for a non-CUDA device, the - answer is False: constructors use this as an early check so the registry falls back to the - PyTorch reference, and ``det_linear_forward`` repeats it for the device of the actual tensors. + ``device`` defaults to the current CUDA device. Without CUDA, for a non-CUDA device, for an + index this host does not have, or when the CUDA runtime cannot be queried, the answer is False: + constructors use this as an early check so the registry falls back to the PyTorch reference, + and ``det_linear_forward`` repeats it for the device of the actual tensors. """ if not torch.cuda.is_available(): return False - if device is not None and torch.device(device).type != "cuda": + try: + if device is not None and torch.device(device).type != "cuda": + return False + return torch.cuda.get_device_capability(device)[0] >= 8 + except (AssertionError, RuntimeError, TypeError, ValueError): + # torch raises AssertionError for an index outside the host's devices and RuntimeError for + # a device string it cannot parse or a CUDA runtime that fails to start. return False - return torch.cuda.get_device_capability(device)[0] >= 8 def _require() -> None: diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py index a6b9fd714..62ed98c90 100644 --- a/rl_engine/reference/minimax_h3/adaln_projection.py +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -34,6 +34,8 @@ def adaln_hidden_size(weight: torch.Tensor) -> int: 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}" ) + if weight.shape[0] == 0: + raise ValueError(f"AdaLN weight must have a nonzero row count, got {tuple(weight.shape)}") return weight.shape[0] // rows_per_hidden diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 23aff80f6..c464fd28b 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -107,6 +107,14 @@ def test_rejects_bad_shapes(self): with pytest.raises(TypeError): op.forward(_temb(2, 16, device="cpu"), weight, bias.float()) + def test_rejects_a_weight_without_rows_instead_of_hidden_size_zero(self): + """A (0, D) weight divides by 18 and would give H = 0 and six empty outputs.""" + + temb = _temb(2, 16, device="cpu") + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(ValueError, match="nonzero row count"): + NativeH3AdaLNProjectionOp().forward(temb, weight[:0], bias[:0]) + def test_rejects_a_zero_width_input_instead_of_returning_the_bias(self): weight, bias = _synthetic(8, 16, device="cpu") with pytest.raises(ValueError, match="nonzero input width"): @@ -295,6 +303,45 @@ def test_rows_invariant_across_tensor_core_row_tiles(self, num): for s_out, f_out in zip(single, full, strict=True): assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) + @pytest.mark.parametrize( + "frozen", [("bias",), ("weight",), ("temb",), ("weight", "bias"), ("temb", "bias")] + ) + def test_frozen_inputs_get_no_gradient_and_do_not_change_the_others(self, frozen, monkeypatch): + """Only a bias that needs a gradient is folded; the other gradients stay bitwise equal.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3 import adaln_projection + + op = _cuda_op() + weight, bias = _synthetic(64, 2688, seed=11) + temb = _temb(3, seed=11) + grads = [torch.randn(9, 64, device="cuda").bfloat16() for _ in range(6)] + names = ("temb", "weight", "bias") + + def run(frozen_names): + leaves = [t.detach().clone() for t in (temb, weight, bias)] + for name, leaf in zip(names, leaves, strict=True): + leaf.requires_grad_(name not in frozen_names) + torch.autograd.backward(list(op(*leaves)), grads) + return leaves + + full = run(()) + calls = [] + real = adaln_projection.det_linear_backward_weight + + def recording(*args, **kwargs): + calls.append(kwargs.get("with_bias", True)) + return real(*args, **kwargs) + + monkeypatch.setattr(adaln_projection, "det_linear_backward_weight", recording) + part = run(frozen) + for name, a, b in zip(names, full, part, strict=True): + if name in frozen: + assert b.grad is None, name + else: + assert torch.equal(a.grad, b.grad), name + wants_params = "weight" not in frozen or "bias" not in frozen + assert calls == ([("bias" not in frozen)] if wants_params else []), (frozen, calls) + def test_rejects_cpu(self): """Reject CPU tensors at the CUDA projection boundary.""" diff --git a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py index ffb87bd65..ae039698b 100644 --- a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -48,6 +48,22 @@ def test_a_case_the_reference_never_saved_is_unavailable_and_does_not_abort_the_ assert "unavailable" in missing and _cell(missing) is None +def test_a_different_number_of_outputs_is_a_difference_not_a_shorter_comparison(): + import torch + + ones = torch.ones(3) + run = {"out": [ones, ones], "grads": {"x": torch.zeros(2)}} + ref = {"out": [ones], "grads": {"x": torch.zeros(2)}} + done = _compare_case(run, ref) + # the extra output is not dropped (a zip would call this case bitwise equal) ... + assert done["out"] == [{"equal": False, "note": "2 vs 1 outputs"}] + assert _cell(done) == 1.0 + # ... and the gradients are still compared. + assert done["grads"]["x"]["equal"] + same = _compare_case(run, {"out": [ones, ones], "grads": {"x": torch.zeros(2)}}) + assert [o["equal"] for o in same["out"]] == [True, True] + + def test_missing_comparisons_are_drawn_grey_not_as_bitwise_equal(): import matplotlib.pyplot as plt diff --git a/tests/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py index f1a2e3175..1bd78a42b 100644 --- a/tests/models/minimax_h3/test_h3_det_linear.py +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -5,6 +5,10 @@ from __future__ import annotations +import shutil +import subprocess +from pathlib import Path + import pytest import torch @@ -88,6 +92,26 @@ def capability(device=None): assert seen == [torch.device("cuda:0"), torch.device("cuda:1")] +@pytest.mark.parametrize("device", [torch.cuda.device_count(), 99, "cuda:99", "bogus", 1.5]) +def test_a_device_this_host_does_not_have_is_unsupported_not_an_error(device): + """The probe answers False where torch raises, so a constructor falls back, not fails.""" + + assert det_linear_device_supported(device) is False + + +def test_an_integer_index_names_a_cuda_device(): + assert det_linear_device_supported(0) == det_linear_device_supported(torch.device("cuda", 0)) + assert det_linear_device_supported("cuda:0") == det_linear_device_supported(0) + + +def test_a_cuda_runtime_that_cannot_be_queried_is_unsupported(monkeypatch): + def broken(device=None): + raise RuntimeError("CUDA driver initialization failed") + + monkeypatch.setattr(torch.cuda, "get_device_capability", broken) + assert det_linear_device_supported() is False + + def test_no_cuda_means_no_support(monkeypatch): monkeypatch.setattr(torch.cuda, "is_available", lambda: False) assert not det_linear_device_supported() @@ -104,3 +128,39 @@ def test_bf16_forward_is_refused_on_a_pre_sm80_device_with_a_clear_error(monkeyp # FP32 does not use mma.sync and is not refused. out = det_linear_forward(x.float(), w.float(), None)[0] assert out.shape == (1, 8) + + +@pytest.mark.skipif(shutil.which("g++") is None, reason="needs a C++ preprocessor") +@pytest.mark.parametrize( + ("defines", "builds"), + [ + ([], True), + (["-DH3_FWD_F32_COLS=2", "-DH3_FWD_F32_AHEAD=6"], True), + (["-DH3_FWD_F32_COLS=4"], False), + (["-DH3_FWD_F32_AHEAD=4"], False), + ], +) +def test_the_fp32_launch_shape_knobs_are_overridden_together(tmp_path, defines, builds): + """One knob alone would leave the compiled case and the run-time shape apart: say so. + + Only the preprocessor runs, on empty stand-ins for the torch and CUDA headers. + """ + + for header in ( + "torch/extension.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAException.h", + "cuda_bf16.h", + ): + (tmp_path / header).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / header).write_text("") + source = Path(__file__).resolve().parents[3] / "csrc/cuda/h3/det_linear.cu" + done = subprocess.run( + ["g++", "-E", "-x", "c++", "-std=c++17", f"-I{tmp_path}", *defines, str(source)], + capture_output=True, + text=True, + ) + assert (done.returncode == 0) is builds, done.stderr[-400:] + if not builds: + assert "define H3_FWD_F32_COLS and H3_FWD_F32_AHEAD together" in done.stderr diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py index 81a80b75c..134c82a07 100644 --- a/tools/validation/models/h3_bitwise_reuse.py +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -172,8 +172,16 @@ def _compare_case(run: dict[str, Any], ref_run: dict[str, Any] | None) -> dict[s if ref_run is None: return {"unavailable": "the reference run has no saved result for this case"} + if len(run["out"]) != len(ref_run["out"]): + # Outputs that have no counterpart are a difference, not something to drop: a zip would + # stop at the shorter list and could report the case as bitwise equal. + outs: list[dict[str, Any] | None] = [ + {"equal": False, "note": f"{len(run['out'])} vs {len(ref_run['out'])} outputs"} + ] + else: + outs = [_same(a, b) for a, b in zip(run["out"], ref_run["out"], strict=True)] return { - "out": [_same(a, b) for a, b in zip(run["out"], ref_run["out"])], + "out": outs, "grads": {g: _same(run["grads"].get(g), ref_run["grads"].get(g)) for g in run["grads"]}, } From 65e86208d96b1d8a89fdc0d0c67c0076461f027f Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:21:18 +0000 Subject: [PATCH 133/139] fix(h3): check that the gather segments cut the tile list exactly, and show the device asserts survive a release build * The tile-metadata pre-pass also requires `seg_first_tile[0] == 0` and `seg_first_tile[R] == tiles`. The old bound (`<= tiles`) kept every read inside `partial`, but a tile that no segment owns, or a first segment that starts after tile 0, would drop contributions from a row's sum without an error. Two more corrupted-metadata cases run in the child-process test. The launch uses at least one block, so an empty list is a no-op instead of an invalid configuration. * `CUDA_KERNEL_ASSERT` is not removed by NDEBUG on CUDA (c10 expands it to `if (!cond) __assert_fail(...)`; only Android, Apple and FreeBSD compile it out). The header says so, and `test_h3_assert_in_release_build.py` compiles a one-kernel program with `-O3 -DNDEBUG` against the torch headers in use and requires the "device-side assert triggered" stop (and a clean run for a valid value) so the claim is checked for the torch this runs against. * `test_backward_writes_the_requested_dtype_whatever_the_gradient_dtype` pins that the gather backward allocates the table gradient in `out_dtype` for FP32, BF16 and FP16 gradients (exact integer data, torch.equal against the FP64 sum). * `_C.pyi`: contracts for `h3_adaln_row_gather_forward` / `_backward`; the modulation tests zip strictly (Ruff B905). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/tile_metadata.cuh | 13 ++- rl_engine/_C.pyi | 24 ++++- .../minimax_h3/test_h3_adaln_modulation.py | 10 +- .../minimax_h3/test_h3_adaln_row_gather.py | 33 +++++++ .../test_h3_assert_in_release_build.py | 91 +++++++++++++++++++ 5 files changed, 163 insertions(+), 8 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_assert_in_release_build.py diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index a959054e2..96318225d 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -14,6 +14,12 @@ // pass on purpose: an assert in or before the loops of the partial and fold kernels changes how the // compiler schedules their loads and slowed them by 6-12% (and by several times inside the loop), // while this pass costs one small launch (about 4 us). +// +// The checks are CUDA_KERNEL_ASSERTs. Unlike assert(), that macro is not removed by NDEBUG on CUDA +// (c10 expands it to a plain `if (!cond) __assert_fail(...)`; only Android, Apple and FreeBSD +// compile it out, and ROCm aborts), so they stay in the release build this extension ships as +// well; tests/models/minimax_h3/test_h3_assert_in_release_build.py compiles with -DNDEBUG -O3 to +// keep that true for the torch it runs against. #pragma once #include @@ -43,6 +49,10 @@ static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ so for (int64_t i = first_thread; i <= segments; i += stride) { CUDA_KERNEL_ASSERT(seg_first_tile[i] >= 0 && seg_first_tile[i] <= tiles); CUDA_KERNEL_ASSERT(i == segments || seg_first_tile[i] <= seg_first_tile[i + 1]); + // The segments cut the tile list exactly. A tile that no segment owns would be missing from + // every row's sum, silently. + CUDA_KERNEL_ASSERT(i != 0 || seg_first_tile[0] == 0); + CUDA_KERNEL_ASSERT(i != segments || seg_first_tile[segments] == tiles); } } @@ -51,7 +61,8 @@ static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ so inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, int64_t positions, const int64_t* tile_begin, const int64_t* tile_end, int64_t tiles, const int64_t* seg_first_tile, int64_t segments) { - const int64_t items = std::max({positions, tiles, segments + 1}); + // At least one block, so an empty list (segments + 1 == 0) is a no-op and not an invalid launch. + const int64_t items = std::max({positions, tiles, segments + 1, int64_t{1}}); const unsigned blocks = static_cast(std::min((items + 255) / 256, 1024)); check_tile_metadata_kernel<<>>(sorted_pos, positions, tile_begin, tile_end, tiles, seg_first_tile, segments); diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 35fee7775..cef79ab51 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -392,7 +392,18 @@ def h3_adaln_row_gather_forward( token_tags: torch.Tensor, chunks: int = 6, modality_num: int = 3, -) -> torch.Tensor: ... +) -> torch.Tensor: + """Gather AdaLN table rows into packed per-position chunks; return ``(chunks, S, H)``. + + ``rows`` is a 2-D CUDA ``(R, chunks * H)`` table with unit column stride and ``R`` a multiple of + ``modality_num``; ``timestep_indices`` and ``token_tags`` are ``(S,)`` contiguous integer + tensors (int64 or int32, one dtype) on its device, ``S > 0``. Position ``s`` reads table row + ``timestep_indices[s] * modality_num + token_tags[s]`` and receives its ``chunks`` slices of + width ``H`` in ``rows``' dtype (a pure copy). Out-of-range indices stop the kernel with a + device-side assert. + """ + ... + def h3_adaln_row_gather_backward( grad: torch.Tensor, sorted_pos: torch.Tensor, @@ -400,4 +411,13 @@ def h3_adaln_row_gather_backward( tile_end: torch.Tensor, seg_first_tile: torch.Tensor, out_dtype: torch.dtype, -) -> torch.Tensor: ... +) -> torch.Tensor: + """Sum a ``(chunks, S, H)`` gradient into table rows; return ``(R, chunks * H)`` in ``out_dtype``. + + ``sorted_pos`` holds the ``S`` positions sorted stably by table row, ``[tile_begin, tile_end)`` + cut that order into fixed-size tiles, and ``seg_first_tile`` (``R + 1`` entries) gives the + first tile of each row (see ``_segment_tiles``). Tiles are summed in FP32 and folded in order, + with no atomics. The metadata is checked on the device first: positions inside ``[0, S)``, tiles + inside ``sorted_pos``, segments that cut the tile list exactly and in order. + """ + ... diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py index 9adc6b4aa..cc3424a04 100644 --- a/tests/models/minimax_h3/test_h3_adaln_modulation.py +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -81,7 +81,7 @@ def test_forward_bitwise_equal_to_separate_ops(self): 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)) + assert all(torch.equal(a, b) for a, b in zip(ours, separate, strict=True)) def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): fused, proj, gather = _ops() @@ -121,7 +121,7 @@ def rel(grad, golden): 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)): + for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad), strict=True): assert (grad == golden.to(torch.bfloat16)).float().mean() > 0.9999 @pytest.mark.parametrize("frozen", ["bias", "weight", "temb", "bias_and_temb"]) @@ -134,10 +134,10 @@ def test_frozen_inputs_get_no_gradient_and_the_others_are_unchanged(self, frozen leaves = [t.detach().clone() for t in (temb, weight, bias)] names = ("temb", "weight", "bias") - for leaf, name in zip(leaves, names): + for leaf, name in zip(leaves, names, strict=True): leaf.requires_grad_(name not in frozen) torch.autograd.backward(list(fused(*leaves, ti, tags)), grads) - for leaf, name, expected in zip(leaves, names, full): + for leaf, name, expected in zip(leaves, names, full, strict=True): if name in frozen: assert leaf.grad is None else: @@ -150,7 +150,7 @@ def test_backward_repeat_bitwise(self): 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)) + assert all(torch.equal(a, b) for a, b in zip(first, again, strict=True)) def test_rejects_bad_indices(self): fused, _, _ = _ops() 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 ff9662f75..220311269 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -364,6 +364,10 @@ def _backward_probe(mutation: str): sorted_pos = sorted_pos.clone(); sorted_pos[0] = seq elif mutation == "segments_past_tiles": first = first.clone(); first[-1] = begin.numel() + 3 + elif mutation == "first_segment_not_at_tile_zero": + first = first.clone(); first[0] = 1 + elif mutation == "last_tile_owned_by_no_segment": + first = first.clone(); first[-1] = first[-1] - 1 elif mutation == "no_tiles": begin = begin[:0]; end = end[:0] _C.h3_adaln_row_gather_backward(grad, sorted_pos, begin, end, first, torch.float32) @@ -387,12 +391,41 @@ def test_the_backward_stops_on_metadata_that_does_not_describe_the_tiles(self): "tile_end_past_seq", "position_past_seq", "segments_past_tiles", + "first_segment_not_at_tile_zero", + "last_tile_owned_by_no_segment", "no_tiles", ): done = self._backward_probe(mutation) assert done.returncode != 0, mutation assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) + @pytest.mark.parametrize("grad_dtype", [torch.float32, torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) + def test_backward_writes_the_requested_dtype_whatever_the_gradient_dtype( + self, grad_dtype, out_dtype + ): + """The table gradient has ``out_dtype``, not the dtype of the incoming gradient.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + from rl_engine.backends.extension import _C + + _cuda_op() + seq, hidden = 40, 8 + index = torch.arange(seq, device="cuda") % 3 + generator = torch.Generator(device="cpu").manual_seed(0) + # Small integers: every sum is exact in all three dtypes, so the result can be compared + # with torch.equal. + grad = torch.randint(-3, 4, (6, seq, hidden), generator=generator).to("cuda", grad_dtype) + tiles = _segment_tiles(index, 3) + got = _C.h3_adaln_row_gather_backward(grad, *tiles, out_dtype) + assert got.dtype == out_dtype and got.shape == (3, 6 * hidden) + expected = torch.stack( + [grad.double()[:, index == r, :].sum(1).reshape(-1) for r in range(3)] + ).to(out_dtype) + assert torch.equal(got, expected) + def test_backward_with_no_table_rows_returns_an_empty_table_and_rejects_no_boundaries(self): from rl_engine.backends.extension import _C diff --git a/tests/models/minimax_h3/test_h3_assert_in_release_build.py b/tests/models/minimax_h3/test_h3_assert_in_release_build.py new file mode 100644 index 000000000..57043d30f --- /dev/null +++ b/tests/models/minimax_h3/test_h3_assert_in_release_build.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The H3 kernels' CUDA_KERNEL_ASSERT bounds checks survive a release (-DNDEBUG) build. + +The native entry points check caller-built metadata and index values with ``CUDA_KERNEL_ASSERT`` +on the device (a host-side check would add a synchronisation to every call). Unlike ``assert()``, +c10 keeps that macro under ``NDEBUG`` on CUDA, so a release build still stops on a bad value +instead of reading out of bounds. This compiles a one-kernel program with the release flags and +the torch headers this repository builds against, and checks that the assertion fires. The tests +that run the built extension on bad metadata and bad indices live next to each op +(``test_native_bounds_assertions``, ``test_the_backward_stops_on_metadata...``). +""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import textwrap +from pathlib import Path + +import pytest +import torch +from torch.utils import cpp_extension + +SOURCE = textwrap.dedent( + """ + #include + #include + #include + + __global__ void check(const long* v, int n) { + const int i = threadIdx.x; + if (i < n) CUDA_KERNEL_ASSERT(v[i] >= 0 && v[i] < 100); + } + + int main(int argc, char** argv) { + long host[4] = {0, 1, 2, 3}; + if (argc > 1) host[2] = std::atol(argv[1]); + long* device = nullptr; + cudaMalloc(&device, sizeof(host)); + cudaMemcpy(device, host, sizeof(host), cudaMemcpyHostToDevice); + check<<<1, 4>>>(device, 4); + const cudaError_t status = cudaDeviceSynchronize(); + std::printf("%s\\n", cudaGetErrorString(status)); + return status == cudaSuccess ? 0 : 2; + } + """ +) + + +def _nvcc() -> str | None: + for root in (os.environ.get("CUDA_HOME"), "/usr/local/cuda"): + if root and (Path(root) / "bin" / "nvcc").exists(): + return str(Path(root) / "bin" / "nvcc") + return shutil.which("nvcc") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.skipif(_nvcc() is None, reason="needs nvcc") +def test_cuda_kernel_assert_stays_active_with_ndebug(tmp_path): + major, minor = torch.cuda.get_device_capability() + arch = f"{major}{minor}" + (tmp_path / "check.cu").write_text(SOURCE) + build = subprocess.run( + [ + _nvcc(), + "-O3", + "-DNDEBUG", + "-std=c++17", + f"-gencode=arch=compute_{arch},code=sm_{arch}", + *(f"-I{path}" for path in cpp_extension.include_paths()), + str(tmp_path / "check.cu"), + "-o", + str(tmp_path / "check"), + ], + capture_output=True, + text=True, + timeout=300, + ) + assert build.returncode == 0, build.stderr[-1500:] + + ok = subprocess.run([str(tmp_path / "check"), "2"], capture_output=True, text=True, timeout=60) + assert ok.returncode == 0, ok.stdout + ok.stderr # the control: a valid value passes + for bad in ("-1", "100", "123456"): + done = subprocess.run( + [str(tmp_path / "check"), bad], capture_output=True, text=True, timeout=60 + ) + assert done.returncode != 0, bad + assert "device-side assert triggered" in done.stdout, (bad, done.stdout, done.stderr) From 560111898846b358238e1dbcbadc950178c73d3c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:33:52 +0000 Subject: [PATCH 134/139] fix(h3): bound the sizes the rmsnorm kernels take as int, check the modulation index once before the dweight loop * `check_xw`, called by every rmsnorm entry point, rejects N or M above INT_MAX (the kernels take N as `int` and launch one block per row), and the modulated backward bounds the table rows and tiles it launches as blocks. A test allocates a (1, 2**31) FP16 row and checks forward and backward. * The per-row `CUDA_KERNEL_ASSERT` on the modulation index inside `rmsnorm_dweight_partial_kernel` cost about 11% of the modulated backward (A/B, interleaved min of 300, B200: 1770 -> 1586 us at S = 32768, 6346 -> 5694 us at S = 131072, 378 -> 371 us at S = 4097 once removed). The loop now reads the index unchecked and the entry point runs `h3_tiles::check_index_range` over the S index values first, a 4 us launch like the tile-metadata pre-pass. The forward and dx kernels keep their once-per-row assert. A child-process test requires the device-side assert for negative, past-the-table and far indices in both forward and backward, with a valid-index control. * The backward header states the exact rounding chain of dx (coef = ((rstd*rstd)*rstd)*dot/N, one cast, FMA contraction allowed with fixed flags) and that the backward is measured against FP64, not claimed bitwise to autograd; the empty `Modulation` initialiser says why `seq` is 1. * `_C.pyi`: contracts for `h3_rmsnorm_forward` / `h3_rmsnorm_backward`. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/rmsnorm_modulate.cu | 36 ++++++++-- csrc/cuda/h3/tile_metadata.cuh | 20 ++++++ tests/models/minimax_h3/test_h3_rmsnorm.py | 79 ++++++++++++++++++++++ 3 files changed, 131 insertions(+), 4 deletions(-) diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 3f352f417..6d0146408 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -18,6 +18,12 @@ // 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) +// evaluated as coef = ((rstd * rstd) * rstd) * dot / N and +// dx = cast(rstd * w * d_n - x * coef), with dot a fixed block tree over fmaf terms. +// nvcc may contract the last line to an FMA; the build flags are fixed, so the result +// is the same on every run. That is the contract of the backward: it is measured +// against an FP64 golden (the *_against_fp64 tests), and not claimed bitwise equal to +// autograd's RMSNorm backward, which forms rstd^3 differently. // 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) @@ -112,13 +118,21 @@ struct Modulation { }; // Table row of x row `r`. check_range=False callers skip the Python range check, so the native side -// asserts the index before it is used as an offset (as the row-gather kernel does). +// asserts the index before it is used as an offset (as the row-gather kernel does). Kernels that read +// it once per row use this; the dweight partials read it inside their row loop, where an assert +// costs about 11% of the backward, so they use the unchecked read and the entry points run +// h3_tiles::check_index_range over the S index values before launching them. __device__ __forceinline__ int64_t modulation_row(const Modulation& mod, int64_t r) { const int64_t i = mod.index[r % mod.seq]; CUDA_KERNEL_ASSERT(i >= 0 && i < mod.rows); return i; } +__device__ __forceinline__ int64_t modulation_row_checked_by_entry_point(const Modulation& mod, + int64_t r) { + return mod.index[r % mod.seq]; +} + template __global__ void __launch_bounds__(32 * kWarps) rmsnorm_modulate_fwd_kernel(const T* __restrict__ x, const T* __restrict__ w, T* __restrict__ y, @@ -204,9 +218,10 @@ __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) + - modulation_row(mod, r) * mod.row_stride - : nullptr; + const T* scale = mod.index != nullptr + ? static_cast(mod.scale) + + modulation_row_checked_by_entry_point(mod, r) * 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; @@ -277,6 +292,8 @@ Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, int64_t n) { + // No modulation: kernels test `index == nullptr`. seq = 1 (not 0) keeps `r % seq` defined for + // any code that computes it before looking at the index. 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"); @@ -321,6 +338,11 @@ void check_xw(const torch::Tensor& x, const torch::Tensor& w) { 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)"); + // The kernels take N as `int` and launch one block per row on grid.x. + TORCH_CHECK(x.size(1) <= std::numeric_limits::max(), "N=", x.size(1), + " is too large for the kernels (at most ", std::numeric_limits::max(), ")"); + TORCH_CHECK(x.size(0) <= std::numeric_limits::max(), "M=", x.size(0), + " rows do not fit one grid.x launch (at most ", std::numeric_limits::max(), ")"); } // float32 / float16 / bfloat16 only (float64 has no PyTorch-matching vector path here). @@ -396,6 +418,8 @@ std::vector h3_rmsnorm_backward( } const c10::cuda::CUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); + // The dweight partials read the index inside their row loop: check its values first. + if (mod.index != nullptr) h3_tiles::check_index_range(stream, mod.index, mod.seq, mod.rows); 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)); @@ -424,6 +448,10 @@ std::vector h3_rmsnorm_backward( const int64_t seg_tiles = tile_begin->numel(); const int64_t width = 2 * n; TORCH_CHECK(width <= std::numeric_limits::max(), "N=", n, " is too large for the kernels"); + TORCH_CHECK(segments <= std::numeric_limits::max() && + seg_tiles <= std::numeric_limits::max(), + "table rows (", segments, ") and tiles (", seg_tiles, + ") must each fit one grid.x launch"); TORCH_CHECK((width + threads - 1) / threads <= 65535, "N=", n, " needs more column blocks than the grid.y limit of 65535"); auto seg_partial = torch::empty({std::max(seg_tiles, 1), width}, x.options().dtype(at::kFloat)); diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index 96318225d..09c6ad4be 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -69,4 +69,24 @@ inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, C10_CUDA_KERNEL_LAUNCH_CHECK(); } +// Table-row indices read inside a kernel's loop: `n` entries, each a row of a table with `limit` +// rows. The assert would cost the loop its scheduling (the same 6-12% as above, measured on the +// rmsnorm dweight partials), so the entry point checks the S index values once, here, before it +// launches the kernels that read them. +static __global__ void check_index_range_kernel(const int64_t* __restrict__ index, int64_t n, + int64_t limit) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n; i += stride) { + CUDA_KERNEL_ASSERT(index[i] >= 0 && index[i] < limit); + } +} + +inline void check_index_range(cudaStream_t stream, const int64_t* index, int64_t n, + int64_t limit) { + const unsigned blocks = + static_cast(std::max(1, std::min((n + 255) / 256, 1024))); + check_index_range_kernel<<>>(index, n, limit); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + } // namespace h3_tiles diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py index 6d91f60bb..be899893b 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm.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 @@ -206,6 +212,22 @@ def test_rejects_unsupported(self): with pytest.raises(ValueError): _cuda_op()(_x((2, 8)).cpu(), torch.ones(8).bfloat16()) + def test_rejects_a_row_wider_than_the_int_the_kernels_count_with(self): + """N goes to the kernels as ``int``: forward and backward both refuse it up front.""" + + from rl_engine.backends.extension import _C + + _cuda_op() + n = 2**31 # a multiple of 4 above INT_MAX + if torch.cuda.mem_get_info()[0] < 10 * 2**30: + pytest.skip("needs about 8 GiB of free device memory for the two operands") + x = torch.empty(1, n, device="cuda", dtype=torch.float16) + w = torch.empty(n, device="cuda", dtype=torch.float16) + with pytest.raises(RuntimeError, match="too large"): + _C.h3_rmsnorm_forward(x, w, 1e-6) + with pytest.raises(RuntimeError, match="too large"): + _C.h3_rmsnorm_backward(x, x, w, torch.empty(1, device="cuda")) + @requires_cuda class TestCudaBackward: @@ -300,3 +322,60 @@ def test_registry_dispatches_cuda(self): assert ( type(KernelRegistry().get_op("h3_rmsnorm", device="cuda")).__name__ == "H3RMSNormCudaOp" ) + + +def _index_probe(entry: str, bad: int): + """Run one native entry point with a table-row index outside the table, in a child process.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + entry, bad = sys.argv[1], int(sys.argv[2]) + rows, hidden, table_rows = 600, 64, 4 + x = torch.randn(rows, hidden, device="cuda").bfloat16() + w = torch.ones(hidden, device="cuda").bfloat16() + table = torch.randn(table_rows, 2 * hidden, device="cuda").bfloat16() + shift, scale = table.chunk(2, dim=-1) + index = torch.arange(rows, device="cuda") % table_rows + index[rows - 3] = bad # the last tile of rows, past the first dx blocks + if entry == "forward": + _C.h3_rmsnorm_forward(x, w, 1e-5, shift, scale, index) + else: + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + valid = index.clamp(0, table_rows - 1) + y, rstd = _C.h3_rmsnorm_forward(x, w, 1e-5, shift, scale, valid) + tiles = _segment_tiles(valid, table_rows) + _C.h3_rmsnorm_backward(torch.ones_like(x), x, w, rstd, shift, scale, index, *tiles) + torch.cuda.synchronize() + """ + ) + return subprocess.run( + [sys.executable, "-c", probe, entry, str(bad)], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=120, + ) + + +@requires_cuda +@pytest.mark.parametrize("entry", ["forward", "backward"]) +@pytest.mark.parametrize("bad", [-1, 4, 10**9], ids=["negative", "past-table", "far"]) +def test_a_modulation_index_outside_the_table_stops_the_kernels(entry, bad): + """check_range=False and direct calls get a device-side assert, not an out-of-bounds read. + + The dweight partials read the index inside their row loop, where an assert costs about 11% of + the backward; the entry point checks the index values once before launching them. + """ + + _cuda_op() + assert _index_probe(entry, 1).returncode == 0 # the control: the same call with a valid index + done = _index_probe(entry, bad) + assert done.returncode != 0, entry + assert "device-side assert triggered" in done.stderr, done.stderr From ab407b39ccc7f09066ce7300a00c8ff17c9c4401 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:38:27 +0000 Subject: [PATCH 135/139] fix(h3): reject gate tables whose rows overlap, and state the residual gradient shape * `make_gate` refuses a table whose row stride is below N (an expanded table would make every row alias the same memory), as the modulation tables already did; `H3GateResidualCudaOp` copies such a table to dense rows first, so an expanded gate still works and gives the bytes of its dense copy. Views of a wider AdaLN table (unit column stride, larger row stride) go to the kernel unchanged. * The residual's VJP is the incoming cotangent, which already has the residual's shape because forward returns `out.view(residual.shape)`; `ctx.residual_shape` was saved and unused, which read like a missing reshape. The return is now `grad.reshape(ctx.residual_shape)` (a no-op for equal shapes, and it accepts a non-contiguous gradient where `view` would not), and a test feeds a (B, S, H) residual with a transposed cotangent and compares every gradient with the eager VJP bitwise. * `_C.pyi`: contracts for `h3_gate_residual_forward` / `_backward`. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 4 ++ .../minimax_h3/gate_residual.py | 10 ++++- .../minimax_h3/test_h3_adaln_gate_residual.py | 37 +++++++++++++++++++ 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 123b434fa..81051fcc6 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -149,6 +149,10 @@ 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"); + // Rows are read as dense N-element spans at `stride(0)`: a view of a wider table is fine, an + // expanded or overlapping one (stride(0) < N) is not a table of rows. + TORCH_CHECK(gate.size(0) <= 1 || gate.stride(0) >= gate.size(1), "gate row stride ", + gate.stride(0), " is smaller than N=", gate.size(1)); 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()); 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 04f6bdf7d..da7319bfb 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 @@ -51,7 +51,9 @@ def backward(ctx, grad): record_backward( "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) - return grad, dy.view(ctx.y_shape), dgate, None + # The cotangent has the output's shape, which is the residual's; the residual's VJP is that + # cotangent itself. reshape (not view) so a non-contiguous gradient is accepted. + return grad.reshape(ctx.residual_shape), dy.view(ctx.y_shape), dgate, None class H3GateResidualCudaOp: @@ -80,7 +82,11 @@ 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: + # A view of a wider table (``table.view(R, 6, H)[:, 2]``: unit column stride, row stride 6H) + # goes to the kernel as it is, which indexes rows by the stride. Anything that is not a + # table of dense rows (non-unit columns, or rows that overlap because the stride is below + # H, as for an expanded table) is copied first. + if gate.stride(1) != 1 or (gate.shape[0] > 1 and gate.stride(0) < gate.shape[1]): gate = gate.contiguous() return _H3GateResidualCuda.apply(residual, y, gate, index.contiguous()) diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index b26ab8a54..193abce60 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -125,6 +125,43 @@ def grads(fn, dtype=None): 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_residual_gradient_keeps_the_residual_shape_for_batched_and_strided_gradients(self): + """A (B, S, H) residual gets a (B, S, H) gradient even when the incoming one is a view.""" + + table, residual, y, index = _case(300, hidden=64, seed=5) + leaves = [t.detach().clone().requires_grad_(True) for t in (residual, y, table)] + out = _cuda_op()(leaves[0], leaves[1], _gate(leaves[2]), index) + # a non-contiguous cotangent: the transpose of a (B, H, S) tensor + grad = torch.randn(2, 64, 300, device="cuda").bfloat16().transpose(1, 2) + assert not grad.is_contiguous() and grad.shape == out.shape + out.backward(grad) + ref = [t.detach().clone().requires_grad_(True) for t in (residual, y, table)] + provider_gate_residual(ref[0], _gate(ref[2]), index, ref[1]).backward(grad) + assert leaves[0].grad.shape == residual.shape and leaves[1].grad.shape == y.shape + assert torch.equal(leaves[0].grad, ref[0].grad) + assert torch.equal(leaves[1].grad, ref[1].grad) + + def test_a_gate_table_whose_rows_overlap_is_copied_not_read_in_place(self): + """An expanded table (row stride 0) is a legal gate: it is copied to dense rows first.""" + + from rl_engine.backends.extension import _C + + hidden = 64 + row = (torch.randn(1, hidden, device="cuda") * 0.5).bfloat16() + expanded = row.expand(9, hidden) # stride (0, 1): all nine rows alias one + assert expanded.stride(0) == 0 + _, residual, y, index = _case(100, hidden=hidden, seed=6) + index = index % 9 + dense = expanded.contiguous() + got = _cuda_op()(residual, y, expanded, index) + assert torch.equal(got, _cuda_op()(residual, y, dense, index)) + assert torch.equal(got, provider_gate_residual(residual, dense, index, y)) + # the native entry point refuses the layout instead of reading overlapping rows + res2, y2 = residual.reshape(-1, hidden), y.reshape(-1, hidden) + with pytest.raises(RuntimeError, match="row stride"): + _C.h3_gate_residual_forward(res2, y2, expanded, index) + _C.h3_gate_residual_forward(res2, y2, dense, index) # the dense copy is accepted + def test_rejects_cpu_and_bad_index(self): table, residual, y, index = _case(10, hidden=16) with pytest.raises(IndexError): From 562d3fa30902ddd9f8060cb1986926b46d66a6eb Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:55:49 +0000 Subject: [PATCH 136/139] fix(h3): read quoted options-file names as nvcc does, fail closed on unknown flag values, bound the options-file read The relaxed-math check split every shlex word again on whitespace and ';', so --options-file "my opts.txt" was read as the file 'my' (and a decoy named 'my' hid the real flags). Quoted words now stay whole for file names and are also scanned piece by piece for relaxed flags. A value nvcc does not take (--ftz=yes, --prec-div=no) is reported instead of read as the exact spelling. Options files are read only if they are regular files of at most 1 MiB, so /dev/zero or a FIFO cannot hang the configure step. The unused _external_relaxed_math helper is removed (get_extensions uses _relaxed_math_reason). The nvcc --dryrun oracle test now also asserts that the three cicc options are present and which flag sets relax. H3_ADALN_CHUNKS is the length of H3_ADALN_CHUNK_NAMES, and the sinusoid docstring states the divisor as half. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- build_tools/extensions.py | 62 +++++++++---- rl_engine/reference/minimax_h3/__init__.py | 2 +- .../reference/minimax_h3/timestep_sinusoid.py | 6 +- .../build/test_build_platform_collectives.py | 90 +++++++++++++++++-- .../minimax_h3/test_h3_timestep_sinusoid.py | 15 ++++ 5 files changed, 148 insertions(+), 27 deletions(-) diff --git a/build_tools/extensions.py b/build_tools/extensions.py index cbd64b53e..098779c80 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -4,6 +4,7 @@ import importlib.util import os import shlex +import stat import sysconfig import warnings from distutils.errors import CompileError @@ -34,6 +35,12 @@ "prec-div": {"false", "0"}, "prec-sqrt": {"false", "0"}, } +# The spelling that keeps the math exact. A value that is neither this nor in the set above is one +# nvcc does not know; it is reported as unreadable, not taken for the safe one. +_EXACT_MATH_VALUES = {"ftz": "false", "prec-div": "true", "prec-sqrt": "true"} +# An options file is read in full before anything checks it; a device node or a huge file named by +# a stray variable would hang or exhaust memory the configure step. +_OPTIONS_FILE_MAX_BYTES = 1 << 20 # nvcc reads more options from the files named by ``--options-file`` / ``-optf`` (a comma-separated # list, as ``--options-file=a,b`` or ``--options-file a,b``), also when the option comes from # NVCC_*_FLAGS. nvcc 13.0 does not expand ``@file`` (it takes it for an input file and stops); that @@ -46,16 +53,46 @@ class _UncheckableFlags(Exception): """The flags cannot be read well enough to say whether they relax math.""" +def _nvcc_flag_words(text: str) -> list[str]: + """Split like nvcc does for the double-quoted values it accepts, and on ``;`` as well. + + ``;`` separates the flags of TORCH_NVCC_FLAGS; here it separates words outside quotes only, so + ``--options-file "my opts;1.txt"`` keeps its file name in one piece, as nvcc reads it. + """ + + lexer = shlex.shlex(text, posix=True) + lexer.whitespace_split = True + lexer.whitespace += ";" + lexer.commenters = "" + return list(lexer) + + +def _read_options_file(name: str) -> str: + path = Path(name) + try: + info = path.stat() # a missing file is an OSError, reported as unreadable + if not stat.S_ISREG(info.st_mode): + raise _UncheckableFlags(f"options file {name!r} is not a regular file") + if info.st_size > _OPTIONS_FILE_MAX_BYTES: + raise _UncheckableFlags( + f"options file {name!r} is larger than {_OPTIONS_FILE_MAX_BYTES} bytes" + ) + return path.read_text() + except OSError as exc: + raise _UncheckableFlags(f"cannot read options file {name!r}: {exc.strerror}") from exc + + def _nvcc_flag_tokens(text: str, *, depth: int = 0) -> list[str]: """Split an NVCC_*_FLAGS value into tokens, expanding options files. - Tokens are split like a shell would, then on whitespace and ``;`` (TORCH_NVCC_FLAGS is - ';'-separated), so ``"-O2 --use_fast_math"`` in one quoted word is still two options. Raises + A word that holds whitespace or ``;`` came out of quotes; nvcc takes it as one argument (an + option it then rejects, or a file name). It is kept whole for that reading, and its pieces + follow it, so ``"-O2 --use_fast_math"`` in one quoted word still reads as relaxed. Raises _UncheckableFlags for an options file that cannot be read or an option nvcc rejects for lack of its file name, and ValueError for a string that cannot be split. """ - words = [part for word in shlex.split(text) for part in word.replace(";", " ").split()] + words = _nvcc_flag_words(text) tokens: list[str] = [] i = 0 while i < len(words): @@ -74,17 +111,14 @@ def _nvcc_flag_tokens(text: str, *, depth: int = 0) -> list[str]: files = value.split(",") else: tokens.append(token) + pieces = token.replace(";", " ").split() + if len(pieces) > 1: + tokens.extend(pieces) continue if depth >= _OPTIONS_FILE_DEPTH: raise _UncheckableFlags(f"options files nest deeper than {_OPTIONS_FILE_DEPTH} levels") for name in files: - try: - content = Path(name).read_text() - except OSError as exc: - raise _UncheckableFlags( - f"cannot read options file {name!r}: {exc.strerror}" - ) from exc - tokens.extend(_nvcc_flag_tokens(content, depth=depth + 1)) + tokens.extend(_nvcc_flag_tokens(_read_options_file(name), depth=depth + 1)) return tokens @@ -111,6 +145,8 @@ def _tokens_relax_math(tokens: list[str]) -> str | None: spelled = f"{token} {nxt}" if value.lower() in relaxed_values: return spelled + if value.lower() != _EXACT_MATH_VALUES[name]: + return f"{spelled} (a value nvcc does not take)" return None @@ -133,12 +169,6 @@ def _relaxed_math_reason() -> str | None: return None -def _external_relaxed_math() -> bool: - """Return whether NVCC_*_FLAGS / TORCH_NVCC_FLAGS relax FP32 math or cannot be checked.""" - - return _relaxed_math_reason() is not None - - def _resolve_package_data_source( filename: str | os.PathLike[str], *, diff --git a/rl_engine/reference/minimax_h3/__init__.py b/rl_engine/reference/minimax_h3/__init__.py index a01f9218b..e54c60b77 100644 --- a/rl_engine/reference/minimax_h3/__init__.py +++ b/rl_engine/reference/minimax_h3/__init__.py @@ -19,7 +19,6 @@ # 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", @@ -28,3 +27,4 @@ "scale_mlp", "gate_mlp", ) +H3_ADALN_CHUNKS = len(H3_ADALN_CHUNK_NAMES) diff --git a/rl_engine/reference/minimax_h3/timestep_sinusoid.py b/rl_engine/reference/minimax_h3/timestep_sinusoid.py index 030159836..2f374c81a 100644 --- a/rl_engine/reference/minimax_h3/timestep_sinusoid.py +++ b/rl_engine/reference/minimax_h3/timestep_sinusoid.py @@ -7,9 +7,11 @@ 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 + freq[k] = exp(-ln(10000) * k / half) k = 0..half-1, half = num_channels // 2, FP32 arg[t, k] = t * freq[k] FP32 - out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256) FP32 + out[t] = [cos(arg[t]) | sin(arg[t])] (T, num_channels) FP32 + +H3 pins ``num_channels = 256`` (``half = 128``). The output is always FP32: ``time_embedder`` is a ``_keep_in_fp32_modules`` module in the checkpoint, so the features never pass through BF16. diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index ce2abea84..dfbcb29b5 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -17,9 +17,9 @@ from torch.utils import cpp_extension from build_tools.extensions import ( + _OPTIONS_FILE_MAX_BYTES, H3_BITWISE_SOURCES, H3_KERNELS_DEFINE, - _external_relaxed_math, _relaxed_math_reason, _resolve_package_data_source, _tokens_relax_math, @@ -29,6 +29,10 @@ _NVCC_ENV = ("NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS", "TORCH_NVCC_FLAGS") +def _external_relaxed_math() -> bool: + return _relaxed_math_reason() is not None + + def _clear_nvcc_env(monkeypatch) -> None: # The build reads these variables; a runner that sets one must not change what a test sees. monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) @@ -153,6 +157,7 @@ def test_relaxed_math_nvcc_flags_from_the_environment_exclude_bitwise_h3_kernels # The flags are appended after the ones built here, so a build with # NVCC_APPEND_FLAGS='--use_fast_math' would otherwise ship H3 kernels with approximated math. _clear_nvcc_env(monkeypatch) + assert _relaxed_math_reason() is None # a leaked variable would fail the control, not the flags control = _load_extension_config(monkeypatch, hip=None) assert set(H3_BITWISE_SOURCES) <= set(control["sources"]) # the flags are the only difference @@ -201,10 +206,9 @@ def test_predicate_sees_every_spelling_of_relaxed_math(monkeypatch, flags) -> No "--fmad=false", "--ftz=false", "--ftz false", - "-ftz=0", + "--ftz=FALSE", "--prec-div=true", - "--prec-div true -O2", - "-prec-sqrt 1", + "--prec-div TRUE -O2", "-O2 -gencode arch=compute_90,code=sm_90", "--use_fast_mathematically", ], @@ -239,6 +243,14 @@ def test_predicate_reads_response_files_and_fails_closed(monkeypatch, tmp_path) "sqrt.opts": "--prec-sqrt=false", "plain.opts": "-lineinfo\n--fmad=false\n", "nested.opts": "-optf fast.opts\n", + # Names with a space or ';' are one word when double-quoted, as nvcc reads them. The decoys + # ``my`` and ``a`` are what a parser that split the quoted word would read instead. + "my opts.txt": "--use_fast_math\n", + "my": "-lineinfo\n", + "a;b.txt": "--ftz=true\n", + "a": "-lineinfo\n", + "plain opts.txt": "-lineinfo\n", + "plain": "--use_fast_math\n", } @@ -256,6 +268,10 @@ def test_predicate_reads_response_files_and_fails_closed(monkeypatch, tmp_path) ("--options-file=plain.opts,plain.opts", False), ("@fast.opts", True), ("@plain.opts", False), + ('--options-file "my opts.txt"', True), + ('-optf="my opts.txt"', True), + ('-optf "a;b.txt"', True), + ('-optf "plain opts.txt"', False), ], ) def test_predicate_reads_the_options_files_nvcc_reads( @@ -281,6 +297,40 @@ def test_options_files_nvcc_cannot_open_count_as_relaxed(monkeypatch, tmp_path, assert _external_relaxed_math() +@pytest.mark.parametrize( + "flags", + ["--ftz=yes", "--ftz=0", "-ftz 1", "--prec-div=no", "--prec-div maybe", "-prec-sqrt=1"], +) +def test_values_nvcc_does_not_take_are_not_read_as_the_exact_spelling(monkeypatch, flags) -> None: + # nvcc stops on these; the build then fails, but the predicate must not call them safe. + _clear_nvcc_env(monkeypatch) + monkeypatch.setenv("NVCC_APPEND_FLAGS", flags) + reason = _relaxed_math_reason() + assert reason is not None, flags + assert "NVCC_APPEND_FLAGS contains" in reason, reason + + +def test_options_files_that_are_not_ordinary_small_files_are_not_read( + monkeypatch, tmp_path +) -> None: + # An options file is read whole before anything checks it. A device node (``/dev/zero`` never + # ends) or a FIFO would hang the configure step, and a huge file would exhaust memory. + _clear_nvcc_env(monkeypatch) + big = tmp_path / "big.opts" + big.write_text("-lineinfo\n" * (_OPTIONS_FILE_MAX_BYTES // 8)) + fifo = tmp_path / "fifo.opts" + os.mkfifo(fifo) + for name, why in [ + ("/dev/zero", "not a regular file"), + (str(fifo), "not a regular file"), + (str(tmp_path), "not a regular file"), + (str(big), "larger than"), + ]: + monkeypatch.setenv("NVCC_APPEND_FLAGS", f"-optf {name}") + reason = _relaxed_math_reason() + assert reason is not None and why in reason, (name, reason) + + def test_the_build_warning_names_the_variable_and_the_reason(monkeypatch, tmp_path) -> None: _clear_nvcc_env(monkeypatch) missing = str(tmp_path / "missing.opts") @@ -306,6 +356,23 @@ def _nvcc() -> str | None: return shutil.which("nvcc") +_FLAG_SETS_THAT_RELAX = { + "--use_fast_math", + "-use_fast_math", + "--ftz=true", + "-ftz true", + "--prec-div=false", + "-prec-div false", + "--prec-sqrt=false", + "-optf fast.opts", + "--options-file=ftz.opts", + "--options-file plain.opts,sqrt.opts", + "-optf nested.opts", + '--options-file "my opts.txt"', + '-optf "a;b.txt"', +} + + @pytest.mark.skipif(_nvcc() is None, reason="needs nvcc") @pytest.mark.parametrize( "flags", @@ -328,6 +395,10 @@ def _nvcc() -> str | None: "--options-file plain.opts,sqrt.opts", "-optf plain.opts", "-optf nested.opts", + '--options-file "my opts.txt"', + '-optf "a;b.txt"', + '-optf "plain opts.txt"', + "-O2 --ftz=false --prec-div=true --prec-sqrt=true", ], ) @pytest.mark.parametrize("variable", ["NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS"]) @@ -354,10 +425,13 @@ def test_predicate_agrees_with_what_nvcc_compiles_with( assert done.returncode == 0, done.stderr cicc = [line for line in done.stderr.splitlines() if "cicc" in line] assert cicc, done.stderr - nvcc_relaxes = any( - f" {opt} " in f" {cicc[0]} " for opt in ("-ftz=1", "-prec_div=0", "-prec_sqrt=0") - ) - assert _external_relaxed_math() is nvcc_relaxes, (variable, flags, cicc[0]) + line = f" {cicc[0]} " + # The three options must be in the line at all, so a format change fails here and does not make + # every flag set look unrelaxed; and the sets that relax must be seen to relax. + assert all(f" {opt}=" in line for opt in ("-ftz", "-prec_div", "-prec_sqrt")), line + nvcc_relaxes = any(f" {opt} " in line for opt in ("-ftz=1", "-prec_div=0", "-prec_sqrt=0")) + assert nvcc_relaxes is (flags in _FLAG_SETS_THAT_RELAX), (flags, line) + assert _external_relaxed_math() is nvcc_relaxes, (variable, flags, line) def test_no_flag_built_here_relaxes_math_while_the_h3_define_is_set(monkeypatch) -> None: diff --git a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py index 5ccca25ef..6adbdd208 100644 --- a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py +++ b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py @@ -258,3 +258,18 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda") assert type(op).__name__ == "H3TimestepSinusoidCudaOp" + + +def test_h3_layout_constants_agree_with_each_other(): + from rl_engine.reference.minimax_h3 import ( + H3_ADALN_CHUNK_NAMES, + H3_ADALN_CHUNKS, + H3_FREQ_DIM, + H3_MODALITY_NUM, + ) + + # The chunk count is the length of the name list, and the pinned layout is six chunks of three + # modalities (diffusers order) with 256 sinusoid channels. + assert H3_ADALN_CHUNKS == len(H3_ADALN_CHUNK_NAMES) == 6 + assert H3_ADALN_CHUNK_NAMES[:3] == ("shift_msa", "scale_msa", "gate_msa") + assert H3_MODALITY_NUM == 3 and H3_FREQ_DIM == 256 From e366ecef628030ab33a4cdbcb1a7491e2b958370 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:00:14 +0000 Subject: [PATCH 137/139] fix(h3): require the segment tiles to partition the positions in order check_tile_metadata bounded every tile and cut the tile list exactly by segment, but a gap between tiles (positions dropped from every row's sum) or an overlap (positions counted twice) passed all of it. Each tile must now start where the one before it ends, the first at 0 and the last at the number of positions, and a tile list that is empty over a non-empty sequence is refused on the host. Mutation tests cover a gap, an overlap, a first tile not at 0, a last tile that stops short, and the empty list with and without matching segment bounds. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/tile_metadata.cuh | 12 ++++++++++- .../minimax_h3/test_h3_adaln_row_gather.py | 21 ++++++++++++++++++- 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index 96318225d..ca11e47e6 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -10,7 +10,9 @@ // fold kernels read the partial rows of a segment's tiles, so a stale or short list would read out // of bounds. // -// check_tile_metadata() verifies the list on the device before any of that happens. It is a separate +// check_tile_metadata() verifies the list on the device before any of that happens: indices in +// range, the tiles an exact in-order partition of the positions, and the segments an exact cut of +// the tile list. It is a separate // pass on purpose: an assert in or before the loops of the partial and fold kernels changes how the // compiler schedules their loads and slowed them by 6-12% (and by several times inside the loop), // while this pass costs one small launch (about 4 us). @@ -45,6 +47,11 @@ static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ so for (int64_t i = first_thread; i < tiles; i += stride) { CUDA_KERNEL_ASSERT(tile_begin[i] >= 0 && tile_begin[i] <= tile_end[i] && tile_end[i] <= positions); + // The tiles partition [0, positions) in order: each starts where the one before it ends, the + // first at 0 and the last at `positions`. A gap would drop positions from every row's sum and + // an overlap would count them twice, both without leaving the bounds checked above. + CUDA_KERNEL_ASSERT(tile_begin[i] == (i == 0 ? int64_t{0} : tile_end[i - 1])); + CUDA_KERNEL_ASSERT(i != tiles - 1 || tile_end[i] == positions); } for (int64_t i = first_thread; i <= segments; i += stride) { CUDA_KERNEL_ASSERT(seg_first_tile[i] >= 0 && seg_first_tile[i] <= tiles); @@ -61,6 +68,9 @@ static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ so inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, int64_t positions, const int64_t* tile_begin, const int64_t* tile_end, int64_t tiles, const int64_t* seg_first_tile, int64_t segments) { + // No tiles over a non-empty sequence would leave every row's sum at zero. + TORCH_CHECK(tiles > 0 || positions == 0, "the tile list is empty but there are ", positions, + " positions to sum"); // At least one block, so an empty list (segments + 1 == 0) is a no-op and not an invalid launch. const int64_t items = std::max({positions, tiles, segments + 1, int64_t{1}}); const unsigned blocks = static_cast(std::min((items + 255) / 256, 1024)); 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 220311269..892e86c85 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -370,6 +370,16 @@ def _backward_probe(mutation: str): first = first.clone(); first[-1] = first[-1] - 1 elif mutation == "no_tiles": begin = begin[:0]; end = end[:0] + elif mutation == "no_tiles_and_empty_segments": + begin = begin[:0]; end = end[:0]; first = torch.zeros_like(first) + elif mutation == "gap_between_tiles": + begin = begin.clone(); begin[1] += 1 + elif mutation == "overlapping_tiles": + begin = begin.clone(); begin[1] -= 1 + elif mutation == "first_tile_not_at_zero": + begin = begin.clone(); begin[0] += 1 + elif mutation == "last_tile_stops_short": + end = end.clone(); end[-1] -= 1 _C.h3_adaln_row_gather_backward(grad, sorted_pos, begin, end, first, torch.float32) torch.cuda.synchronize() """ @@ -393,11 +403,20 @@ def test_the_backward_stops_on_metadata_that_does_not_describe_the_tiles(self): "segments_past_tiles", "first_segment_not_at_tile_zero", "last_tile_owned_by_no_segment", - "no_tiles", + "gap_between_tiles", + "overlapping_tiles", + "first_tile_not_at_zero", + "last_tile_stops_short", ): done = self._backward_probe(mutation) assert done.returncode != 0, mutation assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) + # An empty tile list over a non-empty sequence is refused on the host, before any launch, + # whether or not the segment boundaries agree with it. + for mutation in ("no_tiles", "no_tiles_and_empty_segments"): + done = self._backward_probe(mutation) + assert done.returncode != 0, mutation + assert "tile list is empty" in done.stderr, (mutation, done.stderr) @pytest.mark.parametrize("grad_dtype", [torch.float32, torch.bfloat16, torch.float16]) @pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) From 2164211d8b3ec7dee0f4245cbcbdf578683a0dd0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:00:59 +0000 Subject: [PATCH 138/139] fix(h3): reject a missing shift, scale or index in the modulated RMSNorm validate_h3_modulation read index.dtype first, so forward_modulated(..., index=None) failed with an AttributeError. All three tensors are required for the modulated form; a missing one now raises a TypeError that names it. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/reference/minimax_h3/rmsnorm.py | 5 +++++ tests/models/minimax_h3/test_h3_rmsnorm.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/rl_engine/reference/minimax_h3/rmsnorm.py b/rl_engine/reference/minimax_h3/rmsnorm.py index ddbdaf6bf..dda2e1352 100644 --- a/rl_engine/reference/minimax_h3/rmsnorm.py +++ b/rl_engine/reference/minimax_h3/rmsnorm.py @@ -56,6 +56,11 @@ def validate_h3_modulation( same_dtype: bool = True, ) -> None: hidden = x.shape[-1] + # The modulated form always gathers rows: a missing table or index is an error, not the plain + # norm (which is ``forward``, without shift and scale). + for name, tensor in (("shift", shift), ("scale", scale), ("index", index)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") for name, tensor in (("shift", shift), ("scale", scale)): if tensor.device != x.device: raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py index be899893b..309775562 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm.py @@ -379,3 +379,21 @@ def test_a_modulation_index_outside_the_table_stops_the_kernels(entry, bad): done = _index_probe(entry, bad) assert done.returncode != 0, entry assert "device-side assert triggered" in done.stderr, done.stderr + + +@pytest.mark.parametrize("missing", ["shift", "scale", "index"]) +def test_the_modulated_form_rejects_a_missing_table_or_index(missing): + """``index=None`` must not fall back to the plain norm (or fail on a ``None`` attribute).""" + + x = torch.randn(4, 8, dtype=torch.bfloat16) + weight = torch.ones(8, dtype=torch.bfloat16) + args = { + "shift": torch.randn(2, 8, dtype=torch.bfloat16), + "scale": torch.randn(2, 8, dtype=torch.bfloat16), + "index": torch.tensor([0, 1, 1, 0]), + } + args[missing] = None + with pytest.raises(TypeError, match=f"{missing} must be a torch.Tensor"): + NativeH3RMSNormOp().forward_modulated( + x, weight, args["shift"], args["scale"], args["index"] + ) From 60e7785c754a2f08b4f2e207977430063a181dd1 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:05:10 +0000 Subject: [PATCH 139/139] fix(h3): return no residual gradient for a residual that needs none The gate residual backward handed back the cotangent as the residual's gradient whether or not the residual required one. It now follows needs_input_grad[0], like the other ports of this group. The tile-metadata mutation test gains a gap, an overlap and a short last tile, and expects the empty tile list to be refused on the host. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/gate_residual.py | 6 ++-- .../minimax_h3/test_h3_adaln_gate_residual.py | 31 ++++++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) 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 da7319bfb..f39cf351b 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 @@ -52,8 +52,10 @@ def backward(ctx, grad): "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) # The cotangent has the output's shape, which is the residual's; the residual's VJP is that - # cotangent itself. reshape (not view) so a non-contiguous gradient is accepted. - return grad.reshape(ctx.residual_shape), dy.view(ctx.y_shape), dgate, None + # cotangent itself. reshape (not view) so a non-contiguous gradient is accepted. A residual + # that needs no gradient (a frozen skip path) gets none. + d_residual = grad.reshape(ctx.residual_shape) if ctx.needs_input_grad[0] else None + return d_residual, dy.view(ctx.y_shape), dgate, None class H3GateResidualCudaOp: diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index 193abce60..f93561263 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -141,6 +141,26 @@ def test_residual_gradient_keeps_the_residual_shape_for_batched_and_strided_grad assert torch.equal(leaves[0].grad, ref[0].grad) assert torch.equal(leaves[1].grad, ref[1].grad) + def test_no_residual_gradient_is_returned_for_a_residual_that_needs_none(self): + """The backward node hands back ``None`` for a frozen residual, the cotangent otherwise.""" + + table, residual, y, index = _case(100, hidden=64, seed=7) + grad = torch.randn_like(residual) + for needs_residual in (False, True): + leaves = [ + residual.detach().clone().requires_grad_(needs_residual), + y.detach().clone().requires_grad_(True), + table.detach().clone().requires_grad_(True), + ] + out = _cuda_op()(leaves[0], leaves[1], _gate(leaves[2]), index) + # grad_fn.apply returns what backward returned, before autograd drops the unused ones + d_residual, d_y, *_ = out.grad_fn.apply(grad) + if needs_residual: + assert torch.equal(d_residual, grad) + else: + assert d_residual is None + assert d_y.shape == y.shape + def test_a_gate_table_whose_rows_overlap_is_copied_not_read_in_place(self): """An expanded table (row stride 0) is a legal gate: it is copied to dense rows first.""" @@ -248,6 +268,12 @@ def _gate_backward_in_child(mutation: str) -> subprocess.CompletedProcess: first = torch.tensor([0, 1, 5], **i64) elif mutation == "no_tiles": begin, end = begin[:0], end[:0] + elif mutation == "gap_between_tiles": + begin = torch.tensor([0, 3], **i64) + elif mutation == "overlapping_tiles": + begin = torch.tensor([0, 1], **i64) + elif mutation == "last_tile_stops_short": + end = torch.tensor([2, 3], **i64) elif mutation == "short_sorted_pos": sorted_pos = sorted_pos[:3] elif mutation == "no_boundary": @@ -277,13 +303,16 @@ def test_the_gate_backward_stops_on_tile_metadata_that_does_not_describe_the_row "tile_end_past_rows", "position_past_rows", "segments_past_tiles", - "no_tiles", + "gap_between_tiles", + "overlapping_tiles", + "last_tile_stops_short", ): done = _gate_backward_in_child(mutation) assert done.returncode != 0, mutation assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) # Lengths are checked on the host, before any launch. for mutation, message in ( + ("no_tiles", "the tile list is empty but there are 4 positions to sum"), ("short_sorted_pos", "sorted_pos must have M entries"), ("no_boundary", "seg_first_tile must hold R + 1 entries"), ):