H3 sigma schedule

.github/workflows/ci.yml

@@ -71,6 +71,10 @@ jobs:
       - name: Check refactor packaging and compatibility
         run: python ci/run.py layout -q
 
+      - name: Run H3 Sigma Schedule Tests (CPU-safe)
+        run: |
+          python -m pytest tests/reference/minimax_h3 tests/runtime/test_h3_sigma_schedule_dispatch.py -q
+
       - name: Run Mocked Hardware Discovery Tests
         run: |
           python -m pytest tests/runtime/test_dispatch.py -v

benchmarks/models/benchmark_h3_sigma_schedule.py

@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""Benchmark the h3_sigma_schedule operator across its backends (RFC #420).
+
+Operator-only microbenchmark: builds sigma grids for N in {25, 50, 1000,
+1000000} x shift in {3.0, 12.0} on three paths:
+
+* path A — CPU reference, grid built on CPU, outputs stay on CPU;
+* path B — CPU reference, grid built on CPU, sigmas and timesteps moved to
+  the target GPU via the operator's ``device`` argument (the registry default
+  path on every platform: build + host-to-device copy);
+* path C — Triton opt-in backend (semantic ``requested_backend`` or direct
+  construction), grid built on the device.
+
+B vs C is the comparison that matters when the consumer lives on the GPU; A
+is an auxiliary CPU baseline and A/C ratios must not be presented as a
+speedup under equal output conditions.
+
+Timing semantics (deliberately conservative): **per-call end-to-end host
+wall latency**. Each call is timed individually with ``time.perf_counter``
+around the whole ``op()`` call — including Python dispatch, host
+allocations, kernel launches, and the device path's single host
+synchronization (the deduplicated-length read) — and, for the GPU paths,
+one ``torch.cuda.synchronize()`` before starting each sample (so prior
+async work cannot leak into the window) and one after the call (so the
+call's device work is complete before the clock stops). Reported per case:
+raw per-call samples, mean, median, standard deviation, min, and max. The
+median is the headline number; mean and raw samples are recorded so
+outliers stay visible. No CUDA events are used, because they cannot see the
+host-side part of the call. Triton kernel compilation happens inside the
+pre-timing correctness gate (the gate's calls are this process's first
+Triton invocations): the gate's total wall time is reported once as
+``correctness_gate_wall_ms`` and no later call is labeled "first" or
+"compile". Timed samples start only after per-case warmup, so steady-state
+numbers exclude compilation by construction. The benchmark calls the op
+classes directly (registry bypassed for timing); registry-entry identity
+and correctness are pinned by tests/runtime/ and tests/backends/cuda/.
+
+Before any timing on a GPU host where the Triton backend is available, path C
+is verified bitwise against path A for every measured case (sigmas,
+timesteps, and num_model_evals); a mismatch aborts with exit code 2 and no
+performance numbers are produced.
+
+Exit codes: 0 on success; 2 on the pre-timing correctness gate; nonzero on
+invalid arguments, and on any Triton compile, launch, or runtime failure
+after the backend was confirmed available. Path B needs only torch and a
+device, so it is measured on every GPU host. Path C (and its gate) is
+skipped only for a *confirmed* absent backend — a missing Triton import on a
+GPU host — or recorded as unsupported on a HIP/ROCm build, where the
+NVIDIA-only kernels (inline PTX) refuse construction. Without any CUDA/HIP
+device (``torch.cuda.is_available()`` false) both B and C are skipped. The
+script runs from an uninstalled checkout (it bootstraps the repository root
+onto ``sys.path``).
+
+The grid is tiny (4 KB at N=1000), so GPU launch overhead may exceed the
+arithmetic — the numbers are what they are, and this benchmark does not
+assume the device path wins. N=1000000 is a dedup stress case, not a
+typical workload. No model, checkpoint, or tokenizer is loaded.
+
+Example:
+
+    python benchmarks/models/benchmark_h3_sigma_schedule.py --samples 100
+    CUDA_VISIBLE_DEVICES=0 python benchmarks/models/benchmark_h3_sigma_schedule.py \
+        --samples 100 --json /tmp/h3_sigma_bench.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import sys
+import time
+from functools import partial
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))  # uninstalled checkout
+
+import torch  # noqa: E402
+
+_SIZES = [25, 50, 1000, 1000000]
+_SHIFTS = [3.0, 12.0]
+_SPEEDUP_NOTE = "speedup_b_over_c = latency(B) / latency(C); > 1 means C (Triton) is faster"
+
+
+def _positive_int(value: str) -> int:
+    try:
+        parsed = int(value)
+    except ValueError as exc:
+        raise argparse.ArgumentTypeError(f"expected an integer, got {value!r}") from exc
+    if parsed < 1:
+        raise argparse.ArgumentTypeError(f"expected a positive integer, got {value!r}")
+    return parsed
+
+
+def _int32_bits(tensor: torch.Tensor) -> list[int]:
+    return tensor.detach().contiguous().view(torch.int32).tolist()
+
+
+def _time_calls(call, samples: int, *, device_sync: bool) -> dict:
+    """Per-call end-to-end host wall latency; raw samples plus summary stats."""
+
+    raw_ms: list[float] = []
+    for _ in range(samples):
+        if device_sync:
+            torch.cuda.synchronize()  # no prior async work leaks into this window
+        start = time.perf_counter()
+        call()
+        if device_sync:
+            torch.cuda.synchronize()  # the call's device work is done before stop
+        raw_ms.append((time.perf_counter() - start) * 1e3)
+    return {
+        "samples_raw_ms": [round(value, 6) for value in raw_ms],
+        "mean_ms": statistics.fmean(raw_ms),
+        "median_ms": statistics.median(raw_ms),
+        "stdev_ms": statistics.stdev(raw_ms) if len(raw_ms) > 1 else 0.0,
+        "min_ms": min(raw_ms),
+        "max_ms": max(raw_ms),
+    }
+
+
+def _reference_op():
+    from rl_engine.reference.minimax_h3.sigma_schedule import NativeH3SigmaScheduleOp
+
+    return NativeH3SigmaScheduleOp()
+
+
+def _time_reference(samples: int, warmup: int, *, device) -> dict[str, dict]:
+    # Path A when device is None (outputs stay on CPU); path B when device is
+    # "cuda" (same CPU build, then sigmas/timesteps moved by the operator).
+    op = _reference_op()
+    results = {}
+    for n in _SIZES:
+        for shift in _SHIFTS:
+            for _ in range(warmup):
+                op(n, shift=shift, device=device)
+            if device is not None:
+                torch.cuda.synchronize()
+            call = partial(op, n, shift=shift, device=device)
+            results[f"N={n},shift={shift:g}"] = _time_calls(
+                call, samples, device_sync=device is not None
+            )
+    return results
+
+
+def _verify_triton_matches_reference() -> float:
+    """Bitwise gate: path C must equal path A before any timing happens.
+
+    This gate's calls are the FIRST Triton invocations of the process — the
+    kernels are compiled here, not later. Returns the gate's wall time in ms
+    (compilation + all case comparisons) so the cold-start cost stays on the
+    record instead of being mislabeled somewhere else.
+    """
+
+    from rl_engine.backends.cuda.model_specific.minimax_h3.sigma_schedule import (
+        H3SigmaScheduleCudaOp,
+    )
+
+    cpu_op = _reference_op()
+    gpu_op = H3SigmaScheduleCudaOp()
+    start = time.perf_counter()
+    for n in _SIZES:
+        for shift in _SHIFTS:
+            expected = cpu_op(n, shift=shift)
+            actual = gpu_op(n, shift=shift)
+            if (
+                _int32_bits(actual.sigmas.cpu()) != _int32_bits(expected.sigmas)
+                or _int32_bits(actual.timesteps.cpu()) != _int32_bits(expected.timesteps)
+                or actual.num_model_evals != expected.num_model_evals
+            ):
+                first_sigma = next(
+                    (
+                        idx
+                        for idx, (a, e) in enumerate(
+                            zip(_int32_bits(actual.sigmas.cpu()), _int32_bits(expected.sigmas))
+                        )
+                        if a != e
+                    ),
+                    None,
+                )
+                print(
+                    f"correctness gate FAILED at N={n}, shift={shift}: Triton result "
+                    f"differs from the CPU reference (first mismatching sigma index: "
+                    f"{first_sigma}); refusing to benchmark"
+                )
+                sys.exit(2)  # documented gate exit code; SystemExit(str) would exit 1
+    return (time.perf_counter() - start) * 1e3
+
+
+def _time_triton(samples: int, warmup: int) -> dict[str, dict]:
+    # The caller ran the correctness gate first, so every kernel has already
+    # been compiled and executed; no "first call" exists here to re-time and
+    # nothing in this function is labeled as compilation. Any failure below
+    # (launch, host read) propagates and fails the command.
+    from rl_engine.backends.cuda.model_specific.minimax_h3.sigma_schedule import (
+        H3SigmaScheduleCudaOp,
+    )
+
+    op = H3SigmaScheduleCudaOp()
+
+    results = {}
+    for n in _SIZES:
+        for shift in _SHIFTS:
+            for _ in range(warmup):
+                op(n, shift=shift)
+            torch.cuda.synchronize()
+            call = partial(op, n, shift=shift)
+            results[f"N={n},shift={shift:g}"] = _time_calls(call, samples, device_sync=True)
+    return results
+
+
+def _summarize(report: dict) -> dict:
+    """Speedup B/C per case from medians and means; A medians for context."""
+
+    summary = {"definition": _SPEEDUP_NOTE, "cases": {}}
+    path_b = report.get("path_b_reference_cpu_build_then_copy_to_gpu")
+    path_c = report.get("path_c_triton_device_build")
+    # A path that was skipped or recorded as unsupported carries no per-case
+    # timings, so there is no ratio to compute (B alone is still reported).
+    if (
+        not path_b
+        or not path_c
+        or "skipped" in path_b
+        or "unsupported" in path_b
+        or "skipped" in path_c
+        or "unsupported" in path_c
+    ):
+        return summary
+    for key, case_b in path_b.items():
+        case_c = path_c[key]
+        summary["cases"][key] = {
+            "b_median_ms": case_b["median_ms"],
+            "c_median_ms": case_c["median_ms"],
+            "speedup_b_over_c_median": case_b["median_ms"] / case_c["median_ms"],
+            "speedup_b_over_c_mean": case_b["mean_ms"] / case_c["mean_ms"],
+            "separated_from_dispersion": abs(case_b["median_ms"] - case_c["median_ms"])
+            > (case_b["stdev_ms"] + case_c["stdev_ms"]),
+        }
+    return summary
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--samples", type=_positive_int, default=100)
+    parser.add_argument("--warmup", type=_positive_int, default=5)
+    parser.add_argument("--json", default=None, help="optional output path")
+    args = parser.parse_args(argv)
+
+    has_gpu = torch.cuda.is_available()
+    report: dict = {
+        "operator": "h3_sigma_schedule",
+        "timing": (
+            "per-call end-to-end host wall latency; GPU paths sync before each "
+            "sample and after each call; Triton first-call compile excluded"
+        ),
+        "entry": "direct op construction; registry bypassed for timing (identity pinned by tests)",
+        "torch": torch.__version__,
+        "torch_cuda_runtime": torch.version.cuda,
+        "hip": torch.version.hip,
+        "triton": None,
+        "device": torch.cuda.get_device_name(0) if has_gpu else "cpu-only",
+        "n_gpu_devices": torch.cuda.device_count() if has_gpu else 0,
+        "samples": args.samples,
+        "warmup": args.warmup,
+        "sizes": _SIZES,
+        "shifts": _SHIFTS,
+        "stress_case_note": "N=1000000 is a dedup stress case, not a typical H3 workload",
+    }
+
+    print(
+        f"path A (reference, cpu build, cpu outputs): {args.samples} calls x {_SIZES} x {_SHIFTS}"
+    )
+    report["path_a_reference_cpu"] = _time_reference(args.samples, args.warmup, device=None)
+
+    if not has_gpu:
+        skip = "no CUDA/HIP device (torch.cuda.is_available() is False)"
+        report["path_b_reference_cpu_build_then_copy_to_gpu"] = {"skipped": skip}
+        report["path_c_triton_device_build"] = {"skipped": skip}
+        print("paths B and C: skipped — no CUDA/HIP device in this process")
+    else:
+        try:
+            import triton
+
+            report["triton"] = triton.__version__
+        except ImportError as exc:
+            # Path B is the registry default path and needs only torch plus a
+            # device, so a missing Triton skips path C alone.
+            skip = f"triton import failed on a GPU host: {exc}"
+            print(f"path B (reference, cpu build, outputs moved to gpu): {args.samples} calls")
+            report["path_b_reference_cpu_build_then_copy_to_gpu"] = _time_reference(
+                args.samples, args.warmup, device="cuda"
+            )
+            report["path_c_triton_device_build"] = {"skipped": skip}
+            print(f"path C: skipped — triton missing ({exc})")
+        else:
+            if torch.version.hip is not None:
+                # Triton may import on a HIP host, but the backend is NVIDIA-only
+                # (inline PTX) and refuses construction: record that instead of
+                # letting the ImportError escape from inside the correctness gate.
+                unsupported = (
+                    "HIP/ROCm build: the Triton backend is NVIDIA-only (inline PTX "
+                    "div.rn.f32) and refuses construction; the bitwise correctness "
+                    "gate was not run"
+                )
+                print(f"path B (reference, cpu build, outputs moved to gpu): {args.samples} calls")
+                report["path_b_reference_cpu_build_then_copy_to_gpu"] = _time_reference(
+                    args.samples, args.warmup, device="cuda"
+                )
+                report["path_c_triton_device_build"] = {"unsupported": unsupported}
+                print("path C: unsupported — NVIDIA-only backend on a HIP/ROCm build")
+            else:
+                print("verifying path C against path A bitwise before timing...")
+                gate_ms = _verify_triton_matches_reference()
+                report["correctness_gate_wall_ms"] = gate_ms
+                print(
+                    "correctness gate passed "
+                    f"(includes Triton kernel compilation: {gate_ms:.1f} ms wall, "
+                    "reported once, excluded from all timed samples)"
+                )
+
+                print(f"path B (reference, cpu build, outputs moved to gpu): {args.samples} calls")
+                report["path_b_reference_cpu_build_then_copy_to_gpu"] = _time_reference(
+                    args.samples, args.warmup, device="cuda"
+                )
+                print(f"path C (triton, device build): {args.samples} calls")
+                report["path_c_triton_device_build"] = _time_triton(args.samples, args.warmup)
+
+    report["summary"] = _summarize(report)
+
+    for section in (
+        "path_a_reference_cpu",
+        "path_b_reference_cpu_build_then_copy_to_gpu",
+        "path_c_triton_device_build",
+    ):
+        for key, value in report[section].items():
+            if isinstance(value, dict) and "median_ms" in value:
+                print(f"{section:>44s} {key:>20s}: {value['median_ms']:.4f} ms/call (median)")
+    for key, case in report["summary"]["cases"].items():
+        separated = case["separated_from_dispersion"]
+        note = "separated from dispersion" if separated else "within dispersion"
+        print(f"speedup B/C {key:>20s}: {case['speedup_b_over_c_median']:.2f}x ({note})")
+
+    if args.json:
+        with open(args.json, "w", encoding="utf-8") as handle:
+            json.dump(report, handle, indent=2)
+        print(f"wrote {args.json}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

docs/.nav.yml

@@ -32,6 +32,7 @@ nav:
       - operators/det-gemm.md
       - operators/embedding.md
       - MiniMax-H3:
+        - operators/h3-sigma-schedule.md
         - operators/h3-timestep-sinusoid.md
         - operators/h3-timestep-mlp.md
         - operators/h3-adaln-projection.md

docs/architecture/runtime-dispatch.md

@@ -37,6 +37,15 @@ the legacy SM90 LogP backend when `RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP=1` is
 set. The fused linear logp SM90 backend is gated separately and remains the
 default linear logp backend when the extension is built on Hopper.
 
+## H3 Sigma Schedule
+
+`sigma_schedule_h3` lists only the PyTorch reference on every platform, CUDA
+included; the NVIDIA Triton backend is opt-in through the semantic
+`requested_backend="rlkernel.h3_sigma_schedule.triton"` or direct construction,
+because its device-side build measured 2.5–3.2x slower than CPU construction
+plus transfer at typical H3 sizes (PR #566 evidence). See
+[H3 Sigma Schedule](../operators/h3-sigma-schedule.md).
+
 ## Relevant Files
 
 - `rl_engine/runtime/registry.py`