From 0522865203d6bfaffd9c6cd084c903f4e5864c8e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 6 Oct 2026 23:22:54 +0000 Subject: [PATCH 01/24] [WS1][CUDA][MiniMax-H3] timestep_sinusoid_h3: FP32 timestep features + shared H3 harness RFC #420 step 3, row `timestep_sinusoid_h3`. - CUDA kernel: one thread per (t, k), the provider's exact FP32 op sequence (__fmul_rn, precise expf/sinf/cosf); bitwise equal to diffusers' get_timestep_embedding on CUDA. Analytic row-local backward. Non-1-D, empty, integer, non-finite or out-of-[0, 1] timesteps fail closed (RFC probe H10). - PyTorch reference (provider replay + FP64 golden), registry, gtest. - Shared H3 harness used by the following rows: * h3_manifest.json + prepare_h3_weights.py: pinned revision and sha256, only shard 1 of 14 is downloaded; * h3_provider.py: diffusers replayed op for op; * h3_chain.py: stage-wise chain replay, each stage declaring its promises (provider-bitwise, golden tolerance); used by scripts/h3_chain_replay.py and tests/h3/test_h3_conditioning_e2e.py; * h3_report.py + scripts/h3_evidence.py + scripts/plot_h3_evidence.py: per-op performance/accuracy JSON and figure; * benchmarks/benchmark_h3_conditioning.py. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- benchmarks/benchmark_h3_conditioning.py | 60 ++++++ csrc/cuda/h3/timestep_sinusoid.cu | 83 +++++++++ csrc/ops.cpp | 14 ++ docs/.nav.yml | 2 + docs/operators/README.md | 1 + docs/operators/h3-timestep-sinusoid.md | 102 ++++++++++ rl_engine/_C.pyi | 3 + rl_engine/kernels/gtest/operator_inputs.py | 29 +++ rl_engine/kernels/gtest/operator_specs.py | 16 ++ rl_engine/kernels/ops/cuda/h3/__init__.py | 2 + .../kernels/ops/cuda/h3/timestep_sinusoid.py | 94 ++++++++++ rl_engine/kernels/ops/pytorch/h3/__init__.py | 30 +++ .../kernels/ops/pytorch/h3/fixed_order.py | 46 +++++ .../ops/pytorch/h3/timestep_sinusoid.py | 114 ++++++++++++ rl_engine/kernels/registry.py | 15 ++ rl_engine/testing/h3_cases.py | 35 ++++ rl_engine/testing/h3_chain.py | 170 +++++++++++++++++ rl_engine/testing/h3_manifest.json | 59 ++++++ rl_engine/testing/h3_provider.py | 33 ++++ rl_engine/testing/h3_report.py | 119 ++++++++++++ rl_engine/testing/h3_weights.py | 69 +++++++ rl_engine/tests/test_dispatch.py | 15 ++ scripts/h3_chain_replay.py | 97 ++++++++++ scripts/h3_evidence.py | 64 +++++++ scripts/plot_h3_evidence.py | 176 ++++++++++++++++++ scripts/prepare_h3_weights.py | 95 ++++++++++ setup.py | 2 + tests/h3/conftest.py | 25 +++ tests/h3/test_h3_conditioning_e2e.py | 55 ++++++ tests/h3/test_h3_timestep_sinusoid.py | 171 +++++++++++++++++ tests/test_ws1_gtest_gpu.py | 1 + 31 files changed, 1797 insertions(+) create mode 100644 benchmarks/benchmark_h3_conditioning.py create mode 100644 csrc/cuda/h3/timestep_sinusoid.cu create mode 100644 docs/operators/h3-timestep-sinusoid.md create mode 100644 rl_engine/kernels/ops/cuda/h3/__init__.py create mode 100644 rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/__init__.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/fixed_order.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py create mode 100644 rl_engine/testing/h3_cases.py create mode 100644 rl_engine/testing/h3_chain.py create mode 100644 rl_engine/testing/h3_manifest.json create mode 100644 rl_engine/testing/h3_provider.py create mode 100644 rl_engine/testing/h3_report.py create mode 100644 rl_engine/testing/h3_weights.py create mode 100644 scripts/h3_chain_replay.py create mode 100644 scripts/h3_evidence.py create mode 100644 scripts/plot_h3_evidence.py create mode 100644 scripts/prepare_h3_weights.py create mode 100644 tests/h3/conftest.py create mode 100644 tests/h3/test_h3_conditioning_e2e.py create mode 100644 tests/h3/test_h3_timestep_sinusoid.py diff --git a/benchmarks/benchmark_h3_conditioning.py b/benchmarks/benchmark_h3_conditioning.py new file mode 100644 index 000000000..de1aaceb7 --- /dev/null +++ b/benchmarks/benchmark_h3_conditioning.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark the MiniMax-H3 conditioning ops (RFC #420) against the provider path. + +CUDA-event medians, bandwidth and peak memory per case, plus the backend +the registry dispatched. Cases live in ``rl_engine/testing/h3_report.py``. + + python benchmarks/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 + python benchmarks/benchmark_h3_conditioning.py --op all --json out.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from rl_engine.kernels.registry import KernelRegistry # noqa: E402 +from rl_engine.testing.h3_chain import environment # noqa: E402 +from rl_engine.testing.h3_report import PERF_CASES, TIMED_KEYS, measure # noqa: E402 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", default="all", choices=["all", *PERF_CASES]) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--json", type=Path, default=None) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + torch.backends.cuda.matmul.allow_tf32 = False + + registry = KernelRegistry() + env = environment() + print(f"device={env['gpu']} torch={env['torch']} cuda={env['cuda']}") + results = [] + for name in list(PERF_CASES) if args.op == "all" else [args.op]: + for case in PERF_CASES[name](registry): + row = measure(case, args.warmup, args.iters) + results.append(row) + parts = [ + f"{key}={row[f'{key}_us']:.2f}us ({row[f'{key}_gbps']:.1f} GB/s, " + f"peak {row[f'{key}_peak_mib']:.2f} MiB)" + for key in TIMED_KEYS + if f"{key}_us" in row + ] + print(f"{name} {row['case']} [{row['backend']}]: " + "; ".join(parts)) + if args.json is not None: + args.json.write_text(json.dumps({"environment": env, "results": results}, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/csrc/cuda/h3/timestep_sinusoid.cu b/csrc/cuda/h3/timestep_sinusoid.cu new file mode 100644 index 000000000..778b78965 --- /dev/null +++ b/csrc/cuda/h3/timestep_sinusoid.cu @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 sinusoidal timestep features (RFC #420 `timestep_sinusoid_h3`). +// +// freq[k] = expf((c * k) * (1 / half)) c = (float)(-ln(max_period)) +// arg[t, k] = t[t] * freq[k] +// out[t] = [cosf(arg) | sinf(arg)] (T, 2 * half) FP32 +// +// Every value is produced by the same FP32 operation sequence as diffusers' +// get_timestep_embedding on CUDA (scalar * arange, multiply by the reciprocal +// of the CPU-scalar divisor, exp, multiply, cos/sin), so the output is meant +// to be bitwise equal to that path. Each element depends only on (t, k): +// batch-, position- and repeat-invariant by construction. Build without +// --use_fast_math: expf/sinf/cosf must stay the precise libdevice functions. + +#include +#include +#include +#include + +#include +#include + +namespace { + +__global__ void h3_timestep_sinusoid_kernel( + const float* __restrict__ timestep, + float* __restrict__ out, + int64_t num_timesteps, + int64_t half, + float neg_log_max_period, + float inv_half) { + const int64_t total = num_timesteps * half; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t row = idx / half; + const int64_t k = idx - row * half; + // __fmul_rn keeps each product a separately rounded FP32 multiply (no FMA + // contraction), matching the two elementwise torch kernels it replays. + const float exponent = __fmul_rn(__fmul_rn(neg_log_max_period, static_cast(k)), inv_half); + const float freq = expf(exponent); + const float arg = __fmul_rn(timestep[row], freq); + float* out_row = out + row * 2 * half; + out_row[k] = cosf(arg); + out_row[half + k] = sinf(arg); + } +} + +} // namespace + +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, + double max_period) { + TORCH_CHECK(timestep.is_cuda(), "timestep must be a CUDA tensor"); + TORCH_CHECK(timestep.scalar_type() == at::kFloat, "timestep must be float32, got ", + timestep.scalar_type()); + TORCH_CHECK(timestep.dim() == 1, "timestep must be 1-D (num_timesteps,)"); + TORCH_CHECK(timestep.numel() > 0, "timestep must hold at least one timestep"); + TORCH_CHECK(num_channels > 0 && num_channels % 2 == 0, + "num_channels must be a positive even number, got ", num_channels); + TORCH_CHECK(max_period > 0.0, "max_period must be positive"); + + const c10::cuda::CUDAGuard device_guard(timestep.device()); + auto t = timestep.contiguous(); + const int64_t num_timesteps = t.size(0); + const int64_t half = num_channels / 2; + auto out = torch::empty({num_timesteps, num_channels}, t.options()); + + // torch evaluates `-math.log(max_period) * arange` with the Python double + // rounded to the FP32 opmath type, and `x / half` as x * (1 / half) in FP32. + const float neg_log_max_period = static_cast(-std::log(max_period)); + const float inv_half = 1.0f / static_cast(half); + + const int threads = 256; + const int64_t total = num_timesteps * half; + const int64_t blocks = std::min((total + threads - 1) / threads, 65535); + auto stream = at::cuda::getCurrentCUDAStream(); + h3_timestep_sinusoid_kernel<<(blocks), threads, 0, stream>>>( + t.data_ptr(), out.data_ptr(), num_timesteps, half, neg_log_max_period, + inv_half); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 11fda1159..71b5b5502 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -497,6 +497,13 @@ at::Tensor prefix_shared_attention( #endif #endif +// MiniMax-H3 (RFC #420) conditioning-path declarations. CUDA only. +#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \ + (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, + double max_period); +#endif + // PyBind11 Module Registration PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "RL-Kernel High-Performance Operator Extension Library"; @@ -752,4 +759,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Deterministic GPT-NeoX token-major RoPE apply for ROCm"); #endif #endif + +#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \ + (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) + m.def("h3_timestep_sinusoid_forward", &h3_timestep_sinusoid_forward, + "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", + py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0); +#endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index 7d71230a3..d4c01f6c5 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -29,6 +29,8 @@ nav: - operators/sampling.md - operators/det-gemm.md - operators/embedding.md + - MiniMax-H3: + - operators/h3-timestep-sinusoid.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 00f4cbb45..dee8796e4 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -32,4 +32,5 @@ Every operator page should include: - [Matmul](matmul.md) - [Sampling](sampling.md) - [Token Embedding](embedding.md) +- [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md new file mode 100644 index 000000000..ce45133c3 --- /dev/null +++ b/docs/operators/h3-timestep-sinusoid.md @@ -0,0 +1,102 @@ +# MiniMax-H3 Timestep Sinusoid + +## Summary + +`timestep_sinusoid_h3` produces the FP32 sinusoidal timestep features that feed the +MiniMax-H3 timestep MLP (RFC #420, WS1 step 3). H3 builds +`Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)` and calls it +on the *distinct* timesteps of the packed sequence, unscaled in `[0, 1]` +(`t = 1 - sigma`): + +```text +freq[k] = exp(-ln(10000) * k / 128) k = 0..127 +arg[t, k] = t * freq[k] +out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256), float32 +``` + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227` (`rl_engine/testing/h3_manifest.json`). +Provider reference: `huggingface/diffusers@f53d552`, `get_timestep_embedding`. + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("timestep_sinusoid_h3", device="cuda") +features = op(timestep) # timestep: (T,) in [0, 1] -> (T, 256) float32 +``` + +## Backends + +| Backend | Wrapper | Native symbol | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.kernels.ops.cuda.h3.timestep_sinusoid.H3TimestepSinusoidCudaOp` | `rl_engine._C.h3_timestep_sinusoid_forward` | Bitwise equal to the provider path | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid.NativeH3TimestepSinusoidOp` | n/a | `forward`: provider replay; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `timestep` | `(T,)`, `T >= 1` | fp32 / bf16 / fp16 | Finite, in `[0, 1]`; any stride; CUDA for the CUDA backend | +| `num_channels` | scalar | int | Positive and even; H3 uses 256 | +| output | `(T, num_channels)` | float32 | `[cos | sin]` order | + +Inputs that break the contract fail closed: non-1-D, empty, integer or out-of-range +timesteps raise. `t * 1000` callers (RFC probe H10) therefore fail instead of +silently producing a different embedding. The range check costs one host read-back; +pass `check_range=False` to `forward` when the caller has already validated `t`. + +## Numerics + +- **Operation order.** The kernel evaluates exactly the FP32 operations the provider + runs on CUDA: `(float)(-ln 10000) * k`, multiplied by the FP32 reciprocal of `128` + (PyTorch divides by a CPU scalar this way), `expf`, `t * freq`, `cosf` and `sinf`. + The multiplies use `__fmul_rn`, so they are never fused into an FMA. The extension + is built without `--use_fast_math`, so `expf`, `sinf` and `cosf` are the precise + libdevice functions. +- **Cast points.** bf16 and fp16 timesteps are upcast to FP32 before anything else, + matching `timesteps[:, None].float()`. The output is never rounded below FP32. +- **Invariance.** Each output element depends only on `(t, k)`. Results are + bitwise identical across batch size, position, permutation and repeated runs. +- **Backward.** `d/dt` is the analytic row-local VJP + (`-f * sin(t f)` for the cos half, `f * cos(t f)` for the sin half). It is summed over + the 256 channels with a fixed pairwise tree (`tree_sum_lastdim_fp32`), so it does not + depend on how many timesteps share the launch. + +Measured on a B200 (torch 2.13.0+cu130): + +| Comparison | Result | +| --- | --- | +| CUDA vs provider path, T in {1, 2, 3, 7, 64, 1000, 4097} | bitwise equal | +| CUDA vs FP64 golden | max abs 1.19e-7 (contract: FP32 elementwise atol 1e-5) | +| gradient vs FP64 golden (gtest, T = 7) | max abs 2.4e-7 | + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 +``` + +On a B200 this is a single launch of about 15 µs, independent of `T` for `T <= 64`. With +the range check it takes about 51 µs. The provider path takes about 64 µs, because it is +six eager kernels plus two concatenations. Either way the op is launch-bound: it moves +about 1 KB per timestep. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_timestep_sinusoid.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end, pinned weights +python scripts/check_operator.py --op timestep_sinusoid_h3 --candidate cuda --device cuda \ + --dtype fp32 --batch 7 --check-grad +python scripts/h3_evidence.py --op timestep_sinusoid_h3 \ + --out docs/usage/evidence/h3-timestep-sinusoid-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-timestep-sinusoid-b200/report.json +``` + +## Known Limitations + +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- `max_period` is fixed at 10000, the only value H3 uses. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fb3c4379c..970df8515 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -325,3 +325,6 @@ def deterministic_collective_rocm_ipc_all_gather_input( input: torch.Tensor, output: torch.Tensor, ) -> None: ... +def h3_timestep_sinusoid_forward( + timestep: torch.Tensor, num_channels: int = 256, max_period: float = 10000.0 +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index ca3b7c120..4562f2aad 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -16,6 +16,7 @@ DEFAULT_VOCAB = 151936 DEFAULT_ROPE_THETA = 1.0e6 DEFAULT_RMS_EPS = 1.0e-6 +H3_FREQ_DIM = 256 def make_operator_inputs( @@ -42,6 +43,7 @@ def make_operator_inputs( "embedding": _make_embedding_inputs, "lm_head": _make_lm_head_inputs, "kv_cache_attention": _make_kv_cache_attention_inputs, + "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, } try: return builders[op_name](args, dtype, device) @@ -72,6 +74,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", + "timestep_sinusoid_h3": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}", } try: return names[op_name] @@ -331,6 +334,32 @@ def _make_kv_cache_attention_inputs( } +def _h3_num_timesteps(args: argparse.Namespace) -> int: + # H3 packs a handful of distinct timesteps; reuse --batch as their count. + return _arg_int(args, "num_timesteps", _arg_int(args, "batch", 2)) + + +def _h3_timesteps( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + num = _h3_num_timesteps(args) + mode = _arg_str(args, "input_mode", "random") + if mode == "constant": + t = torch.full((num,), 0.5, device=device) + else: + t = torch.rand((num,), generator=_generator(args, device, offset=7), device=device) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(dtype) + + +def _make_timestep_sinusoid_h3_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + return {"timestep": _h3_timesteps(args, dtype, device)} + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 4b86c2b7e..cc740d0ff 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -245,6 +245,22 @@ def _load_object(path: str) -> Any: }, grad_input_names=("logits",), ), + # MiniMax-H3 (RFC #420) conditioning path. The op is FP32 end to end, so + # its FP32 rows are the declared contract; other dtypes only change the + # timestep input's storage dtype. + "timestep_sinusoid_h3": OperatorSpec( + name="timestep_sinusoid_h3", + op_class="elementwise", + gold_path="rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid.NativeH3TimestepSinusoidOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ), + "cuda": "rl_engine.kernels.ops.cuda.h3.timestep_sinusoid.H3TimestepSinusoidCudaOp", + }, + grad_input_names=("timestep",), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/__init__.py b/rl_engine/kernels/ops/cuda/h3/__init__.py new file mode 100644 index 000000000..86cf4c9da --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py new file mode 100644 index 000000000..a779941b2 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 sinusoidal timestep features (RFC #420 ``timestep_sinusoid_h3``).""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.pytorch.h3 import H3_FREQ_DIM, H3_MAX_PERIOD +from rl_engine.kernels.ops.pytorch.h3.fixed_order import tree_sum_lastdim_fp32 +from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import ( + NativeH3TimestepSinusoidOp, + validate_h3_timesteps, +) + +KERNEL_ID = "rl_engine._C.h3_timestep_sinusoid_forward" +BACKWARD_IMPL = "row_local_fp32_analytic_tree_sum" + + +def h3_sinusoid_cuda_available() -> bool: + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) + + +class _H3TimestepSinusoidCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, timestep: torch.Tensor, num_channels: int): + t32 = timestep.detach().float().contiguous() + out = _C.h3_timestep_sinusoid_forward(t32, int(num_channels), float(H3_MAX_PERIOD)) + ctx.save_for_backward(out) + ctx.num_channels = int(num_channels) + ctx.timestep_dtype = timestep.dtype + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + (out,) = ctx.saved_tensors + half = ctx.num_channels // 2 + freq = NativeH3TimestepSinusoidOp.frequencies_fp32(ctx.num_channels, out.device) + cos_part, sin_part = out[:, :half], out[:, half:] + g = grad_out.float() + # d cos(t f)/dt = -f sin(t f); d sin(t f)/dt = f cos(t f). Row-local, + # summed over channels with a fixed pairwise tree. + contrib = torch.cat( + [-(g[:, :half] * sin_part) * freq, (g[:, half:] * cos_part) * freq], dim=-1 + ) + grad_t = tree_sum_lastdim_fp32(contrib).to(ctx.timestep_dtype) + record_backward( + "timestep_sinusoid_h3", + kernel_id="rl_engine.kernels.ops.pytorch.h3.fixed_order.tree_sum_lastdim_fp32", + impl=BACKWARD_IMPL, + family="cuda", + ) + return grad_t, None + + +class H3TimestepSinusoidCudaOp: + """CUDA candidate: one thread per (timestep, frequency), FP32 output. + + Construction fails when the extension lacks the symbol, so the registry + falls back to the PyTorch reference instead of silently mis-dispatching. + """ + + op_class = "elementwise" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_sinusoid_cuda_available(): + raise RuntimeError( + "rl_engine._C.h3_timestep_sinusoid_forward is unavailable; rebuild the " + "CUDA extension with csrc/cuda/h3/timestep_sinusoid.cu" + ) + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + if not timestep.is_cuda: + raise ValueError("H3TimestepSinusoidCudaOp needs a CUDA timestep tensor") + return _H3TimestepSinusoidCuda.apply(timestep, num_channels) + + def forward_fp32(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + # The op is FP32 end to end; the FP32 path is the op itself. + return self.forward(timestep, num_channels=num_channels) diff --git a/rl_engine/kernels/ops/pytorch/h3/__init__.py b/rl_engine/kernels/ops/pytorch/h3/__init__.py new file mode 100644 index 000000000..56a3e01ef --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 (RFC #420) PyTorch goldens. + +Shape constants are pinned to ``MiniMaxAI/MiniMax-H3@42ed227`` (see +``rl_engine/testing/h3_manifest.json``). The ops accept other sizes so that +tiny synthetic shapes can be tested, but the H3 layout rules (three modality +rows per timestep, six modulation chunks) are fixed. +""" + +from __future__ import annotations + +H3_FREQ_DIM = 256 +H3_MAX_PERIOD = 10000 +H3_TIME_EMBED_HIDDEN_DIM = 5376 +H3_TIME_EMBED_DIM = 2688 +H3_HIDDEN_SIZE = 5376 +# Modality tags of the packed sequence: 0 video, 1 text, 2 audio. +H3_MODALITY_NUM = 3 +# shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp +H3_ADALN_CHUNKS = 6 +H3_ADALN_CHUNK_NAMES = ( + "shift_msa", + "scale_msa", + "gate_msa", + "shift_mlp", + "scale_mlp", + "gate_mlp", +) diff --git a/rl_engine/kernels/ops/pytorch/h3/fixed_order.py b/rl_engine/kernels/ops/pytorch/h3/fixed_order.py new file mode 100644 index 000000000..e95a66054 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/fixed_order.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shape-independent FP32 reductions shared by the H3 goldens and backwards. + +``torch.sum`` picks its reduction tree from the tensor shape, so a row's +result can change with the number of rows next to it. These helpers spell +the order out with elementwise ops only, which keeps every row's bytes a +function of that row alone. +""" + +from __future__ import annotations + +import torch + + +def tree_sum_lastdim_fp32(x: torch.Tensor) -> torch.Tensor: + """Pairwise tree over the last dim: pad to a power of two, halve until 1. + + Level ``l`` adds element ``i`` and ``i + width/2`` for every ``i``. The + tree depends only on the last-dim size. + """ + + x = x.float() + width = x.shape[-1] + if width == 0: + return x.new_zeros(x.shape[:-1]) + padded = 1 << (width - 1).bit_length() + if padded != width: + x = torch.nn.functional.pad(x, (0, padded - width)) + while x.shape[-1] > 1: + half = x.shape[-1] // 2 + x = x[..., :half] + x[..., half:] + return x[..., 0] + + +def fold_rows_fp32(rows: torch.Tensor) -> torch.Tensor: + """Left fold over dim 0 in ascending row order: ((r0 + r1) + r2) + ...""" + + rows = rows.float() + if rows.shape[0] == 0: + return rows.new_zeros(rows.shape[1:]) + acc = rows[0].clone() + for index in range(1, rows.shape[0]): + acc = acc + rows[index] + return acc diff --git a/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py b/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py new file mode 100644 index 000000000..8ca0c0ccb --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 sinusoidal timestep features (RFC #420 row ``timestep_sinusoid_h3``). + +H3 builds ``Timesteps(num_channels=256, flip_sin_to_cos=True, +downscale_freq_shift=0)`` and feeds it the *distinct* timestep values of the +packed sequence, unscaled in ``[0, 1]`` (``t = 1 - sigma``):: + + freq[k] = exp(-ln(10000) * k / 128) k = 0..127, FP32 + arg[t, k] = t * freq[k] FP32 + out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256) FP32 + +The output is always FP32: ``time_embedder`` is a ``_keep_in_fp32_modules`` +module in the checkpoint, so the features never pass through BF16. +""" + +from __future__ import annotations + +import math + +import torch + +from rl_engine.kernels.ops.pytorch.h3 import H3_FREQ_DIM, H3_MAX_PERIOD + +_TIMESTEP_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_timesteps( + timestep: torch.Tensor, num_channels: int, *, check_range: bool = True +) -> None: + """Fail closed on inputs the H3 convention does not define. + + ``check_range`` rejects values outside ``[0, 1]``: H3 consumes + ``t = 1 - sigma`` unscaled, so a ``t * 1000`` caller (RFC probe H10) is a + convention error, not a different valid input. It reads the tensor back + to the host, which is cheap for the handful of distinct timesteps. + """ + + if not isinstance(timestep, torch.Tensor): + raise TypeError("timestep must be a torch.Tensor") + if timestep.dim() != 1: + raise ValueError(f"timestep must be 1-D (num_timesteps,), got {tuple(timestep.shape)}") + if timestep.numel() == 0: + raise ValueError("timestep must hold at least one timestep") + if timestep.dtype not in _TIMESTEP_DTYPES: + raise TypeError(f"timestep must be fp32, bf16 or fp16, got {timestep.dtype}") + if num_channels <= 0 or num_channels % 2 != 0: + raise ValueError(f"num_channels must be a positive even number, got {num_channels}") + if check_range: + t32 = timestep.detach().float() + # One host sync: NaN fails both comparisons, so this also rejects it. + if not bool(((t32 >= 0) & (t32 <= 1)).all()): + raise ValueError( + "timestep must be finite and lie in [0, 1]: H3 consumes t = 1 - sigma " + f"unscaled (got {t32.cpu().tolist()[:8]})" + ) + + +class NativeH3TimestepSinusoidOp: + """PyTorch reference for the H3 sinusoidal timestep features. + + ``forward`` replays diffusers' ``get_timestep_embedding`` op for op on the + input's device (the provider path). ``forward_fp32`` is the independent + golden: the same formula evaluated in FP64 and rounded once to FP32. + """ + + op_class = "elementwise" + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + # Same op sequence and dtypes as diffusers get_timestep_embedding with + # flip_sin_to_cos=True, downscale_freq_shift=0, scale=1. + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / half + freq = torch.exp(exponent) + arg = timestep[:, None].float() * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1) + + def forward_fp32( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + k = torch.arange(half, dtype=torch.float64, device=timestep.device) + freq = torch.exp(-math.log(H3_MAX_PERIOD) * k / half) + arg = timestep.double()[:, None] * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1).float() + + @staticmethod + def frequencies_fp32(num_channels: int, device: torch.device | str) -> torch.Tensor: + """The FP32 frequency table the provider path multiplies by.""" + + half = num_channels // 2 + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=device + ) + return torch.exp(exponent / half) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 9eeb1c42b..533536f57 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -200,6 +200,14 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp" CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" + # MiniMax-H3 (RFC #420) conditioning path + PYTORCH_H3_TIMESTEP_SINUSOID = ( + "rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ) + CUDA_H3_TIMESTEP_SINUSOID = ( + "rl_engine.kernels.ops.cuda.h3.timestep_sinusoid.H3TimestepSinusoidCudaOp" + ) + def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: return ( @@ -609,6 +617,10 @@ def __init__(self): OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE, ], + "timestep_sinusoid_h3": [ + OpBackend.CUDA_H3_TIMESTEP_SINUSOID, + OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID, + ], }, "rocm": { "logp": [ @@ -649,6 +661,7 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -693,6 +706,7 @@ def __init__(self): ], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -718,6 +732,7 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_cases.py b/rl_engine/testing/h3_cases.py new file mode 100644 index 000000000..527a820ec --- /dev/null +++ b/rl_engine/testing/h3_cases.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic MiniMax-H3 (RFC #420) test inputs shared by tests and scripts.""" + +from __future__ import annotations + +import torch + + +def h3_timesteps(num: int, *, seed: int = 0, device: str = "cuda") -> torch.Tensor: + """Distinct-looking timesteps in [0, 1] with both endpoints present.""" + + generator = torch.Generator(device="cpu").manual_seed(seed) + t = torch.rand(num, generator=generator) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(device) + + +def h3_packed_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """(timestep_indices, token_tags) for a packed sequence with every modality. + + Rows are video, text and audio blocks in random order, each block mapped + to a random timestep, like H3's conditioning/target packing. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + token_tags = torch.randint(0, 3, (seq_len,), generator=generator) + timestep_indices = torch.randint(0, num_timesteps, (seq_len,), generator=generator) + token_tags[: min(3, seq_len)] = torch.arange(min(3, seq_len)) + return timestep_indices.to(device), token_tags.to(device) diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py new file mode 100644 index 000000000..fb1f10943 --- /dev/null +++ b/rl_engine/testing/h3_chain.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Stage-by-stage replay of the MiniMax-H3 conditioning chain (RFC #420). + +Each stage runs three ways on the same device: + +* provider: the diffusers op sequence (``rl_engine.testing.h3_provider``); +* candidate: the dispatched RL-Kernel op, fed the previous candidate output + (chained) and, separately, the previous provider output (isolated); +* golden: the FP32/FP64 reference, fed the previous golden output. + +A stage also declares what it promises, so the end-to-end test and the +evidence JSON check the same thing: whether its isolated output is bitwise +equal to the provider, and its tolerance against the golden (rows of +``tolerance_contract.json``). ``first_drift`` is the first stage whose chained +output differs from the provider; ``first_isolated_drift`` the first stage +that differs even on provider inputs. +""" + +from __future__ import annotations + +import platform +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable + +import torch + +from rl_engine.kernels.registry import KernelRegistry +from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps +from rl_engine.testing.h3_provider import provider_time_proj + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +@dataclass(frozen=True) +class Stage: + name: str + op_type: str + # (op, ctx, upstream) -> outputs; ``upstream`` is the previous stage's output. + candidate: Callable[[Any, dict, Any], Any] + provider: Callable[[dict, Any], Any] + golden: Callable[[Any, dict, Any], Any] + # Promises checked by tests/h3/test_h3_conditioning_e2e.py. + provider_bitwise_isolated: bool + golden_atol: float + golden_rtol: float + + +STAGES: list[Stage] = [ + Stage( + name="timestep_sinusoid_h3", + op_type="timestep_sinusoid_h3", + candidate=lambda op, ctx, _up: op(ctx["timestep"]), + provider=lambda ctx, _up: provider_time_proj(ctx["timestep"]), + golden=lambda op, ctx, _up: op.forward_fp32(ctx["timestep"]), + provider_bitwise_isolated=True, + golden_atol=1e-5, # elementwise / float32 + golden_rtol=1e-5, + ), +] + + +def _flat(value: Any) -> list[torch.Tensor]: + if isinstance(value, torch.Tensor): + return [value] + return [tensor for item in value for tensor in _flat(item)] + + +def compare(lhs: Any, rhs: Any, atol: float = 0.0, rtol: float = 0.0) -> dict[str, Any]: + lhs_t, rhs_t = _flat(lhs), _flat(rhs) + bitwise = all(torch.equal(a, b) for a, b in zip(lhs_t, rhs_t, strict=True)) + max_abs = 0.0 + within = True + for a, b in zip(lhs_t, rhs_t, strict=True): + diff = (a.float() - b.float()).abs() + if diff.numel(): + max_abs = max(max_abs, float(diff.max())) + within = within and bool((diff <= atol + rtol * b.float().abs()).all()) + return {"bitwise_equal": bitwise, "max_abs": max_abs, "within_tolerance": within} + + +def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dict[str, Any]: + timestep_indices, token_tags = h3_packed_layout(seq_len, num_timesteps, seed=seed) + return { + "timestep": h3_timesteps(num_timesteps, seed=seed), + "timestep_indices": timestep_indices, + "token_tags": token_tags, + "weights": weights, + } + + +def golden_op(registry: KernelRegistry, op_type: str): + """The PyTorch reference: the last entry of the CPU priority list.""" + + return registry._get_or_create_backend(registry._priority_map["cpu"][op_type][-1]) + + +def run_case( + registry: KernelRegistry, + weights, + *, + num_timesteps: int, + seq_len: int, + seed: int = 0, + stages: list[Stage] | None = None, +) -> dict[str, Any]: + stages = STAGES if stages is None else stages + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + report: dict[str, Any] = { + "num_timesteps": num_timesteps, + "seq_len": seq_len, + "seed": seed, + "stages": [], + } + chained = provider = golden = None + first_drift = first_isolated = None + for stage in stages: + op = registry.get_op(stage.op_type, device="cuda") + gold = golden_op(registry, stage.op_type) + with torch.no_grad(): + isolated = stage.candidate(op, ctx, provider) + chained_input = chained + chained = stage.candidate(op, ctx, chained_input) + repeat = stage.candidate(op, ctx, chained_input) + golden = stage.golden(gold, ctx, golden) + provider = stage.provider(ctx, provider) + entry = { + "stage": stage.name, + "backend": type(op).__name__, + "kernel_id": getattr(op, "kernel_id", type(op).__name__), + "repeat_bitwise_equal": compare(repeat, chained)["bitwise_equal"], + "chained_vs_provider": compare(chained, provider), + "isolated_vs_provider": compare(isolated, provider), + "chained_vs_golden": compare(chained, golden, stage.golden_atol, stage.golden_rtol), + "provider_vs_golden": compare(provider, golden, stage.golden_atol, stage.golden_rtol), + } + if first_drift is None and not entry["chained_vs_provider"]["bitwise_equal"]: + first_drift = stage.name + if first_isolated is None and not entry["isolated_vs_provider"]["bitwise_equal"]: + first_isolated = stage.name + report["stages"].append(entry) + report["first_drift"] = first_drift + report["first_isolated_drift"] = first_isolated + return report + + +def git_state() -> dict[str, Any]: + def run(*args: str) -> str: + return subprocess.check_output(["git", *args], cwd=REPO_ROOT, text=True).strip() + + try: + sha = run("rev-parse", "HEAD") + dirty = bool(run("status", "--porcelain", "--untracked-files=no")) + except Exception: # noqa: BLE001 + sha, dirty = "unknown", True + return {"rl_kernel_commit": sha, "tracked_tree_dirty": dirty} + + +def environment() -> dict[str, Any]: + return { + "gpu": torch.cuda.get_device_name(), + "capability": list(torch.cuda.get_device_capability()), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "python": platform.python_version(), + "tf32": bool(torch.backends.cuda.matmul.allow_tf32), + } diff --git a/rl_engine/testing/h3_manifest.json b/rl_engine/testing/h3_manifest.json new file mode 100644 index 000000000..52b1989b1 --- /dev/null +++ b/rl_engine/testing/h3_manifest.json @@ -0,0 +1,59 @@ +{ + "version": "h3-conditioning-v1", + "rfc": "RL-Align/RL-Kernel#420", + "model_identity": { + "hf_repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "subfolder": "transformer", + "architecture": "MiniMaxH3Transformer3DModel", + "config_sha256": "74c11bff524336576096993cbfcdcdc2ef4fa2fa4409df693bdcbc6c666282ae", + "index_file": "diffusion_pytorch_model.safetensors.index.json", + "index_sha256": "ac30a3b58963f2e735d493475fbb81853a5735ec947619648b3e045acda6783e", + "config_fingerprint": { + "hidden_size": 5376, + "num_attention_heads": 56, + "attention_head_dim": 128, + "num_layers": 50, + "num_refiner_layers": 2, + "ffn_dim": 14336, + "freq_dim": 256, + "time_embed_hidden_dim": 5376, + "time_embed_dim": 2688, + "norm_eps": 1e-05, + "final_norm_eps": 1e-05 + } + }, + "weight_shards": { + "diffusion_pytorch_model-00001-of-00014.safetensors": { + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028", + "size_bytes": 4825958704 + } + }, + "tensors": { + "time_embedder.linear_1.weight": {"dtype": "float32", "shape": [5376, 256]}, + "time_embedder.linear_1.bias": {"dtype": "float32", "shape": [5376]}, + "time_embedder.linear_2.weight": {"dtype": "float32", "shape": [2688, 5376]}, + "time_embedder.linear_2.bias": {"dtype": "float32", "shape": [2688]}, + "transformer_blocks.0.adaln_proj.linear.weight": {"dtype": "bfloat16", "shape": [96768, 2688]}, + "transformer_blocks.0.adaln_proj.linear.bias": {"dtype": "bfloat16", "shape": [96768]} + }, + "reference_implementation": { + "repo": "huggingface/diffusers", + "commit": "f53d552036a0d1bd5570782a39cd40cfabf112bc", + "file": "src/diffusers/models/transformers/transformer_minimax_h3.py", + "file_sha256": "435873d9bc26b8770342961e18cb20324aeeaa4d045e85e0941638d4857d7db3", + "timestep_embedding_file": "src/diffusers/models/embeddings.py" + }, + "conventions": { + "timestep": "distinct timesteps t = 1 - sigma, unscaled in [0, 1], shape (T,)", + "modality_tags": {"video": 0, "text": 1, "audio": 2}, + "adaln_row": "timestep_index * 3 + token_tag", + "adaln_table_rows": "[t0_mod0, t0_mod1, t0_mod2, t1_mod0, ...]", + "adaln_chunk_order": ["shift_msa", "scale_msa", "gate_msa", "shift_mlp", "scale_mlp", "gate_mlp"], + "precision": { + "timestep_sinusoid": "float32", + "timestep_mlp": "float32 weights, float32 activations", + "adaln_projection": "SiLU in float32, cast to bfloat16, bfloat16 weight/bias/output" + } + } +} diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py new file mode 100644 index 000000000..ff4718c00 --- /dev/null +++ b/rl_engine/testing/h3_provider.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The diffusers MiniMax-H3 conditioning path, op for op, without diffusers. + +Each function replays the tensor ops of the pinned reference +(``huggingface/diffusers@f53d552``, see ``h3_manifest.json``) in the same +order and dtypes, so it runs the same kernels the provider does on a given +device. This is the *provider* side of the RFC #420 comparisons; the FP32/FP64 +goldens live in ``rl_engine.kernels.ops.pytorch.h3``. +""" + +from __future__ import annotations + +import math + +import torch + + +def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch.Tensor: + """``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0)``.""" + + half_dim = num_channels // 2 + exponent = -math.log(10000) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / (half_dim - 0) + emb = torch.exp(exponent) + emb = timestep[:, None].float() * emb[None, :] + emb = 1 * emb + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + return emb diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py new file mode 100644 index 000000000..bd5f9c9a6 --- /dev/null +++ b/rl_engine/testing/h3_report.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Per-operator performance and accuracy measurements for the H3 evidence. + +``PERF_CASES[op]`` builds the timed cases that ``benchmarks/benchmark_h3_conditioning.py`` +prints and ``scripts/h3_evidence.py`` stores. ``ACCURACY[op]`` measures the +op against the provider path and its golden. Each RFC #420 row registers both. +""" + +from __future__ import annotations + +import statistics +from typing import Any, Callable + +import torch + +from rl_engine.kernels.registry import KernelRegistry +from rl_engine.testing.h3_cases import h3_timesteps +from rl_engine.testing.h3_provider import provider_time_proj + +# Keys of a perf case that hold a timed callable, in display order. +TIMED_KEYS = ( + "candidate", + "candidate_checked", + "provider", + "candidate_backward", + "provider_backward", +) + + +def time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 200) -> float: + """Median CUDA-event time of ``fn`` in microseconds (includes the Python wrapper).""" + + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1e3) + return statistics.median(samples) + + +def peak_mib(fn: Callable[[], Any]) -> float: + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + fn() + torch.cuda.synchronize() + return (torch.cuda.max_memory_allocated() - base) / 2**20 + + +def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} + for key in TIMED_KEYS: + if key in case: + us = time_us(case[key], warmup, iters) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key]) + return row + + +# --------------------------------------------------------------------------- # +# timestep_sinusoid_h3 +# --------------------------------------------------------------------------- # + + +def _sinusoid_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + cases = [] + for num in (1, 2, 4, 64): + t = h3_timesteps(num) + cases.append( + { + "op": "timestep_sinusoid_h3", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": num * 4 + num * 256 * 4, + "candidate": lambda t=t: op.forward(t, check_range=False), + "candidate_checked": lambda t=t: op.forward(t), + "provider": lambda t=t: provider_time_proj(t), + } + ) + return cases + + +def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["timestep_sinusoid_h3"][-1] + ) + rows = [] + for num in (1, 2, 3, 7, 64, 1000, 4097): + t = h3_timesteps(num, seed=num) + out, ref, gold = op(t), provider_time_proj(t), golden.forward_fp32(t) + rows.append( + { + "num_timesteps": num, + "bitwise_equal_to_provider": bool(torch.equal(out, ref)), + "max_abs_vs_fp64": float((out - gold).abs().max()), + "provider_max_abs_vs_fp64": float((ref - gold).abs().max()), + } + ) + return {"contract_atol": 1e-5, "cases": rows} + + +PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { + "timestep_sinusoid_h3": _sinusoid_perf, +} +ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { + "timestep_sinusoid_h3": _sinusoid_accuracy, +} diff --git a/rl_engine/testing/h3_weights.py b/rl_engine/testing/h3_weights.py new file mode 100644 index 000000000..8090c6e83 --- /dev/null +++ b/rl_engine/testing/h3_weights.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Pinned MiniMax-H3 conditioning weights for the RFC #420 tests. + +Only the tensors listed in ``h3_manifest.json`` are used. They all live in +shard 1 of 14, so ``scripts/prepare_h3_weights.py`` downloads that shard, +checks it against the manifest and extracts the tensors into +``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import torch + +MANIFEST_PATH = Path(__file__).with_name("h3_manifest.json") +WEIGHTS_ENV = "RL_KERNEL_H3_WEIGHTS" +EXTRACTED_FILE = "h3_conditioning.safetensors" + + +def load_h3_manifest() -> dict[str, Any]: + return json.loads(MANIFEST_PATH.read_text()) + + +def h3_weights_dir() -> Path | None: + value = os.environ.get(WEIGHTS_ENV, "").strip() + return Path(value).expanduser() if value else None + + +def sha256_file(path: Path, chunk: int = 1 << 24) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(chunk): + digest.update(block) + return digest.hexdigest() + + +def load_h3_conditioning_weights( + device: torch.device | str = "cpu", names: list[str] | None = None +) -> dict[str, torch.Tensor]: + """Load the extracted pinned tensors; raise if they are missing or off-manifest.""" + + from safetensors.torch import load_file + + root = h3_weights_dir() + if root is None: + raise FileNotFoundError(f"set {WEIGHTS_ENV} to the prepare_h3_weights.py output dir") + path = root / EXTRACTED_FILE + if not path.is_file(): + raise FileNotFoundError(f"{path} missing; run scripts/prepare_h3_weights.py") + manifest = load_h3_manifest() + tensors = load_file(str(path), device=str(device)) + wanted = names if names is not None else list(manifest["tensors"]) + out: dict[str, torch.Tensor] = {} + for name in wanted: + spec = manifest["tensors"][name] + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise ValueError(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise ValueError(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + out[name] = tensor + return out diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index b63a67882..f416295d0 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -4,6 +4,7 @@ import sys from types import ModuleType +import pytest import torch import rl_engine.platforms.device as device_module @@ -225,6 +226,20 @@ def fake_load_backend(backend): ] +H3_OPS = ("timestep_sinusoid_h3",) + + +@pytest.mark.parametrize("op_name", H3_OPS) +def test_h3_ops_prefer_cuda_and_fall_back_to_pytorch_reference(op_name): + registry = KernelRegistry() + cuda_chain = registry._priority_map["cuda"][op_name] + assert cuda_chain[0].name.startswith("CUDA_H3_") + assert cuda_chain[-1].name.startswith("PYTORCH_H3_") + for platform in ("rocm", "musa", "cpu", "npu"): + chain = registry._priority_map[platform][op_name] + assert [backend.name for backend in chain] == [cuda_chain[-1].name] + + def test_npu_available_handles_runtime_failure(monkeypatch): class BrokenNPU: @staticmethod diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py new file mode 100644 index 000000000..e262ca053 --- /dev/null +++ b/scripts/h3_chain_replay.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay the MiniMax-H3 conditioning chain and write JSON evidence (RFC #420). + +See ``rl_engine/testing/h3_chain.py`` for what is compared. + + export RL_KERNEL_H3_WEIGHTS= + python scripts/h3_chain_replay.py --out chain_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from rl_engine.kernels.registry import KernelRegistry # noqa: E402 +from rl_engine.testing import h3_chain # noqa: E402 +from rl_engine.testing.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) + + +def _ints(text: str) -> list[int]: + return [int(v) for v in text.split(",")] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--stages", default=",".join(s.name for s in h3_chain.STAGES)) + parser.add_argument("--timesteps", type=_ints, default=[1, 2, 4], help="comma list of T") + parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("the chain replay needs a CUDA device") + wanted = args.stages.split(",") + unknown = sorted(set(wanted) - {s.name for s in h3_chain.STAGES}) + if unknown: + raise SystemExit(f"unknown stages: {unknown}") + stages = [s for s in h3_chain.STAGES if s.name in wanted] + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + weights = load_h3_conditioning_weights("cuda") + registry = KernelRegistry() + manifest = load_h3_manifest() + + cases = [] + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_case( + registry, + weights, + num_timesteps=num_timesteps, + seq_len=seq_len, + seed=args.seed, + stages=stages, + ) + cases.append(case) + summary = ", ".join( + f"{e['stage']}=" + f"{'eq' if e['chained_vs_provider']['bitwise_equal'] else 'drift'}" + f"(gold {e['chained_vs_golden']['max_abs']:.3e})" + for e in case["stages"] + ) + print(f"T={num_timesteps} S={seq_len}: {summary}; first_drift={case['first_drift']}") + + evidence = { + "kind": "h3_conditioning_chain_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "stages": [s.name for s in stages], + "cases": cases, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/h3_evidence.py b/scripts/h3_evidence.py new file mode 100644 index 000000000..c20778234 --- /dev/null +++ b/scripts/h3_evidence.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write one RFC #420 operator's evidence report (performance + accuracy) as JSON. + +Run it on a clean tree; the report records the commit and whether the tree +was dirty. ``scripts/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python scripts/h3_evidence.py --op timestep_sinusoid_h3 \\ + --out docs/usage/evidence/h3-timestep-sinusoid-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from rl_engine.kernels.registry import KernelRegistry # noqa: E402 +from rl_engine.testing.h3_chain import environment, git_state # noqa: E402 +from rl_engine.testing.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 +from rl_engine.testing.h3_weights import load_h3_manifest # noqa: E402 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + + registry = KernelRegistry() + manifest = load_h3_manifest() + report = { + "kind": "h3_operator_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "reference_commit": manifest["reference_implementation"]["commit"], + **git_state(), + "environment": environment(), + "accuracy": ACCURACY[args.op](registry), + "perf": [measure(case) for case in PERF_CASES[args.op](registry)], + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py new file mode 100644 index 000000000..b0bcc8eb1 --- /dev/null +++ b/scripts/plot_h3_evidence.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render an RFC #420 operator report (``scripts/h3_evidence.py``) as a PNG. + +Needs only the report JSON and matplotlib (not torch or a GPU), so it can run +anywhere: + + python scripts/plot_h3_evidence.py docs/usage/evidence/h3-timestep-sinusoid-b200/report.json + +writes ``figure.png`` next to the report. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any, Callable + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + +# Fixed series identity across every H3 figure (categorical slots 1-3). +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +THIRD = "#1baf7a" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" + + +def _style(ax, title: str, ylabel: str) -> None: + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) + ax.set_ylabel(ylabel, color=MUTED, fontsize=9) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_axisbelow(True) + for side in ("top", "right", "left"): + ax.spines[side].set_visible(False) + ax.spines["bottom"].set_color(GRID) + + +def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt: str) -> None: + """Grouped bars with a 2px-style gap and value labels on top.""" + + n = len(series) + width = 0.8 / n + for i, (name, color, values) in enumerate(series): + xs = [x + (i - (n - 1) / 2) * width for x in range(len(labels))] + bars = ax.bar(xs, values, width * 0.92, color=color, label=name, zorder=2) + for bar, value in zip(bars, values): + ax.annotate( + fmt.format(value), + (bar.get_x() + bar.get_width() / 2, bar.get_height()), + xytext=(0, 2), + textcoords="offset points", + ha="center", + fontsize=7, + color=INK, + ) + ax.set_xticks(range(len(labels)), labels) + top = max(v for _, _, values in series for v in values) + ax.set_ylim(0, top * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=n) + + +# Exact zeros cannot sit on a log axis; draw them on this floor and label them. +LOG_FLOOR = 1e-9 + + +def _log_values(values: list[float]) -> list[float]: + return [v if v > 0 else LOG_FLOOR for v in values] + + +def _figure(report: dict[str, Any], suptitle: str): + fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), facecolor=SURFACE) + env = report["environment"] + fig.suptitle(suptitle, x=0.01, ha="left", fontsize=12, color=INK, fontweight="bold") + fig.text( + 0.01, + 0.01, + f"{env['gpu']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · MiniMax-H3@{report['model_revision'][:7]}", + fontsize=7, + color=MUTED, + ) + return fig, axes + + +def plot_sinusoid(report: dict[str, Any]): + fig, (left, right) = _figure(report, "timestep_sinusoid_h3") + perf = report["perf"] + _style(left, "Latency per call (lower is better)", "µs") + _bars( + left, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + ("CUDA + range check", THIRD, [row["candidate_checked_us"] for row in perf]), + ("diffusers", PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + cases = report["accuracy"]["cases"] + xs = [c["num_timesteps"] for c in cases] + _style(right, "Max |error| vs FP64 golden (log)", "max abs error") + cuda = [c["max_abs_vs_fp64"] for c in cases] + provider = [c["provider_max_abs_vs_fp64"] for c in cases] + right.plot(xs, _log_values(cuda), "o-", color=RL_KERNEL, lw=2, ms=8, label="RL-Kernel CUDA") + right.plot(xs, _log_values(provider), "s--", color=PROVIDER, lw=2, ms=5, label="diffusers") + atol = report["accuracy"]["contract_atol"] + right.axhline(atol, color=MUTED, lw=1, ls=":") + right.annotate( + "contract atol 1e-5", + (xs[-1], atol), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + for x, value in zip(xs, cuda): + if value == 0: + right.annotate( + "exact", + (x, LOG_FLOOR), + xytext=(0, 8), + textcoords="offset points", + fontsize=7, + color=INK, + ha="center", + ) + right.set_xscale("log") + right.set_yscale("log") + right.set_ylim(LOG_FLOOR / 3, atol * 30) + right.set_xlabel("number of timesteps T", color=MUTED, fontsize=9) + equal = sum(c["bitwise_equal_to_provider"] for c in cases) + right.legend( + frameon=False, + fontsize=8, + labelcolor=INK, + loc="upper left", + ncol=2, + title=f"bitwise equal to diffusers at {equal}/{len(cases)} sizes", + title_fontsize=8, + alignment="left", + ) + return fig + + +PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { + "timestep_sinusoid_h3": plot_sinusoid, +} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + fig = PLOTS[report["op"]](report) + fig.tight_layout(rect=(0, 0.04, 1, 0.95)) + out = args.out or args.report.with_name("figure.png") + fig.savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_h3_weights.py b/scripts/prepare_h3_weights.py new file mode 100644 index 000000000..dc2ed209e --- /dev/null +++ b/scripts/prepare_h3_weights.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fetch and pin the MiniMax-H3 conditioning tensors used by the RFC #420 tests. + +Downloads only the config, the shard index and the shards that hold the +manifest tensors (shard 1 of 14, ~4.8 GB) at the pinned revision, verifies +every sha256 against ``rl_engine/testing/h3_manifest.json`` and writes the +tensors to ``/h3_conditioning.safetensors``. Point the tests at it with +``export RL_KERNEL_H3_WEIGHTS=``. + + python scripts/prepare_h3_weights.py --out ~/.cache/rl-kernel/h3 +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.testing.h3_weights import EXTRACTED_FILE, load_h3_manifest, sha256_file # noqa: E402 + + +def _check(path: Path, expected: str) -> None: + actual = sha256_file(path) + if actual != expected: + raise SystemExit(f"sha256 mismatch for {path.name}: {actual} != {expected}") + print(f"ok {path.name} {actual}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, required=True, help="output directory") + parser.add_argument( + "--download-dir", + type=Path, + default=None, + help="where the raw shard goes (default: /hf)", + ) + args = parser.parse_args() + + from huggingface_hub import hf_hub_download + from safetensors import safe_open + from safetensors.torch import save_file + + manifest = load_h3_manifest() + identity = manifest["model_identity"] + out_dir: Path = args.out.expanduser() + download_dir = (args.download_dir or out_dir / "hf").expanduser() + out_dir.mkdir(parents=True, exist_ok=True) + + def fetch(filename: str) -> Path: + return Path( + hf_hub_download( + identity["hf_repo"], + f"{identity['subfolder']}/{filename}", + revision=identity["revision"], + local_dir=str(download_dir), + ) + ) + + _check(fetch("config.json"), identity["config_sha256"]) + index_path = fetch(identity["index_file"]) + _check(index_path, identity["index_sha256"]) + weight_map = json.loads(index_path.read_text())["weight_map"] + + tensors = {} + for shard in sorted({weight_map[name] for name in manifest["tensors"]}): + if shard not in manifest["weight_shards"]: + raise SystemExit(f"shard {shard} is not pinned in the manifest") + shard_path = fetch(shard) + _check(shard_path, manifest["weight_shards"][shard]["sha256"]) + with safe_open(str(shard_path), framework="pt") as handle: + for name in manifest["tensors"]: + if weight_map[name] == shard: + tensors[name] = handle.get_tensor(name).contiguous() + + for name, spec in manifest["tensors"].items(): + tensor = tensors[name] + assert str(tensor.dtype).removeprefix("torch.") == spec["dtype"], name + assert list(tensor.shape) == spec["shape"], name + + target = out_dir / EXTRACTED_FILE + save_file(tensors, str(target), metadata={"revision": identity["revision"]}) + print(f"wrote {len(tensors)} tensors to {target}") + print(f"export RL_KERNEL_H3_WEIGHTS={out_dir}") + + +if __name__ == "__main__": + main() diff --git a/setup.py b/setup.py index 9831c7022..29f870cc9 100644 --- a/setup.py +++ b/setup.py @@ -156,6 +156,8 @@ def get_extensions(): # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. cuda_sources.append("csrc/cuda/distributed/deterministic_collective.cu") + # MiniMax-H3 (RFC #420) conditioning path; plain-arch, runs on SM90 and SM100. + cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/conftest.py b/tests/h3/conftest.py new file mode 100644 index 000000000..44c8fe117 --- /dev/null +++ b/tests/h3/conftest.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared fixtures for the MiniMax-H3 (RFC #420) operator tests. + +Real-weight tests read the pinned tensors from ``$RL_KERNEL_H3_WEIGHTS`` +(written by ``scripts/prepare_h3_weights.py``) and skip when it is unset. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.testing.h3_weights import WEIGHTS_ENV, h3_weights_dir, load_h3_conditioning_weights + + +@pytest.fixture(scope="session") +def h3_weights_cpu() -> dict[str, torch.Tensor]: + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run scripts/prepare_h3_weights.py") + try: + return load_h3_conditioning_weights("cpu") + except FileNotFoundError as exc: + pytest.skip(str(exc)) diff --git a/tests/h3/test_h3_conditioning_e2e.py b/tests/h3/test_h3_conditioning_e2e.py new file mode 100644 index 000000000..edd75a3b2 --- /dev/null +++ b/tests/h3/test_h3_conditioning_e2e.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""End to end: the H3 conditioning chain on the pinned checkpoint, through the registry. + +Runs every stage registered in ``rl_engine.testing.h3_chain.STAGES`` (each +RFC #420 row adds its own), chained exactly as the model calls them, and +checks every stage's promises: + +* the registry dispatched the CUDA backend; +* a repeated run is bitwise equal; +* the chained output is within the stage's contract tolerance of the golden; +* stages that promise it are bitwise equal to diffusers on diffusers' inputs, + so the first isolated drift can only be a stage that declares a reduction. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.registry import KernelRegistry +from rl_engine.testing import h3_chain + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + +# (T distinct timesteps, S packed rows): single row, short packed, realistic-size packed. +CASES = [(1, 3), (3, 257), (4, 4097)] + + +@pytest.fixture(scope="module") +def chain_weights(h3_weights_cpu): + return {name: tensor.cuda() for name, tensor in h3_weights_cpu.items()} + + +@pytest.fixture(scope="module") +def registry(): + return KernelRegistry() + + +@pytest.mark.parametrize("num_timesteps, seq_len", CASES) +def test_chain_forward(registry, chain_weights, num_timesteps, seq_len): + report = h3_chain.run_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + stages = {stage.name: stage for stage in h3_chain.STAGES} + for entry in report["stages"]: + stage = stages[entry["stage"]] + assert entry["backend"].endswith("CudaOp"), entry + assert entry["repeat_bitwise_equal"], entry["stage"] + assert entry["chained_vs_golden"]["within_tolerance"], entry + if stage.provider_bitwise_isolated: + assert entry["isolated_vs_provider"]["bitwise_equal"], entry["stage"] + drift = report["first_isolated_drift"] + assert drift is None or not stages[drift].provider_bitwise_isolated diff --git a/tests/h3/test_h3_timestep_sinusoid.py b/tests/h3/test_h3_timestep_sinusoid.py new file mode 100644 index 000000000..8575041a3 --- /dev/null +++ b/tests/h3/test_h3_timestep_sinusoid.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_sinusoid_h3``: FP32 [cos | sin] timestep features. + +Three properties, each checked separately: + +* provider parity: the CUDA kernel is bitwise equal to the diffusers + ``get_timestep_embedding`` replay on the same device (an elementwise op + with the same FP32 operation sequence); +* accuracy: within the FP32 elementwise contract of an FP64 golden; +* invariance: a row's bytes do not depend on how many timesteps share the + launch, their order, or repeated runs. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.testing.h3_cases import h3_timesteps +from rl_engine.testing.h3_provider import provider_time_proj + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.timestep_sinusoid import ( + H3TimestepSinusoidCudaOp, + h3_sinusoid_cuda_available, + ) + + if not h3_sinusoid_cuda_available(): + pytest.skip("rl_engine._C lacks h3_timestep_sinusoid_forward") + return H3TimestepSinusoidCudaOp() + + +class TestReference: + def test_layout_is_cos_then_sin(self): + t = torch.tensor([0.0, 1.0]) + out = NativeH3TimestepSinusoidOp().forward(t) + assert out.shape == (2, 256) and out.dtype == torch.float32 + assert torch.equal(out[0, :128], torch.ones(128)) + assert torch.equal(out[0, 128:], torch.zeros(128)) + # k = 0 has frequency 1: cos(1), sin(1) at t = 1. + assert out[1, 0].item() == pytest.approx(math.cos(1.0), abs=1e-7) + assert out[1, 128].item() == pytest.approx(math.sin(1.0), abs=1e-7) + + def test_provider_path_matches_diffusers_cpu(self): + t = h3_timesteps(9, device="cpu") + ours = NativeH3TimestepSinusoidOp().forward(t) + assert torch.equal(ours, provider_time_proj(t, 256)) + + def test_golden_close_to_provider_path(self): + t = h3_timesteps(33, device="cpu") + op = NativeH3TimestepSinusoidOp() + torch.testing.assert_close(op.forward(t), op.forward_fp32(t), atol=1e-6, rtol=0) + + @pytest.mark.parametrize( + "bad, error", + [ + (torch.tensor([[0.5]]), ValueError), # not 1-D + (torch.empty(0), ValueError), # empty + (torch.tensor([1, 0]), TypeError), # integer + (torch.tensor([0.5, 1.5]), ValueError), # H10: t outside [0, 1] + (torch.tensor([500.0]), ValueError), # H10: t * 1000 convention + (torch.tensor([float("nan")]), ValueError), + ], + ) + def test_rejects_invalid_timesteps(self, bad, error): + with pytest.raises(error): + NativeH3TimestepSinusoidOp().forward(bad) + + def test_rejects_odd_channels(self): + with pytest.raises(ValueError): + NativeH3TimestepSinusoidOp().forward(torch.tensor([0.5]), num_channels=255) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("num", [1, 2, 3, 7, 64, 1000, 4097]) + def test_bitwise_equal_to_diffusers_cuda_path(self, num): + t = h3_timesteps(num, seed=num) + out = _cuda_op()(t) + assert out.dtype == torch.float32 and out.shape == (num, 256) + assert torch.equal(out, provider_time_proj(t, 256)) + assert torch.equal(out, NativeH3TimestepSinusoidOp().forward(t)) + + @pytest.mark.parametrize("num_channels", [2, 8, 96, 256, 320]) + def test_other_channel_counts_match_reference(self, num_channels): + t = h3_timesteps(5) + out = _cuda_op()(t, num_channels=num_channels) + ref = NativeH3TimestepSinusoidOp().forward(t, num_channels=num_channels) + assert torch.equal(out, ref) + + def test_fp32_contract_against_fp64_golden(self): + t = h3_timesteps(257) + out = _cuda_op()(t) + gold = NativeH3TimestepSinusoidOp().forward_fp32(t) + # tolerance_contract.json forward_accuracy / elementwise / float32. + torch.testing.assert_close(out, gold, atol=1e-5, rtol=1e-5) + assert (out - gold).abs().max().item() <= 2.0**-23 + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + def test_low_precision_timesteps_are_upcast_first(self, dtype): + t = h3_timesteps(6).to(dtype) + assert torch.equal(_cuda_op()(t), provider_time_proj(t, 256)) + + def test_batch_position_and_repeat_invariance(self): + op = _cuda_op() + base = torch.tensor(0.3141592, device="cuda") + single = op(base.view(1)) + for num in (2, 5, 64, 1023): + for pos in (0, num // 2, num - 1): + t = h3_timesteps(num, seed=num) + t[pos] = base + assert torch.equal(op(t)[pos], single[0]), (num, pos) + t = h3_timesteps(129, seed=11) + first = op(t) + for _ in range(3): + assert torch.equal(op(t), first) + perm = torch.randperm(129, generator=torch.Generator().manual_seed(0)).cuda() + assert torch.equal(op(t[perm]), first[perm]) + + def test_non_contiguous_input(self): + t = h3_timesteps(20) + strided = t[::2] + assert not strided.is_contiguous() + assert torch.equal(_cuda_op()(strided), _cuda_op()(strided.contiguous())) + + def test_rejects_cpu_input(self): + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.5])) + + def test_rejects_out_of_range(self): + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.0, 999.0], device="cuda")) + + def test_backward_matches_autograd_reference(self): + t = h3_timesteps(17).requires_grad_(True) + t_ref = t.detach().clone().requires_grad_(True) + grad = torch.randn(17, 256, device="cuda", generator=torch.Generator("cuda").manual_seed(1)) + _cuda_op()(t).backward(grad) + NativeH3TimestepSinusoidOp().forward_fp32(t_ref).backward(grad) + torch.testing.assert_close(t.grad, t_ref.grad, atol=1e-4, rtol=1e-4) + + def test_backward_is_row_invariant(self): + op = _cuda_op() + grad_row = torch.randn(256, device="cuda", generator=torch.Generator("cuda").manual_seed(2)) + + def grad_of(t, pos): + t = t.clone().requires_grad_(True) + grad = torch.zeros(t.shape[0], 256, device="cuda") + grad[pos] = grad_row + op(t).backward(grad) + return t.grad[pos] + + single = grad_of(torch.tensor([0.7], device="cuda"), 0) + t = h3_timesteps(40) + t[13] = 0.7 + assert torch.equal(grad_of(t, 13), single) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda") + assert type(op).__name__ == "H3TimestepSinusoidCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 2f4ec8648..c688cd3d2 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -49,6 +49,7 @@ def test_all_ws1_single_ops_are_registered(): "pack", "linear_logp", "prefix_shared_attention", + "timestep_sinusoid_h3", } <= names From 23a530e267c9be17d4a27d3ea96db921e84dd04e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:28:41 +0000 Subject: [PATCH 02/24] [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 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 03/24] 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 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 04/24] 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 05/24] 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 06/24] 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 07/24] 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 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 08/24] 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 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 09/24] 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 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 10/24] 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 11/24] 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 12/24] 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 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 13/24] 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 14/24] 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 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 15/24] 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 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 16/24] 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 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 17/24] 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 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 18/24] 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 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 19/24] 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 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 20/24] 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 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 21/24] 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 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 22/24] 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 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 23/24] 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 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 24/24] 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