[WS2][CUDA][MiniMax-H3] sp_norm_adaln: sequence-parallel norm/AdaLN/gated residual, byte-equal to WS1 (RFC #420 stack 9/10)

csrc/bindings/ops.cpp

@@ -540,6 +540,25 @@ std::vector<torch::Tensor> h3_gate_residual_backward(torch::Tensor grad, torch::
                                                      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
@@ -848,5 +867,28 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           "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() &&
+    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());
+                "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/gate_residual.cu

@@ -16,6 +16,8 @@
 // 16-bit values is exact in FP32, so this equals the eager VJP bitwise);
 // d_gate[i] = sum over positions mapped to i of grad * y, an FP32 segmented
 // sum over positions sorted stably by row, in fixed tiles folded in order.
+// The tile partials and the fold are also bound separately for sequence
+// parallel callers (each tile computed where its rows live, folded in order).
 
 #include <torch/extension.h>
 #include <ATen/cuda/CUDAContext.h>
@@ -262,3 +264,96 @@ std::vector<torch::Tensor> h3_gate_residual_backward(torch::Tensor grad, torch::
   C10_CUDA_KERNEL_LAUNCH_CHECK();
   return {dy, dgate};
 }
+
+// Sequence parallel: the WS1 d_gate tile partials over explicit row lists of
+// (grad, y) rows -> (tiles, N) FP32.
+torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch::Tensor rows,
+                                    torch::Tensor tile_begin, torch::Tensor tile_end) {
+  TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(),
+              "grad must be a contiguous (M, N) CUDA tensor");
+  check_act(y, "y", grad);
+  for (const auto* t : {&rows, &tile_begin, &tile_end}) {
+    TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->scalar_type() == at::kLong &&
+                    t->dim() == 1 && t->is_contiguous(),
+                "tile metadata must be contiguous int64 CUDA tensors on grad's device");
+  }
+  TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin and tile_end differ in length");
+  const int64_t n = grad.size(1);
+  const int64_t tiles = tile_begin.numel();
+  const c10::cuda::CUDAGuard guard(grad.device());
+  auto partial = torch::empty({tiles, n}, grad.options().dtype(at::kFloat));
+  if (tiles == 0) return partial;
+  const int threads = 256;
+  const int64_t col_blocks64 = (n + threads - 1) / threads;
+  TORCH_CHECK(col_blocks64 <= 65535, "N=", n, " needs ", col_blocks64,
+              " column blocks, more than the grid.y limit of 65535");
+  const unsigned col_blocks = static_cast<unsigned>(col_blocks64);
+  // The row and tile lists are caller-built: every row must be a row of grad and every tile must
+  // lie inside the list before the kernel reads through them.
+  h3_tiles::check_lists(at::cuda::getCurrentCUDAStream(), rows.data_ptr<int64_t>(), rows.numel(),
+                        grad.size(0), tile_begin.data_ptr<int64_t>(), tile_end.data_ptr<int64_t>(),
+                        tiles, nullptr, 0, 0);
+  H3_DISPATCH(grad.scalar_type(), "h3_gate_grad_partials", [&] {
+    gate_grad_partial_kernel<T><<<dim3(static_cast<unsigned>(tiles), col_blocks), threads, 0,
+                                  at::cuda::getCurrentCUDAStream()>>>(
+        reinterpret_cast<const T*>(grad.data_ptr()), reinterpret_cast<const T*>(y.data_ptr()),
+        rows.data_ptr<int64_t>(), tile_begin.data_ptr<int64_t>(), tile_end.data_ptr<int64_t>(),
+        partial.data_ptr<float>(), n);
+  });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return partial;
+}
+
+// The WS1 d_gate fold: each segment's tiles in ascending order, cast once to dtype.
+torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_tile,
+                                c10::ScalarType dtype) {
+  TORCH_CHECK(partial.is_cuda() && partial.dim() == 2 && partial.is_contiguous() &&
+                  partial.scalar_type() == at::kFloat,
+              "partial must be a contiguous float32 (tiles, N) CUDA tensor");
+  TORCH_CHECK(seg_first_tile.is_cuda() && seg_first_tile.device() == partial.device() &&
+                  seg_first_tile.scalar_type() == at::kLong && seg_first_tile.dim() == 1 &&
+                  seg_first_tile.is_contiguous(),
+              "seg_first_tile must be a contiguous int64 CUDA tensor on partial's device");
+  TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries");
+  const int64_t n = partial.size(1);
+  const int64_t segments = seg_first_tile.numel() - 1;
+  const c10::cuda::CUDAGuard guard(partial.device());
+  auto out = torch::empty({segments, n}, partial.options().dtype(dtype));
+  if (segments == 0) return out;  // a zero-sized grid is not a valid launch
+  const int threads = 256;
+  const int64_t col_blocks64 = (n + threads - 1) / threads;
+  TORCH_CHECK(col_blocks64 <= 65535, "N=", n, " needs ", col_blocks64,
+              " column blocks, more than the grid.y limit of 65535");
+  const unsigned col_blocks = static_cast<unsigned>(col_blocks64);
+  h3_tiles::check_lists(at::cuda::getCurrentCUDAStream(), nullptr, 0, 0, nullptr, nullptr, 0,
+                        seg_first_tile.data_ptr<int64_t>(), segments, partial.size(0));
+  H3_DISPATCH(dtype, "h3_gate_grad_fold", [&] {
+    gate_grad_fold_kernel<T><<<dim3(static_cast<unsigned>(segments), col_blocks), threads, 0,
+                               at::cuda::getCurrentCUDAStream()>>>(
+        partial.data_ptr<float>(), seg_first_tile.data_ptr<int64_t>(),
+        reinterpret_cast<T*>(out.data_ptr()), n);
+  });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return out;
+}
+
+// Row-local part of the WS1 backward only: d_y = cast(grad * gate[index]).
+torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate,
+                                           torch::Tensor index) {
+  TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(),
+              "grad must be a contiguous (M, N) CUDA tensor");
+  const Gate g = make_gate(gate, index, grad);
+  const int64_t rows = grad.size(0);
+  const int64_t n = grad.size(1);
+  const c10::cuda::CUDAGuard guard(grad.device());
+  auto dy = torch::empty_like(grad);
+  H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward_dy", [&] {
+    const bool vec = can_vectorize(gate, n, {&grad, &dy});
+    gate_residual_rows_kernel<T, true><<<row_blocks(rows), 256, 0,
+                                         at::cuda::getCurrentCUDAStream()>>>(
+        reinterpret_cast<const T*>(grad.data_ptr()), nullptr, reinterpret_cast<T*>(dy.data_ptr()),
+        rows, n, g, vec);
+  });
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return dy;
+}