[WS1][CUDA][MiniMax-H3] token_refiner_block: two refiner blocks over the strict operator rows (RFC #420 stack 4/5)

.github/workflows/ci.yml

@@ -5,9 +5,9 @@ name: CI-Pipeline
 
 on:
   push:
-    branches: [ main, test ]
+    branches: [ main, test, test-h3 ]
   pull_request:
-    branches: [ main, test ]
+    branches: [ main, test, test-h3 ]
 
 permissions:
   contents: read
@@ -75,6 +75,9 @@ jobs:
         run: |
           python -m pytest tests/reference/minimax_h3 tests/runtime/test_h3_sigma_schedule_dispatch.py -q
 
+      - name: Run full H3 binding contracts
+        run: python -m pytest -q tests/models/minimax_h3/test_h3_full_model.py
+
       - name: Run Mocked Hardware Discovery Tests
         run: |
           python -m pytest tests/runtime/test_dispatch.py -v

build_tools/extensions.py

@@ -26,6 +26,7 @@
     "csrc/cuda/h3/det_linear.cu",
     "csrc/cuda/h3/adaln_row_gather.cu",
     "csrc/cuda/h3/rmsnorm_modulate.cu",
+    "csrc/cuda/h3/gate_residual.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

csrc/bindings/ops.cpp

@@ -510,6 +510,8 @@ std::vector<torch::Tensor> h3_det_linear_forward(torch::Tensor x, torch::Tensor
                                                  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);
@@ -530,6 +532,33 @@ std::vector<torch::Tensor> h3_rmsnorm_backward(
     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
@@ -803,6 +832,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     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);
@@ -825,5 +860,35 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           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
 }

csrc/cuda/h3/adaln_row_gather.cu

@@ -212,9 +212,10 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor
   TORCH_CHECK(grad.is_cuda() && grad.dim() == 3 && grad.is_contiguous(),
               "grad must be a contiguous (C, S, H) CUDA tensor");
   for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) {
-    TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->dim() == 1 &&
-                    t->is_contiguous() && t->scalar_type() == at::kLong,
-                "tile metadata must be contiguous 1-D int64 tensors on grad's device");
+    TORCH_CHECK(t->is_cuda() && t->device() == grad.device() &&
+                    t->dim() == 1 && t->is_contiguous() &&
+                    t->scalar_type() == at::kLong,
+                "tile metadata must be contiguous 1-D int64 tensors on grad device");
   }
   TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries");
   TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch");
@@ -225,9 +226,10 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor
   const int64_t num_rows = seg_first_tile.numel() - 1;
   const int64_t tiles = tile_begin.numel();
   TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries");
+  TORCH_CHECK(chunks > 0 && seq > 0 && hidden > 0 && num_rows > 0,
+              "gradient dimensions and table row count must be positive");
   const c10::cuda::CUDAGuard device_guard(grad.device());
   auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype));
-  if (num_rows == 0) return out;  // a zero-sized grid is not a valid launch
   auto partial = torch::empty({std::max<int64_t>(tiles, 1), width},
                               grad.options().dtype(at::kFloat));
   auto stream = at::cuda::getCurrentCUDAStream();

csrc/cuda/h3/det_linear.cu

@@ -26,6 +26,9 @@
 //   * N is split into fixed chunks of 64 rows; each chunk is an ascending
 //     fmaf chain into FP32 from 0; the chunk partials are then left-folded in
 //     ascending chunk order and cast once.
+//     The partials and the fold are also bound separately: a contiguous N
+//     slice starting on a chunk boundary yields exactly the full call's
+//     partials for those chunks, which is what tensor-parallel shards gather.
 // d_weight dW[n, k] = sum_t g[t, n] * x[t, k];  d_bias db[n] = sum_t g[t, n]
 //   * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one
 //     cross-row reduction; its order is the logical row order. db is the same
@@ -509,25 +512,22 @@ std::vector<torch::Tensor> h3_det_linear_forward(torch::Tensor x, torch::Tensor
   return {out};
 }
 
-// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype.
-torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight,
-                                           c10::ScalarType out_dtype) {
+// grad [T, N] (float32), weight [N, K] -> partial [ceil(N / 64), T, K] float32:
+// the per-chunk sums of contract h3-det-linear-v1, before the ascending fold.
+// A contiguous slice of N that starts on a chunk boundary yields exactly the
+// matching rows of the full call (tensor-parallel shards rely on this).
+torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight) {
   TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() &&
                   grad.scalar_type() == at::kFloat,
               "grad must be a contiguous 2-D float32 CUDA tensor");
   check_matrix(weight, "weight");
   TORCH_CHECK(grad.device() == weight.device(), "grad and weight must be on the same device");
   TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N");
-  TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16,
-              "out_dtype must be float32 or bfloat16");
   const int64_t rows = grad.size(0);
   TORCH_CHECK(rows > 0, "grad must have at least one row");
   const int64_t n_out = weight.size(0);
   const int64_t k_in = weight.size(1);
   const c10::cuda::CUDAGuard device_guard(grad.device());
-  auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype));
-  if (k_in == 0) return out;
-  if (n_out == 0) return out.zero_();
   const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk;
   // chunks is blockIdx.y and rows * k_in / threads is blockIdx.x of the fold: fail with a clear message
   // instead of "invalid configuration argument" when a launch dimension would overflow.
@@ -536,6 +536,7 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei
   TORCH_CHECK((rows * k_in + 255) / 256 <= std::numeric_limits<int32_t>::max(), "rows * K = ",
               rows * k_in, " is too large for one input-gradient fold launch");
   auto partial = torch::empty({chunks, rows, k_in}, grad.options());
+  if (partial.numel() == 0) return partial;
   auto stream = at::cuda::getCurrentCUDAStream();
   const int threads = 256;
   dim3 grid(static_cast<unsigned>((k_in + threads - 1) / threads), static_cast<unsigned>(chunks));
@@ -549,7 +550,27 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei
         n_out, k_in);
   }
   C10_CUDA_KERNEL_LAUNCH_CHECK();
-  const int64_t elems = rows * k_in;
+  return partial;
+}
+
+// partial [C, T, K] float32 -> left fold over C in ascending order, cast once.
+torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype) {
+  TORCH_CHECK(partial.is_cuda() && partial.dim() == 3 && partial.is_contiguous() &&
+                  partial.scalar_type() == at::kFloat && partial.size(0) > 0,
+              "partial must be a non-empty contiguous 3-D float32 CUDA tensor");
+  TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16,
+              "out_dtype must be float32 or bfloat16");
+  const c10::cuda::CUDAGuard device_guard(partial.device());
+  const int64_t chunks = partial.size(0);
+  auto out = torch::empty({partial.size(1), partial.size(2)}, partial.options().dtype(out_dtype));
+  const int64_t elems = out.numel();
+  if (elems == 0) return out;
+  auto stream = at::cuda::getCurrentCUDAStream();
+  const int threads = 256;
+  // Same bound as h3_det_linear_backward_input: grid.x holds at most INT32_MAX blocks. (That is 5e11
+  // elements, beyond any tensor that fits in memory; the check keeps the two entry points alike.)
+  TORCH_CHECK((elems + threads - 1) / threads <= std::numeric_limits<int32_t>::max(),
+              "rows * K = ", elems, " needs more blocks than one grid.x launch holds");
   const unsigned fold_blocks = static_cast<unsigned>((elems + threads - 1) / threads);
   if (out_dtype == at::kFloat) {
     det_fold_chunks_kernel<float><<<fold_blocks, threads, 0, stream>>>(
@@ -562,6 +583,18 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei
   return out;
 }
 
+// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype.
+torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight,
+                                           c10::ScalarType out_dtype) {
+  TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16,
+              "out_dtype must be float32 or bfloat16");
+  auto partial = h3_det_linear_backward_input_partials(grad, weight);
+  if (partial.size(0) == 0) {  // N == 0: no output column contributes
+    return torch::zeros({partial.size(1), partial.size(2)}, grad.options().dtype(out_dtype));
+  }
+  return h3_det_linear_fold_chunks(partial, out_dtype);
+}
+
 // grad [T, N] float32, x [T, K] (float32 or bfloat16) -> (dW [N, K], dbias [N]) in w_dtype.
 std::vector<torch::Tensor> h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x,
                                                          c10::ScalarType w_dtype,