[WS1][CUDA][MiniMax-H3] ws1_one_h3_block: one real block forward/backward with first-drift report (RFC #420 stack 10/10)

.coderabbit.yaml

@@ -0,0 +1,6 @@
+# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json
+inheritance: true
+reviews:
+  auto_review:
+    base_branches:
+      - "^test-h3$"

benchmarks/models/benchmark_h3_conditioning.py

@@ -0,0 +1,71 @@
+# 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. Timings alternate candidate/provider execution order;
+backward timings exclude forward setup. 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:
+    """Benchmark selected CUDA operators, print summaries, and optionally save JSON."""
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--op", default="all", choices=["all", *PERF_CASES])
+    parser.add_argument("--warmup", type=int, default=20)
+    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']}")
+    print("timing order alternates each iteration; backward timings exclude forward setup")
+    results = []
+    for name in list(PERF_CASES) if args.op == "all" else [args.op]:
+        for case in PERF_CASES[name](registry):
+            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
+            ]
+            order = row["execution_order"]
+            phases = [" -> ".join(order[f"iteration_{i}"]) for i in (0, 1)]
+            print(
+                f"{name} {row['case']} [{row['backend']}]: "
+                + "; ".join(parts)
+                + "; alternating order: "
+                + " / ".join(phases)
+            )
+    if args.json is not None:
+        args.json.write_text(json.dumps({"environment": env, "results": results}, indent=2) + "\n")
+
+
+if __name__ == "__main__":
+    main()

build_tools/extensions.py

@@ -236,6 +236,11 @@ def get_extensions():
             # CUDA IPC and the fixed-tree collective implementation are not
             # part of the ROCm extension.
             cuda_sources.append("csrc/cuda/collectives/deterministic_collective.cu")
+            cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu")
+            cuda_sources.append("csrc/cuda/h3/det_linear.cu")
+            cuda_sources.append("csrc/cuda/h3/adaln_row_gather.cu")
+            cuda_sources.append("csrc/cuda/h3/rmsnorm_modulate.cu")
+            cuda_sources.append("csrc/cuda/h3/gate_residual.cu")
             # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync).
             # The ROCm dispatcher falls back to PyTorch SDPA for this operator.
             cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu")

csrc/bindings/ops.cpp

