[WS2][CUDA][MiniMax-H3] tp_adaln_3mod: column-parallel AdaLN projection, byte-equal to WS1 (RFC #420 stack 8/10)

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,14 @@ 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);
 #endif
 
 // PyBind11 Module Registration
@@ -803,6 +813,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 +841,12 @@ 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"));
 #endif
 }

csrc/cuda/h3/adaln_row_gather.cu

@@ -212,9 +212,9 @@ 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 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");

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,