diff --git a/benchmarks/models/benchmark_h3_conditioning.py b/benchmarks/models/benchmark_h3_conditioning.py new file mode 100644 index 000000000..b739f34ab --- /dev/null +++ b/benchmarks/models/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/validation/models/h3_report.py``. + + python benchmarks/models/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 + python benchmarks/models/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[2])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models.h3_chain import environment # noqa: E402 +from rl_engine.validation.models.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/build_tools/extensions.py b/build_tools/extensions.py index 670fdb072..098779c80 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -3,6 +3,8 @@ import importlib.util import os +import shlex +import stat import sysconfig import warnings from distutils.errors import CompileError @@ -14,6 +16,158 @@ _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" +# 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. 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"}, +} +# 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 +# spelling is read as well, so a driver that does expand it is covered. +_OPTIONS_FILE_FLAGS = {"optf", "options-file"} +_OPTIONS_FILE_DEPTH = 4 + + +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. + + 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 = _nvcc_flag_words(text) + tokens: list[str] = [] + 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) + 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: + tokens.extend(_nvcc_flag_tokens(_read_options_file(name), depth=depth + 1)) + return tokens + + +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 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 f"{token} without a value" + value = nxt + 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 + + +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 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: + 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 _resolve_package_data_source( filename: str | os.PathLike[str], @@ -219,6 +373,20 @@ 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_reason = _relaxed_math_reason() + relaxed_math_env = relaxed_reason is not None + if relaxed_math_env and not is_rocm: + warnings.warn( + 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, + ) + 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", @@ -241,12 +409,14 @@ 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 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. 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() @@ -333,6 +503,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 h3_enabled: + 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 0c7c3e7f4..265937e10 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -500,6 +500,13 @@ at::Tensor prefix_shared_attention( #endif #endif +// MiniMax-H3 (RFC #420) conditioning-path declarations. CUDA only. +// 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 + // PyBind11 Module Registration PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "RL-Kernel High-Performance Operator Extension Library"; @@ -758,4 +765,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Deterministic GPT-NeoX token-major RoPE apply for ROCm"); #endif #endif + +#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, + py::arg("check_range") = true); +#endif } diff --git a/csrc/cuda/h3/timestep_sinusoid.cu b/csrc/cuda/h3/timestep_sinusoid.cu new file mode 100644 index 000000000..c618a0ad3 --- /dev/null +++ b/csrc/cuda/h3/timestep_sinusoid.cu @@ -0,0 +1,94 @@ +// 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. +// 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 +#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, 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()); + 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(); + // 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()); + + // 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/docs/.nav.yml b/docs/.nav.yml index 1037edce5..b82106501 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -30,6 +30,8 @@ nav: - operators/sampling.md - operators/det-gemm.md - operators/embedding.md + - MiniMax-H3: + - operators/h3-timestep-sinusoid.md - Developer Guide: - contributing/README.md - Contributor Guide: contributing/contributor-guide.md diff --git a/docs/operators/README.md b/docs/operators/README.md index 184d1d0e4..18b46f229 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -32,5 +32,6 @@ Every operator page should include: - [Matmul](matmul.md) - [Sampling](sampling.md) - [Token Embedding](embedding.md) +- [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) - [H3 Flash Attention (optional fast backend)](h3-flash-attention.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..03b9b489d --- /dev/null +++ b/docs/operators/h3-timestep-sinusoid.md @@ -0,0 +1,166 @@ +# 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/validation/models/h3_manifest.json`). +Provider reference: `huggingface/diffusers@f53d552`, `get_timestep_embedding`. + +## Entry Point + +```python +from rl_engine.runtime.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.backends.cuda.model_specific.minimax_h3.timestep_sinusoid.H3TimestepSinusoidCudaOp` | `rl_engine._C.h3_timestep_sinusoid_forward` | Bitwise equal to the provider path | +| PyTorch reference | `rl_engine.reference.minimax_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. +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 + +- **Operation order.** The kernel evaluates exactly the FP32 operations the provider + 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. +- **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 `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): + +| 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/models/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 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. + +## Evidence + +![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`](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](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 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 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 `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 +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/timestep_sinusoid.json +``` + +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](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`](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` (`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 \ + --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 +``` + +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 +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_timestep_sinusoid.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end, pinned weights +python tools/validation/operators/check_operator.py --op timestep_sinusoid_h3 --candidate cuda --device cuda \ + --dtype fp32 --batch 7 --check-grad +python tools/validation/models/h3_evidence.py --op timestep_sinusoid_h3 \ + --out reports/experiments/h3-timestep-sinusoid-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/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. +- `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. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fb3c4379c..3371735f8 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -325,3 +325,16 @@ 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, + max_period: float = 10000.0, + check_range: bool = True, +) -> torch.Tensor: ... diff --git a/rl_engine/backends/cuda/model_specific/__init__.py b/rl_engine/backends/cuda/model_specific/__init__.py new file mode 100644 index 000000000..3b09d3c11 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA implementations with model-specific arithmetic contracts.""" diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py b/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py new file mode 100644 index 000000000..86cf4c9da --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors 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 new file mode 100644 index 000000000..86b838403 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py @@ -0,0 +1,115 @@ +# 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.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 ( + 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, check_range: bool): + t32 = timestep.detach().float().contiguous() + 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 + return out + + @staticmethod + 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() + # 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 + ) + # 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", + kernel_id="rl_engine.reference.minimax_h3.fixed_order.tree_sum_lastdim_fp32", + impl=BACKWARD_IMPL, + family="cuda", + ) + return grad_t, None, 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 (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): + 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: + # 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, + 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/__init__.py b/rl_engine/reference/minimax_h3/__init__.py new file mode 100644 index 000000000..e54c60b77 --- /dev/null +++ b/rl_engine/reference/minimax_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/validation/models/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_CHUNK_NAMES = ( + "shift_msa", + "scale_msa", + "gate_msa", + "shift_mlp", + "scale_mlp", + "gate_mlp", +) +H3_ADALN_CHUNKS = len(H3_ADALN_CHUNK_NAMES) diff --git a/rl_engine/reference/minimax_h3/fixed_order.py b/rl_engine/reference/minimax_h3/fixed_order.py new file mode 100644 index 000000000..e95a66054 --- /dev/null +++ b/rl_engine/reference/minimax_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/reference/minimax_h3/timestep_sinusoid.py b/rl_engine/reference/minimax_h3/timestep_sinusoid.py new file mode 100644 index 000000000..2f374c81a --- /dev/null +++ b/rl_engine/reference/minimax_h3/timestep_sinusoid.py @@ -0,0 +1,123 @@ +# 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 / 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, 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. +""" + +from __future__ import annotations + +import math + +import torch + +from rl_engine.reference.minimax_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 (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): + 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. + + 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" + + 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.""" + + 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 + ) + return torch.exp(exponent / half) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 528c77b98..5dc3eb558 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -204,6 +204,15 @@ 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.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ) + CUDA_H3_TIMESTEP_SINUSOID = ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_sinusoid.H3TimestepSinusoidCudaOp" + ) + def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: return ( @@ -673,6 +682,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": [ @@ -714,6 +727,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": [ @@ -763,6 +777,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], @@ -789,6 +804,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/validation/models/h3_cases.py b/rl_engine/validation/models/h3_cases.py new file mode 100644 index 000000000..527a820ec --- /dev/null +++ b/rl_engine/validation/models/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/validation/models/h3_chain.py b/rl_engine/validation/models/h3_chain.py new file mode 100644 index 000000000..31f9718b6 --- /dev/null +++ b/rl_engine/validation/models/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.validation.models.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.runtime.registry import KernelRegistry +from rl_engine.validation.models.h3_cases import h3_packed_layout, h3_timesteps +from rl_engine.validation.models.h3_provider import provider_time_proj + +REPO_ROOT = Path(__file__).resolve().parents[3] + + +@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/models/minimax_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/validation/models/h3_manifest.json b/rl_engine/validation/models/h3_manifest.json new file mode 100644 index 000000000..52b1989b1 --- /dev/null +++ b/rl_engine/validation/models/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/validation/models/h3_provider.py b/rl_engine/validation/models/h3_provider.py new file mode 100644 index 000000000..c2a738f70 --- /dev/null +++ b/rl_engine/validation/models/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.reference.minimax_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/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py new file mode 100644 index 000000000..3d640f9e2 --- /dev/null +++ b/rl_engine/validation/models/h3_report.py @@ -0,0 +1,169 @@ +# 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/models/benchmark_h3_conditioning.py`` +prints and ``tools/validation/models/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.runtime.registry import KernelRegistry +from rl_engine.validation.models.h3_cases import h3_timesteps +from rl_engine.validation.models.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 _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): + _prepared_call(fn, setup)() + torch.cuda.synchronize() + 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() + run() + torch.cuda.synchronize() + return (torch.cuda.max_memory_allocated() - base) / 2**20 + + +def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + """Interleave timed keys, reversing their execution order on alternate iterations. + + Record raw CUDA-event samples, executed orders, median microseconds, + decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. + """ + + row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} + keys = [key for key in TIMED_KEYS if key in case] + if not keys: + return row + orders = (keys, list(reversed(keys))) + row["execution_order"] = { + "policy": "alternating", + "iteration_0": orders[0], + "iteration_1": orders[1], + } + if any(key.endswith("_backward") for key in keys): + row["backward_timing_scope"] = "backward_only" + for iteration in range(warmup): + for key in orders[iteration % 2]: + _prepared_call(case[key], case.get(f"{key}_setup"))() + torch.cuda.synchronize() + samples = {key: [] for key in keys} + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) + row["timing_samples_us"] = samples + 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], setup=case.get(f"{key}_setup")) + 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/validation/models/h3_weights.py b/rl_engine/validation/models/h3_weights.py new file mode 100644 index 000000000..261b4385e --- /dev/null +++ b/rl_engine/validation/models/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 ``tools/weights/prepare_h3_weights.py`` downloads that shard, +checks it against the manifest and extracts the tensors into +``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. +""" + +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 tools/weights/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/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index aaebc279f..a0397b96b 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/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( @@ -44,6 +45,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) @@ -78,6 +80,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] @@ -370,6 +373,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/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 92995b938..1c9d5ad71 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -272,6 +272,25 @@ 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.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_sinusoid.H3TimestepSinusoidCudaOp" + ), + }, + grad_input_names=("timestep",), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index b813506b9..dfbcb29b5 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -3,14 +3,41 @@ from __future__ import annotations +import os +import re import runpy +import shutil +import subprocess +from pathlib import Path from typing import Any +import pytest import setuptools import torch from torch.utils import cpp_extension -from build_tools.extensions import _resolve_package_data_source, get_cmdclass +from build_tools.extensions import ( + _OPTIONS_FILE_MAX_BYTES, + H3_BITWISE_SOURCES, + H3_KERNELS_DEFINE, + _relaxed_math_reason, + _resolve_package_data_source, + _tokens_relax_math, + get_cmdclass, +) + +_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) + for name in _NVCC_ENV: + monkeypatch.delenv(name, raising=False) def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: @@ -83,3 +110,346 @@ 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() + + +def test_cuda_build_compiles_h3_kernels_without_fast_math(monkeypatch) -> None: + _clear_nvcc_env(monkeypatch) + extension = _load_extension_config(monkeypatch, hip=None) + + 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_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. + _clear_nvcc_env(monkeypatch) + 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 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"] + + +@pytest.mark.parametrize("variable", _NVCC_ENV) +@pytest.mark.parametrize( + "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, 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. + _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 + + monkeypatch.setenv(variable, flags) + 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"] + + +@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=FALSE", + "--prec-div=true", + "--prec-div TRUE -O2", + "-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() + + +_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", + # 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", +} + + +@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), + ('--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( + 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() + + +@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") + 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") + + +_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", + [ + "", + "-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", + '--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"]) +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 + 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: + # 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") + 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"] diff --git a/tests/models/minimax_h3/conftest.py b/tests/models/minimax_h3/conftest.py new file mode 100644 index 000000000..c1c9296a1 --- /dev/null +++ b/tests/models/minimax_h3/conftest.py @@ -0,0 +1,41 @@ +# 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 ``tools/weights/prepare_h3_weights.py``) and skip when it is unset. +""" + +from __future__ import annotations + +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, + 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 tools/weights/prepare_h3_weights.py") + try: + 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 new file mode 100644 index 000000000..104090d10 --- /dev/null +++ b/tests/models/minimax_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.validation.models.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.runtime.registry import KernelRegistry +from rl_engine.validation.models 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(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 + ) + 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/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py new file mode 100644 index 000000000..b79973eb1 --- /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, **kwargs: 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 new file mode 100644 index 000000000..853680a67 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""``measure`` on a real CUDA device with real kernels. + +``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 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +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, + "candidate": lambda: small @ small, + "provider": lambda: large @ large, + } + + +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"] diff --git a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py new file mode 100644 index 000000000..6adbdd208 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py @@ -0,0 +1,275 @@ +# 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.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.validation.models.h3_cases import h3_timesteps +from rl_engine.validation.models.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.backends.cuda.model_specific.minimax_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() + + +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]) + 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) + + @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: + @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)) + + 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) + 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")) + + @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.backends.extension 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.backends.extension 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) + 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) + + @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)) + + 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.runtime.registry import KernelRegistry + + _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 diff --git a/tests/models/minimax_h3/test_prepare_h3_weights.py b/tests/models/minimax_h3/test_prepare_h3_weights.py new file mode 100644 index 000000000..0c470e8d2 --- /dev/null +++ b/tests/models/minimax_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.validation.models.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[3] / "tools" / "weights" / "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) diff --git a/tests/runtime/test_dispatch.py b/tests/runtime/test_dispatch.py index 3c47bbbcf..6dfb5a6ce 100644 --- a/tests/runtime/test_dispatch.py +++ b/tests/runtime/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 @@ -231,6 +232,42 @@ 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_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 + + # 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) + + 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): class BrokenNPU: @staticmethod diff --git a/tests/validation/operators/test_ws1_gtest_gpu.py b/tests/validation/operators/test_ws1_gtest_gpu.py index bcaebc2f5..3c5fa33a4 100644 --- a/tests/validation/operators/test_ws1_gtest_gpu.py +++ b/tests/validation/operators/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 diff --git a/tools/migration/README.md b/tools/migration/README.md index 4a6cee300..48491e383 100644 --- a/tools/migration/README.md +++ b/tools/migration/README.md @@ -8,3 +8,6 @@ Historical workload/operator identifiers and archived evidence intentionally retain their original values. New source imports should use the mapped canonical module. Remove compatibility paths only in a separately announced migration once downstream launchers, patches, imports and saved objects have migrated. + +The MiniMax-H3 entries also record the open RFC #420 PR stack migration to +`test-h3`; those H3 files were not part of the original main baseline. diff --git a/tools/migration/layout.json b/tools/migration/layout.json index cb8fe469a..784cb93ba 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -347,6 +347,7 @@ "benchmarks/benchmark_deterministic_attention.py": "benchmarks/operators/attention/benchmark_deterministic_attention.py", "benchmarks/benchmark_grpo_loss.py": "benchmarks/operators/loss/benchmark_grpo_loss.py", "benchmarks/benchmark_grpo_op.py": "benchmarks/operators/loss/benchmark_grpo_op.py", + "benchmarks/benchmark_h3_conditioning.py": "benchmarks/models/benchmark_h3_conditioning.py", "benchmarks/benchmark_linear_logp.py": "benchmarks/operators/logprob/benchmark_linear_logp.py", "benchmarks/benchmark_pack.py": "benchmarks/operators/packing/benchmark_pack.py", "benchmarks/benchmark_qwen_ffn_layout.py": "benchmarks/layers/benchmark_qwen_ffn_layout.py", @@ -466,6 +467,10 @@ "docs/design/ws2-attention-single-gpu-harness.md": "docs/archive/design/ws2-attention-single-gpu-harness.md", "docs/design/ws2-cp-attention-contract.md": "docs/contracts/ws2-cp-attention-contract.md", "docs/design/ws2-module-debug-matrix.md": "docs/archive/design/ws2-module-debug-matrix.md", + "docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.json": "reports/experiments/h3-prior-art-b200/timestep_sinusoid.json", + "docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.png": "reports/experiments/h3-prior-art-b200/timestep_sinusoid.png", + "docs/usage/evidence/h3-timestep-sinusoid-b200/figure.png": "reports/experiments/h3-timestep-sinusoid-b200/figure.png", + "docs/usage/evidence/h3-timestep-sinusoid-b200/report.json": "reports/experiments/h3-timestep-sinusoid-b200/report.json", "envs.py": "build_tools/envs.py", "examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json": "configs/experiments/cross_config/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json", "examples/cross_config_s0_cpu_smoke.json": "configs/experiments/cross_config/cross_config_s0_cpu_smoke.json", @@ -675,6 +680,8 @@ "rl_engine/kernels/ops/cuda/attention/prefix_shared_attn.py": "rl_engine/backends/cuda/attention/prefix_shared_attn.py", "rl_engine/kernels/ops/cuda/attention/strict_runtime.py": "rl_engine/backends/cuda/attention/strict_runtime.py", "rl_engine/kernels/ops/cuda/ffn.py": "rl_engine/backends/cuda/ffn/ffn.py", + "rl_engine/kernels/ops/cuda/h3/__init__.py": "rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py", + "rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py": "rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py", "rl_engine/kernels/ops/cuda/linear/__init__.py": "rl_engine/backends/cuda/embedding/__init__.py", "rl_engine/kernels/ops/cuda/linear/embedding.py": "rl_engine/backends/cuda/embedding/embedding.py", "rl_engine/kernels/ops/cuda/linear/lm_head.py": "rl_engine/backends/cuda/embedding/lm_head.py", @@ -705,6 +712,9 @@ "rl_engine/kernels/ops/pytorch/attention/stateful_kv.py": "rl_engine/reference/attention/stateful_kv.py", "rl_engine/kernels/ops/pytorch/ffn/__init__.py": "rl_engine/ops/ffn/__init__.py", "rl_engine/kernels/ops/pytorch/ffn/ffn.py": "rl_engine/ops/ffn/qwen3.py", + "rl_engine/kernels/ops/pytorch/h3/__init__.py": "rl_engine/reference/minimax_h3/__init__.py", + "rl_engine/kernels/ops/pytorch/h3/fixed_order.py": "rl_engine/reference/minimax_h3/fixed_order.py", + "rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py": "rl_engine/reference/minimax_h3/timestep_sinusoid.py", "rl_engine/kernels/ops/pytorch/linear/__init__.py": "rl_engine/reference/embedding/__init__.py", "rl_engine/kernels/ops/pytorch/linear/embedding.py": "rl_engine/reference/embedding/embedding.py", "rl_engine/kernels/ops/pytorch/linear/lm_head.py": "rl_engine/reference/embedding/lm_head.py", @@ -776,6 +786,12 @@ "rl_engine/testing/__init__.py": "rl_engine/validation/common/__init__.py", "rl_engine/testing/attention_comparison.py": "rl_engine/validation/common/attention_comparison.py", "rl_engine/testing/distributed_logprob_comparison.py": "rl_engine/validation/common/distributed_logprob_comparison.py", + "rl_engine/testing/h3_cases.py": "rl_engine/validation/models/h3_cases.py", + "rl_engine/testing/h3_chain.py": "rl_engine/validation/models/h3_chain.py", + "rl_engine/testing/h3_manifest.json": "rl_engine/validation/models/h3_manifest.json", + "rl_engine/testing/h3_provider.py": "rl_engine/validation/models/h3_provider.py", + "rl_engine/testing/h3_report.py": "rl_engine/validation/models/h3_report.py", + "rl_engine/testing/h3_weights.py": "rl_engine/validation/models/h3_weights.py", "rl_engine/testing/logprob_comparison.py": "rl_engine/validation/common/logprob_comparison.py", "rl_engine/testing/logprob_drift.py": "rl_engine/validation/common/logprob_drift.py", "rl_engine/testing/reference_ops.py": "rl_engine/reference/definitions.py", @@ -790,6 +806,12 @@ "scripts/check_rocm_env.py": "tools/env/check_rocm_env.py", "scripts/check_stateful_kv.py": "tools/validation/operators/check_stateful_kv.py", "scripts/ci_smoke.py": "tools/checks/ci_smoke.py", + "scripts/h3_chain_replay.py": "tools/validation/models/h3_chain_replay.py", + "scripts/h3_evidence.py": "tools/validation/models/h3_evidence.py", + "scripts/h3_prior_art.py": "tools/validation/models/h3_prior_art.py", + "scripts/plot_h3_evidence.py": "tools/validation/models/plot_h3_evidence.py", + "scripts/plot_h3_prior_art.py": "tools/validation/models/plot_h3_prior_art.py", + "scripts/prepare_h3_weights.py": "tools/weights/prepare_h3_weights.py", "scripts/prepare_ws1_weights.py": "tools/weights/prepare_ws1_weights.py", "scripts/run_perf.py": "tools/benchmarking/run_perf.py", "scripts/run_profile_suite.py": "tools/benchmarking/run_profile_suite.py", @@ -820,6 +842,10 @@ "tests/distributed/test_rocm_attention_transport.py": "tests/distributed/cp/test_rocm_attention_transport.py", "tests/distributed/test_rocm_strict_attention_cp.py": "tests/distributed/cp/test_rocm_strict_attention_cp.py", "tests/distributed/test_transport_deterministic_collective.py": "tests/distributed/collectives/test_transport_deterministic_collective.py", + "tests/h3/conftest.py": "tests/models/minimax_h3/conftest.py", + "tests/h3/test_h3_conditioning_e2e.py": "tests/models/minimax_h3/test_h3_conditioning_e2e.py", + "tests/h3/test_h3_timestep_sinusoid.py": "tests/models/minimax_h3/test_h3_timestep_sinusoid.py", + "tests/h3/test_prepare_h3_weights.py": "tests/models/minimax_h3/test_prepare_h3_weights.py", "tests/linear_logp_tp.py": "tests/distributed/tp/linear_logp_tp.py", "tests/test_alignment_model_wrappers.py": "tests/models/qwen3/test_alignment_model_wrappers.py", "tests/test_alignment_wrapper_interfaces.py": "tests/models/qwen3/test_alignment_wrapper_interfaces.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/h3_chain_replay.py b/tools/validation/models/h3_chain_replay.py new file mode 100644 index 000000000..217db85e1 --- /dev/null +++ b/tools/validation/models/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/validation/models/h3_chain.py`` for what is compared. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/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[3])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models import h3_chain # noqa: E402 +from rl_engine.validation.models.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/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py new file mode 100644 index 000000000..b4852d61f --- /dev/null +++ b/tools/validation/models/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. ``tools/validation/models/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_evidence.py --op timestep_sinusoid_h3 \\ + --out reports/experiments/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[3])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models.h3_chain import environment, git_state # noqa: E402 +from rl_engine.validation.models.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 +from rl_engine.validation.models.h3_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/tools/validation/models/h3_prior_art.py b/tools/validation/models/h3_prior_art.py new file mode 100644 index 000000000..9012542e2 --- /dev/null +++ b/tools/validation/models/h3_prior_art.py @@ -0,0 +1,946 @@ +#!/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 tools/validation/models/h3_prior_art.py --op adaln_row_gather \\ + --out reports/experiments/h3-prior-art-b200/adaln_row_gather.json \\ + [--megatron-src /path/to/Megatron-LM] [--quick] + python tools/validation/models/plot_h3_prior_art.py \\ + reports/experiments/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 +from pathlib import Path +from typing import Any, Callable + +REPO_ROOT = Path(__file__).resolve().parents[3] +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.backends.cuda.model_specific.minimax_h3.timestep_sinusoid", + "timestep_mlp": "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp", + "adaln_projection": "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection", + "adaln_row_gather": "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather", + "norm_modulate": "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm", + "gate_residual": "rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual", + "final_adaln_out": "rl_engine.backends.cuda.model_specific.minimax_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.validation.models.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.validation.models 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.backends.cuda.model_specific.minimax_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.backends.cuda.model_specific.minimax_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.backends.cuda.model_specific.minimax_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.backends.cuda.model_specific.minimax_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.backends.cuda.model_specific.minimax_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.backends.cuda.model_specific.minimax_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.backends.cuda.model_specific.minimax_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.backends.cuda.model_specific.minimax_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.validation.models.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) + # 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 = ( + 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/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() diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py new file mode 100644 index 000000000..74107ec0d --- /dev/null +++ b/tools/validation/models/plot_h3_evidence.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render an RFC #420 operator report (``tools/validation/models/h3_evidence.py``) as a PNG. + +Needs only the report JSON and matplotlib (not torch or a GPU), so it can run +anywhere: + + python tools/validation/models/plot_h3_evidence.py \\ + reports/experiments/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, strict=True): + 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, strict=True): + 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/tools/validation/models/plot_h3_prior_art.py b/tools/validation/models/plot_h3_prior_art.py new file mode 100644 index 000000000..893612e87 --- /dev/null +++ b/tools/validation/models/plot_h3_prior_art.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plot H3 prior-art latency, accuracy, and batch invariance. + + python tools/validation/models/plot_h3_prior_art.py \\ + reports/experiments/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 +from matplotlib import transforms # 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") + 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( + 0.98, + y, + 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_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, loc="lower right") + ax.grid(True, axis="x", alpha=0.3) + + ax = axes[2] + for y, (_, e) in zip(ys, entries): + bi = e["batch_invariance"] + rep = bi["params_repeatable"] + 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( + 0.02, + y, + " | ".join(lines), + va="center", + fontsize=8, + color="#3aa676" if bi["batch_invariant"] else "#d14b4b", + ) + ax.set_ylim(len(entries) - 0.5, -0.5) + ax.set_xlim(0, 1) + ax.axis("off") + ax.set_title("batch invariance (bitwise)") + 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", + ) + + out = path.with_suffix(".png") + fig.savefig(out, dpi=120) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py new file mode 100644 index 000000000..c21f56c2d --- /dev/null +++ b/tools/weights/prepare_h3_weights.py @@ -0,0 +1,101 @@ +#!/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/validation/models/h3_manifest.json`` and writes the +tensors to ``/h3_conditioning.safetensors``. Point the tests at it with +``export RL_KERNEL_H3_WEIGHTS=``. + + python tools/weights/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[2] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + EXTRACTED_FILE, + load_h3_manifest, + sha256_file, +) + + +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] + 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"]}) + print(f"wrote {len(tensors)} tensors to {target}") + print(f"export RL_KERNEL_H3_WEIGHTS={out_dir}") + + +if __name__ == "__main__": + main()