@@ -500,6 +500,67 @@ at::Tensor prefix_shared_attention(
 #endif
 #endif
 
+// MiniMax-H3 (RFC #420) conditioning-path declarations. CUDA only.
+#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \
+    (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA))
+torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels,
+                                           double max_period, bool check_range);
+std::vector<torch::Tensor> h3_det_linear_forward(torch::Tensor x, torch::Tensor weight,
+                                                 c10::optional<torch::Tensor> bias,
+                                                 int64_t activation, bool save_pre_activation);
+torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight,
+                                           c10::ScalarType out_dtype);
+torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight);
+torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype);
+std::vector<torch::Tensor> h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x,
+                                                         c10::ScalarType w_dtype,
+                                                         bool with_bias);
+torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor timestep_indices,
+                                          torch::Tensor token_tags, int64_t chunks,
+                                          int64_t modality_num);
+torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sorted_pos,
+                                           torch::Tensor tile_begin, torch::Tensor tile_end,
+                                           torch::Tensor seg_first_tile,
+                                           c10::ScalarType out_dtype);
+std::vector<torch::Tensor> h3_rmsnorm_forward(torch::Tensor x, torch::Tensor weight, double eps,
+                                              c10::optional<torch::Tensor> shift,
+                                              c10::optional<torch::Tensor> scale,
+                                              c10::optional<torch::Tensor> index);
+std::vector<torch::Tensor> h3_rmsnorm_backward(
+    torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd,
+    c10::optional<torch::Tensor> shift, c10::optional<torch::Tensor> scale,
+    c10::optional<torch::Tensor> index, c10::optional<torch::Tensor> sorted_pos,
+    c10::optional<torch::Tensor> tile_begin, c10::optional<torch::Tensor> tile_end,
+    c10::optional<torch::Tensor> seg_first_tile);
+torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate,
+                                       torch::Tensor index);
+std::vector<torch::Tensor> h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y,
+                                                     torch::Tensor gate, torch::Tensor index,
+                                                     torch::Tensor sorted_pos,
+                                                     torch::Tensor tile_begin,
+                                                     torch::Tensor tile_end,
+                                                     torch::Tensor seg_first_tile);
+std::vector<torch::Tensor> h3_rmsnorm_backward_partials(
+    torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd,
+    c10::optional<torch::Tensor> shift, c10::optional<torch::Tensor> scale,
+    c10::optional<torch::Tensor> index, torch::Tensor dw_rows, torch::Tensor dw_begin,
+    torch::Tensor dw_end, c10::optional<torch::Tensor> seg_rows,
+    c10::optional<torch::Tensor> seg_begin, c10::optional<torch::Tensor> seg_end);
+std::vector<torch::Tensor> h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight,
+                                                    c10::optional<torch::Tensor> seg_partial,
+                                                    c10::optional<torch::Tensor> seg_first_tile);
+torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch::Tensor rows,
+                                    torch::Tensor tile_begin, torch::Tensor tile_end);
+torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_tile,
+                                c10::ScalarType dtype);
+torch::Tensor h3_rmsnorm_backward_dx(torch::Tensor grad, torch::Tensor x, torch::Tensor weight,
+                                     torch::Tensor rstd, c10::optional<torch::Tensor> shift,
+                                     c10::optional<torch::Tensor> scale,
+                                     c10::optional<torch::Tensor> index);
+torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate,
+                                           torch::Tensor index);
+#endif
+
 // PyBind11 Module Registration
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.doc() = "RL-Kernel High-Performance Operator Extension Library";
@@ -758,4 +819,77 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         "Deterministic GPT-NeoX token-major RoPE apply for ROCm");
 #endif
 #endif
