[WS1][CUDA][MiniMax-H3] timestep_sinusoid_h3: FP32 timestep features + shared H3 harness (RFC #420 stack 1/10)

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()

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 = [

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
 }

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 <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+
+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<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < total;
+       idx += static_cast<int64_t>(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<float>(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<bool>(),
+                      "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<float>(-std::log(max_period));
+  const float inv_half = 1.0f / static_cast<float>(half);
+
+  const int threads = 256;
+  const int64_t total = num_timesteps * half;
+  const int64_t blocks = std::min<int64_t>((total + threads - 1) / threads, 65535);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  h3_timestep_sinusoid_kernel<<<static_cast<unsigned int>(blocks), threads, 0, stream>>>(
+      t.data_ptr<float>(), out.data_ptr<float>(), num_timesteps, half, neg_log_max_period,
+      inv_half);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return out;
+}