+
+#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) && \
+    (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA))
+    m.def("h3_timestep_sinusoid_forward", 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);
+    m.def("h3_det_linear_forward", &h3_det_linear_forward,
+          "Batch-invariant warp-per-column linear (contract h3-det-linear-v1), optional SiLU",
+          py::arg("x"), py::arg("weight"), py::arg("bias") = py::none(),
+          py::arg("activation") = 0, py::arg("save_pre_activation") = false);
+    m.def("h3_det_linear_backward_input", &h3_det_linear_backward_input,
+          "Deterministic grad @ weight with fixed 64-row N chunks folded in order",
+          py::arg("grad"), py::arg("weight"), py::arg("out_dtype"));
+    m.def("h3_det_linear_backward_input_partials", &h3_det_linear_backward_input_partials,
+          "Per-64-row-chunk FP32 partials of grad @ weight, before the ascending fold",
+          py::arg("grad"), py::arg("weight"));
+    m.def("h3_det_linear_fold_chunks", &h3_det_linear_fold_chunks,
+          "Ascending left fold of h3_det_linear_backward_input_partials, cast once",
+          py::arg("partial"), py::arg("out_dtype"));
+    m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight,
+          "Deterministic dW/dbias as ascending-row FP32 folds",
+          py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true);
+    m.def("h3_adaln_row_gather_forward", &h3_adaln_row_gather_forward,
+          "Fused six-way AdaLN row gather by timestep_index * 3 + token_tag (pure copy)",
+          py::arg("rows"), py::arg("timestep_indices"), py::arg("token_tags"),
+          py::arg("chunks") = 6, py::arg("modality_num") = 3);
+    m.def("h3_adaln_row_gather_backward", &h3_adaln_row_gather_backward,
+          "Deterministic segmented sum (sorted tiles folded in order) for the row gather",
+          py::arg("grad"), py::arg("sorted_pos"), py::arg("tile_begin"), py::arg("tile_end"),
+          py::arg("seg_first_tile"), py::arg("out_dtype"));
+    m.def("h3_rmsnorm_forward", &h3_rmsnorm_forward,
+          "RMSNorm replaying PyTorch's reduction order, with optional fused AdaLN modulation",
+          py::arg("x"), py::arg("weight"), py::arg("eps"), py::arg("shift") = py::none(),
+          py::arg("scale") = py::none(), py::arg("index") = py::none());
+    m.def("h3_rmsnorm_backward", &h3_rmsnorm_backward,
+          "Deterministic RMSNorm(+modulation) backward: row-local dx, tiled dweight, sorted table grads",
+          py::arg("grad"), py::arg("x"), py::arg("weight"), py::arg("rstd"),
+          py::arg("shift") = py::none(), py::arg("scale") = py::none(),
+          py::arg("index") = py::none(), py::arg("sorted_pos") = py::none(),
+          py::arg("tile_begin") = py::none(), py::arg("tile_end") = py::none(),
+          py::arg("seg_first_tile") = py::none());
+    m.def("h3_gate_residual_forward", &h3_gate_residual_forward,
+          "residual + gate[index] * y with the gate row gathered in-kernel (eager rounding order)",
+          py::arg("residual"), py::arg("y"), py::arg("gate"), py::arg("index"));
+    m.def("h3_gate_residual_backward", &h3_gate_residual_backward,
+          "Gated-residual backward: exact dy, deterministic sorted segment sum for dgate",
+          py::arg("grad"), py::arg("y"), py::arg("gate"), py::arg("index"), py::arg("sorted_pos"),
+          py::arg("tile_begin"), py::arg("tile_end"), py::arg("seg_first_tile"));
+    m.def("h3_rmsnorm_backward_partials", &h3_rmsnorm_backward_partials,
+          "WS1 RMSNorm dweight / table-gradient tile partials over explicit row lists (SP)",
+          py::arg("grad"), py::arg("x"), py::arg("weight"), py::arg("rstd"), py::arg("shift"),
+          py::arg("scale"), py::arg("index"), py::arg("dw_rows"), py::arg("dw_begin"),
+          py::arg("dw_end"), py::arg("seg_rows") = py::none(), py::arg("seg_begin") = py::none(),
+          py::arg("seg_end") = py::none());
+    m.def("h3_rmsnorm_fold_partials", &h3_rmsnorm_fold_partials,
+          "WS1 ascending folds of RMSNorm dweight and per-segment table-gradient partials",
+          py::arg("dw_partial"), py::arg("weight"), py::arg("seg_partial") = py::none(),
+          py::arg("seg_first_tile") = py::none());
+    m.def("h3_gate_grad_partials", &h3_gate_grad_partials,
+          "WS1 d_gate tile partials over explicit row lists (SP)", py::arg("grad"), py::arg("y"),
+          py::arg("rows"), py::arg("tile_begin"), py::arg("tile_end"));
+    m.def("h3_gate_grad_fold", &h3_gate_grad_fold,
+          "WS1 per-segment ascending fold of d_gate partials, cast once", py::arg("partial"),
+          py::arg("seg_first_tile"), py::arg("dtype"));
+    m.def("h3_rmsnorm_backward_dx", &h3_rmsnorm_backward_dx,
+          "Row-local dx of the WS1 RMSNorm(+modulation) backward", py::arg("grad"), py::arg("x"),
+          py::arg("weight"), py::arg("rstd"), py::arg("shift") = py::none(),
+          py::arg("scale") = py::none(), py::arg("index") = py::none());
+    m.def("h3_gate_residual_backward_dy", &h3_gate_residual_backward_dy,
+          "Row-local d_y = grad * gate[index] of the WS1 gated-residual backward",
+          py::arg("grad"), py::arg("gate"), py::arg("index"));
+#endif
 }