From 507c9ad86540dd0aaab9376996d85d8c1c5f60a8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 10:39:17 +0000 Subject: [PATCH 01/19] [WS2][CUDA][MiniMax-H3] tp_adaln_3mod: column-parallel AdaLN projection, byte-equal to WS1 RFC #420 WS2 row `tp_adaln_3mod` (stacked on final_adaln_out). - Ownership: rank r owns the contiguous output columns [r*N/tp, (r+1)*N/tp) of the 3 x 6 x 5376 table (the same weight/bias rows); N/tp must be a multiple of the 64-row d_input chunk, so a shard never splits a WS1 partial. Unsupported TP sizes fail closed. - Every floating-point op is the WS1 kernel: the local tensor-core GEMV gives the WS1 columns, a rank-ordered all-gather (a copy) rebuilds the table, dW/db are shard-local, and d_temb all-gathers the WS1 chunk partials and runs the WS1 ascending fold on every rank. No collective reduction arithmetic. - det_linear: d_input split into h3_det_linear_backward_input_partials and h3_det_linear_fold_chunks (same kernels; backward_input composes them). - ws2_comm.gather_rows: rank-ordered row gather that pads outputs of at most 256 KiB past DeterministicCollective's single-block byte-copy path (~4 us/KiB); keeps the T = 1 forward at ~0.12 ms. - Readback: tp, rank, owned columns and (modality, chunk) slots, collective backend, kernel ids, fallback. - Tests: ownership, single-GPU shard arithmetic (columns and chunk partials equal the WS1 slice, the fold equals WS1 d_input), fail-closed inputs, and real NCCL TP2/4/8 on the pinned block-0 weights. - scripts/h3_ws2_evidence.py + figure; operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 41 +++- csrc/ops.cpp | 8 + docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-tp-adaln-3mod.md | 94 ++++++++ rl_engine/_C.pyi | 4 + rl_engine/kernels/ops/cuda/h3/det_linear.py | 17 ++ .../ops/cuda/h3/tp_adaln_projection.py | 211 ++++++++++++++++++ rl_engine/kernels/ops/cuda/h3/ws2_comm.py | 33 +++ rl_engine/testing/h3_ws2.py | 187 ++++++++++++++++ scripts/h3_ws2_evidence.py | 90 ++++++++ scripts/plot_h3_evidence.py | 96 +++++++- tests/h3/test_h3_tp_adaln.py | 148 ++++++++++++ 13 files changed, 923 insertions(+), 8 deletions(-) create mode 100644 docs/operators/h3-tp-adaln-3mod.md create mode 100644 rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py create mode 100644 rl_engine/kernels/ops/cuda/h3/ws2_comm.py create mode 100644 rl_engine/testing/h3_ws2.py create mode 100644 scripts/h3_ws2_evidence.py create mode 100644 tests/h3/test_h3_tp_adaln.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 95e16b234..66d62931d 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/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. @@ -487,24 +490,24 @@ std::vector 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); const int64_t n_out = weight.size(0); const int64_t k_in = weight.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; auto partial = torch::empty({chunks, rows, k_in}, grad.options()); - auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); + if (partial.numel() == 0) return partial; auto stream = at::cuda::getCurrentCUDAStream(); const int threads = 256; dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); @@ -518,7 +521,23 @@ 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; const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); if (out_dtype == at::kFloat) { det_fold_chunks_kernel<<>>( @@ -531,6 +550,14 @@ 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"); + return h3_det_linear_fold_chunks(h3_det_linear_backward_input_partials(grad, weight), out_dtype); +} + // grad [T, N] float32, x [T, K] (float32 or bfloat16) -> (dW [N, K], dbias [N]) in w_dtype. std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, c10::ScalarType w_dtype, diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 52ece3f13..f3dc5cca2 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -507,6 +507,8 @@ std::vector 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 h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, c10::ScalarType w_dtype, bool with_bias); @@ -806,6 +808,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); diff --git a/docs/.nav.yml b/docs/.nav.yml index 1ef4fa017..48f12fe3b 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -37,6 +37,7 @@ nav: - operators/h3-rmsnorm.md - operators/h3-adaln-gate-residual.md - operators/h3-final-adaln-out.md + - operators/h3-tp-adaln-3mod.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 27ffce7d0..99e02cf0c 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -39,4 +39,5 @@ Every operator page should include: - [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) - [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) - [MiniMax-H3 Final AdaLN Output](h3-final-adaln-out.md) +- [MiniMax-H3 Tensor-Parallel AdaLN Projection](h3-tp-adaln-3mod.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md new file mode 100644 index 000000000..da7b0ae42 --- /dev/null +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -0,0 +1,94 @@ +# MiniMax-H3 Tensor-Parallel AdaLN Projection + +## Summary + +`tp_adaln_3mod` shards the [AdaLN projection](h3-adaln-projection.md) (`2688 -> 96768`, 520 MB of +BF16 weight per block) across tensor-parallel ranks (RFC #420, WS2). Every rank ends with the +**same bytes as the single-GPU (WS1) op**: the six modulation tensors, the `3T` modality rows, +`d_temb`, and its shard of `dW`/`db`. + +## Entry Point + +```python +from rl_engine.distributed.collectives import DeterministicCollective +from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, +) + +collective = DeterministicCollective(group=tp_group, device=local_rank) +op = H3TPAdaLNProjectionCudaOp(collective, n_total=96768) +w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, w_shard, b_shard) +op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, fallback +``` + +## Ownership + +Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the +`N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. +`N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard +never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 qualify. Anything else fails closed when +the op is built. + +| TP | Columns per rank | Rank 0 owns | +| --- | --- | --- | +| 2 | 48384 | video: all six chunks; text: shift/scale/gate_msa | +| 4 | 24192 | video: shift_msa .. shift_mlp, half of scale_mlp | +| 8 | 12096 | video: shift_msa, scale_msa, 1344 channels of gate_msa | + +`AdaLNColumnShard.slots()` lists each rank's `(modality, chunk, h_begin, h_end)` pieces. The +all-gathered table is the WS1 table, so the six tensors and the `t * 3 + modality` rows are views +of it exactly as in WS1. Nothing downstream needs to know the TP size. + +## Numerics + +Every floating-point operation is the WS1 kernel itself. The only collective is a rank-ordered +all-gather, which is a copy. No collective reduction arithmetic is used. + +- **Forward.** An output column depends only on `silu(temb)` and its own weight row, so the local + tensor-core GEMV yields exactly the WS1 columns. The all-gather rebuilds the `(T, N)` table. +- **`dW`, `db`.** These are rows of the shard, computed locally with the WS1 ascending-`t` fold. +- **`d_temb`.** WS1 sums `N` in 64-row chunks and then left-folds the chunk partials in ascending + order. Each rank computes the partials of its own chunks + (`h3_det_linear_backward_input_partials`), the all-gather puts them in global chunk order, and + every rank runs the WS1 fold (`h3_det_linear_fold_chunks`) and the FP32 SiLU VJP. + +The backward reads its own columns of the table gradient. It therefore requires that gradient to +be identical on every TP rank, as it is when every rank applies the modulation to the same rows. + +## Communication + +| Direction | Payload per rank | At T = 3, TP = 8 | +| --- | --- | --- | +| forward | `T x N / tp` BF16 | 72.6 KB | +| backward | `(N / 64 / tp) x T x 2688` FP32 | 6.1 MB | + +`DeterministicCollective` copies all-gathers whose output is at most 256 KiB with a single-block +byte loop, at about 4 µs per KiB on B200. `ws2_comm.gather_rows` pads such shards just past that +size, where the multi-block path takes about 45 µs. This matters for the forward at small `T`. + +## Evidence + +The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 +node, one GPU per rank. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_tp_adaln.py -v # ownership, shard arithmetic, NCCL TP2/4/8 +python scripts/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \ + --out docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json +``` + +The NCCL tests need as many GPUs as ranks and skip otherwise. The shard-arithmetic tests check +the same property on one GPU: each rank's columns and chunk partials equal the matching slice of +the WS1 call, and the fold of the gathered partials equals WS1's `d_input`. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- `norm_out.linear` (`2688 -> 10752`) is not sharded here. It is small and stays replicated. +- There is no ROCm path. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 89843f979..672bf6049 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -341,6 +341,10 @@ def h3_det_linear_forward( def h3_det_linear_backward_input( grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype ) -> torch.Tensor: ... +def h3_det_linear_backward_input_partials( + grad: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: ... +def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: ... def h3_det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True ) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/ops/cuda/h3/det_linear.py b/rl_engine/kernels/ops/cuda/h3/det_linear.py index f5dab854e..17e7bce0f 100644 --- a/rl_engine/kernels/ops/cuda/h3/det_linear.py +++ b/rl_engine/kernels/ops/cuda/h3/det_linear.py @@ -16,10 +16,13 @@ ACT_NONE = 0 ACT_SILU = 1 CONTRACT = "h3-det-linear-v1" +DINPUT_CHUNK = 64 # N rows per d_input partial (kDInputChunk) _SYMBOLS = ( "h3_det_linear_forward", "h3_det_linear_backward_input", "h3_det_linear_backward_weight", + "h3_det_linear_backward_input_partials", + "h3_det_linear_fold_chunks", ) @@ -62,6 +65,20 @@ def det_linear_backward_input( ) +def det_linear_backward_input_partials(grad: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``(ceil(N / 64), T, K)`` FP32 chunk sums of ``grad @ weight``, before the fold.""" + + _require() + return _C.h3_det_linear_backward_input_partials(grad.float().contiguous(), weight.contiguous()) + + +def det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Ascending left fold of the chunk partials, cast once (the second half of d_input).""" + + _require() + return _C.h3_det_linear_fold_chunks(partial.contiguous(), out_dtype) + + def det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True ) -> list[torch.Tensor]: diff --git a/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py new file mode 100644 index 000000000..0eb5a4709 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tensor-parallel H3 AdaLN projection (RFC #420 ``tp_adaln_3mod``). + +Ownership: rank ``r`` of ``tp`` owns the contiguous output columns +``[r * N / tp, (r + 1) * N / tp)`` of the ``N = 3 * 6 * H`` projection, i.e. the +same rows of ``adaln_proj.linear.weight``/``bias``. ``N / tp`` must be a +multiple of the 64-row ``d_input`` chunk of contract ``h3-det-linear-v1``. + +Every floating-point operation is the WS1 kernel itself, so every rank ends +with the WS1 bytes: + +* forward: each column depends only on ``x`` and its own weight row, so the + local GEMV produces the WS1 columns; a rank-ordered all-gather (a copy) + rebuilds the ``(T, N)`` table, and the six tensors and ``3T`` modality rows + are views of it exactly as in WS1; +* ``dW``/``db``: rows of the shard, computed locally (no cross-rank sum); +* ``d_temb``: each rank computes the WS1 chunk partials of its columns, the + all-gather puts them in global chunk order, and every rank runs the WS1 + ascending fold. No collective reduction arithmetic is used. + +The backward expects the table gradient to be identical on every TP rank +(the modulation consumes the full replicated table); it reads its own columns. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + CONTRACT, + DINPUT_CHUNK, + det_linear_available, + det_linear_backward_input_partials, + det_linear_backward_weight, + det_linear_fold_chunks, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.cuda.h3.ws2_comm import gather_rows +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.kernels.ops.pytorch.h3.adaln_projection import split_adaln_table + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[column_shard]+all_gather[rank_order]" +BACKWARD_IMPL = "local_dweight+all_gather_chunk_partials+ascending_fold" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input_partials" + "+all_gather+rl_engine._C.h3_det_linear_fold_chunks" +) + + +@dataclass(frozen=True) +class AdaLNColumnShard: + """The output columns a TP rank owns, and the WS1 objects they map to.""" + + tp: int + rank: int + n_total: int + begin: int + end: int + + @property + def hidden(self) -> int: + return self.n_total // (H3_MODALITY_NUM * H3_ADALN_CHUNKS) + + def slots(self) -> list[tuple[int, int, int, int]]: + """``(modality, chunk, h_begin, h_end)`` pieces covered by this shard.""" + + out, col = [], self.begin + while col < self.end: + slot, h0 = divmod(col, self.hidden) + h1 = min(self.hidden, h0 + self.end - col) + out.append((*divmod(slot, H3_ADALN_CHUNKS), h0, h1)) + col += h1 - h0 + return out + + +def adaln_column_shard(n_total: int, tp: int, rank: int) -> AdaLNColumnShard: + if tp < 1 or not 0 <= rank < tp: + raise ValueError(f"need 0 <= rank < tp, got rank={rank}, tp={tp}") + if n_total % (H3_MODALITY_NUM * H3_ADALN_CHUNKS): + raise ValueError(f"N={n_total} is not 3 modalities x 6 chunks x H") + if n_total % (tp * DINPUT_CHUNK): + # A shard boundary inside a d_input chunk would split one WS1 partial. + raise ValueError( + f"tp={tp} does not split N={n_total} on {DINPUT_CHUNK}-column chunk boundaries" + ) + width = n_total // tp + return AdaLNColumnShard(tp, rank, n_total, rank * width, (rank + 1) * width) + + +def shard_adaln_projection(weight: torch.Tensor, bias: torch.Tensor, tp: int, rank: int): + """This rank's rows of the full projection weight and bias (views).""" + + shard = adaln_column_shard(weight.shape[0], tp, rank) + return weight[shard.begin : shard.end], bias[shard.begin : shard.end] + + +def _gather_columns(collective, local: torch.Tensor) -> torch.Tensor: + """``(T, N / tp)`` column shards -> the ``(T, N)`` table, in rank order (a copy).""" + + num_t, width = local.shape + rows = gather_rows(collective, local) # (tp * T, N / tp), rank-major + return rows.view(-1, num_t, width).transpose(0, 1).reshape(num_t, -1) + + +def _validate_shard_inputs(temb, weight, bias, shard: AdaLNColumnShard) -> None: + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if not tensor.is_cuda or tensor.device != temb.device: + raise ValueError(f"{name} must be on this rank's CUDA device, got {tensor.device}") + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in (torch.bfloat16, torch.float32) or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share bfloat16 or float32, got {weight.dtype}/{bias.dtype}" + ) + width = shard.end - shard.begin + if weight.dim() != 2 or weight.shape[0] != width or bias.shape != (width,): + raise ValueError( + f"rank {shard.rank} of {shard.tp} owns {width} rows, got weight " + f"{tuple(weight.shape)} and bias {tuple(bias.shape)}" + ) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + + +class _TPAdaLNProjection(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight_shard, bias_shard, collective, shard): + temb = temb.contiguous() + act = F.silu(temb).to(weight_shard.dtype) + (local,) = det_linear_forward(act, weight_shard, bias_shard) + ctx.save_for_backward(temb, act, weight_shard) + ctx.collective, ctx.shard = collective, shard + return _gather_columns(collective, local) + + @staticmethod + def backward(ctx, grad_table): + temb, act, weight = ctx.saved_tensors + shard = ctx.shard + grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() + dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) + partial = det_linear_backward_input_partials(grad_local, weight) + d_act = det_linear_fold_chunks(gather_rows(ctx.collective, partial), torch.float32) + record_backward( + "tp_adaln_3mod", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return silu_backward_fp32(d_act, temb), dw, db, None, None + + +class H3TPAdaLNProjectionCudaOp: + """Column-parallel ``adaln_projection_3mod``, byte-equal to WS1 on every rank. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` + and ``backend_id`` (e.g. ``DeterministicCollective``). Inputs are the replicated + FP32 ``temb`` and this rank's weight/bias rows. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self, collective, n_total: int) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + self.collective = collective + self.shard = adaln_column_shard(n_total, collective.world_size, collective.rank) + + def __call__(self, temb, weight_shard, bias_shard): + return self.forward(temb, weight_shard, bias_shard) + + def forward(self, temb, weight_shard, bias_shard) -> tuple[torch.Tensor, ...]: + table = self.forward_table(temb, weight_shard, bias_shard) + return split_adaln_table(table, self.shard.hidden) + + def forward_table(self, temb, weight_shard, bias_shard) -> torch.Tensor: + _validate_shard_inputs(temb, weight_shard, bias_shard, self.shard) + return _TPAdaLNProjection.apply(temb, weight_shard, bias_shard, self.collective, self.shard) + + def readback(self) -> dict: + """Runtime identity of this rank's projection, for evidence and strict traces.""" + + return { + "op": "tp_adaln_3mod", + "kernel_id": KERNEL_ID, + "backward_kernel_id": BACKWARD_KERNEL_ID, + "contract": CONTRACT, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_ascending_chunk_fold_after_rank_ordered_gather", + "tp": self.shard.tp, + "rank": self.shard.rank, + "columns": [self.shard.begin, self.shard.end], + "slots": [list(s) for s in self.shard.slots()], + "fallback": None, + } diff --git a/rl_engine/kernels/ops/cuda/h3/ws2_comm.py b/rl_engine/kernels/ops/cuda/h3/ws2_comm.py new file mode 100644 index 000000000..f03597b7d --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/ws2_comm.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rank-ordered row gather shared by the H3 WS2 ops (a copy: no arithmetic).""" + +from __future__ import annotations + +import torch + +# DeterministicCollective sends all-gathers whose output is at most this many bytes +# through a single-block byte-copy kernel (csrc/cuda/distributed/ +# deterministic_collective.cu, kSingleBlockFastPathMaxBytes) that costs about 4 us per +# KiB on B200; larger outputs take the multi-block path (~45 us). Padding the shard +# past the threshold is cheaper for every non-trivial size. +_SINGLE_BLOCK_GATHER_MAX_BYTES = 256 * 1024 + + +def gather_rows(collective, shard: torch.Tensor) -> torch.Tensor: + """``(world * rows, ...)``: every rank's ``(rows, ...)`` shard, in rank order.""" + + shard = shard.contiguous() + rows, world = shard.shape[0], collective.world_size + row_bytes = shard[0].numel() * shard.element_size() if rows else 0 + if world == 1: + return shard + if row_bytes == 0: + return collective.all_gather(shard) + padded = max(rows, _SINGLE_BLOCK_GATHER_MAX_BYTES // (world * row_bytes) + 1) + if padded == rows: + return collective.all_gather(shard) + staged = torch.cat([shard, shard.new_zeros((padded - rows, *shard.shape[1:]))]) + out = collective.all_gather(staged).view(world, padded, *shard.shape[1:]) + return out[:, :rows].reshape(world * rows, *shard.shape[1:]) diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py new file mode 100644 index 000000000..a0dc116f6 --- /dev/null +++ b/rl_engine/testing/h3_ws2.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS1 vs WS2 comparisons for the H3 conditioning path (tests and evidence). + +The per-rank helpers run inside real multi-process groups (tests and +``scripts/h3_ws2_evidence.py``); the WS1 side runs on one GPU. +""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def ws1_projection(temb, weight, bias, grad) -> dict[str, torch.Tensor]: + """The single-GPU ``adaln_projection_3mod`` table and its three gradients.""" + + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + table = H3AdaLNProjectionCudaOp().forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight": leaves[1].grad, + "d_bias": leaves[2].grad, + } + + +def tp_projection_rank(collective, temb, weight, bias, grad) -> dict[str, Any]: + """One TP rank: shard the full weight, run the TP op, return its view of the results.""" + + from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + op = H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]) + w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, w_shard, b_shard)] + table = op.forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight_shard": leaves[1].grad, + "d_bias_shard": leaves[2].grad, + "readback": op.readback(), + } + + +def tp_matches_ws1(ws1: dict[str, torch.Tensor], ranks: list[dict[str, Any]]) -> dict[str, bool]: + """Byte equality of every rank's replicated outputs and of the concatenated shards.""" + + return { + "table": all(torch.equal(r["table"], ws1["table"]) for r in ranks), + "d_temb": all(torch.equal(r["d_temb"], ws1["d_temb"]) for r in ranks), + "d_weight": torch.equal(torch.cat([r["d_weight_shard"] for r in ranks]), ws1["d_weight"]), + "d_bias": torch.equal(torch.cat([r["d_bias_shard"] for r in ranks]), ws1["d_bias"]), + } + + +def tp_projection_sweep(collective, temb, weight, bias, grad, iters: int = 100) -> dict[str, Any]: + """Evidence for one rank: per-T byte equality against WS1 on this GPU, timing, readback. + + ``temb``/``grad`` hold the largest T; every T in ``1..T`` uses their first + rows. Each rank compares its own replicated outputs and its own shard of + ``dW``/``db`` with the WS1 result computed here, so nothing large is shipped. + """ + + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + adaln_column_shard, + ) + from rl_engine.testing.h3_report import time_us + + shard = adaln_column_shard(weight.shape[0], collective.world_size, collective.rank) + cols = slice(shard.begin, shard.end) + tp_op, ws1_op = ( + H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]), + H3AdaLNProjectionCudaOp(), + ) + equality, perf = [], [] + for num_t in range(1, temb.shape[0] + 1): + t, g = temb[:num_t], grad[:num_t] + ws1 = ws1_projection(t, weight, bias, g) + mine = tp_projection_rank(collective, t, weight, bias, g) + equality.append( + { + "num_timesteps": num_t, + "table": torch.equal(mine["table"], ws1["table"]), + "d_temb": torch.equal(mine["d_temb"], ws1["d_temb"]), + "d_weight_shard": torch.equal(mine["d_weight_shard"], ws1["d_weight"][cols]), + "d_bias_shard": torch.equal(mine["d_bias_shard"], ws1["d_bias"][cols]), + } + ) + leaves = [x.detach().clone().requires_grad_(True) for x in (t, weight[cols], bias[cols])] + full = [x.detach().clone().requires_grad_(True) for x in (t, weight, bias)] + + def step(op, args): + for x in args: + x.grad = None + op.forward_table(*args).backward(g) + + perf.append( + { + "num_timesteps": num_t, + "tp_forward_us": time_us( + lambda: tp_op.forward_table(t, weight[cols], bias[cols]), 30, iters + ), + "tp_forward_backward_us": time_us(lambda: step(tp_op, leaves), 10, iters), + "ws1_forward_us": time_us(lambda: ws1_op.forward_table(t, weight, bias), 30, iters), + "ws1_forward_backward_us": time_us(lambda: step(ws1_op, full), 10, iters), + } + ) + return {"equality": equality, "perf": perf, "readback": tp_op.readback()} + + +# --- real multi-process groups ------------------------------------------------- + + +def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): + import traceback + + import torch.distributed as dist + + from rl_engine.distributed.collectives import DeterministicCollective + + try: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", init_method=init_method, rank=rank, world_size=world, device_id=rank + ) + inputs = {k: v.cuda() for k, v in torch.load(inputs_path).items()} + collective = DeterministicCollective(device=rank, max_size_bytes=256 << 20) + out = target(collective, **inputs, **kwargs) + # Results go through a file: tensors sent over the queue are shared by + # file descriptor and vanish if this process exits before the parent reads. + torch.save( + {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, + inputs_path.with_name(f"rank{rank}.pt"), + ) + queue.put({"rank": rank}) + collective.close() + dist.destroy_process_group() + except Exception: # forwarded to the parent + queue.put({"rank": rank, "error": traceback.format_exc()}) + + +def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: float = 900, **kwargs): + """Run ``target(collective, **inputs, **kwargs)`` on ranks ``0..world-1`` (one GPU each). + + ``target`` must be importable (it is pickled into spawned processes). Each + rank gets a NCCL group and a ``DeterministicCollective``; results come back + on the CPU, ordered by rank. Any rank's exception is re-raised here. + """ + + import tempfile + from pathlib import Path + + import torch.multiprocessing as mp + + if torch.cuda.device_count() < world: + raise RuntimeError(f"needs {world} GPUs, found {torch.cuda.device_count()}") + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "inputs.pt" + torch.save({k: v.cpu() for k, v in inputs.items()}, path) + queue = ctx.Queue() + init = (Path(tmp) / "init").as_uri() + procs = [ + ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) + for r in range(world) + ] + for proc in procs: + proc.start() + status = sorted((queue.get(timeout=timeout) for _ in procs), key=lambda r: r["rank"]) + for proc in procs: + proc.join(timeout=120) + for result in status: + if "error" in result: + raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") + return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] diff --git a/scripts/h3_ws2_evidence.py b/scripts/h3_ws2_evidence.py new file mode 100644 index 000000000..c27332f77 --- /dev/null +++ b/scripts/h3_ws2_evidence.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write an RFC #420 WS2 evidence report (byte equality vs WS1, timing, readback). + +Each TP size runs as real NCCL processes, one GPU per rank, with the +deterministic CUDA collective. Run it on a clean tree on one 8-GPU node; +``scripts/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python scripts/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \\ + --out docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from rl_engine.testing.h3_chain import environment, git_state # noqa: E402 +from rl_engine.testing.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) +from rl_engine.testing.h3_ws2 import run_world, tp_projection_sweep # noqa: E402 + +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _tp_adaln(worlds: list[int], max_t: int) -> dict: + weights = load_h3_conditioning_weights("cpu") + weight, bias = weights[WEIGHT], weights[BIAS] + g = torch.Generator(device="cpu").manual_seed(0) + inputs = { + "temb": torch.randn(max_t, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(max_t, weight.shape[0], generator=g).bfloat16(), + } + runs = [] + for world in worlds: + ranks = run_world(world, tp_projection_sweep, inputs) + runs.append({"tp": world, "ranks": ranks}) + equal = all( + v for r in ranks for e in r["equality"] for k, v in e.items() if k != "num_timesteps" + ) + print(f"tp={world}: byte-equal to WS1 on every rank and T: {equal}") + return {"weights": [WEIGHT, BIAS], "runs": runs} + + +OPS = {"tp_adaln_3mod": _tp_adaln} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--worlds", default="1,2,4,8") + parser.add_argument("--max-timesteps", type=int, default=4) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + worlds = [int(w) for w in args.worlds.split(",")] + if torch.cuda.device_count() < max(worlds): + raise SystemExit(f"needs {max(worlds)} GPUs, found {torch.cuda.device_count()}") + + manifest = load_h3_manifest() + report = { + "kind": "h3_ws2_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + **git_state(), + "environment": {**environment(), "gpus": torch.cuda.device_count()}, + **OPS[args.op](worlds, args.max_timesteps), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + commit, dirty = report["rl_kernel_commit"][:7], report["tracked_tree_dirty"] + print(f"wrote {args.out} (commit {commit}, dirty={dirty})") + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 3159c7379..ba4b3b4ef 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -2,7 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Render an RFC #420 operator report (``scripts/h3_evidence.py``) as a PNG. +"""Render an RFC #420 operator or WS2 report (``scripts/h3_evidence.py``, +``scripts/h3_ws2_evidence.py``) as a PNG. Needs only the report JSON and matplotlib (not torch or a GPU), so it can run anywhere: @@ -22,6 +23,7 @@ import matplotlib matplotlib.use("Agg") +import matplotlib.patches # noqa: E402 import matplotlib.pyplot as plt # noqa: E402 # Fixed series identity across every H3 figure (categorical slots 1-3). @@ -496,6 +498,97 @@ def plot_final(report: dict[str, Any]): return fig +TP_TENSORS = ( + ("table", "table (6 x 3T rows)"), + ("d_temb", "d_temb"), + ("d_weight_shard", "dW shard"), + ("d_bias_shard", "db shard"), +) + + +def plot_tp_adaln(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "tp_adaln_3mod (column-parallel AdaLN projection, 2688 -> 96768, NCCL)" + ) + runs = report["runs"] + # Left: every rank, every T, every tensor byte-equal to WS1? + left.set_facecolor(SURFACE) + left.set_title( + "Byte-equal to WS1 (every rank, every T)", loc="left", fontsize=11, color=INK, pad=10 + ) + for y, run in enumerate(runs): + checks = [e for r in run["ranks"] for e in r["equality"]] + for x, (key, _) in enumerate(TP_TENSORS): + ok = sum(bool(e[key]) for e in checks) + good = ok == len(checks) + left.add_patch( + matplotlib.patches.FancyBboxPatch( + (x + 0.06, y + 0.08), + 0.88, + 0.84, + boxstyle="round,pad=0,rounding_size=0.06", + facecolor=RL_KERNEL if good else PROVIDER, + edgecolor=SURFACE, + linewidth=2, + ) + ) + left.text( + x + 0.5, + y + 0.5, + f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} rank x T", + ha="center", + va="center", + fontsize=8, + color="white", + ) + left.set_xlim(0, len(TP_TENSORS)) + left.set_ylim(len(runs), 0) + left.set_xticks([i + 0.5 for i in range(len(TP_TENSORS))], [label for _, label in TP_TENSORS]) + left.set_yticks([i + 0.5 for i in range(len(runs))], [f"TP{run['tp']}" for run in runs]) + left.tick_params(colors=MUTED, labelsize=8, length=0) + for side in left.spines.values(): + side.set_visible(False) + # Right: slowest rank's time at the largest T; WS1 on one GPU as reference lines. + num_t = max(p["num_timesteps"] for p in runs[0]["ranks"][0]["perf"]) + + def worst(run, key): + perf = [p for r in run["ranks"] for p in r["perf"] if p["num_timesteps"] == num_t] + return max(p[key] for p in perf) / 1e3 + + tp_runs = [run for run in runs if run["tp"] > 1] + _style(right, f"Time per call, T = {num_t} (slowest rank; lower is better)", "ms") + _bars( + right, + [f"TP{run['tp']}" for run in tp_runs], + [ + ("TP fwd", RL_KERNEL, [worst(run, "tp_forward_us") for run in tp_runs]), + ( + "TP fwd+bwd", + RL_KERNEL, + [worst(run, "tp_forward_backward_us") for run in tp_runs], + "////", + ), + ], + "{:.2f}", + log=True, + ) + for key, label, style in ( + ("ws1_forward_us", "WS1 fwd", ":"), + ("ws1_forward_backward_us", "WS1 fwd+bwd", "--"), + ): + value = worst(runs[0], key) + right.axhline( + value, + color=PROVIDER, + linestyle=style, + linewidth=1.4, + zorder=1, + label=f"{label}, 1 GPU: {value:.2f}", + ) + right.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, @@ -504,6 +597,7 @@ def plot_final(report: dict[str, Any]): "h3_rmsnorm": plot_rmsnorm, "adaln_gate_residual": plot_gate_residual, "final_adaln_out": plot_final, + "tp_adaln_3mod": plot_tp_adaln, } diff --git a/tests/h3/test_h3_tp_adaln.py b/tests/h3/test_h3_tp_adaln.py new file mode 100644 index 000000000..8b341c8d3 --- /dev/null +++ b/tests/h3/test_h3_tp_adaln.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``tp_adaln_3mod``: column-parallel AdaLN projection, byte-equal to WS1. + +* ownership: contiguous 64-aligned column shards cover the 3 x 6 x H table + exactly once; unsupported TP sizes and wrong shards fail closed; +* shard arithmetic (one GPU): a rank's forward columns and ``d_input`` chunk + partials are bitwise the matching slice of the full WS1 call, and folding + the rank-ordered partials reproduces WS1's ``d_input``; +* real NCCL (2/4/8 GPUs): every rank's table and ``d_temb`` and the + concatenated ``dW``/``db`` shards are bitwise equal to WS1. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import adaln_column_shard +from rl_engine.testing.h3_ws2 import run_world, tp_matches_ws1, tp_projection_rank, ws1_projection + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" +N_H3 = 3 * 6 * 5376 +ALL_TRUE = dict.fromkeys(("table", "d_temb", "d_weight", "d_bias"), True) + + +def _require_ext(): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + + +def _inputs(num_t, n_out, k, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + temb = (torch.randn(num_t, k, generator=g) * 2).cuda() + weight = (torch.randn(n_out, k, generator=g) / k**0.5).to(dtype).cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + grad = torch.randn(num_t, n_out, generator=g).to(dtype).cuda() + return temb, weight, bias, grad + + +class TestOwnership: + @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8]) + def test_shards_tile_the_table_once(self, tp): + shards = [adaln_column_shard(N_H3, tp, r) for r in range(tp)] + assert [s.begin for s in shards] == [0, *[s.end for s in shards[:-1]]] + assert shards[-1].end == N_H3 + assert sum(h1 - h0 for s in shards for (_, _, h0, h1) in s.slots()) == N_H3 + + def test_tp2_slots(self): + # rank 0: video's six chunks and text's first three; rank 1: the rest. + assert adaln_column_shard(N_H3, 2, 0).slots() == [ + *[(0, c, 0, 5376) for c in range(6)], + *[(1, c, 0, 5376) for c in range(3)], + ] + + @pytest.mark.parametrize("tp, rank", [(5, 0), (16, 0), (0, 0), (2, 2), (2, -1)]) + def test_rejects_unsupported(self, tp, rank): + with pytest.raises(ValueError): + adaln_column_shard(N_H3, tp, rank) + + +@requires_cuda +class TestShardArithmetic: + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_pinned_shards_reproduce_ws1(self, h3_weights_cpu, tp): + _require_ext() + from rl_engine.kernels.ops.cuda.h3 import det_linear as dl + + weight, bias = h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + g = torch.Generator(device="cpu").manual_seed(tp) + act = F.silu(torch.randn(3, weight.shape[1], generator=g).cuda() * 2).bfloat16() + grad = torch.randn(3, weight.shape[0], generator=g).cuda() + (full,) = dl.det_linear_forward(act, weight, bias) + full_partials = dl.det_linear_backward_input_partials(grad, weight) + pieces = [] + for rank in range(tp): + s = adaln_column_shard(weight.shape[0], tp, rank) + (cols,) = dl.det_linear_forward(act, weight[s.begin : s.end], bias[s.begin : s.end]) + assert torch.equal(cols, full[:, s.begin : s.end]) + part = dl.det_linear_backward_input_partials( + grad[:, s.begin : s.end].contiguous(), weight[s.begin : s.end] + ) + c0 = s.begin // dl.DINPUT_CHUNK + assert torch.equal(part, full_partials[c0 : c0 + part.shape[0]]) + pieces.append(part) + folded = dl.det_linear_fold_chunks(torch.cat(pieces), torch.float32) + assert torch.equal(folded, dl.det_linear_backward_input(grad, weight, torch.float32)) + + def test_tp1_op_is_ws1(self): + _require_ext() + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(3, 3 * 6 * 256, 136) + ws1 = ws1_projection(temb, weight, bias, grad) + assert tp_matches_ws1(ws1, [tp_projection_rank(one, temb, weight, bias, grad)]) == ALL_TRUE + + def test_fails_closed(self): + _require_ext() + from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + temb, weight, bias, _ = _inputs(2, 3 * 6 * 256, 136) + rank1 = SimpleNamespace(rank=1, world_size=2, backend_id="unused", all_gather=None) + op = H3TPAdaLNProjectionCudaOp(rank1, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 2, 1) + with pytest.raises(TypeError): # RFC probe H7: cast before the SiLU + op(temb.bfloat16(), w, b) + with pytest.raises(ValueError): # the full weight instead of this rank's rows + op(temb, weight, bias) + with pytest.raises(ValueError): + op(temb.cpu(), w, b) + with pytest.raises(ValueError): # 4608 / 16 is not a multiple of 64 + H3TPAdaLNProjectionCudaOp(SimpleNamespace(rank=0, world_size=16), weight.shape[0]) + assert op.readback()["columns"] == [2304, 4608] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_byte_equal_to_ws1(h3_weights_cpu, world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + weight, bias = h3_weights_cpu[WEIGHT], h3_weights_cpu[BIAS] + g = torch.Generator(device="cpu").manual_seed(world) + inputs = { + "temb": torch.randn(3, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(3, weight.shape[0], generator=g).bfloat16(), + } + ws1 = { + k: v.cpu() for k, v in ws1_projection(**{k: v.cuda() for k, v in inputs.items()}).items() + } + torch.cuda.empty_cache() + ranks = run_world(world, tp_projection_rank, inputs) + assert tp_matches_ws1(ws1, ranks) == ALL_TRUE + assert [r["readback"]["rank"] for r in ranks] == list(range(world)) + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} From ecde5bb76a19bdc47457b67dad19e5423225fe55 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 10:58:54 +0000 Subject: [PATCH 02/19] [WS2][CUDA][MiniMax-H3] sp_norm_adaln: sequence-parallel norm/AdaLN/gated residual, byte-equal to WS1 RFC #420 WS2 row `sp_norm_adaln` (stacked on tp_adaln_3mod). - Ownership: rank r holds packed positions [r*S//sp, (r+1)*S//sp) of every batch item with the matching metadata; any S >= sp. The (S,) row index is replicated and every rank derives the same reduction plan from it. - Forward, dx, d_sublayer and d_residual: the WS1 kernels on local rows. - d_norm_weight, d_shift/d_scale, d_gate keep the WS1 two-level order (256-element tiles, then an ascending fold): each tile is computed by the rank holding its first row, after all-gathering the tile's rows from other ranks (only tiles crossing a shard boundary move rows); the partials are all-gathered into WS1 tile order and every rank runs the WS1 fold. Rank-ordered all-gathers only; no collective arithmetic. - New bindings around the unchanged WS1 kernels: h3_rmsnorm_backward_{dx,partials}, h3_rmsnorm_fold_partials, h3_gate_residual_backward_dy, h3_gate_grad_{partials,fold}; the WS1 dweight tile kernel also takes explicit row lists (same arithmetic). - The plan is built once per index object and shared by every norm and gated residual that uses it. - Tests: layout, per-rank backward vs WS1 for SP 2-8 on one GPU (odd S, B = 2, shards smaller than a tile, block and interleaved packing, bf16/fp32), fail-closed inputs, and the autograd block region over real NCCL SP2/4/8. - h3_block_layout (H3's block packing), evidence sweep with rows exchanged, timing and a naive-SP comparison; figure; operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 80 ++++ csrc/cuda/h3/rmsnorm_modulate.cu | 165 ++++++- csrc/ops.cpp | 42 ++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-sp-norm-adaln.md | 96 +++++ rl_engine/_C.pyi | 43 ++ .../kernels/ops/cuda/h3/sp_norm_adaln.py | 404 ++++++++++++++++++ rl_engine/testing/h3_cases.py | 24 ++ rl_engine/testing/h3_ws2.py | 181 ++++++++ scripts/h3_ws2_evidence.py | 42 +- scripts/plot_h3_evidence.py | 183 ++++++-- tests/h3/test_h3_sp_norm_adaln.py | 210 +++++++++ 13 files changed, 1428 insertions(+), 44 deletions(-) create mode 100644 docs/operators/h3-sp-norm-adaln.md create mode 100644 rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py create mode 100644 tests/h3/test_h3_sp_norm_adaln.py diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 18ede9e72..9a03660f1 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/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 #include @@ -240,3 +242,81 @@ std::vector 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 unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(grad.scalar_type(), "h3_gate_grad_partials", [&] { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(y.data_ptr()), + rows.data_ptr(), tile_begin.data_ptr(), tile_end.data_ptr(), + partial.data_ptr(), 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"); + 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)); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(dtype, "h3_gate_grad_fold", [&] { + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(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<<>>( + reinterpret_cast(grad.data_ptr()), nullptr, reinterpret_cast(dy.data_ptr()), + rows, n, g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dy; +} diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index d42fe3525..0c7a1bbfe 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -20,6 +20,10 @@ // dx = rstd * w * d_n - x * rstd^3 * sum(w * d_n * x) / N (row-local) // dweight = sum_r d_n * x * rstd over fixed 256-row tiles, tiles folded in order // dshift = sum_{r -> i} g, dscale = sum_{r -> i} g * n (sorted tiles) +// +// The tile partials and the folds are also bound separately (`*_partials`, +// `*_fold_*`) so a sequence-parallel caller can compute each WS1 tile on the +// rank that holds its rows and fold the gathered partials in WS1 order. #include #include @@ -180,18 +184,25 @@ __global__ void __launch_bounds__(32 * kWarps) } } -// partial[tile, j] = sum over rows of the tile (ascending) of d_n * x * rstd +// partial[tile, j] = sum over rows of the tile (ascending) of d_n * x * rstd. +// Tiles are the fixed 256-row blocks, or (sequence parallel) explicit row lists +// rows[tile_begin[tile] .. tile_end[tile]) that hold the same rows in the same order. template __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* __restrict__ x, const float* __restrict__ rstd, float* __restrict__ partial, int64_t rows, int n, - Modulation mod) { + Modulation mod, + const int64_t* __restrict__ row_list = nullptr, + const int64_t* __restrict__ tile_begin = nullptr, + const int64_t* __restrict__ tile_end = nullptr) { const int j = blockIdx.y * blockDim.x + threadIdx.x; if (j >= n) return; - const int64_t r0 = static_cast(blockIdx.x) * kRowTile; - const int64_t r1 = min(r0 + kRowTile, rows); + const int64_t tile = blockIdx.x; + const int64_t p0 = row_list != nullptr ? tile_begin[tile] : tile * kRowTile; + const int64_t p1 = row_list != nullptr ? tile_end[tile] : min(p0 + kRowTile, rows); float acc = 0.0f; - for (int64_t r = r0; r < r1; ++r) { + for (int64_t p = p0; p < p1; ++p) { + const int64_t r = row_list != nullptr ? row_list[p] : p; const T* scale = mod.index != nullptr ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride : nullptr; @@ -403,3 +414,147 @@ std::vector h3_rmsnorm_backward( C10_CUDA_KERNEL_LAUNCH_CHECK(); return {dx, dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; } + +namespace { + +void check_tiles(const torch::Tensor& rows, const torch::Tensor& begin, const torch::Tensor& end, + const torch::Tensor& like) { + for (const auto* t : {&rows, &begin, &end}) { + TORCH_CHECK(t->is_cuda() && t->device() == like.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors on x's device"); + } + TORCH_CHECK(begin.numel() == end.numel(), "tile_begin and tile_end differ in length"); +} + +} // namespace + +// Sequence parallel: the dweight and (with modulation) table-gradient tile +// partials of the WS1 backward, for explicit tiles over the rows of x. `index` +// gives each row's table row (one entry per row of x). Returns +// {dweight_partial (tiles, N)} or {dweight_partial, table_partial (seg_tiles, 2N)}. +std::vector h3_rmsnorm_backward_partials( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, torch::Tensor dw_rows, torch::Tensor dw_begin, + torch::Tensor dw_end, c10::optional seg_rows, + c10::optional seg_begin, c10::optional seg_end) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics on x's device"); + check_tiles(dw_rows, dw_begin, dw_end, x); + const int64_t rows = x.size(0); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + TORCH_CHECK(mod.index == nullptr || mod.seq == rows, "index needs one entry per row of x"); + const c10::cuda::CUDAGuard guard(x.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + const int64_t tiles = dw_begin.numel(); + auto dw_partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + if (tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_partials", [&] { + rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, + 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + rstd.data_ptr(), dw_partial.data_ptr(), rows, static_cast(n), mod, + dw_rows.data_ptr(), dw_begin.data_ptr(), dw_end.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + if (mod.index == nullptr) return {dw_partial}; + + TORCH_CHECK(seg_rows.has_value() && seg_begin.has_value() && seg_end.has_value(), + "modulated partials need the segment tiles"); + check_tiles(*seg_rows, *seg_begin, *seg_end, x); + const int64_t seg_tiles = seg_begin->numel(); + const int64_t width = 2 * n; + auto seg_partial = torch::empty({seg_tiles, width}, x.options().dtype(at::kFloat)); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + if (seg_tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_modulation_partials", [&] { + modulation_grad_partial_kernel + <<(seg_tiles), wide_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + seg_rows->data_ptr(), seg_begin->data_ptr(), + seg_end->data_ptr(), seg_partial.data_ptr(), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + return {dw_partial, seg_partial}; +} + +// The WS1 folds: dweight = ascending fold of all tiles, cast to weight's dtype; +// table rows = ascending fold of each segment's tiles, FP32 (R, 2N). +std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight, + c10::optional seg_partial, + c10::optional seg_first_tile) { + TORCH_CHECK(dw_partial.is_cuda() && dw_partial.dim() == 2 && dw_partial.is_contiguous() && + dw_partial.scalar_type() == at::kFloat, + "dw_partial must be a contiguous float32 (tiles, N) CUDA tensor"); + TORCH_CHECK(weight.is_cuda() && weight.device() == dw_partial.device() && weight.dim() == 1 && + weight.size(0) == dw_partial.size(1), + "weight must be (N,) on dw_partial's device"); + const int64_t n = dw_partial.size(1); + const c10::cuda::CUDAGuard guard(dw_partial.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + auto dweight = torch::empty_like(weight); + H3_DISPATCH(weight.scalar_type(), "h3_rmsnorm_fold_partials", [&] { + fold_tiles_kernel<<>>( + dw_partial.data_ptr(), reinterpret_cast(dweight.data_ptr()), + dw_partial.size(0), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!seg_partial.has_value()) return {dweight}; + + TORCH_CHECK(seg_first_tile.has_value() && seg_first_tile->is_cuda() && + seg_first_tile->device() == dw_partial.device() && + seg_first_tile->scalar_type() == at::kLong && seg_first_tile->is_contiguous(), + "segment fold needs seg_first_tile (int64, on dw_partial's device)"); + TORCH_CHECK(seg_partial->is_cuda() && seg_partial->device() == dw_partial.device() && + seg_partial->dim() == 2 && seg_partial->is_contiguous() && + seg_partial->scalar_type() == at::kFloat && seg_partial->size(1) == 2 * n, + "seg_partial must be a contiguous float32 (tiles, 2N) CUDA tensor"); + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t width = 2 * n; + auto table_grad = torch::empty({segments, width}, dw_partial.options()); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + fold_segments_kernel<<(segments), wide_blocks), threads, 0, stream>>>( + seg_partial->data_ptr(), seg_first_tile->data_ptr(), + table_grad.data_ptr(), static_cast(width)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; +} + +// Row-local part of the WS1 backward only: dx (sequence-parallel callers +// compute the cross-row reductions with h3_rmsnorm_backward_partials). +torch::Tensor h3_rmsnorm_backward_dx(torch::Tensor grad, torch::Tensor x, torch::Tensor weight, + torch::Tensor rstd, c10::optional shift, + c10::optional scale, + c10::optional index) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics on x's device"); + const Modulation mod = make_modulation(shift, scale, index, x, x.size(1)); + const c10::cuda::CUDAGuard guard(x.device()); + auto dx = torch::empty_like(x); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_dx", [&] { + rmsnorm_modulate_dx_kernel<<(x.size(0)), dim3(32, kWarps), 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + reinterpret_cast(dx.data_ptr()), static_cast(x.size(1)), mod); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dx; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index f3dc5cca2..342d09516 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -537,6 +537,25 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: torch::Tensor tile_begin, torch::Tensor tile_end, torch::Tensor seg_first_tile); +std::vector h3_rmsnorm_backward_partials( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, torch::Tensor dw_rows, torch::Tensor dw_begin, + torch::Tensor dw_end, c10::optional seg_rows, + c10::optional seg_begin, c10::optional seg_end); +std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight, + c10::optional seg_partial, + c10::optional 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 shift, + c10::optional scale, + c10::optional index); +torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate, + torch::Tensor index); #endif // PyBind11 Module Registration @@ -843,5 +862,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 } diff --git a/docs/.nav.yml b/docs/.nav.yml index 48f12fe3b..cfc833fe2 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -38,6 +38,7 @@ nav: - operators/h3-adaln-gate-residual.md - operators/h3-final-adaln-out.md - operators/h3-tp-adaln-3mod.md + - operators/h3-sp-norm-adaln.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 99e02cf0c..065d754b6 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -40,4 +40,5 @@ Every operator page should include: - [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) - [MiniMax-H3 Final AdaLN Output](h3-final-adaln-out.md) - [MiniMax-H3 Tensor-Parallel AdaLN Projection](h3-tp-adaln-3mod.md) +- [MiniMax-H3 Sequence-Parallel Norm and AdaLN](h3-sp-norm-adaln.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md new file mode 100644 index 000000000..167b63f7c --- /dev/null +++ b/docs/operators/h3-sp-norm-adaln.md @@ -0,0 +1,96 @@ +# MiniMax-H3 Sequence-Parallel Norm, AdaLN Modulation and Gated Residual + +## Summary + +`sp_norm_adaln` runs the row-wise part of an H3 block with the packed sequence split across +sequence-parallel ranks (RFC #420, WS2). It covers the [RMSNorm and AdaLN modulation](h3-rmsnorm.md) +and the [gated residual](h3-adaln-gate-residual.md). Every rank's output and row gradients are the +**WS1 bytes of its rows**. The gradients of the replicated tensors (`d_norm_weight`, `d_shift`, +`d_scale`, `d_gate`) are **the WS1 bytes, identical on every rank**. + +## Entry Point + +```python +from rl_engine.distributed.collectives import DeterministicCollective +from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + +collective = DeterministicCollective(group=sp_group, device=local_rank) +op = H3SPNormAdaLNCudaOp(collective, seq_len=S, batch=B) +rows = slice(op.layout.lo, op.layout.hi) # this rank's packed positions +hidden = op.gate_residual(residual[:, rows], attn_out_local, gate_msa, adaln_indices) +normed = op.norm_modulated(hidden, norm2_weight, shift_mlp, scale_mlp, adaln_indices) +final = op.norm_modulated(h, norm_out_weight, shift, scale, timestep_indices) # norm_out +op.readback() # sp, rank, positions, collective backend, reduction order, fallback +``` + +The table views, the norm weight and the full `(S,)` row index are replicated. Activations are +the rank's `(B, S_local, H)` rows. + +## Ownership + +Rank `r` of `sp` holds packed positions `[r * S // sp, (r + 1) * S // sp)` of every batch item, +together with the matching rows of the metadata. Any `S >= sp` works, including `S` not divisible +by `sp`. Shapes that do not match the rank's slice, a sharded index, an out-of-range index, CPU +tensors and malformed tables all fail closed. + +## Numerics + +- **Forward and row gradients.** `out`, `dx`, `d_sublayer` and `d_residual` are the WS1 kernels on + the local rows, with the local slice of the index. Rows are independent, so they are byte-equal + to WS1. +- **Cross-row gradients.** In WS1, `d_norm_weight` sums `d_n * x * rstd` over fixed 256-row tiles + of the `B x S` rows. `d_shift`/`d_scale`/`d_gate` sum over 256-element tiles of each table row's + positions, sorted stably. In both cases the tile partials are then folded in ascending order. SP + keeps exactly that two-level order: + 1. Every rank builds the global WS1 tile list from the replicated index (`sp_plan`). + 2. Each tile is computed by the rank that holds its first row. Rows of that tile that live on + other ranks are all-gathered first. Only tiles that cross a shard boundary move rows. + 3. The tile partials (`h3_rmsnorm_backward_partials`, `h3_gate_grad_partials`, the WS1 partial + kernels over explicit row lists) are all-gathered and put back in WS1 tile order. + 4. Every rank runs the WS1 folds (`h3_rmsnorm_fold_partials`, `h3_gate_grad_fold`). + + The only collectives are rank-ordered all-gathers, which are copies. No collective reduction + arithmetic is used. + +A naive SP backward runs the WS1 backward on each rank's rows and sums the per-rank results. It +changes the reduction tree, so it is not WS1. The evidence measures how often it differs. + +The plan depends only on the layout and the index. `H3SPNormAdaLNCudaOp` builds it once per index +object and reuses it for every norm and gated residual that uses the same index, which in a forward +pass is every block. + +## Rows Exchanged + +How many rows move depends on how the table rows interleave along the sequence: + +| Packing | Rows sent per rank (S = 32768, SP8, 4096 rows per rank) | +| --- | --- | +| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 174 | +| interleaved (modality and timestep random per position; stress case) | 0 – 1536 | + +## Evidence + +The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 +node, one GPU per rank. + +## Tests + +```bash +python -m pytest tests/h3/test_h3_sp_norm_adaln.py -v # layout, rank backward, NCCL SP2/4/8 +python scripts/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \ + --out docs/usage/evidence/h3-sp-norm-adaln-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-sp-norm-adaln-b200/report.json +``` + +The rank-backward tests run SP 2–8 on one GPU, with ranks as threads around the plain backward +functions. They cover odd `S`, `B = 2`, shards smaller than a tile, block and interleaved packing, +and bf16 and fp32, and compare every rank against WS1. The NCCL tests run the autograd region end +to end and need as many GPUs as ranks. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- With interleaved packing, up to a full shard of rows can move in the backward. The result is + still exact. +- Attention and FFN are other rows. Until they land, the region uses a seeded stand-in for the + sublayer output. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 672bf6049..d9948d98d 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -397,3 +397,46 @@ def h3_gate_residual_backward( tile_end: torch.Tensor, seg_first_tile: torch.Tensor, ) -> list[torch.Tensor]: ... +def h3_rmsnorm_backward_partials( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + index: torch.Tensor | None, + dw_rows: torch.Tensor, + dw_begin: torch.Tensor, + dw_end: torch.Tensor, + seg_rows: torch.Tensor | None = None, + seg_begin: torch.Tensor | None = None, + seg_end: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_fold_partials( + dw_partial: torch.Tensor, + weight: torch.Tensor, + seg_partial: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_gate_grad_partials( + grad: torch.Tensor, + y: torch.Tensor, + rows: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, +) -> torch.Tensor: ... +def h3_gate_grad_fold( + partial: torch.Tensor, seg_first_tile: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: ... +def h3_rmsnorm_backward_dx( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> torch.Tensor: ... +def h3_gate_residual_backward_dy( + grad: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py new file mode 100644 index 000000000..19a867b3e --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py @@ -0,0 +1,404 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sequence-parallel H3 RMSNorm / AdaLN modulation / gated residual (RFC #420 ``sp_norm_adaln``). + +Ownership: rank ``r`` of ``sp`` holds packed positions ``[r * S // sp, (r + 1) * S // sp)`` +of every batch item. The full ``(S,)`` row index is replicated; it is small, and every +rank derives the same reduction plan from it. + +Forward and the row-local gradients (``dx``, ``d_sublayer``, ``d_residual``) are the WS1 +kernels on the local rows. Rows are independent, so they are byte-equal by construction. + +The cross-row reductions (``d_norm_weight``, ``d_shift``/``d_scale``, ``d_gate``) are WS1 +two-level folds: FP32 sums over fixed 256-element tiles, then an ascending fold of the +tiles. Every rank builds the global WS1 tile list. A tile is computed by the rank that +holds its first row; the tile's rows from other ranks are all-gathered beforehand (only +rows of tiles that straddle a shard boundary move). The tile partials are all-gathered, +put back in WS1 tile order, and every rank runs the WS1 fold. The partial and fold kernels +are the WS1 kernels, so every gradient equals WS1 on every rank. The only collectives are +rank-ordered all-gathers (copies). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import BACKWARD_TILE, _segment_tiles +from rl_engine.kernels.ops.cuda.h3.ws2_comm import gather_rows +from rl_engine.kernels.ops.pytorch.h3.gate_residual import validate_h3_gate_residual +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +_SYMBOLS = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward_dx", + "h3_rmsnorm_backward_partials", + "h3_rmsnorm_fold_partials", + "h3_gate_residual_forward", + "h3_gate_residual_backward_dy", + "h3_gate_grad_partials", + "h3_gate_grad_fold", +) +BACKWARD_IMPL = "row_local_ws1+gather_straddling_rows+ws1_tile_partials+gather+ws1_fold" + + +def sp_norm_adaln_available() -> bool: + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +@dataclass(frozen=True) +class SPRowLayout: + """Which packed positions each sequence-parallel rank holds.""" + + seq_len: int + batch: int + sp: int + rank: int + + def bounds(self, rank: int) -> tuple[int, int]: + return rank * self.seq_len // self.sp, (rank + 1) * self.seq_len // self.sp + + @property + def lo(self) -> int: + return self.bounds(self.rank)[0] + + @property + def hi(self) -> int: + return self.bounds(self.rank)[1] + + @property + def local_len(self) -> int: + return self.hi - self.lo + + +def sp_row_layout(seq_len: int, batch: int, sp: int, rank: int) -> SPRowLayout: + if sp < 1 or not 0 <= rank < sp: + raise ValueError(f"need 0 <= rank < sp, got rank={rank}, sp={sp}") + if batch < 1 or seq_len < sp: + raise ValueError(f"every rank needs a row: seq_len={seq_len}, sp={sp}, batch={batch}") + return SPRowLayout(seq_len, batch, sp, rank) + + +def _ordinal_by_group(group: torch.Tensor, groups: int) -> tuple[torch.Tensor, torch.Tensor]: + """Per-group counts, and each element's position within its group (in input order).""" + + counts = torch.bincount(group, minlength=groups) + order = torch.argsort(group, stable=True) + ordinal = torch.empty_like(group) + first = torch.cumsum(counts, 0) - counts + ordinal[order] = torch.arange(group.numel(), device=group.device) - first[group[order]] + return counts, ordinal + + +class _Plan: + """Global WS1 tiles of one backward, and this rank's part in computing them. + + ``families`` are ``(rows, tile_begin, tile_end)`` over global flattened rows + ``b * S + s``; consecutive tiles cover ``rows`` in order. Every rank builds the same + plan from replicated metadata, so collective calls and shapes agree. + """ + + def __init__(self, layout: SPRowLayout, families, device) -> None: + self.layout, sp = layout, layout.sp + self._his = torch.tensor([layout.bounds(r)[1] for r in range(sp)], device=device) + + # Rows of a tile computed by another rank: the only activation rows that move. + tile_owner, shared = [], [] + for rows, begin, end in families: + owner = self.owner(rows[begin]) + tile_owner.append(owner) + shared.append(rows[torch.repeat_interleave(owner, end - begin) != self.owner(rows)]) + shared = torch.unique(torch.cat(shared)) + shared_owner = self.owner(shared) + counts, slot = _ordinal_by_group(shared_owner, sp) + self.max_send = int(counts.max()) if shared.numel() else 0 + + # Buffer row of every global row this rank reads: its own rows first, then + # every rank's send block (padded to max_send) in rank order. + own = self.global_rows(layout.rank) + buffer_index = torch.full((layout.batch * layout.seq_len,), -1, device=device) + buffer_index[own] = torch.arange(own.numel(), device=device) + foreign = shared_owner != layout.rank + buffer_index[shared[foreign]] = ( + own.numel() + shared_owner[foreign] * self.max_send + slot[foreign] + ) + mine = ~foreign + self.send_local = self.local_index(shared[mine][torch.argsort(slot[mine])]) + self.recv_global = torch.zeros(sp * self.max_send, dtype=torch.long, device=device) + self.recv_global[shared_owner * self.max_send + slot] = shared # padding: never read + + self.families = [ + self._my_tiles(rows, begin, end, owner, buffer_index) + for (rows, begin, end), owner in zip(families, tile_owner) + ] + + def owner(self, rows: torch.Tensor) -> torch.Tensor: + return torch.searchsorted(self._his, rows % self.layout.seq_len, right=True) + + def global_rows(self, rank: int) -> torch.Tensor: + lo, hi = self.layout.bounds(rank) + s = torch.arange(lo, hi, device=self._his.device) + b = torch.arange(self.layout.batch, device=self._his.device) + return (b[:, None] * self.layout.seq_len + s[None, :]).reshape(-1) + + def local_index(self, rows: torch.Tensor) -> torch.Tensor: + b, s = rows // self.layout.seq_len, rows % self.layout.seq_len + return b * self.layout.local_len + s - self.layout.lo + + def _my_tiles(self, rows, begin, end, owner, buffer_index): + mine = torch.nonzero(owner == self.layout.rank).flatten() + lengths = (end - begin)[mine] + my_end = torch.cumsum(lengths, 0) + my_begin = my_end - lengths + elems = torch.repeat_interleave(begin[mine] - my_begin, lengths) + elems = elems + torch.arange(elems.numel(), device=rows.device) + counts, ordinal = _ordinal_by_group(owner, self.layout.sp) + width = int(counts.max()) if owner.numel() else 0 + return { + "tiles": (buffer_index[rows[elems]].contiguous(), my_begin, my_end), + "width": width, # tiles per rank in the padded partial gather + "gather_index": owner * width + ordinal, # WS1 tile order within that gather + } + + def exchange(self, collective, *tensors: torch.Tensor) -> list[torch.Tensor]: + """For each ``(M_local, ...)`` tensor: local rows, then the gathered send blocks.""" + + out = [] + for t in tensors: + if self.max_send == 0: + out.append(t) + continue + send = t.new_zeros((self.max_send, *t.shape[1:])) + send[: self.send_local.numel()] = t[self.send_local] + out.append(torch.cat([t, gather_rows(collective, send)])) + return out + + def gather_partials(self, collective, family: int, partial: torch.Tensor) -> torch.Tensor: + """This rank's tile partials -> every tile's partial, in WS1 tile order.""" + + fam = self.families[family] + padded = partial.new_zeros((fam["width"], *partial.shape[1:])) + padded[: partial.shape[0]] = partial + return gather_rows(collective, padded).index_select(0, fam["gather_index"]).contiguous() + + +def _dweight_tiles(total: int, device): + begin = torch.arange(0, total, BACKWARD_TILE, device=device) + return torch.arange(total, device=device), begin, torch.clamp(begin + BACKWARD_TILE, max=total) + + +@dataclass(frozen=True) +class SPPlan: + """Everything an SP backward needs that depends only on the layout and the row index. + + Built once per ``(row index, table rows)`` and shared by every norm and gated + residual that uses that index (every block of a forward pass). + """ + + layout: SPRowLayout + plan: _Plan + local_index: torch.Tensor | None # (S_local,) table row of each local position + index_buf: torch.Tensor | None # table row of each exchange-buffer row + seg_first_tile: torch.Tensor | None + num_rows: int | None + + +def sp_plan(layout: SPRowLayout, index_full=None, num_rows: int | None = None) -> SPPlan: + """The reduction plan for one row index (``None``: the plain norm's ``dweight`` only).""" + + device = index_full.device if index_full is not None else torch.device("cuda") + families = [_dweight_tiles(layout.batch * layout.seq_len, device)] + if index_full is None: + return SPPlan(layout, _Plan(layout, families, device), None, None, None, None) + *seg_tiles, seg_first_tile = _segment_tiles(index_full.repeat(layout.batch), num_rows) + plan = _Plan(layout, [*families, tuple(seg_tiles)], device) + local_index = index_full[layout.lo : layout.hi].contiguous() + received = index_full[plan.recv_global % layout.seq_len] + index_buf = torch.cat([local_index.repeat(layout.batch), received]).contiguous() + return SPPlan(layout, plan, local_index, index_buf, seg_first_tile, num_rows) + + +def sp_rmsnorm_backward(grad, x, weight, rstd, shift, scale, sp: SPPlan, collective): + """One rank's backward on its ``(B * S_local, H)`` rows: ``(dx, dweight, d_shift, d_scale)``. + + ``dx`` is row-local; the other gradients are the full WS1 reductions, identical on + every rank. ``shift``/``scale`` are ``None`` for the plain norm. + """ + + weight, plan, modulated = weight.contiguous(), sp.plan, sp.local_index is not None + dx = _C.h3_rmsnorm_backward_dx(grad, x, weight, rstd, shift, scale, sp.local_index) + g_buf, x_buf, rstd_buf = plan.exchange(collective, grad, x, rstd) + seg = plan.families[1]["tiles"] if modulated else (None, None, None) + partials = _C.h3_rmsnorm_backward_partials( + g_buf, x_buf, weight, rstd_buf, shift, scale, sp.index_buf, + *plan.families[0]["tiles"], *seg, + ) # fmt: skip + gathered = [plan.gather_partials(collective, i, p) for i, p in enumerate(partials)] + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_rmsnorm_backward_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + if not modulated: + (dweight,) = _C.h3_rmsnorm_fold_partials(gathered[0], weight) + return dx, dweight, None, None + dweight, d_shift, d_scale = _C.h3_rmsnorm_fold_partials( + gathered[0], weight, gathered[1], sp.seg_first_tile + ) + return dx, dweight, d_shift.to(shift.dtype), d_scale.to(scale.dtype) + + +def sp_gate_residual_backward(grad, y, gate, sp: SPPlan, collective): + """One rank's backward on its rows: ``(d_sublayer, d_gate)``; ``d_residual`` is ``grad``.""" + + plan = sp.plan + dy = _C.h3_gate_residual_backward_dy(grad, gate, sp.local_index) + g_buf, y_buf = plan.exchange(collective, grad, y) + partial = _C.h3_gate_grad_partials(g_buf, y_buf, *plan.families[1]["tiles"]) + dgate = _C.h3_gate_grad_fold( + plan.gather_partials(collective, 1, partial), sp.seg_first_tile, gate.dtype + ) + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_gate_grad_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + return dy, dgate + + +class _SPRMSNorm(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, sp, collective, eps): + x2 = x.contiguous().view(-1, x.shape[-1]) + y, rstd = _C.h3_rmsnorm_forward( + x2, weight.contiguous(), float(eps), shift, scale, sp.local_index + ) + ctx.save_for_backward(x2, weight, rstd, shift, scale) + ctx.sp, ctx.collective, ctx.x_shape = sp, collective, x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + dx, dweight, d_shift, d_scale = sp_rmsnorm_backward( + g2, x2, weight, rstd, shift, scale, ctx.sp, ctx.collective + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None, None + + +class _SPGateResidual(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, sp, collective): + hidden = residual.shape[-1] + res2 = residual.contiguous().view(-1, hidden) + y2 = y.contiguous().view(-1, hidden) + out = _C.h3_gate_residual_forward(res2, y2, gate, sp.local_index) + ctx.save_for_backward(y2, gate) + ctx.sp, ctx.collective, ctx.shape = sp, collective, residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + dy, dgate = sp_gate_residual_backward(g2, y2, gate, ctx.sp, ctx.collective) + return grad, dy.view(ctx.shape), dgate, None, None + + +class H3SPNormAdaLNCudaOp: + """Sequence-parallel ``h3_rmsnorm`` (+ modulation) and ``adaln_gate_residual``. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` and + ``backend_id`` (e.g. ``DeterministicCollective``). Activations are this rank's + ``(B, S_local, H)`` rows; the norm weight, the table views and the ``(S,)`` row + index are the full, replicated tensors. Gradients of the replicated tensors come + back identical on every rank, equal to WS1. + """ + + op_class = "reduction" + backward_impl = BACKWARD_IMPL + + def __init__(self, collective, seq_len: int, batch: int) -> None: + if not sp_norm_adaln_available(): + raise RuntimeError("rl_engine._C lacks the H3 SP symbols; rebuild the CUDA extension") + self.collective = collective + self.layout = sp_row_layout(seq_len, batch, collective.world_size, collective.rank) + self._plans: list[tuple[object, int, int | None, SPPlan]] = [] + + def plan(self, index: torch.Tensor | None, num_rows: int | None) -> SPPlan: + """The cached plan for this index object (rebuilt if it was modified in place).""" + + version = -1 if index is None else index._version + for obj, ver, rows, plan in self._plans: + if obj is index and ver == version and rows == num_rows: + return plan + plan = sp_plan(self.layout, index, num_rows) + # Holding the index keeps its storage alive, so a new tensor can never alias it. + self._plans = [(index, version, num_rows, plan), *self._plans[:3]] + return plan + + def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | None) -> None: + expected = (self.layout.batch, self.layout.local_len) + if not isinstance(x, torch.Tensor) or x.dim() != 3 or tuple(x.shape[:2]) != expected: + shape = tuple(x.shape) if isinstance(x, torch.Tensor) else type(x).__name__ + raise ValueError(f"activations must be this rank's {expected} + (H,) rows, got {shape}") + if not x.is_cuda: + raise ValueError("H3SPNormAdaLNCudaOp needs CUDA tensors") + if index is None: + return + if index.dim() != 1 or index.shape[0] != self.layout.seq_len: + raise ValueError( + f"index must be the full ({self.layout.seq_len},) row index, " + f"got {tuple(index.shape)}" + ) + if index.numel() and (int(index.min()) < 0 or int(index.max()) >= num_rows): + raise IndexError(f"index must be in [0, {num_rows})") + + def norm(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + self._check(x, None, None) + validate_h3_rmsnorm(x, weight, eps) + return _SPRMSNorm.apply(x, weight, None, None, self.plan(None, None), self.collective, eps) + + def norm_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + self._check(x, index, shift.shape[0]) + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index[self.layout.lo : self.layout.hi]) + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + sp = self.plan(index, shift.shape[0]) + return _SPRMSNorm.apply(x, weight, shift, scale, sp, self.collective, eps) + + def gate_residual(self, residual, y, gate, index) -> torch.Tensor: + self._check(residual, index, gate.shape[0]) + validate_h3_gate_residual(residual, y, gate, index[self.layout.lo : self.layout.hi]) + if gate.stride(1) != 1: + gate = gate.contiguous() + sp = self.plan(index, gate.shape[0]) + return _SPGateResidual.apply(residual, y, gate, sp, self.collective) + + def readback(self) -> dict: + return { + "op": "sp_norm_adaln", + "backward_impl": BACKWARD_IMPL, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_256_element_tiles_then_ascending_fold", + "sp": self.layout.sp, + "rank": self.layout.rank, + "positions": [self.layout.lo, self.layout.hi], + "batch": self.layout.batch, + "fallback": None, + } diff --git a/rl_engine/testing/h3_cases.py b/rl_engine/testing/h3_cases.py index 527a820ec..512d3ca4a 100644 --- a/rl_engine/testing/h3_cases.py +++ b/rl_engine/testing/h3_cases.py @@ -33,3 +33,27 @@ def h3_packed_layout( timestep_indices = torch.randint(0, num_timesteps, (seq_len,), generator=generator) token_tags[: min(3, seq_len)] = torch.arange(min(3, seq_len)) return timestep_indices.to(device), token_tags.to(device) + + +def h3_block_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """(timestep_indices, token_tags) for block-structured packing, as H3 packs requests. + + Each timestep owns one contiguous run of text, video and audio blocks (about + 5% / 80% / 15% of its tokens), so table rows form contiguous position runs. + ``h3_packed_layout`` is the interleaved stress case. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + cuts = torch.sort(torch.randperm(seq_len - 1, generator=generator)[: num_timesteps - 1] + 1) + edges = [0, *cuts.values.tolist(), seq_len] + timestep_indices = torch.empty(seq_len, dtype=torch.long) + token_tags = torch.empty(seq_len, dtype=torch.long) + for t, (lo, hi) in enumerate(zip(edges, edges[1:])): + text, audio = (hi - lo) // 20, (hi - lo) * 3 // 20 + timestep_indices[lo:hi] = t + token_tags[lo:hi] = 0 # video + token_tags[lo : lo + text] = 1 + token_tags[hi - audio : hi] = 2 + return timestep_indices.to(device), token_tags.to(device) diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py index a0dc116f6..b2e3b448f 100644 --- a/rl_engine/testing/h3_ws2.py +++ b/rl_engine/testing/h3_ws2.py @@ -120,6 +120,187 @@ def step(op, args): return {"equality": equality, "perf": perf, "readback": tp_op.readback()} +def _region(gate_fn, norm_fn, residual, y, weight, table): + """One block's MLP-side SP region: ``norm2(residual + gate_msa * y)`` with modulation.""" + + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = table.chunk(6, dim=-1) + hidden = gate_fn(residual, y, gate_msa) + return norm_fn(hidden, weight, shift_mlp, scale_mlp) + + +def ws1_sp_region(residual, y, weight, table, index, grad) -> dict[str, torch.Tensor]: + """The SP region on one GPU with the WS1 ops (the reference for ``sp_norm_adaln``).""" + + from rl_engine.kernels.ops.cuda.h3.gate_residual import H3GateResidualCudaOp + from rl_engine.kernels.ops.cuda.h3.rmsnorm import H3RMSNormCudaOp + + gate_op, norm_op = H3GateResidualCudaOp(), H3RMSNormCudaOp() + leaves = [t.detach().clone().requires_grad_(True) for t in (residual, y, weight, table)] + out = _region( + lambda r, yy, g: gate_op(r, yy, g, index), + lambda h, w, sh, sc: norm_op.forward_modulated(h, w, sh, sc, index), + *leaves, + ) + out.backward(grad) + names = ("d_residual", "d_y", "d_weight", "d_table") + return {"out": out.detach(), **{n: leaf.grad for n, leaf in zip(names, leaves)}} + + +def sp_region_rank(collective, residual, y, weight, table, index, grad) -> dict[str, Any]: + """One SP rank: its rows of the region and its (replicated) parameter gradients.""" + + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + op = H3SPNormAdaLNCudaOp(collective, seq_len=residual.shape[1], batch=residual.shape[0]) + rows = slice(op.layout.lo, op.layout.hi) + local = [t[:, rows].contiguous() for t in (residual, y)] + leaves = [t.detach().clone().requires_grad_(True) for t in (*local, weight, table)] + out = _region( + lambda r, yy, g: op.gate_residual(r, yy, g, index), + lambda h, w, sh, sc: op.norm_modulated(h, w, sh, sc, index), + *leaves, + ) + out.backward(grad[:, rows].contiguous()) + names = ("d_residual", "d_y", "d_weight", "d_table") + return { + "out": out.detach(), + **{n: leaf.grad for n, leaf in zip(names, leaves)}, + "readback": op.readback(), + } + + +def sp_matches_ws1(ws1: dict[str, torch.Tensor], ranks: list[dict[str, Any]]) -> dict[str, bool]: + """Rows: each rank's slice equals WS1's; parameters: every rank's copy equals WS1's.""" + + def rows(key): + return all( + torch.equal(r[key], ws1[key][:, slice(*r["readback"]["positions"])]) for r in ranks + ) + + return { + "out": rows("out"), + "d_residual": rows("d_residual"), + "d_y": rows("d_y"), + "d_weight": all(torch.equal(r["d_weight"], ws1["d_weight"]) for r in ranks), + "d_table": all(torch.equal(r["d_table"], ws1["d_table"]) for r in ranks), + } + + +def sp_case(batch, seq, hidden=5376, num_t=3, *, layout="block", dtype=torch.bfloat16, seed=0): + """Inputs of the SP region; identical on every rank (CPU generator).""" + + from rl_engine.testing.h3_cases import h3_block_layout, h3_packed_layout + + g = torch.Generator(device="cpu").manual_seed(seed) + make = h3_block_layout if layout == "block" else h3_packed_layout + ti, tags = make(seq, num_t, seed=seed) + return { + "residual": torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda(), + "y": (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda(), + "weight": (torch.rand(hidden, generator=g) + 0.5).to(dtype).cuda(), + "table": (torch.randn(3 * num_t, 6 * hidden, generator=g) * 0.5).to(dtype).cuda(), + "index": (ti * 3 + tags).cuda(), + "grad": torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda(), + } + + +def naive_sp_mismatch(case: dict[str, torch.Tensor], sp: int) -> dict[str, float]: + """Fraction of parameter-gradient elements where a naive SP backward differs from WS1. + + Naive: each rank runs the WS1 backward on its own rows, then the per-rank results + are summed in rank order (what a plain all-reduce of local gradients does). + """ + + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import sp_row_layout + + ws1 = ws1_sp_region(**case) + parts = [] + for rank in range(sp): + lay = sp_row_layout(case["residual"].shape[1], case["residual"].shape[0], sp, rank) + part = slice(lay.lo, lay.hi) + local = {k: (v[:, part] if k in ("residual", "y", "grad") else v) for k, v in case.items()} + local["index"] = case["index"][part] + parts.append(ws1_sp_region(**local)) + out = {} + for key in ("d_weight", "d_table"): + total = parts[0][key].float() + for p in parts[1:]: + total = total + p[key].float() + out[key] = (total.to(ws1[key].dtype) != ws1[key]).float().mean().item() + return out + + +def sp_region_sweep(collective, cases: list[dict[str, Any]], iters: int = 20) -> dict[str, Any]: + """Evidence for one SP rank: byte equality, rows exchanged and time, per case.""" + + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import _dweight_tiles, _Plan, sp_row_layout + + results = [] + for spec in cases: + case = sp_case(**spec) + batch, seq = case["residual"].shape[:2] + ws1 = ws1_sp_region(**case) + mine = sp_region_rank(collective, **case) + equal = sp_matches_ws1(ws1, [mine]) + lay = sp_row_layout(seq, batch, collective.world_size, collective.rank) + tiles = _segment_tiles(case["index"].repeat(batch), case["table"].shape[0]) + plan = _Plan(lay, [_dweight_tiles(batch * seq, "cuda"), tuple(tiles[:3])], "cuda") + del ws1, mine + torch.cuda.empty_cache() + results.append( + { + **spec, + "equal": equal, + "local_rows": batch * lay.local_len, + "rows_sent": int(plan.send_local.numel()), + "rows_gathered_per_rank": plan.max_send, + **_region_times(collective, case, lay, iters), + } + ) + return {"cases": results} + + +def _region_times(collective, case, lay, iters) -> dict[str, float]: + """Forward + backward time of the region (leaves prepared once, grads reset per call).""" + + from rl_engine.kernels.ops.cuda.h3.gate_residual import H3GateResidualCudaOp + from rl_engine.kernels.ops.cuda.h3.rmsnorm import H3RMSNormCudaOp + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + from rl_engine.testing.h3_report import time_us + + index, rows = case["index"], slice(lay.lo, lay.hi) + sp_op = H3SPNormAdaLNCudaOp(collective, seq_len=lay.seq_len, batch=lay.batch) + gate_op, norm_op = H3GateResidualCudaOp(), H3RMSNormCudaOp() + + def leaves(local): + acts = [case[k][:, rows] if local else case[k] for k in ("residual", "y")] + return [ + t.detach().contiguous().requires_grad_(True) + for t in (*acts, case["weight"], case["table"]) + ] + + def step(fns, args, grad): + for leaf in args: + leaf.grad = None + _region(*fns, *args).backward(grad) + + sp_fns = ( + lambda r, yy, g: sp_op.gate_residual(r, yy, g, index), + lambda h, w, sh, sc: sp_op.norm_modulated(h, w, sh, sc, index), + ) + ws1_fns = ( + lambda r, yy, g: gate_op(r, yy, g, index), + lambda h, w, sh, sc: norm_op.forward_modulated(h, w, sh, sc, index), + ) + sp_args, ws1_args = leaves(True), leaves(False) + sp_grad = case["grad"][:, rows].contiguous() + return { + "sp_forward_backward_us": time_us(lambda: step(sp_fns, sp_args, sp_grad), 3, iters), + "ws1_forward_backward_us": time_us(lambda: step(ws1_fns, ws1_args, case["grad"]), 3, iters), + } + + # --- real multi-process groups ------------------------------------------------- diff --git a/scripts/h3_ws2_evidence.py b/scripts/h3_ws2_evidence.py index c27332f77..021cfc1f8 100644 --- a/scripts/h3_ws2_evidence.py +++ b/scripts/h3_ws2_evidence.py @@ -11,6 +11,8 @@ export RL_KERNEL_H3_WEIGHTS= python scripts/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \\ --out docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json + python scripts/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \\ + --out docs/usage/evidence/h3-sp-norm-adaln-b200/report.json """ from __future__ import annotations @@ -29,7 +31,13 @@ load_h3_conditioning_weights, load_h3_manifest, ) -from rl_engine.testing.h3_ws2 import run_world, tp_projection_sweep # noqa: E402 +from rl_engine.testing.h3_ws2 import ( # noqa: E402 + naive_sp_mismatch, + run_world, + sp_case, + sp_region_sweep, + tp_projection_sweep, +) WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" BIAS = "transformer_blocks.0.adaln_proj.linear.bias" @@ -56,7 +64,37 @@ def _tp_adaln(worlds: list[int], max_t: int) -> dict: return {"weights": [WEIGHT, BIAS], "runs": runs} -OPS = {"tp_adaln_3mod": _tp_adaln} +SP_CASES = [ + {"batch": 1, "seq": 4097, "layout": "block", "seed": 1}, + {"batch": 1, "seq": 4097, "layout": "interleaved", "seed": 2}, + {"batch": 2, "seq": 4097, "layout": "block", "seed": 3}, + {"batch": 1, "seq": 32768, "layout": "block", "seed": 4}, + {"batch": 1, "seq": 32768, "layout": "interleaved", "seed": 5}, + {"batch": 1, "seq": 131072, "layout": "block", "seed": 6}, +] + + +def _sp_norm(worlds: list[int], _max_t: int) -> dict: + runs = [] + for world in worlds: + ranks = run_world(world, sp_region_sweep, {}, cases=SP_CASES) + runs.append({"sp": world, "ranks": ranks}) + equal = all(all(c["equal"].values()) for r in ranks for c in r["cases"]) + print(f"sp={world}: byte-equal to WS1 on every rank and case: {equal}") + naive = { + layout: naive_sp_mismatch(sp_case(1, 32768, layout=layout, seed=4), max(worlds)) + for layout in ("block", "interleaved") + } + region = "norm2(residual + gate_msa[row] * y) * (1 + scale_mlp[row]) + shift_mlp[row]" + return { + "region": region, + "hidden": 5376, + "runs": runs, + "naive_sp": {"sp": max(worlds), **naive}, + } + + +OPS = {"tp_adaln_3mod": _tp_adaln, "sp_norm_adaln": _sp_norm} def main() -> None: diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index ba4b3b4ef..6ee3de114 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -511,43 +511,13 @@ def plot_tp_adaln(report: dict[str, Any]): report, "tp_adaln_3mod (column-parallel AdaLN projection, 2688 -> 96768, NCCL)" ) runs = report["runs"] - # Left: every rank, every T, every tensor byte-equal to WS1? - left.set_facecolor(SURFACE) - left.set_title( - "Byte-equal to WS1 (every rank, every T)", loc="left", fontsize=11, color=INK, pad=10 - ) - for y, run in enumerate(runs): - checks = [e for r in run["ranks"] for e in r["equality"]] - for x, (key, _) in enumerate(TP_TENSORS): - ok = sum(bool(e[key]) for e in checks) - good = ok == len(checks) - left.add_patch( - matplotlib.patches.FancyBboxPatch( - (x + 0.06, y + 0.08), - 0.88, - 0.84, - boxstyle="round,pad=0,rounding_size=0.06", - facecolor=RL_KERNEL if good else PROVIDER, - edgecolor=SURFACE, - linewidth=2, - ) - ) - left.text( - x + 0.5, - y + 0.5, - f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} rank x T", - ha="center", - va="center", - fontsize=8, - color="white", - ) - left.set_xlim(0, len(TP_TENSORS)) - left.set_ylim(len(runs), 0) - left.set_xticks([i + 0.5 for i in range(len(TP_TENSORS))], [label for _, label in TP_TENSORS]) - left.set_yticks([i + 0.5 for i in range(len(runs))], [f"TP{run['tp']}" for run in runs]) - left.tick_params(colors=MUTED, labelsize=8, length=0) - for side in left.spines.values(): - side.set_visible(False) + _equality_grid( + left, + "Byte-equal to WS1 (every rank, every T)", + [(f"TP{run['tp']}", [e for r in run["ranks"] for e in r["equality"]]) for run in runs], + TP_TENSORS, + "rank x T", + ) # Right: slowest rank's time at the largest T; WS1 on one GPU as reference lines. num_t = max(p["num_timesteps"] for p in runs[0]["ranks"][0]["perf"]) @@ -589,6 +559,144 @@ def worst(run, key): return fig +def _equality_grid(ax, title, rows, tensors, unit) -> None: + """One cell per (configuration, tensor): how many checks were byte-equal to WS1.""" + + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) + for y, (_, checks) in enumerate(rows): + for x, (key, _) in enumerate(tensors): + ok = sum(bool(e[key]) for e in checks) + good = ok == len(checks) + ax.add_patch( + matplotlib.patches.FancyBboxPatch( + (x + 0.06, y + 0.08), + 0.88, + 0.84, + boxstyle="round,pad=0,rounding_size=0.06", + facecolor=RL_KERNEL if good else PROVIDER, + edgecolor=SURFACE, + linewidth=2, + ) + ) + label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} {unit}" + ax.text(x + 0.5, y + 0.5, label, ha="center", va="center", fontsize=8, color="white") + ax.set_xlim(0, len(tensors)) + ax.set_ylim(len(rows), 0) + ax.set_xticks([i + 0.5 for i in range(len(tensors))], [label for _, label in tensors]) + ax.set_yticks([i + 0.5 for i in range(len(rows))], [name for name, _ in rows]) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + for side in ax.spines.values(): + side.set_visible(False) + + +SP_TENSORS = ( + ("out", "output rows"), + ("d_residual", "d_residual"), + ("d_y", "d_sublayer"), + ("d_weight", "d_norm_w"), + ("d_table", "d_table"), +) + + +def plot_sp_norm(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(15, 3.9), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.35, 1, 0.9]} + ) + env = report["environment"] + fig.suptitle( + "sp_norm_adaln (sequence-parallel norm2(residual + gate * y) with AdaLN modulation, " + "H = 5376, NCCL)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.01, + f"{env['gpu']} x {env['gpus']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · bf16", + fontsize=7, + color=MUTED, + ) + runs = report["runs"] + _equality_grid( + axes[0], + "Byte-equal to WS1 (every rank, every case)", + [ + (f"SP{run['sp']}", [c["equal"] for r in run["ranks"] for c in r["cases"]]) + for run in runs + ], + SP_TENSORS, + "rank x case", + ) + + # Middle: the largest block-layout case, slowest rank. + big = max( + (c for c in runs[0]["ranks"][0]["cases"] if c["layout"] == "block"), + key=lambda c: c["seq"] * c["batch"], + ) + + def worst(run, key): + return ( + max( + c[key] + for r in run["ranks"] + for c in r["cases"] + if (c["seq"], c["batch"], c["layout"]) == (big["seq"], big["batch"], big["layout"]) + ) + / 1e3 + ) + + sp_runs = [run for run in runs if run["sp"] > 1] + ax = axes[1] + _style(ax, f"Region fwd+bwd, S = {big['seq']}, block layout", "ms (slowest rank)") + _bars( + ax, + [f"SP{run['sp']}" for run in sp_runs], + [("SP", RL_KERNEL, [worst(run, "sp_forward_backward_us") for run in sp_runs])], + "{:.2f}", + ) + ws1 = worst(runs[0], "ws1_forward_backward_us") + ax.axhline( + ws1, color=PROVIDER, linestyle="--", linewidth=1.4, zorder=1, label=f"WS1, 1 GPU: {ws1:.2f}" + ) + ax.set_ylim(0, ws1 * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper right") + + # Right: what a naive SP reduction would do instead. + naive = report["naive_sp"] + ax = axes[2] + _style(ax, f"Naive SP{naive['sp']} vs WS1 (S = 32768)", "% of elements that differ") + _bars( + ax, + ["d_norm_w", "d_table"], + [ + ("naive, block", PROVIDER, [100 * naive["block"][k] for k in ("d_weight", "d_table")]), + ( + "naive, interleaved", + PROVIDER, + [100 * naive["interleaved"][k] for k in ("d_weight", "d_table")], + "////", + ), + ("RL-Kernel SP", RL_KERNEL, [0.0, 0.0]), + ], + "{:.1f}", + ) + ax.text( + 0.98, + 0.62, + "naive: each rank's WS1 backward,\nthen a rank-order sum", + transform=ax.transAxes, + ha="right", + fontsize=7, + color=MUTED, + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, @@ -598,6 +706,7 @@ def worst(run, key): "adaln_gate_residual": plot_gate_residual, "final_adaln_out": plot_final, "tp_adaln_3mod": plot_tp_adaln, + "sp_norm_adaln": plot_sp_norm, } diff --git a/tests/h3/test_h3_sp_norm_adaln.py b/tests/h3/test_h3_sp_norm_adaln.py new file mode 100644 index 000000000..04d58395c --- /dev/null +++ b/tests/h3/test_h3_sp_norm_adaln.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``sp_norm_adaln``: sequence-parallel norm / modulation / gated residual. + +* ownership: contiguous position shards per batch item; bad layouts fail closed; +* backward (one GPU, ranks as threads around the plain backward functions): every + rank's row gradients equal WS1's rows, and its ``dweight``/``d_shift``/``d_scale``/ + ``d_gate`` equal WS1's, for SP 2..8, odd S, B > 1, interleaved modality rows and + bf16/fp32; +* end to end (real NCCL, autograd): the MLP-side block region + ``norm2(residual + gate_msa * y)`` with modulation is byte-equal to WS1 on every rank. +""" + +from __future__ import annotations + +import os +import threading +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import SPRowLayout, sp_row_layout +from rl_engine.testing.h3_ws2 import ( + run_world, + sp_case, + sp_matches_ws1, + sp_region_rank, + ws1_sp_region, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +ALL_TRUE = dict.fromkeys(("out", "d_residual", "d_y", "d_weight", "d_table"), True) + + +def _require_ext(): + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import sp_norm_adaln_available + + if not sp_norm_adaln_available(): + pytest.skip("rl_engine._C lacks the H3 SP symbols") + + +class _ThreadGroup: + """Ranks as threads on one GPU; ``all_gather`` is the rank-ordered concatenation. + + Only for the plain backward functions: autograd runs every backward of a device + on one worker thread, so autograd-driven ranks need real processes. + """ + + def __init__(self, world: int) -> None: + self.barrier, self.slots = threading.Barrier(world), [None] * world + + def handle(self, rank: int): + def all_gather(t): + self.slots[rank] = t + self.barrier.wait() + out = torch.cat(self.slots) + self.barrier.wait() + return out + + return SimpleNamespace(rank=rank, world_size=len(self.slots), all_gather=all_gather) + + def run(self, fn): + results, errors = [None] * len(self.slots), [] + + def body(rank): + try: + results[rank] = fn(self.handle(rank)) + except BaseException as exc: # re-raised below + errors.append(exc) + self.barrier.abort() + + threads = [threading.Thread(target=body, args=(r,)) for r in range(len(self.slots))] + for t in threads: + t.start() + for t in threads: + t.join() + if errors: + raise next( + (e for e in errors if not isinstance(e, threading.BrokenBarrierError)), errors[0] + ) + return results + + +class TestLayout: + @pytest.mark.parametrize("seq, sp", [(8, 8), (257, 2), (1000, 3), (4097, 8)]) + def test_shards_cover_every_position_once(self, seq, sp): + bounds = [SPRowLayout(seq, 1, sp, r).bounds(r) for r in range(sp)] + assert bounds[0][0] == 0 and bounds[-1][1] == seq + assert all(a[1] == b[0] and a[0] < a[1] for a, b in zip(bounds, bounds[1:])) + + @pytest.mark.parametrize("seq, batch, sp, rank", [(3, 1, 4, 0), (8, 0, 2, 0), (8, 1, 2, 2)]) + def test_rejects(self, seq, batch, sp, rank): + with pytest.raises(ValueError): + sp_row_layout(seq, batch, sp, rank) + + +@requires_cuda +class TestBackwardRanks: + @pytest.mark.parametrize( + "sp, batch, seq, hidden, layout, dtype", + [ + (2, 1, 4097, 5376, "block", torch.bfloat16), + (2, 1, 4097, 5376, "interleaved", torch.bfloat16), + (3, 2, 1000, 256, "interleaved", torch.bfloat16), + (4, 2, 777, 512, "block", torch.float32), + (5, 1, 1301, 128, "interleaved", torch.bfloat16), + (8, 2, 2049, 256, "block", torch.bfloat16), + (8, 1, 40, 64, "interleaved", torch.float32), # shards far smaller than a tile + ], + ) + def test_ranks_reproduce_ws1(self, sp, batch, seq, hidden, layout, dtype): + _require_ext() + from rl_engine import _C + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import ( + sp_gate_residual_backward, + sp_plan, + sp_rmsnorm_backward, + ) + + c = sp_case(batch, seq, hidden, layout=layout, dtype=dtype, seed=sp * 100 + seq) + chunks = c["table"].chunk(6, dim=-1) + shift, scale, gate = chunks[3], chunks[4], chunks[2] # shift/scale_mlp, gate_msa + x2, g2, y2 = (c[k].view(-1, hidden) for k in ("residual", "grad", "y")) + idx, rows = c["index"], shift.shape[0] + tiles = _segment_tiles(idx.repeat(batch), rows) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5, shift, scale, idx) + ws1_norm = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd, shift, scale, idx, *tiles) + ws1_gate = _C.h3_gate_residual_backward(g2, y2, gate, idx, *tiles) + + def rank(coll): + lay = sp_row_layout(seq, batch, sp, coll.rank) + part = slice(lay.lo, lay.hi) + local = { + k: c[k][:, part].reshape(-1, hidden).contiguous() for k in ("residual", "grad", "y") + } + rs = rstd.view(batch, seq)[:, part].reshape(-1).contiguous() + plan = sp_plan(lay, idx, rows) + norm = sp_rmsnorm_backward( + local["grad"], local["residual"], c["weight"], rs, shift, scale, plan, coll + ) + gate_grads = sp_gate_residual_backward(local["grad"], local["y"], gate, plan, coll) + return lay, norm, gate_grads + + for lay, (dx, dw, dsh, dsc), (dy, dgate) in _ThreadGroup(sp).run(rank): + part = slice(lay.lo, lay.hi) + assert torch.equal( + dx.view(batch, -1, hidden), ws1_norm[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dw, ws1_norm[1]) + assert torch.equal(dsh, ws1_norm[2].to(dtype)) and torch.equal( + dsc, ws1_norm[3].to(dtype) + ) + assert torch.equal( + dy.view(batch, -1, hidden), ws1_gate[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dgate, ws1_gate[1]) + + def test_plain_norm_ranks_reproduce_ws1(self): + _require_ext() + from rl_engine import _C + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import sp_plan, sp_rmsnorm_backward + + c = sp_case(2, 1500, 256, seed=5) + x2, g2 = c["residual"].view(-1, 256), c["grad"].view(-1, 256) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5) + ws1 = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd) + + def rank(coll): + lay = sp_row_layout(1500, 2, 4, coll.rank) + part = slice(lay.lo, lay.hi) + loc = [c[k][:, part].reshape(-1, 256).contiguous() for k in ("grad", "residual")] + rs = rstd.view(2, 1500)[:, part].reshape(-1).contiguous() + return sp_rmsnorm_backward(*loc, c["weight"], rs, None, None, sp_plan(lay), coll) + + assert all(torch.equal(dw, ws1[1]) for _, dw, _, _ in _ThreadGroup(4).run(rank)) + + def test_fails_closed(self): + _require_ext() + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + c = sp_case(1, 100, 64) + shift, scale = c["table"].chunk(6, dim=-1)[3:5] + op = H3SPNormAdaLNCudaOp(SimpleNamespace(rank=1, world_size=2), seq_len=100, batch=1) + local = c["residual"][:, 50:] + with pytest.raises(ValueError): # the full sequence instead of this rank's rows + op.norm_modulated(c["residual"], c["weight"], shift, scale, c["index"]) + with pytest.raises(ValueError): # a shard of the index instead of the full index + op.norm_modulated(local, c["weight"], shift, scale, c["index"][50:]) + with pytest.raises(IndexError): + op.gate_residual(local, c["y"][:, 50:], shift, c["index"] + 9) + with pytest.raises(ValueError): + op.norm(local.cpu(), c["weight"].cpu()) + assert op.readback()["positions"] == [50, 100] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_region_byte_equal_to_ws1(world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + case = sp_case(2, 4097, layout="interleaved", seed=world) + ws1 = {k: v.cpu() for k, v in ws1_sp_region(**case).items()} + torch.cuda.empty_cache() + ranks = run_world(world, sp_region_rank, {k: v.cpu() for k, v in case.items()}) + assert sp_matches_ws1(ws1, ranks) == ALL_TRUE + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} From ef0c9ee24bf2924af0cbdea5e6d64022118a5366 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 19:43:01 +0000 Subject: [PATCH 03/19] [WS2][MiniMax-H3] tp_adaln_3mod: 8 x B200 NCCL evidence report and figure Produced by scripts/h3_ws2_evidence.py --worlds 1,2,4,8 at 4947996 on a clean tree (real NCCL, one GPU per rank); rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-tp-adaln-3mod.md | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index da7b0ae42..6f525efe6 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -70,8 +70,24 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ## Evidence -The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 -node, one GPU per rank. +![tp_adaln_3mod on 8 x B200: byte equality and time](../usage/evidence/h3-tp-adaln-3mod-b200/figure.png) + +The report is written by `scripts/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For +TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are +byte-equal to WS1 computed on the same GPU. + +| T | | WS1 (1 GPU) | TP2 | TP4 | TP8 | +| --- | --- | --- | --- | --- | --- | +| 1 | forward | 0.10 ms | 0.12 ms | 0.11 ms | 0.12 ms | +| 1 | forward + backward | 1.90 ms | 1.49 ms | 1.19 ms | 1.04 ms | +| 3 | forward | 0.10 ms | 0.11 ms | 0.11 ms | 0.12 ms | +| 3 | forward + backward | 2.25 ms | 1.77 ms | 1.48 ms | 1.32 ms | + +Times are for the slowest rank. The WS1 forward already streams the 520 MB weight at about +5 TB/s, so the TP forward is bounded by the all-gather. `DeterministicCollective` synchronises the +host on every call, which costs about 0.1 ms. The backward is dominated by the shard-local `dW` +and gains 1.7-1.8x at TP8. Each rank also holds only `1 / tp` of the weight and its gradient. ## Tests From 63b8caa226e2af3aa49b7e9a2a913f8f235692be Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 19:43:48 +0000 Subject: [PATCH 04/19] [WS2][MiniMax-H3] sp_norm_adaln: 8 x B200 NCCL evidence report and figure Produced by scripts/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 at 6c0d380 on a clean tree (real NCCL, one GPU per rank); rendered with scripts/plot_h3_evidence.py (equality-grid label moved to the panel title) and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-sp-norm-adaln.md | 25 +++++++++++++++++++++++-- scripts/plot_h3_evidence.py | 12 ++++++------ 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 167b63f7c..8506060ea 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -70,8 +70,29 @@ How many rows move depends on how the table rows interleave along the sequence: ## Evidence -The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 -node, one GPU per rank. +![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](../usage/evidence/h3-sp-norm-adaln-b200/figure.png) + +The report is written by `scripts/h3_ws2_evidence.py` from a clean tree at commit `6c0d380`, with +real NCCL processes on one 8 x B200 node, one GPU per rank. The region is +`norm2(residual + gate_msa[row] * y)` with `shift_mlp`/`scale_mlp` modulation, H = 5376, BF16. It +covers six cases: S = 4097 (block, interleaved, B = 2), 32768 (block, interleaved) and 131072 +(block). For SP 1, 2, 4 and 8, every rank's output rows, `d_residual` and `d_sublayer`, and its +`d_norm_w` and `d_table`, are byte-equal to WS1 computed on the same GPU. + +Forward + backward of the region (slowest rank): + +| S, packing | WS1 (1 GPU) | SP2 | SP4 | SP8 | +| --- | --- | --- | --- | --- | +| 4097, block | 2.04 ms | 1.73 ms | 1.88 ms | 2.43 ms | +| 32768, block | 3.72 ms | 2.85 ms | 2.40 ms | 2.72 ms | +| 32768, interleaved | 3.72 ms | 3.45 ms | 3.88 ms | 6.82 ms | +| 131072, block | 10.73 ms | 7.27 ms | 4.66 ms | 3.92 ms | + +At long sequences the rows split and the gain grows: 2.74x at SP8 for S = 131072. At short +sequences, each backward's handful of host-synchronising all-gathers (about 0.1 ms each) +dominates. Interleaved packing also moves more rows. A naive SP8 backward (each rank's WS1 +backward, then a rank-order sum) differs from WS1 on 40% of `d_norm_w` elements and on 6% (block) +or 21% (interleaved) of `d_table` elements. ## Tests diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 6ee3de114..c4f7c05c6 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -513,10 +513,10 @@ def plot_tp_adaln(report: dict[str, Any]): runs = report["runs"] _equality_grid( left, - "Byte-equal to WS1 (every rank, every T)", + "Byte-equal to WS1", [(f"TP{run['tp']}", [e for r in run["ranks"] for e in r["equality"]]) for run in runs], TP_TENSORS, - "rank x T", + "ranks x T", ) # Right: slowest rank's time at the largest T; WS1 on one GPU as reference lines. num_t = max(p["num_timesteps"] for p in runs[0]["ranks"][0]["perf"]) @@ -563,7 +563,7 @@ def _equality_grid(ax, title, rows, tensors, unit) -> None: """One cell per (configuration, tensor): how many checks were byte-equal to WS1.""" ax.set_facecolor(SURFACE) - ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) + ax.set_title(f"{title}; checks = {unit}", loc="left", fontsize=11, color=INK, pad=10) for y, (_, checks) in enumerate(rows): for x, (key, _) in enumerate(tensors): ok = sum(bool(e[key]) for e in checks) @@ -579,7 +579,7 @@ def _equality_grid(ax, title, rows, tensors, unit) -> None: linewidth=2, ) ) - label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} {unit}" + label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)}" ax.text(x + 0.5, y + 0.5, label, ha="center", va="center", fontsize=8, color="white") ax.set_xlim(0, len(tensors)) ax.set_ylim(len(rows), 0) @@ -624,13 +624,13 @@ def plot_sp_norm(report: dict[str, Any]): runs = report["runs"] _equality_grid( axes[0], - "Byte-equal to WS1 (every rank, every case)", + "Byte-equal to WS1", [ (f"SP{run['sp']}", [c["equal"] for r in run["ranks"] for c in r["cases"]]) for run in runs ], SP_TENSORS, - "rank x case", + "ranks x cases", ) # Middle: the largest block-layout case, slowest rank. From a9ed1267326f3efb50a9cf1d2a27188fdf52664a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:11:51 +0000 Subject: [PATCH 05/19] Fix H3 native validation, backward timing, and rank cleanup Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 4 +- csrc/cuda/h3/gate_residual.cu | 7 +++- csrc/cuda/h3/rmsnorm_modulate.cu | 23 +++++++---- docs/operators/h3-adaln-projection.md | 5 +++ docs/operators/h3-tp-adaln-3mod.md | 1 + rl_engine/testing/h3_report.py | 28 ++++++++++---- rl_engine/testing/h3_ws2.py | 55 ++++++++++++++++++++++----- tests/h3/test_h3_benchmark_timing.py | 50 +++++++++++++++++++++++- tests/h3/test_h3_native_validation.py | 51 ++++++++++++++++++++++++- tests/h3/test_h3_ws2_processes.py | 47 +++++++++++++++++++++++ 10 files changed, 239 insertions(+), 32 deletions(-) create mode 100644 tests/h3/test_h3_ws2_processes.py diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..9f3d00973 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -207,9 +207,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->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 1-D int64 CUDA tensors"); + "tile metadata must be contiguous int64 tensors on ", grad.device()); } TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); const int64_t chunks = grad.size(0); diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index c2bd7921a..6e0b5a85b 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -50,6 +50,7 @@ struct Gate { int64_t stride; const int64_t* index; // (S,) int64_t seq; + int64_t num_rows; }; // One block per row (grid-strided); 16-byte vectors when the row allows it. @@ -60,7 +61,9 @@ __global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __re Gate gate, bool vectorized) { constexpr int kVecElems = 16 / sizeof(T); for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { - const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; const T* ar = a + r * n; const T* br = kDy ? nullptr : b + r * n; T* orow = out + r * n; @@ -148,7 +151,7 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); - return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; } void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index e2c05d2d0..19700d35a 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -106,6 +106,7 @@ struct Modulation { int64_t row_stride; const int64_t* index; // (S,) table row per sequence position int64_t seq; // S: row r of x uses index[r % S] + int64_t num_rows; }; template @@ -120,6 +121,7 @@ __global__ void __launch_bounds__(32 * kWarps) const T* scale = nullptr; if (mod.index != nullptr) { const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); shift = static_cast(mod.shift) + i * mod.row_stride; scale = static_cast(mod.scale) + i * mod.row_stride; } @@ -156,9 +158,12 @@ __global__ void __launch_bounds__(32 * kWarps) const int64_t r = blockIdx.x; const T* xr = x + r * n; const T* gr = g + r * n; - const T* scale = mod.index != nullptr - ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } const int numx = blockDim.x * blockDim.y; const int thrx = threadIdx.x + threadIdx.y * blockDim.x; float s = 0.0f; @@ -193,9 +198,12 @@ __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* const int64_t r1 = min(r0 + kRowTile, rows); float acc = 0.0f; for (int64_t r = r0; r < r1; ++r) { - const T* scale = mod.index != nullptr ? static_cast(mod.scale) + - mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); } partial[static_cast(blockIdx.x) * n + j] = acc; @@ -260,7 +268,7 @@ Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, int64_t n) { - Modulation mod{nullptr, nullptr, 0, nullptr, 1}; + Modulation mod{nullptr, nullptr, 0, nullptr, 1, 0}; if (!index.has_value()) { TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); return mod; @@ -287,6 +295,7 @@ Modulation make_modulation(const c10::optional& shift, mod.row_stride = sh.stride(0); mod.index = ix.data_ptr(); mod.seq = ix.size(0); + mod.num_rows = sh.size(0); return mod; } diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 910165a3b..2cd2cbe0d 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -135,6 +135,11 @@ or in a batch of 9. ![adaln_projection vs existing implementations](../usage/evidence/h3-prior-art-b200/adaln_projection.png) +The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` +streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional +weight passes and are outside its intended T ≤ 8 range. This accounts for the +steep increase in the reported CUDA forward + backward times (77.4 and 627 ms). + | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| | diffusers AdaLN projection, BF16 F.linear [plain] | **no** (14214 rows, 3080 sub-batches, 2232 sweep cases) | 3.6e-03 / 3.9e-03 / 716 µs | 3.4e-03 / 3.2e-03 / 728 µs | 3.2e-03 / 3.5e-03 / 2346 µs | diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 6f525efe6..7a33f6d26 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -83,6 +83,7 @@ byte-equal to WS1 computed on the same GPU. | 1 | forward + backward | 1.90 ms | 1.49 ms | 1.19 ms | 1.04 ms | | 3 | forward | 0.10 ms | 0.11 ms | 0.11 ms | 0.12 ms | | 3 | forward + backward | 2.25 ms | 1.77 ms | 1.48 ms | 1.32 ms | +| 4 | forward + backward | 2.09 ms | 1.83 ms | 1.57 ms | 1.45 ms | Times are for the slowest rank. The WS1 forward already streams the 520 MB weight at about 5 TB/s, so the TP forward is bounded by the all-gather. `DeterministicCollective` synchronises the diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 407906a88..3bd094314 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -590,9 +590,12 @@ def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: gate = _gate_view(table) grad = torch.randn_like(residual) - def backward(fn, residual=residual, y=y, table=table, grad=grad): + def prepare_backward(fn, residual=residual, y=y, table=table): leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] - fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + return fn(leaves[0], leaves[1], _gate_view(leaves[2])) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -606,10 +609,12 @@ def backward(fn, residual=residual, y=y, table=table, grad=grad): "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( r, g, i, y ), - "candidate_backward": lambda b=backward, i=index: b( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: provider_gate_residual(r, g, i, y_) ), } @@ -695,9 +700,12 @@ def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: x, nw, temb, w, b, ti = _final_inputs(seq) grad = torch.randn_like(x) - def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + def prepare_backward(fn, tensors=(x, nw, temb, w, b)): leaves = [t.detach().requires_grad_(True) for t in tensors] - fn(*leaves).backward(grad) + return fn(*leaves) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -708,8 +716,12 @@ def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), - "candidate_backward": lambda bw=backward, i=ti: bw(lambda *t: op(*t, i)), - "provider_backward": lambda bw=backward, i=ti: bw( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( lambda *t: provider_final_adaln_out(*t, i) ), } diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py index a0dc116f6..ad9041cf9 100644 --- a/rl_engine/testing/h3_ws2.py +++ b/rl_engine/testing/h3_ws2.py @@ -144,9 +144,9 @@ def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, inputs_path.with_name(f"rank{rank}.pt"), ) - queue.put({"rank": rank}) collective.close() dist.destroy_process_group() + queue.put({"rank": rank}) except Exception: # forwarded to the parent queue.put({"rank": rank, "error": traceback.format_exc()}) @@ -160,7 +160,9 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa """ import tempfile + import time from pathlib import Path + from queue import Empty import torch.multiprocessing as mp @@ -176,12 +178,45 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) for r in range(world) ] - for proc in procs: - proc.start() - status = sorted((queue.get(timeout=timeout) for _ in procs), key=lambda r: r["rank"]) - for proc in procs: - proc.join(timeout=120) - for result in status: - if "error" in result: - raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") - return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + started = [] + try: + for proc in procs: + proc.start() + started.append(proc) + deadline = time.monotonic() + timeout + pending = set(range(world)) + while pending: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"timed out waiting for ranks {sorted(pending)}") + try: + result = queue.get(timeout=min(0.1, remaining)) + except Empty: + for rank in pending: + if procs[rank].exitcode is not None: + raise RuntimeError( + f"rank {rank} exited with code {procs[rank].exitcode} " + "without reporting a result" + ) from None + continue + if "error" in result: + raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") + pending.remove(result["rank"]) + for rank, proc in enumerate(procs): + proc.join(timeout=max(0, deadline - time.monotonic())) + if proc.is_alive(): + raise TimeoutError(f"rank {rank} did not exit after reporting a result") + if proc.exitcode != 0: + raise RuntimeError(f"rank {rank} exited with code {proc.exitcode}") + return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + finally: + for proc in started: + if proc.is_alive(): + proc.terminate() + for proc in started: + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join() + queue.close() + queue.join_thread() diff --git a/tests/h3/test_h3_benchmark_timing.py b/tests/h3/test_h3_benchmark_timing.py index e6517bd70..47d5575b1 100644 --- a/tests/h3/test_h3_benchmark_timing.py +++ b/tests/h3/test_h3_benchmark_timing.py @@ -114,7 +114,7 @@ def run(key=key, duration=position): assert not any(key.endswith("_setup") for key in row) -@pytest.mark.parametrize("operator", ["norm", "gather"]) +@pytest.mark.parametrize("operator", ["norm", "gather", "gate", "final"]) def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( monkeypatch, cuda_clock, operator ): @@ -136,7 +136,9 @@ def backward_hook(grad): def track(fn, leaves, *args): if leaves[0].requires_grad: assert not cuda_clock.active - prepared_leaves.append(leaves) + prepared_leaves.append( + [leaf._base if leaf._base is not None else leaf for leaf in leaves] + ) cuda_clock.time += 1.0 outputs = fn(*leaves, *args) return track_backward(outputs) if leaves[0].requires_grad else outputs @@ -162,6 +164,50 @@ def forward(x, weight, shift, scale, index, **kwargs): monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) op = SimpleNamespace(forward_modulated=forward) build = h3_report._norm_perf + elif operator == "gate": + + def inputs(seq): + return ( + torch.randn(3, 6 * 4), + torch.randn(1, seq, 4), + torch.randn(1, seq, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_gate_residual + + def forward(residual, y, gate, index, **kwargs): + return track(lambda r, y_, g, i: provider(r, g, i, y_), [residual, y, gate], index) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_gate_inputs", inputs) + monkeypatch.setattr( + h3_report, "provider_gate_residual", lambda r, g, i, y: forward(r, y, g, i) + ) + op = SimpleNamespace(forward=forward) + build = h3_report._gate_perf + elif operator == "final": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 2), + torch.randn(8, 2), + torch.randn(8), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_final_adaln_out + + def forward(x, nw, temb, w, b, ti): + return track(provider, [x, nw, temb, w, b], ti) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_final_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", forward) + op = forward + build = h3_report._final_perf else: randn = torch.randn diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py index 8dd919028..8c71eba08 100644 --- a/tests/h3/test_h3_native_validation.py +++ b/tests/h3/test_h3_native_validation.py @@ -16,6 +16,7 @@ def native_extension(): from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE names = ( + "h3_adaln_row_gather_backward", "h3_rmsnorm_forward", "h3_rmsnorm_backward", "h3_gate_residual_forward", @@ -49,6 +50,10 @@ def native_case(): def _native_call(extension, operation, boundary, case): x, index = case["x"], case["index"] tiles = tuple(case[name] for name in ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile")) + if operation == "gather": + return extension.h3_adaln_row_gather_backward( + case["grad"].unsqueeze(0), *tiles, case["grad"].dtype + ) if operation == "rmsnorm": modulation = (case["shift"], case["scale"], index) if boundary == "forward": @@ -86,7 +91,7 @@ def test_native_rejects_invalid_row_index( _native_call(native_extension, operation, boundary, native_case) -@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual", "gather"]) @pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) @pytest.mark.parametrize("device", ["cpu", "other_cuda"]) def test_native_rejects_mixed_device_tiles( @@ -122,3 +127,47 @@ def test_native_accepts_same_device_inputs(native_extension, native_case, operat actual_grads = _native_call(native_extension, operation, "backward", native_case) for actual, expected in zip(actual_grads, expected_grads, strict=True): torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("boundary", ["forward", "backward"]) +@pytest.mark.parametrize("bad_index", [-1, 2]) +def test_native_rejects_out_of_range_index_in_subprocess( + native_extension, operation, boundary, bad_index +): + # A device assertion poisons the CUDA context, so isolate each invalid launch. + import os + import subprocess + import sys + + code = f""" +import torch +from rl_engine import _C +x = torch.ones(2, 8, device='cuda') +w = torch.ones(8, device='cuda') +index = torch.tensor([0, {bad_index}], device='cuda', dtype=torch.int64) +pos = torch.arange(2, device='cuda', dtype=torch.int64) +ends = pos + 1 +seg = torch.arange(3, device='cuda', dtype=torch.int64) +if {operation!r} == 'gate_residual': + if {boundary!r} == 'forward': + _C.h3_gate_residual_forward(x, x, x, index) + else: + _C.h3_gate_residual_backward(x, x, x, index, pos, pos, ends, seg) +else: + if {boundary!r} == 'forward': + _C.h3_rmsnorm_forward(x, w, 1e-5, x, x, index) + else: + rstd = torch.ones(2, device='cuda') + _C.h3_rmsnorm_backward(x, x, w, rstd, x, x, index, pos, pos, ends, seg) +torch.cuda.synchronize() +""" + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + timeout=60, + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + ) + assert result.returncode != 0 + assert "device-side assert" in result.stderr diff --git a/tests/h3/test_h3_ws2_processes.py b/tests/h3/test_h3_ws2_processes.py new file mode 100644 index 000000000..fb69430cd --- /dev/null +++ b/tests/h3/test_h3_ws2_processes.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU process lifecycle regressions for the multi-GPU evidence runner.""" + +import os +import time + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.testing import h3_ws2 + + +def _worker(rank, world, init, queue, path, target, kwargs): + mode = kwargs["mode"] + if mode == "success": + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + elif rank == 0 and mode == "error": + queue.put({"rank": rank, "error": "original rank traceback"}) + elif rank == 0 and mode == "crash": + os._exit(7) + else: + time.sleep(120) + + +@pytest.mark.parametrize("mode", ["success", "error", "crash", "timeout"]) +def test_run_world_reaps_workers(monkeypatch, mode): + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + if mode == "success": + ranks = h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert [r["rank"] for r in ranks] == [0, 1] + assert [r["value"].item() for r in ranks] == [0, 1] + elif mode == "timeout": + with pytest.raises(TimeoutError, match="timed out waiting for ranks"): + h3_ws2.run_world(2, None, {}, timeout=0.2, mode=mode) + else: + message = "original rank traceback" if mode == "error" else "rank 0 exited with code 7" + with pytest.raises(RuntimeError, match=message): + h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert time.monotonic() - start < 25 + assert {p.pid for p in mp.active_children()} == before From 72ed9702872be0ae4a5518ebe59579c264f54ead Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:14:09 +0000 Subject: [PATCH 06/19] Regenerate gate and final AdaLN backward-only B200 evidence Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 10 +++++----- docs/operators/h3-final-adaln-out.md | 13 +++++++------ 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6b923d740..8d3b4b380 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -61,13 +61,13 @@ hidden = op(residual, attn_output, gate_msa, adaln_indices) python benchmarks/benchmark_h3_conditioning.py --op adaln_gate_residual ``` -B200, B = 1, H = 5376, BF16: +B200, B = 1, H = 5376, BF16. Backward timings exclude input and forward-graph setup: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | -| 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | -| 131072 | 0.92 ms | 1.59 ms | 3.29 ms | 10.6 ms | +| 4097 | 0.07 ms | 0.09 ms | 1.12 ms | 0.58 ms | +| 32768 | 0.24 ms | 0.42 ms | 1.12 ms | 2.24 ms | +| 131072 | 0.88 ms | 1.61 ms | 2.22 ms | 8.92 ms | The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and @@ -78,7 +78,7 @@ tile setup. ![adaln_gate_residual on B200: latency and backward accuracy](../usage/evidence/h3-adaln-gate-residual-b200/figure.png) The data is in [`report.json`](../usage/evidence/h3-adaln-gate-residual-b200/report.json), -written by `scripts/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +written by `scripts/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index e7bac8c59..2c2e215c6 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -40,7 +40,7 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste summation tree versus cuBLAS. Rows are batch- and position-invariant. - **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into the projection backward without an extra BF16 rounding: - - `d_temb`, `dW` and `db` are 22–48× closer to FP64 in the recorded run than diffusers, + - `d_temb`, `dW` and `db` are 21–37× closer to FP64 in the recorded run than diffusers, whose error varies from run to run because it rounds the table gradient to BF16 and accumulates it with `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; @@ -58,13 +58,14 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste python benchmarks/benchmark_h3_conditioning.py --op final_adaln_out ``` -B200, B = 1, T = 3, pinned `norm_out` weights: +B200, B = 1, T = 3, pinned `norm_out` weights. Backward timings exclude input +and forward-graph setup: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | -| 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | -| 131072 | 1.24 ms | 2.94 ms | 7.75 ms | 20.6 ms | +| 4097 | 0.19 ms | 0.18 ms | 1.40 ms | 0.85 ms | +| 32768 | 0.46 ms | 0.80 ms | 2.59 ms | 4.49 ms | +| 131072 | 1.32 ms | 3.02 ms | 7.05 ms | 17.69 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -77,7 +78,7 @@ There are two data files: - [`report.json`](../usage/evidence/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `001684d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + tree at `bde1d8d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../usage/evidence/h3-final-adaln-out-b200/chain_replay.json): the From d9600066ccd459f6f71e2daf4e071aee6d28854a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:14:40 +0000 Subject: [PATCH 07/19] test(h3): consolidate native validation tests for review file limit Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/h3/test_h3_native_validation.py | 153 +++++++++++++++++ tests/h3/test_h3_rmsnorm_native_validation.py | 158 ------------------ 2 files changed, 153 insertions(+), 158 deletions(-) delete mode 100644 tests/h3/test_h3_rmsnorm_native_validation.py diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py index 8dd919028..9f2a06a07 100644 --- a/tests/h3/test_h3_native_validation.py +++ b/tests/h3/test_h3_native_validation.py @@ -122,3 +122,156 @@ def test_native_accepts_same_device_inputs(native_extension, native_case, operat actual_grads = _native_call(native_extension, operation, "backward", native_case) for actual, expected in zip(actual_grads, expected_grads, strict=True): torch.testing.assert_close(actual, expected) + + +TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") + + +@pytest.fixture +def native(): + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not all( + hasattr(_C, name) for name in ("h3_rmsnorm_forward", "h3_rmsnorm_backward") + ): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return _C + + +@pytest.fixture +def inputs(native): + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device=x.device) + _, rstd = native.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device=x.device).chunk(2, dim=1) + return { + "grad": torch.ones_like(x), + "x": x, + "weight": weight, + "rstd": rstd, + "shift": shift, + "scale": scale, + "index": torch.tensor([0, 1], dtype=torch.int64, device=x.device), + "sorted_pos": torch.tensor([0, 2, 1, 3], dtype=torch.int64, device=x.device), + "tile_begin": torch.tensor([0, 2], dtype=torch.int64, device=x.device), + "tile_end": torch.tensor([2, 4], dtype=torch.int64, device=x.device), + "seg_first_tile": torch.tensor([0, 1, 2], dtype=torch.int64, device=x.device), + } + + +def _call(native, inputs, entrypoint): + if entrypoint == "forward": + return native.h3_rmsnorm_forward( + inputs["x"], + inputs["weight"], + 1e-5, + inputs["shift"], + inputs["scale"], + inputs["index"], + ) + return native.h3_rmsnorm_backward(**inputs) + + +def _malformed(tensor, defect): + if defect == "cpu": + return tensor.cpu() + if defect == "dtype": + dtype = torch.float16 if tensor.dtype == torch.float32 else torch.int32 + return tensor.to(dtype) + if defect == "rank": + return tensor.unsqueeze(0) + if defect == "stride": + return torch.stack((tensor, tensor), dim=1)[:, 0] + if defect == "length": + return tensor[:-1] + raise ValueError(f"unknown defect: {defect}") + + +@pytest.mark.cuda_only +def test_valid_plain_and_modulated_backward(native, inputs): + plain = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + assert len(native.h3_rmsnorm_backward(**plain)) == 2 + dx, dweight, dshift, dscale = native.h3_rmsnorm_backward(**inputs) + assert dx.shape == inputs["x"].shape + assert dweight.shape == inputs["weight"].shape + assert dscale.shape == inputs["scale"].shape + torch.testing.assert_close(dshift, torch.full_like(inputs["shift"], 2.0)) + torch.cuda.synchronize(inputs["x"].device) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_empty_row_index(native, inputs, entrypoint): + inputs["index"] = inputs["index"][:0] + with pytest.raises(RuntimeError, match="row index must be a non-empty"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_zero_hidden_size(native, inputs, entrypoint): + inputs["x"] = inputs["x"][:, :0].contiguous() + inputs["weight"] = inputs["weight"][:0] + with pytest.raises(RuntimeError, match="at least one column"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +@pytest.mark.parametrize("mode", ["plain", "modulated"]) +def test_rejects_invalid_rstd(native, inputs, defect, mode): + inputs["rstd"] = _malformed(inputs["rstd"], defect) + if mode == "plain": + inputs = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +def test_rejects_invalid_segment_tiles(native, inputs, name, defect): + inputs[name] = _malformed(inputs[name], defect) + message = "tile_begin and tile_end" if defect == "length" and name.startswith("tile_") else name + with pytest.raises(RuntimeError, match=message): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_rejects_missing_segment_tiles(native, inputs, name): + inputs[name] = None + with pytest.raises(RuntimeError, match="needs the sorted segment tiles"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_plain_backward_rejects_segment_tiles(native, inputs, name): + plain = {key: inputs[key] for key in ("grad", "x", "weight", "rstd", name)} + with pytest.raises(RuntimeError, match="sorted segment tiles require modulation"): + native.h3_rmsnorm_backward(**plain) + + +@pytest.mark.cuda_only +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize( + ("entrypoint", "name"), + [ + ("forward", "weight"), + ("forward", "index"), + ("backward", "weight"), + ("backward", "index"), + ("backward", "rstd"), + ("backward", "sorted_pos"), + ("backward", "tile_begin"), + ("backward", "tile_end"), + ("backward", "seg_first_tile"), + ], +) +def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name): + other_device = (inputs["x"].device.index + 1) % torch.cuda.device_count() + inputs[name] = inputs[name].to(f"cuda:{other_device}") + message = "row index" if name == "index" else name + with pytest.raises(RuntimeError, match=message): + _call(native, inputs, entrypoint) diff --git a/tests/h3/test_h3_rmsnorm_native_validation.py b/tests/h3/test_h3_rmsnorm_native_validation.py deleted file mode 100644 index d6f2a0d81..000000000 --- a/tests/h3/test_h3_rmsnorm_native_validation.py +++ /dev/null @@ -1,158 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""Direct native calls reject malformed pointers and layouts before launching CUDA work.""" - -from __future__ import annotations - -import pytest -import torch - -pytestmark = [ - pytest.mark.cuda_only, - pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), -] - -TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") - - -@pytest.fixture -def native(): - from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE - - if not _EXT_AVAILABLE or not all( - hasattr(_C, name) for name in ("h3_rmsnorm_forward", "h3_rmsnorm_backward") - ): - pytest.skip("rl_engine._C lacks h3_rmsnorm_*") - return _C - - -@pytest.fixture -def inputs(native): - x = torch.randn(4, 8, device="cuda") - weight = torch.ones(8, device=x.device) - _, rstd = native.h3_rmsnorm_forward(x, weight, 1e-5) - shift, scale = torch.zeros(2, 16, device=x.device).chunk(2, dim=1) - return { - "grad": torch.ones_like(x), - "x": x, - "weight": weight, - "rstd": rstd, - "shift": shift, - "scale": scale, - "index": torch.tensor([0, 1], dtype=torch.int64, device=x.device), - "sorted_pos": torch.tensor([0, 2, 1, 3], dtype=torch.int64, device=x.device), - "tile_begin": torch.tensor([0, 2], dtype=torch.int64, device=x.device), - "tile_end": torch.tensor([2, 4], dtype=torch.int64, device=x.device), - "seg_first_tile": torch.tensor([0, 1, 2], dtype=torch.int64, device=x.device), - } - - -def _call(native, inputs, entrypoint): - if entrypoint == "forward": - return native.h3_rmsnorm_forward( - inputs["x"], - inputs["weight"], - 1e-5, - inputs["shift"], - inputs["scale"], - inputs["index"], - ) - return native.h3_rmsnorm_backward(**inputs) - - -def _malformed(tensor, defect): - if defect == "cpu": - return tensor.cpu() - if defect == "dtype": - dtype = torch.float16 if tensor.dtype == torch.float32 else torch.int32 - return tensor.to(dtype) - if defect == "rank": - return tensor.unsqueeze(0) - if defect == "stride": - return torch.stack((tensor, tensor), dim=1)[:, 0] - if defect == "length": - return tensor[:-1] - raise ValueError(f"unknown defect: {defect}") - - -def test_valid_plain_and_modulated_backward(native, inputs): - plain = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} - assert len(native.h3_rmsnorm_backward(**plain)) == 2 - dx, dweight, dshift, dscale = native.h3_rmsnorm_backward(**inputs) - assert dx.shape == inputs["x"].shape - assert dweight.shape == inputs["weight"].shape - assert dscale.shape == inputs["scale"].shape - torch.testing.assert_close(dshift, torch.full_like(inputs["shift"], 2.0)) - torch.cuda.synchronize(inputs["x"].device) - - -@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) -def test_rejects_empty_row_index(native, inputs, entrypoint): - inputs["index"] = inputs["index"][:0] - with pytest.raises(RuntimeError, match="row index must be a non-empty"): - _call(native, inputs, entrypoint) - - -@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) -def test_rejects_zero_hidden_size(native, inputs, entrypoint): - inputs["x"] = inputs["x"][:, :0].contiguous() - inputs["weight"] = inputs["weight"][:0] - with pytest.raises(RuntimeError, match="at least one column"): - _call(native, inputs, entrypoint) - - -@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) -@pytest.mark.parametrize("mode", ["plain", "modulated"]) -def test_rejects_invalid_rstd(native, inputs, defect, mode): - inputs["rstd"] = _malformed(inputs["rstd"], defect) - if mode == "plain": - inputs = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} - with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): - native.h3_rmsnorm_backward(**inputs) - - -@pytest.mark.parametrize("name", TILE_NAMES) -@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) -def test_rejects_invalid_segment_tiles(native, inputs, name, defect): - inputs[name] = _malformed(inputs[name], defect) - message = "tile_begin and tile_end" if defect == "length" and name.startswith("tile_") else name - with pytest.raises(RuntimeError, match=message): - native.h3_rmsnorm_backward(**inputs) - - -@pytest.mark.parametrize("name", TILE_NAMES) -def test_rejects_missing_segment_tiles(native, inputs, name): - inputs[name] = None - with pytest.raises(RuntimeError, match="needs the sorted segment tiles"): - native.h3_rmsnorm_backward(**inputs) - - -@pytest.mark.parametrize("name", TILE_NAMES) -def test_plain_backward_rejects_segment_tiles(native, inputs, name): - plain = {key: inputs[key] for key in ("grad", "x", "weight", "rstd", name)} - with pytest.raises(RuntimeError, match="sorted segment tiles require modulation"): - native.h3_rmsnorm_backward(**plain) - - -@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") -@pytest.mark.parametrize( - ("entrypoint", "name"), - [ - ("forward", "weight"), - ("forward", "index"), - ("backward", "weight"), - ("backward", "index"), - ("backward", "rstd"), - ("backward", "sorted_pos"), - ("backward", "tile_begin"), - ("backward", "tile_end"), - ("backward", "seg_first_tile"), - ], -) -def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name): - other_device = (inputs["x"].device.index + 1) % torch.cuda.device_count() - inputs[name] = inputs[name].to(f"cuda:{other_device}") - message = "row index" if name == "index" else name - with pytest.raises(RuntimeError, match=message): - _call(native, inputs, entrypoint) From 86c168f2f5091b8f3b6b947dd3ef5a15ce30b19d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:42:12 +0000 Subject: [PATCH 08/19] fix(h3): address automated review bounds, device, and benchmark findings Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 20 +++- csrc/cuda/h3/gate_residual.cu | 7 +- csrc/cuda/h3/rmsnorm_modulate.cu | 23 ++-- docs/operators/h3-adaln-gate-residual.md | 7 +- docs/operators/h3-final-adaln-out.md | 7 +- docs/operators/h3-sp-norm-adaln.md | 2 +- docs/operators/h3-tp-adaln-3mod.md | 6 +- .../kernels/ops/cuda/h3/sp_norm_adaln.py | 27 +++-- .../ops/cuda/h3/tp_adaln_projection.py | 2 + rl_engine/testing/h3_chain.py | 4 +- rl_engine/testing/h3_report.py | 30 +++-- rl_engine/testing/h3_ws2.py | 62 +++++++++-- scripts/plot_h3_evidence.py | 14 ++- tests/h3/test_h3_benchmark_timing.py | 52 ++++++++- tests/h3/test_h3_native_validation.py | 103 ++++++++++++++++++ tests/h3/test_h3_sp_norm_adaln.py | 103 ++++++++++++++++++ tests/h3/test_h3_tp_adaln.py | 8 +- 17 files changed, 423 insertions(+), 54 deletions(-) diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..0de2d98d5 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -110,7 +110,10 @@ __global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, const int64_t h = j - c * hidden; const g_t* g = grad + c * seq * hidden + h; float acc = 0.0f; - for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t begin = tile_begin[tile], end = tile_end[tile]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= seq); + for (int64_t p = begin; p < end; ++p) { + CUDA_KERNEL_ASSERT(sorted_pos[p] >= 0 && sorted_pos[p] < seq); acc += to_float(g[sorted_pos[p] * hidden]); } partial[tile * width + j] = acc; @@ -119,12 +122,14 @@ __global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, template __global__ void adaln_row_gather_fold_kernel(const float* __restrict__ partial, const int64_t* __restrict__ seg_first_tile, - out_t* __restrict__ out, int64_t width) { + out_t* __restrict__ out, int64_t width, int64_t tiles) { const int64_t r = blockIdx.x; const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; if (j >= width) return; float acc = 0.0f; - for (int64_t tile = seg_first_tile[r]; tile < seg_first_tile[r + 1]; ++tile) { + const int64_t begin = seg_first_tile[r], end = seg_first_tile[r + 1]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= tiles); + for (int64_t tile = begin; tile < end; ++tile) { acc += partial[tile * width + j]; } out[r * width + j] = from_float(acc); @@ -207,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->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 1-D int64 CUDA tensors"); + "tile metadata must be contiguous 1-D int64 tensors on grad device"); } TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); const int64_t chunks = grad.size(0); @@ -218,6 +224,8 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor const int64_t width = chunks * hidden; const int64_t num_rows = seg_first_tile.numel() - 1; const int64_t tiles = tile_begin.numel(); + TORCH_CHECK(chunks > 0 && seq > 0 && hidden > 0 && num_rows > 0, + "gradient dimensions and table row count must be positive"); TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries"); const c10::cuda::CUDAGuard device_guard(grad.device()); auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype)); @@ -244,7 +252,7 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor adaln_row_gather_fold_kernel <<(num_rows), col_blocks), threads, 0, stream>>>( partial.data_ptr(), seg_first_tile.data_ptr(), - reinterpret_cast(out.data_ptr()), width); + reinterpret_cast(out.data_ptr()), width, tiles); }); C10_CUDA_KERNEL_LAUNCH_CHECK(); return out; diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 2d8aea209..eb0b03f50 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -52,6 +52,7 @@ struct Gate { int64_t stride; const int64_t* index; // (S,) int64_t seq; + int64_t num_rows; }; // One block per row (grid-strided); 16-byte vectors when the row allows it. @@ -62,7 +63,9 @@ __global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __re Gate gate, bool vectorized) { constexpr int kVecElems = 16 / sizeof(T); for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { - const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; const T* ar = a + r * n; const T* br = kDy ? nullptr : b + r * n; T* orow = out + r * n; @@ -150,7 +153,7 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); - return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; } void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 42e3ee18a..f9304ab13 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -108,6 +108,7 @@ struct Modulation { const void* shift; // (R, N) rows with stride `row_stride` (elements), or nullptr const void* scale; int64_t row_stride; + int64_t num_rows; const int64_t* index; // (S,) table row per sequence position int64_t seq; // S: row r of x uses index[r % S] }; @@ -124,6 +125,7 @@ __global__ void __launch_bounds__(32 * kWarps) const T* scale = nullptr; if (mod.index != nullptr) { const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); shift = static_cast(mod.shift) + i * mod.row_stride; scale = static_cast(mod.scale) + i * mod.row_stride; } @@ -160,9 +162,12 @@ __global__ void __launch_bounds__(32 * kWarps) const int64_t r = blockIdx.x; const T* xr = x + r * n; const T* gr = g + r * n; - const T* scale = mod.index != nullptr - ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } const int numx = blockDim.x * blockDim.y; const int thrx = threadIdx.x + threadIdx.y * blockDim.x; float s = 0.0f; @@ -204,9 +209,12 @@ __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* float acc = 0.0f; for (int64_t p = p0; p < p1; ++p) { const int64_t r = row_list != nullptr ? row_list[p] : p; - const T* scale = mod.index != nullptr ? static_cast(mod.scale) + - mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); } partial[static_cast(blockIdx.x) * n + j] = acc; @@ -271,7 +279,7 @@ Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, int64_t n) { - Modulation mod{nullptr, nullptr, 0, nullptr, 1}; + Modulation mod{nullptr, nullptr, 0, 0, nullptr, 1}; if (!index.has_value()) { TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); return mod; @@ -296,6 +304,7 @@ Modulation make_modulation(const c10::optional& shift, mod.shift = sh.data_ptr(); mod.scale = sc.data_ptr(); mod.row_stride = sh.stride(0); + mod.num_rows = sh.size(0); mod.index = ix.data_ptr(); mod.seq = ix.size(0); return mod; diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6b923d740..9d5ee7e09 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -63,7 +63,7 @@ python benchmarks/benchmark_h3_conditioning.py --op adaln_gate_residual B200, B = 1, H = 5376, BF16: -| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| S | CUDA fwd | diffusers fwd | CUDA fwd+bwd | diffusers fwd+bwd | | --- | --- | --- | --- | --- | | 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | | 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | @@ -73,6 +73,11 @@ The forward is one pass over `residual` and `sublayer_output`, using 16-byte vec block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and tile setup. +The stored timing evidence includes forward graph setup in the columns labeled +`fwd+bwd`; those values are not backward-only measurements. The current benchmark +prepares the graph outside the timed backward region. Historical timings are retained +with their original commit and corrected scope. + ## Evidence ![adaln_gate_residual on B200: latency and backward accuracy](../usage/evidence/h3-adaln-gate-residual-b200/figure.png) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index e7bac8c59..6022af96a 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -60,7 +60,7 @@ python benchmarks/benchmark_h3_conditioning.py --op final_adaln_out B200, B = 1, T = 3, pinned `norm_out` weights: -| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| S | CUDA fwd | diffusers fwd | CUDA fwd+bwd | diffusers fwd+bwd | | --- | --- | --- | --- | --- | | 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | | 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | @@ -69,6 +69,11 @@ B200, B = 1, T = 3, pinned `norm_out` weights: At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. +The stored timing evidence includes forward graph setup in the columns labeled +`fwd+bwd`; those values are not backward-only measurements. The current benchmark +prepares the graph outside the timed backward region. Historical timings are retained +with their original commit and corrected scope. + ## Evidence ![final_adaln_out on B200: latency and backward accuracy](../usage/evidence/h3-final-adaln-out-b200/figure.png) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 8506060ea..65b02d021 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -65,7 +65,7 @@ How many rows move depends on how the table rows interleave along the sequence: | Packing | Rows sent per rank (S = 32768, SP8, 4096 rows per rank) | | --- | --- | -| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 174 | +| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 224 | | interleaved (modality and timestep random per position; stress case) | 0 – 1536 | ## Evidence diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 6f525efe6..0f6a64a35 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -28,8 +28,10 @@ op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the `N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. `N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard -never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 qualify. Anything else fails closed when -the op is built. +never splits a WS1 partial sum. The shard constructor accepts positive TP sizes satisfying +this alignment, including 1, 2, 3, 4, 6, 7, 8 and 9 (TP=9 gives 10752 columns per rank). +The `DeterministicCollective` backend supports only 1, 2, 4 and 8 ranks; shard alignment +alone does not imply backend support. | TP | Columns per rank | Rank 0 owns | | --- | --- | --- | diff --git a/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py index 19a867b3e..7b5eb20d3 100644 --- a/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py +++ b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py @@ -211,16 +211,18 @@ class SPPlan: num_rows: int | None -def sp_plan(layout: SPRowLayout, index_full=None, num_rows: int | None = None) -> SPPlan: +def sp_plan( + layout: SPRowLayout, index_full=None, num_rows: int | None = None, *, device=None +) -> SPPlan: """The reduction plan for one row index (``None``: the plain norm's ``dweight`` only).""" - device = index_full.device if index_full is not None else torch.device("cuda") + device = index_full.device if index_full is not None else torch.device(device or "cuda") families = [_dweight_tiles(layout.batch * layout.seq_len, device)] if index_full is None: return SPPlan(layout, _Plan(layout, families, device), None, None, None, None) *seg_tiles, seg_first_tile = _segment_tiles(index_full.repeat(layout.batch), num_rows) plan = _Plan(layout, [*families, tuple(seg_tiles)], device) - local_index = index_full[layout.lo : layout.hi].contiguous() + local_index = index_full[layout.lo : layout.hi].clone(memory_format=torch.contiguous_format) received = index_full[plan.recv_global % layout.seq_len] index_buf = torch.cat([local_index.repeat(layout.batch), received]).contiguous() return SPPlan(layout, plan, local_index, index_buf, seg_first_tile, num_rows) @@ -334,18 +336,21 @@ def __init__(self, collective, seq_len: int, batch: int) -> None: raise RuntimeError("rl_engine._C lacks the H3 SP symbols; rebuild the CUDA extension") self.collective = collective self.layout = sp_row_layout(seq_len, batch, collective.world_size, collective.rank) - self._plans: list[tuple[object, int, int | None, SPPlan]] = [] + self._plans: list[tuple[object, int, int | None, torch.device, SPPlan]] = [] - def plan(self, index: torch.Tensor | None, num_rows: int | None) -> SPPlan: + def plan(self, index: torch.Tensor | None, num_rows: int | None, *, device=None) -> SPPlan: """The cached plan for this index object (rebuilt if it was modified in place).""" + device = index.device if index is not None else torch.device(device or "cuda") + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) version = -1 if index is None else index._version - for obj, ver, rows, plan in self._plans: - if obj is index and ver == version and rows == num_rows: + for obj, ver, rows, cached_device, plan in self._plans: + if obj is index and ver == version and rows == num_rows and cached_device == device: return plan - plan = sp_plan(self.layout, index, num_rows) + plan = sp_plan(self.layout, index, num_rows, device=device) # Holding the index keeps its storage alive, so a new tensor can never alias it. - self._plans = [(index, version, num_rows, plan), *self._plans[:3]] + self._plans = [(index, version, num_rows, device, plan), *self._plans[:3]] return plan def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | None) -> None: @@ -368,7 +373,9 @@ def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | No def norm(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: self._check(x, None, None) validate_h3_rmsnorm(x, weight, eps) - return _SPRMSNorm.apply(x, weight, None, None, self.plan(None, None), self.collective, eps) + return _SPRMSNorm.apply( + x, weight, None, None, self.plan(None, None, device=x.device), self.collective, eps + ) def norm_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): self._check(x, index, shift.shape[0]) diff --git a/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py index 0eb5a4709..f5ec080a3 100644 --- a/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py +++ b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py @@ -84,6 +84,8 @@ def slots(self) -> list[tuple[int, int, int, int]]: def adaln_column_shard(n_total: int, tp: int, rank: int) -> AdaLNColumnShard: if tp < 1 or not 0 <= rank < tp: raise ValueError(f"need 0 <= rank < tp, got rank={rank}, tp={tp}") + if n_total <= 0: + raise ValueError(f"N must be positive, got {n_total}") if n_total % (H3_MODALITY_NUM * H3_ADALN_CHUNKS): raise ValueError(f"N={n_total} is not 3 modalities x 6 chunks x H") if n_total % (tp * DINPUT_CHUNK): diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 3af6ffd38..8f41be297 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -378,7 +378,9 @@ def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torc else: value = stage.golden(golden_op(registry, stage.op_type), run_ctx, value) outputs = list(value) - torch.autograd.backward(outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs)]) + torch.autograd.backward( + outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs, strict=True)] + ) return [leaves[name].grad for name in GRAD_LEAVES] diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 407906a88..877a2534d 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -590,9 +590,13 @@ def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: gate = _gate_view(table) grad = torch.randn_like(residual) - def backward(fn, residual=residual, y=y, table=table, grad=grad): + def prepare_backward(fn, residual=residual, y=y, table=table, grad=grad): leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] - fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + return fn(leaves[0], leaves[1], _gate_view(leaves[2])), grad + + def backward(state): + output, grad = state + output.backward(grad) cases.append( { @@ -606,10 +610,12 @@ def backward(fn, residual=residual, y=y, table=table, grad=grad): "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( r, g, i, y ), - "candidate_backward": lambda b=backward, i=index: b( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: provider_gate_residual(r, g, i, y_) ), } @@ -695,9 +701,13 @@ def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: x, nw, temb, w, b, ti = _final_inputs(seq) grad = torch.randn_like(x) - def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + def prepare_backward(fn, tensors=(x, nw, temb, w, b), grad=grad): leaves = [t.detach().requires_grad_(True) for t in tensors] - fn(*leaves).backward(grad) + return fn(*leaves), grad + + def backward(state): + output, grad = state + output.backward(grad) cases.append( { @@ -708,8 +718,12 @@ def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), - "candidate_backward": lambda bw=backward, i=ti: bw(lambda *t: op(*t, i)), - "provider_backward": lambda bw=backward, i=ti: bw( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( lambda *t: provider_final_adaln_out(*t, i) ), } diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py index b2e3b448f..13174f7f6 100644 --- a/rl_engine/testing/h3_ws2.py +++ b/rl_engine/testing/h3_ws2.py @@ -325,9 +325,9 @@ def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, inputs_path.with_name(f"rank{rank}.pt"), ) - queue.put({"rank": rank}) collective.close() dist.destroy_process_group() + queue.put({"rank": rank}) except Exception: # forwarded to the parent queue.put({"rank": rank, "error": traceback.format_exc()}) @@ -341,7 +341,9 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa """ import tempfile + import time from pathlib import Path + from queue import Empty import torch.multiprocessing as mp @@ -357,12 +359,52 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) for r in range(world) ] - for proc in procs: - proc.start() - status = sorted((queue.get(timeout=timeout) for _ in procs), key=lambda r: r["rank"]) - for proc in procs: - proc.join(timeout=120) - for result in status: - if "error" in result: - raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") - return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + started = [] + deadline = time.monotonic() + timeout + status = {} + try: + for proc in procs: + proc.start() + started.append(proc) + while len(status) < world or any(proc.is_alive() for proc in procs): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"rank results timed out after {timeout}s") + try: + result = queue.get(timeout=min(0.1, remaining)) + except Empty: + exited = [ + (rank, proc.exitcode) + for rank, proc in enumerate(procs) + if proc.exitcode is not None and (proc.exitcode != 0 or rank not in status) + ] + if not exited: + continue + # A worker may have flushed its final message between the + # timed read and the exit-code check. + try: + result = queue.get_nowait() + except Empty: + rank, code = exited[0] + raise RuntimeError( + f"rank {rank} exited with code {code} without completing" + ) from None + if "error" in result: + raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") + status[result["rank"]] = result + for rank, proc in enumerate(procs): + proc.join() + if proc.exitcode != 0: + raise RuntimeError(f"rank {rank} exited with code {proc.exitcode}") + return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + finally: + for proc in started: + if proc.is_alive(): + proc.terminate() + for proc in started: + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join() + queue.close() + queue.join_thread() diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 7328207ad..bd2ada020 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -440,6 +440,8 @@ def repeat(mode: str) -> str: def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: + combined = all(row.get("backward_timing_scope") == "forward_and_backward" for row in perf) + backward_label = "fwd+bwd" if combined else "bwd" _style(ax, "Forward and backward time (log, lower is better)", "ms") _bars( ax, @@ -447,9 +449,14 @@ def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: [ ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), (f"{provider} fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), - ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), ( - f"{provider} bwd", + f"CUDA {backward_label}", + RL_KERNEL, + [row["candidate_backward_us"] / 1e3 for row in perf], + "////", + ), + ( + f"{provider} {backward_label}", PROVIDER, [row["provider_backward_us"] / 1e3 for row in perf], "////", @@ -459,6 +466,8 @@ def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: log=True, ) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + def _backward_errors(ax, backward: dict[str, Any], keys: list[str], note: str) -> None: _style(ax, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") @@ -490,6 +499,7 @@ def label(name: str, mode: str) -> str: for text in ax.texts: if text.get_text() == f"{floor:.1e}": text.set_text("exact") + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=1) ax.set_ylim(floor / 3, 1e4) ax.text(0.98, 0.80, note, transform=ax.transAxes, ha="right", fontsize=8, color=INK) diff --git a/tests/h3/test_h3_benchmark_timing.py b/tests/h3/test_h3_benchmark_timing.py index e6517bd70..d709468b3 100644 --- a/tests/h3/test_h3_benchmark_timing.py +++ b/tests/h3/test_h3_benchmark_timing.py @@ -114,7 +114,7 @@ def run(key=key, duration=position): assert not any(key.endswith("_setup") for key in row) -@pytest.mark.parametrize("operator", ["norm", "gather"]) +@pytest.mark.parametrize("operator", ["norm", "gather", "gate", "final"]) def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( monkeypatch, cuda_clock, operator ): @@ -136,7 +136,7 @@ def backward_hook(grad): def track(fn, leaves, *args): if leaves[0].requires_grad: assert not cuda_clock.active - prepared_leaves.append(leaves) + prepared_leaves.append([leaf if leaf.is_leaf else leaf._base for leaf in leaves]) cuda_clock.time += 1.0 outputs = fn(*leaves, *args) return track_backward(outputs) if leaves[0].requires_grad else outputs @@ -162,6 +162,54 @@ def forward(x, weight, shift, scale, index, **kwargs): monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) op = SimpleNamespace(forward_modulated=forward) build = h3_report._norm_perf + elif operator == "gate": + + def inputs(seq): + return ( + torch.randn(3, 24), + torch.randn(1, seq, 4), + torch.randn(1, seq, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_gate_residual + + def forward(residual, y, gate, index, **kwargs): + return track(lambda r, y_, g: provider(r, g, index, y_), [residual, y, gate]) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_gate_inputs", inputs) + monkeypatch.setattr( + h3_report, "provider_gate_residual", lambda r, g, i, y: forward(r, y, g, i) + ) + op = SimpleNamespace(forward=forward) + build = h3_report._gate_perf + elif operator == "final": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 2), + torch.randn(8, 2), + torch.randn(8), + torch.arange(seq) % 3, + ) + + def forward(x, nw, temb, w, b, ti): + def compute(x, nw, temb, w, b): + rows = torch.nn.functional.linear(torch.nn.functional.silu(temb), w, b) + shift, scale = rows.chunk(2, dim=-1) + norm = torch.nn.functional.rms_norm(x, (4,), nw) + return norm * (1 + scale[ti]) + shift[ti] + + return track(compute, [x, nw, temb, w, b]) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_final_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", forward) + op = forward + build = h3_report._final_perf else: randn = torch.randn diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py index 9f2a06a07..4e59cfcce 100644 --- a/tests/h3/test_h3_native_validation.py +++ b/tests/h3/test_h3_native_validation.py @@ -275,3 +275,106 @@ def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name) message = "row index" if name == "index" else name with pytest.raises(RuntimeError, match=message): _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize( + "operation", ["gate_forward", "gate_dy", "norm_forward", "norm_dx", "norm_partials"] +) +@pytest.mark.parametrize("bad_index", [-1, 2]) +def test_native_table_bounds_fail_in_isolated_process(native_extension, operation, bad_index): + import subprocess + import sys + import textwrap + + script = textwrap.dedent(""" + import sys + import torch + from rl_engine import _C + op, bad = sys.argv[1], int(sys.argv[2]) + x = torch.ones(2, 8, device="cuda") + w = torch.ones(8, device="cuda") + table = torch.ones(2, 8, device="cuda") + index = torch.tensor([bad, 0], device="cuda") + rstd = torch.ones(2, device="cuda") + rows = torch.arange(2, device="cuda") + begin = torch.tensor([0], device="cuda") + end = torch.tensor([2], device="cuda") + try: + if op == "gate_forward": + _C.h3_gate_residual_forward(x, x, table, index) + elif op == "gate_dy": + _C.h3_gate_residual_backward_dy(x, table, index) + elif op == "norm_forward": + _C.h3_rmsnorm_forward(x, w, 1e-5, table, table, index) + elif op == "norm_dx": + _C.h3_rmsnorm_backward_dx(x, x, w, rstd, table, table, index) + else: + _C.h3_rmsnorm_backward_partials( + x, x, w, rstd, table, table, index, rows, begin, end, rows, begin, end) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid index reached native table lookup without an assertion") + """) + result = subprocess.run( + [sys.executable, "-c", script, operation, str(bad_index)], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +def test_row_gather_rejects_metadata_on_other_device(native_extension, metadata): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + grad = torch.ones(6, 2, 8, device="cuda:0") + tensors = dict( + sorted_pos=torch.tensor([0, 1], device="cuda:0"), + tile_begin=torch.tensor([0], device="cuda:0"), + tile_end=torch.tensor([2], device="cuda:0"), + seg_first_tile=torch.tensor([0, 1], device="cuda:0"), + ) + tensors[metadata] = tensors[metadata].to("cuda:1") + with pytest.raises(RuntimeError, match="grad device"): + native_extension.h3_adaln_row_gather_backward(grad, **tensors, out_dtype=torch.float32) + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +@pytest.mark.parametrize("bound", ["negative", "past_end"]) +def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metadata, bound): + import subprocess + import sys + import textwrap + + script = textwrap.dedent(""" + import sys + import torch + from rl_engine import _C + name, bound = sys.argv[1:] + grad = torch.ones(6, 2, 8, device="cuda") + metadata = dict(sorted_pos=torch.tensor([0, 1], device="cuda"), + tile_begin=torch.tensor([0], device="cuda"), + tile_end=torch.tensor([2], device="cuda"), + seg_first_tile=torch.tensor([0, 1], device="cuda")) + metadata[name][0] = -1 if bound == "negative" else 3 + try: + _C.h3_adaln_row_gather_backward(grad, **metadata, out_dtype=torch.float32) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid metadata was not rejected") + """) + result = subprocess.run( + [sys.executable, "-c", script, metadata, bound], capture_output=True, text=True, timeout=60 + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout diff --git a/tests/h3/test_h3_sp_norm_adaln.py b/tests/h3/test_h3_sp_norm_adaln.py index 04d58395c..2b91aad64 100644 --- a/tests/h3/test_h3_sp_norm_adaln.py +++ b/tests/h3/test_h3_sp_norm_adaln.py @@ -208,3 +208,106 @@ def test_nccl_region_byte_equal_to_ws1(world): ranks = run_world(world, sp_region_rank, {k: v.cpu() for k, v in case.items()}) assert sp_matches_ws1(ws1, ranks) == ALL_TRUE assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@requires_cuda +def test_plan_snapshots_caller_index(): + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import sp_plan + + index = torch.tensor([0, 1, 0, 1], device="cuda") + plan = sp_plan(sp_row_layout(4, 1, 1, 0), index, 2) + local, exchanged = plan.local_index.clone(), plan.index_buf.clone() + index.fill_(99) + assert torch.equal(plan.local_index, local) + assert torch.equal(plan.index_buf, exchanged) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_plain_norm_uses_activation_device_and_separate_cache(): + _require_ext() + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + collective = SimpleNamespace(rank=0, world_size=1, all_gather=lambda x: x) + op = H3SPNormAdaLNCudaOp(collective, 5, 1) + with torch.cuda.device(0): + for device in ("cuda:1", "cuda:0"): + x = torch.randn(1, 5, 8, device=device, requires_grad=True) + weight = torch.randn(8, device=device, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + ref_w = weight.detach().clone().requires_grad_(True) + result = op.norm(x, weight) + expected = torch.nn.functional.rms_norm(ref_x, (8,), ref_w, 1e-5) + result.sum().backward() + expected.sum().backward() + torch.testing.assert_close(result, expected) + torch.testing.assert_close(x.grad, ref_x.grad) + torch.testing.assert_close(weight.grad, ref_w.grad) + assert len(op._plans) == 2 + assert op.plan(None, None, device="cuda:1").plan.send_local.device.index == 1 + assert op.plan(None, None, device="cuda:0").plan.send_local.device.index == 0 + + +def _worker_exit_without_result(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + os._exit(kwargs["exit_code"]) + time.sleep(60) + + +def _worker_reports_error(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + queue.put({"rank": rank, "error": "intentional worker failure"}) + time.sleep(60) + + +def _worker_never_reports(rank, world, init, queue, path, target, kwargs): + import time + + time.sleep(60) + + +@pytest.mark.parametrize("mode", ["clean_exit", "crash", "reported_error", "timeout"]) +def test_run_world_cleans_up_failed_workers(monkeypatch, mode): + import time + import torch.multiprocessing as mp + from rl_engine.testing import h3_ws2 + + workers = { + "clean_exit": _worker_exit_without_result, + "crash": _worker_exit_without_result, + "reported_error": _worker_reports_error, + "timeout": _worker_never_reports, + } + monkeypatch.setattr(h3_ws2, "_bootstrap", workers[mode]) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + expected = TimeoutError if mode == "timeout" else RuntimeError + with pytest.raises(expected): + run_world( + 2, + None, + {}, + timeout=1 if mode == "timeout" else 30, + exit_code=7 if mode == "crash" else 0, + ) + assert time.monotonic() - start < 20 + assert {p.pid for p in mp.active_children()} <= before + + +def _worker_reports_result(rank, world, init, queue, path, target, kwargs): + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + + +def test_run_world_returns_rank_ordered_results(monkeypatch): + from rl_engine.testing import h3_ws2 + + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker_reports_result) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + results = run_world(2, None, {}, timeout=30) + assert [result["rank"] for result in results] == [0, 1] + assert [result["value"].item() for result in results] == [0, 1] diff --git a/tests/h3/test_h3_tp_adaln.py b/tests/h3/test_h3_tp_adaln.py index 8b341c8d3..91079b691 100644 --- a/tests/h3/test_h3_tp_adaln.py +++ b/tests/h3/test_h3_tp_adaln.py @@ -48,7 +48,7 @@ def _inputs(num_t, n_out, k, dtype=torch.bfloat16, seed=0): class TestOwnership: - @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8]) + @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8, 9]) def test_shards_tile_the_table_once(self, tp): shards = [adaln_column_shard(N_H3, tp, r) for r in range(tp)] assert [s.begin for s in shards] == [0, *[s.end for s in shards[:-1]]] @@ -146,3 +146,9 @@ def test_nccl_byte_equal_to_ws1(h3_weights_cpu, world): assert tp_matches_ws1(ws1, ranks) == ALL_TRUE assert [r["readback"]["rank"] for r in ranks] == list(range(world)) assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@pytest.mark.parametrize("n_total", [0, -1152]) +def test_rejects_nonpositive_projection(n_total): + with pytest.raises(ValueError, match="N must be positive"): + adaln_column_shard(n_total, 1, 0) From 76594af5cb43cb73698faa7b1698cdb85ab9aff8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:28:22 +0000 Subject: [PATCH 09/19] evidence(h3): regenerate this branch's final_adaln_out report after the golden fix This branch carries a newer final_adaln_out report.json than #489. Regenerated by tools/validation/models/h3_evidence.py from a clean tree at 0a2d5589 on an idle B200; chain_replay.json comes from #489's regeneration at 50e913fa. CUDA accuracy and every pass/fail flag are unchanged; the doc's latency table follows the new report. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index b95c906b0..3e884dcf2 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -63,9 +63,9 @@ and forward-graph setup: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.19 ms | 0.18 ms | 1.40 ms | 0.85 ms | -| 32768 | 0.46 ms | 0.80 ms | 2.59 ms | 4.49 ms | -| 131072 | 1.32 ms | 3.02 ms | 7.05 ms | 17.69 ms | +| 4097 | 0.19 ms | 0.21 ms | 1.37 ms | 0.83 ms | +| 32768 | 0.42 ms | 0.82 ms | 2.53 ms | 4.45 ms | +| 131072 | 1.28 ms | 2.97 ms | 7.02 ms | 17.61 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -78,13 +78,16 @@ There are two data files: - [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `bde1d8d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `001684d`. + tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 + before, 0.0671 now). ## Existing implementations (RFC #420 reuse rule) From c469b72becec588256a436e7ed45b4609a529d01 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:20 +0000 Subject: [PATCH 10/19] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 12 ++++++------ docs/operators/h3-adaln-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-final-adaln-out.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ docs/operators/h3-tp-adaln-3mod.md | 2 +- 8 files changed, 45 insertions(+), 45 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6c972994d..9a5d815a5 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -75,15 +75,15 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](../../reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-gate-residual-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) -![gate_residual vs existing implementations](../../reports/experiments/h3-prior-art-b200/gate_residual.png) +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -94,7 +94,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../../reports/experiments/h3-prior-art-b200/gate_residual.json) +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) was written from a clean tree at `4010854` by ```bash @@ -111,10 +111,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a4cb57dd8..e8fc09537 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -124,9 +124,9 @@ dependence on T. ## Evidence -![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also @@ -135,7 +135,7 @@ or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![adaln_projection vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_projection.png) +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional @@ -157,7 +157,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../../reports/experiments/h3-prior-art-b200/adaln_projection.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) was written from a clean tree at `c88691c` by ```bash @@ -174,10 +174,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 88fee5f3d..94e32b711 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -122,13 +122,13 @@ alternates each iteration and is recorded in the report. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../../reports/experiments/h3-adaln-row-gather-b200/figure.png) +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward bitwise checks, op-level backward, and the chain-backward cases plotted above. -- [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. `report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only @@ -141,7 +141,7 @@ model does; every CUDA pass/fail flag is unchanged. ## Existing implementations (RFC #420 reuse rule) -![adaln_row_gather vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.png) +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) was written from a clean tree at `6c900ae` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_row_gather bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 3e884dcf2..7970b5c2d 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -72,17 +72,17 @@ projection backward. ## Evidence -![final_adaln_out on B200: latency and backward accuracy](../../reports/experiments/h3-final-adaln-out-b200/figure.png) +![final_adaln_out on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. -- [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the @@ -91,7 +91,7 @@ There are two data files: ## Existing implementations (RFC #420 reuse rule) -![final_adaln_out vs existing implementations](../../reports/experiments/h3-prior-art-b200/final_adaln_out.png) +![final_adaln_out vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -102,7 +102,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](../../reports/experiments/h3-prior-art-b200/final_adaln_out.json) +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.json) was written from a clean tree at `03da729` by ```bash @@ -119,10 +119,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![final_adaln_out bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) +![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 740f16232..75dd3d477 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -92,9 +92,9 @@ each iteration and is recorded in the report. ## Evidence -![h3_rmsnorm on B200: latency and backward accuracy](../../reports/experiments/h3-rmsnorm-b200/figure.png) +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report.json), written by +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; @@ -103,7 +103,7 @@ The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report. ## Existing implementations (RFC #420 reuse rule) -![norm_modulate vs existing implementations](../../reports/experiments/h3-prior-art-b200/norm_modulate.png) +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -119,7 +119,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](../../reports/experiments/h3-prior-art-b200/norm_modulate.json) +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) was written from a clean tree at `ee83dec` by ```bash @@ -136,10 +136,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![norm_modulate bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 8049a87df..cc5809cee 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -109,15 +109,15 @@ kernel time. ## Evidence -![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../../reports/experiments/h3-timestep-mlp-b200/figure.png) +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-mlp-b200/report.json), written +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![timestep_mlp vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_mlp.png) +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -134,7 +134,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../../reports/experiments/h3-prior-art-b200/timestep_mlp.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) was written from a clean tree at `ffd958a` by ```bash @@ -151,10 +151,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index c5f69a15b..75edf172b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -88,15 +88,15 @@ about 1 KB per timestep. ## Evidence -![timestep_sinusoid_h3 on B200: latency and error vs FP64](../../reports/experiments/h3-timestep-sinusoid-b200/figure.png) +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-sinusoid-b200/report.json). +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). `tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. ## Existing implementations (RFC #420 reuse rule) -![timestep_sinusoid vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -108,7 +108,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) was written from a clean tree at `7a82917` by ```bash @@ -126,11 +126,11 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). -![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. -[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) was written from a clean tree at `0eb5ddd5` by ```bash diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 614be59c5..8354dff8b 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -70,7 +70,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ## Evidence -![tp_adaln_3mod on 8 x B200: byte equality and time](../../reports/experiments/h3-tp-adaln-3mod-b200/figure.png) +![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For From 0033b3a5a513604c72987abbc2cc08bd22d5a67d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:22 +0000 Subject: [PATCH 11/19] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 12 ++++++------ docs/operators/h3-adaln-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-final-adaln-out.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-sp-norm-adaln.md | 2 +- docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ docs/operators/h3-tp-adaln-3mod.md | 2 +- 9 files changed, 46 insertions(+), 46 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6c972994d..9a5d815a5 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -75,15 +75,15 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](../../reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-gate-residual-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) -![gate_residual vs existing implementations](../../reports/experiments/h3-prior-art-b200/gate_residual.png) +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -94,7 +94,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../../reports/experiments/h3-prior-art-b200/gate_residual.json) +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) was written from a clean tree at `4010854` by ```bash @@ -111,10 +111,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a4cb57dd8..e8fc09537 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -124,9 +124,9 @@ dependence on T. ## Evidence -![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also @@ -135,7 +135,7 @@ or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![adaln_projection vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_projection.png) +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional @@ -157,7 +157,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../../reports/experiments/h3-prior-art-b200/adaln_projection.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) was written from a clean tree at `c88691c` by ```bash @@ -174,10 +174,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 88fee5f3d..94e32b711 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -122,13 +122,13 @@ alternates each iteration and is recorded in the report. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../../reports/experiments/h3-adaln-row-gather-b200/figure.png) +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward bitwise checks, op-level backward, and the chain-backward cases plotted above. -- [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. `report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only @@ -141,7 +141,7 @@ model does; every CUDA pass/fail flag is unchanged. ## Existing implementations (RFC #420 reuse rule) -![adaln_row_gather vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.png) +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) was written from a clean tree at `6c900ae` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_row_gather bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 93143b858..d40a49f5f 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -77,17 +77,17 @@ with their original commit and corrected scope. ## Evidence -![final_adaln_out on B200: latency and backward accuracy](../../reports/experiments/h3-final-adaln-out-b200/figure.png) +![final_adaln_out on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. -- [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the @@ -96,7 +96,7 @@ There are two data files: ## Existing implementations (RFC #420 reuse rule) -![final_adaln_out vs existing implementations](../../reports/experiments/h3-prior-art-b200/final_adaln_out.png) +![final_adaln_out vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -107,7 +107,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](../../reports/experiments/h3-prior-art-b200/final_adaln_out.json) +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.json) was written from a clean tree at `03da729` by ```bash @@ -124,10 +124,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![final_adaln_out bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) +![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 740f16232..75dd3d477 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -92,9 +92,9 @@ each iteration and is recorded in the report. ## Evidence -![h3_rmsnorm on B200: latency and backward accuracy](../../reports/experiments/h3-rmsnorm-b200/figure.png) +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report.json), written by +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; @@ -103,7 +103,7 @@ The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report. ## Existing implementations (RFC #420 reuse rule) -![norm_modulate vs existing implementations](../../reports/experiments/h3-prior-art-b200/norm_modulate.png) +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -119,7 +119,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](../../reports/experiments/h3-prior-art-b200/norm_modulate.json) +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) was written from a clean tree at `ee83dec` by ```bash @@ -136,10 +136,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![norm_modulate bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 1f19d0444..8eb91cc38 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -70,7 +70,7 @@ How many rows move depends on how the table rows interleave along the sequence: ## Evidence -![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](../../reports/experiments/h3-sp-norm-adaln-b200/figure.png) +![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr494/reports/experiments/h3-sp-norm-adaln-b200/figure.png) The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `6c0d380`, with real NCCL processes on one 8 x B200 node, one GPU per rank. The region is diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 8049a87df..cc5809cee 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -109,15 +109,15 @@ kernel time. ## Evidence -![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../../reports/experiments/h3-timestep-mlp-b200/figure.png) +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-mlp-b200/report.json), written +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![timestep_mlp vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_mlp.png) +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -134,7 +134,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../../reports/experiments/h3-prior-art-b200/timestep_mlp.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) was written from a clean tree at `ffd958a` by ```bash @@ -151,10 +151,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index c5f69a15b..75edf172b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -88,15 +88,15 @@ about 1 KB per timestep. ## Evidence -![timestep_sinusoid_h3 on B200: latency and error vs FP64](../../reports/experiments/h3-timestep-sinusoid-b200/figure.png) +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-sinusoid-b200/report.json). +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). `tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. ## Existing implementations (RFC #420 reuse rule) -![timestep_sinusoid vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -108,7 +108,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) was written from a clean tree at `7a82917` by ```bash @@ -126,11 +126,11 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). -![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. -[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) was written from a clean tree at `0eb5ddd5` by ```bash diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index c348bce35..318dbd2d0 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -72,7 +72,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ## Evidence -![tp_adaln_3mod on 8 x B200: byte equality and time](../../reports/experiments/h3-tp-adaln-3mod-b200/figure.png) +![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For From 4530173e0ccd2f02dbc05535f6c21e98734d3613 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:06:02 +0000 Subject: [PATCH 12/19] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 2 +- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-final-adaln-out.md | 6 +++--- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- docs/operators/h3-tp-adaln-3mod.md | 2 +- 8 files changed, 16 insertions(+), 16 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 9a5d815a5..ffd200b0a 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -114,7 +114,7 @@ gradient with diffusers' path and with this kernel. ![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `a746bff0` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index e8fc09537..526e2b4cf 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -127,7 +127,7 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `579775f1` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone @@ -177,7 +177,7 @@ gradient with diffusers' path and with this kernel. ![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `50d1e924` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 94e32b711..923ad2f77 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,11 +131,11 @@ There are two data files: - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +`report.json` was regenerated from a clean tree at commit `29b07c5d` with backward-only timings and alternating execution order, and with the op-backward check's rows and gradients drawn from a seeded generator. Before, they came from the process's default CUDA generator, -so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). -`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +so the fraction depended on how many draws the runner made first (99.99965% at `d2a6546`). +`chain_replay.json` was written from a clean tree at commit `61d23ade`. Both were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is unchanged. @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `62f8009e` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 7970b5c2d..c4c45ebb5 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -78,14 +78,14 @@ There are two data files: - [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + tree at `ddca8f95` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + tree at `a4bc99f9` (#489). Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 before, 0.0671 now). @@ -122,7 +122,7 @@ gradient with diffusers' path and with this kernel. ![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `8cb4abed` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 75dd3d477..ac625ae35 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -95,7 +95,7 @@ each iteration and is recorded in the report. ![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by -`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -139,7 +139,7 @@ gradient with diffusers' path and with this kernel. ![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `83797833` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index cc5809cee..b65b19fca 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -112,7 +112,7 @@ kernel time. ![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written -by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `418998c`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) @@ -154,7 +154,7 @@ gradient with diffusers' path and with this kernel. ![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `73939f2d` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 75edf172b..ff6579cd8 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -131,7 +131,7 @@ kernel compiled with `--use_fast_math`, which the build excludes (see Known Limi SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) -was written from a clean tree at `0eb5ddd5` by +was written from a clean tree at `be6a94ae` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 8354dff8b..f1b1a2409 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -72,7 +72,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) -The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `507c9ad`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are byte-equal to WS1 computed on the same GPU. From 4ef511c9e05e26073dc4baff6ec0075e5048be24 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:06:03 +0000 Subject: [PATCH 13/19] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 2 +- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-final-adaln-out.md | 6 +++--- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-sp-norm-adaln.md | 2 +- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- docs/operators/h3-tp-adaln-3mod.md | 2 +- 9 files changed, 17 insertions(+), 17 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 9a5d815a5..ffd200b0a 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -114,7 +114,7 @@ gradient with diffusers' path and with this kernel. ![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `a746bff0` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index e8fc09537..526e2b4cf 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -127,7 +127,7 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `579775f1` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone @@ -177,7 +177,7 @@ gradient with diffusers' path and with this kernel. ![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `50d1e924` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 94e32b711..923ad2f77 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,11 +131,11 @@ There are two data files: - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +`report.json` was regenerated from a clean tree at commit `29b07c5d` with backward-only timings and alternating execution order, and with the op-backward check's rows and gradients drawn from a seeded generator. Before, they came from the process's default CUDA generator, -so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). -`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +so the fraction depended on how many draws the runner made first (99.99965% at `d2a6546`). +`chain_replay.json` was written from a clean tree at commit `61d23ade`. Both were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is unchanged. @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `62f8009e` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index d40a49f5f..a46954f06 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -83,14 +83,14 @@ There are two data files: - [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + tree at `ddca8f95` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + tree at `a4bc99f9` (#489). Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 before, 0.0671 now). @@ -127,7 +127,7 @@ gradient with diffusers' path and with this kernel. ![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `8cb4abed` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 75dd3d477..ac625ae35 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -95,7 +95,7 @@ each iteration and is recorded in the report. ![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by -`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -139,7 +139,7 @@ gradient with diffusers' path and with this kernel. ![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `83797833` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 8eb91cc38..eeb90f72f 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -72,7 +72,7 @@ How many rows move depends on how the table rows interleave along the sequence: ![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr494/reports/experiments/h3-sp-norm-adaln-b200/figure.png) -The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `6c0d380`, with +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `ecde5bb`, with real NCCL processes on one 8 x B200 node, one GPU per rank. The region is `norm2(residual + gate_msa[row] * y)` with `shift_mlp`/`scale_mlp` modulation, H = 5376, BF16. It covers six cases: S = 4097 (block, interleaved, B = 2), 32768 (block, interleaved) and 131072 diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index cc5809cee..b65b19fca 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -112,7 +112,7 @@ kernel time. ![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written -by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `418998c`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) @@ -154,7 +154,7 @@ gradient with diffusers' path and with this kernel. ![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `73939f2d` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 75edf172b..ff6579cd8 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -131,7 +131,7 @@ kernel compiled with `--use_fast_math`, which the build excludes (see Known Limi SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) -was written from a clean tree at `0eb5ddd5` by +was written from a clean tree at `be6a94ae` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 318dbd2d0..868df445a 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -74,7 +74,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) -The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `507c9ad`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are byte-equal to WS1 computed on the same GPU. From b3747d90ec5db2e4bd4a20c23cd4cbc9fd9d62b7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:18:32 +0000 Subject: [PATCH 14/19] fix(h3): refuse pre-SM80 before the first collective, check the TP table gradient, handle a TP1-only figure, say which TP sizes run tp_adaln_projection.py: the op refuses a device without SM80 when it is built (so the registry can fall back) and again, for BF16 weights, in the input validation that runs before the first collective, so a rank that cannot run the BF16 forward fails before the other ranks start waiting in gather_rows (a test counts the collective calls). The backward checks that the table gradient is (T, N) before it slices this rank's columns; autograd already checks the shape of a gradient that reaches the node through the graph, so the check guards a direct call of the node's backward, which the test makes. det_linear.cu: h3_det_linear_fold_chunks checks the grid.x block count like h3_det_linear_backward_input (a bound of 5e11 elements, so it cannot be exercised with real memory). plot_h3_evidence.py: a TP evidence report without a run of TP > 1 (--worlds 1) no longer fails in the bar chart on an empty sequence; the right panel says so. Tested for TP1-only, TP1/2/4 and TP2/4 reports. docs/operators/h3-tp-adaln-3mod.md: TP 3, 6 and 7 pass the shard-alignment check but only 1, 2, 4 and 8 can run, because DeterministicCollective rejects the other world sizes when it is constructed. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 4 ++ docs/operators/h3-tp-adaln-3mod.md | 6 ++- .../minimax_h3/tp_adaln_projection.py | 18 +++++++- .../minimax_h3/test_h3_bitwise_reuse_plot.py | 41 ++++++++++++++++++- tests/models/minimax_h3/test_h3_tp_adaln.py | 40 ++++++++++++++++++ tools/validation/models/plot_h3_evidence.py | 13 ++++++ 6 files changed, 118 insertions(+), 4 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index cc84bfa5a..56e6be84f 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -561,6 +561,10 @@ torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType o 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::max(), + "rows * K = ", elems, " needs more blocks than one grid.x launch holds"); const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); if (out_dtype == at::kFloat) { det_fold_chunks_kernel<<>>( diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index f1b1a2409..6487c6a8d 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -28,8 +28,10 @@ op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the `N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. `N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard -never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 qualify. Anything else fails closed when -the op is built. +never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 pass this alignment check and anything +else fails closed when the op is built. Only 1, 2, 4 and 8 can run: `DeterministicCollective` +supports exactly those world sizes and rejects a 3-, 6- or 7-way group when it is constructed, so +those sizes are covered by the ownership tests only. | TP | Columns per rank | Rank 0 owns | | --- | --- | --- | diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py index d262ac030..371f0ae74 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py @@ -31,18 +31,19 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( CONTRACT, DINPUT_CHUNK, det_linear_available, det_linear_backward_input_partials, det_linear_backward_weight, + det_linear_device_supported, det_linear_fold_chunks, det_linear_forward, silu_backward_fp32, ) from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM from rl_engine.reference.minimax_h3.adaln_projection import split_adaln_table @@ -122,6 +123,10 @@ def _validate_shard_inputs(temb, weight, bias, shard: AdaLNColumnShard) -> None: raise TypeError( f"weight/bias must share bfloat16 or float32, got {weight.dtype}/{bias.dtype}" ) + if weight.dtype == torch.bfloat16 and not det_linear_device_supported(weight.device): + # Checked here, before any collective: a rank that fails inside the forward would leave the + # other ranks waiting in gather_rows. + raise RuntimeError(f"the BF16 h3 det_linear forward needs SM80 or newer ({weight.device})") width = shard.end - shard.begin if weight.dim() != 2 or weight.shape[0] != width or bias.shape != (width,): raise ValueError( @@ -148,6 +153,13 @@ def forward(ctx, temb, weight_shard, bias_shard, collective, shard): def backward(ctx, grad_table): temb, act, weight = ctx.saved_tensors shard = ctx.shard + if grad_table.dim() != 2 or grad_table.shape != (temb.shape[0], shard.n_total): + # The backward reads this rank's columns of a replicated (T, N) table gradient; any + # other width would slice the wrong columns without an error. + raise ValueError( + f"table gradient must be ({temb.shape[0]}, {shard.n_total}), " + f"got {tuple(grad_table.shape)}" + ) grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) partial = det_linear_backward_input_partials(grad_local, weight) @@ -176,6 +188,10 @@ def __init__(self, collective, n_total: int) -> None: raise RuntimeError( "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" ) + if not det_linear_device_supported(): + # Best effort at construction so the registry can fall back; forward checks the + # weight's own device before the first collective. + raise RuntimeError("the BF16 h3 det_linear forward needs SM80 or newer") self.collective = collective self.shard = adaln_column_shard(n_total, collective.world_size, collective.rank) diff --git a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py index ffb87bd65..e386f4ec4 100644 --- a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""A heatmap cell counts every forward output of the operator, not only the first.""" +"""A heatmap cell counts every forward output of the operator, not only the first. + +Also: the TP evidence figure for a report that has no TP > 1 run. +""" from __future__ import annotations @@ -71,3 +74,39 @@ def result(frac_differ): texts = [t.get_text() for t in ax.texts] assert "n/a" in texts and "0" in texts plt.close(fig) + + +def _tp_run(tp: int) -> dict: + perf = { + "num_timesteps": 3, + "tp_forward_us": 110.0, + "tp_forward_backward_us": 2000.0, + "ws1_forward_us": 100.0, + "ws1_forward_backward_us": 1900.0, + } + equal = dict.fromkeys(("table", "d_temb", "d_weight_shard", "d_bias_shard"), True) + return { + "tp": tp, + "ranks": [{"rank": 0, "equality": [{"num_timesteps": 3, **equal}], "perf": [perf]}], + } + + +def _tp_report(worlds: list[int]) -> dict: + return { + "op": "tp_adaln_3mod", + "environment": {"gpu": "gpu", "torch": "t", "cuda": "c"}, + "rl_kernel_commit": "0123456789", + "model_revision": "0123456789", + "runs": [_tp_run(tp) for tp in worlds], + } + + +@pytest.mark.parametrize("worlds", [[1], [1, 2, 4], [2, 4]]) +def test_tp_figure_for_reports_with_and_without_tp_above_one(worlds): + import matplotlib.pyplot as plt + from plot_h3_evidence import plot_tp_adaln + + fig = plot_tp_adaln(_tp_report(worlds)) + texts = [t.get_text() for ax in fig.axes for t in ax.texts] + assert ("no run with TP > 1 in this report" in texts) == (worlds == [1]) + plt.close(fig) diff --git a/tests/models/minimax_h3/test_h3_tp_adaln.py b/tests/models/minimax_h3/test_h3_tp_adaln.py index 2a62aa502..6d3737824 100644 --- a/tests/models/minimax_h3/test_h3_tp_adaln.py +++ b/tests/models/minimax_h3/test_h3_tp_adaln.py @@ -128,6 +128,46 @@ def test_fails_closed(self): H3TPAdaLNProjectionCudaOp(SimpleNamespace(rank=0, world_size=16), weight.shape[0]) assert op.readback()["columns"] == [2304, 4608] + def test_pre_sm80_is_refused_at_construction_and_before_the_first_collective(self, monkeypatch): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3 import tp_adaln_projection as module + + calls = [] + collective = SimpleNamespace( + rank=0, world_size=2, backend_id="unused", all_gather=lambda t: calls.append(t) + ) + temb, weight, bias, _ = _inputs(2, 3 * 6 * 256, 136) + op = module.H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]) + w, b = module.shard_adaln_projection(weight, bias, 2, 0) + monkeypatch.setattr(module, "det_linear_device_supported", lambda *args: False) + with pytest.raises(RuntimeError, match="SM80"): + module.H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]) + with pytest.raises(RuntimeError, match="SM80"): + op(temb, w, b) + assert not calls # no rank is left waiting in a collective + + def test_backward_rejects_a_table_gradient_of_the_wrong_width(self): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + _TPAdaLNProjection, + shard_adaln_projection, + ) + + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(2, 3 * 6 * 256, 136) + op = H3TPAdaLNProjectionCudaOp(one, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 1, 0) + leaf = temb.clone().requires_grad_(True) + table = op.forward_table(leaf, w, b) + node = table.grad_fn # the autograd node is the ``ctx`` its backward receives + # Autograd checks the shape of a gradient that reaches this node through the graph, so the + # check guards a direct call of the node's backward, which this makes. + for bad in (grad[:, :-64], grad[:1]): + with pytest.raises(ValueError, match="table gradient must be"): + _TPAdaLNProjection.backward(node, bad.to(table.dtype)) + assert _TPAdaLNProjection.backward(node, grad.to(table.dtype))[0].shape == temb.shape + @pytest.mark.parametrize("world", [2, 4, 8]) @pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py index 564cc3ce6..35688f318 100644 --- a/tools/validation/models/plot_h3_evidence.py +++ b/tools/validation/models/plot_h3_evidence.py @@ -599,6 +599,19 @@ def worst(run, key): tp_runs = [run for run in runs if run["tp"] > 1] _style(right, f"Time per call, T = {num_t} (slowest rank; lower is better)", "ms") + if not tp_runs: + # A TP1-only report (--worlds 1) has nothing to compare with the single-GPU WS1 call. + right.text( + 0.5, + 0.5, + "no run with TP > 1 in this report", + transform=right.transAxes, + ha="center", + va="center", + fontsize=10, + color=MUTED, + ) + return fig _bars( right, [f"TP{run['tp']}" for run in tp_runs], From 5d50d34388e3e79dfd76101d741513d85f5c375d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:27:58 +0000 Subject: [PATCH 15/19] fix(h3): check the lists of the sequence-parallel native entry points, and pin that the exchange padding is never read The SP entry points take caller-built row lists and tile lists, which the pre-pass of the other backward kernels did not cover. h3_gate_grad_partials, h3_gate_grad_fold, h3_rmsnorm_backward_partials and h3_rmsnorm_fold_partials now run a pre-pass (h3_tiles::check_lists in tile_metadata.cuh, one small launch) that asserts on the device that every listed row is a row of the activation, every tile lies inside its list, and the segment boundaries are non-decreasing and inside the partial buffer; they check that the column blocks fit grid.y (65535) and N fits the int indices of the segment kernels, require at least one segment boundary, and skip a fold launch with no segments. One test per entry point runs the corrupted lists in a child process and expects the device assert; the rmsnorm table-layout tests of #487 (overlapping rows rejected, a view of a wider table read by its row pitch) are ported into test_h3_native_validation.py. sp_norm_adaln: no code change. The padding slots of the send blocks and of the partial gathers hold a valid table row but no tile list refers to them; a test now pads both with NaN (sp = 3, 4 and 8, B = 1 and 2, shards smaller than a tile) and the gradients still equal WS1 bitwise, so a read of the padding would show. The plain-norm backward passes None segment tiles because the binding takes optionals; it is covered by test_plain_norm_ranks_reproduce_ws1, and plan.families[1] is only read when the plan was built with an index. The local-slice validator receives x with S_local positions and the matching slice of the index, which is what validate_h3_modulation checks (S = len(index)), against the full tables for the range check. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 5 +- csrc/cuda/h3/gate_residual.cu | 19 ++- csrc/cuda/h3/rmsnorm_modulate.cu | 26 +++- csrc/cuda/h3/tile_metadata.cuh | 41 ++++++ .../minimax_h3/test_h3_adaln_row_gather.py | 6 +- .../minimax_h3/test_h3_native_validation.py | 121 +++++++++++++++++- .../minimax_h3/test_h3_sp_norm_adaln.py | 75 +++++++++++ 7 files changed, 278 insertions(+), 15 deletions(-) diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index dbff808e4..3290eaf5d 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/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"); diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 8439c037b..f9cd6db6f 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -280,7 +280,15 @@ torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch:: auto partial = torch::empty({tiles, n}, grad.options().dtype(at::kFloat)); if (tiles == 0) return partial; const int threads = 256; - const unsigned col_blocks = static_cast((n + threads - 1) / threads); + 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(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(), rows.numel(), + grad.size(0), tile_begin.data_ptr(), tile_end.data_ptr(), + tiles, nullptr, 0, 0); H3_DISPATCH(grad.scalar_type(), "h3_gate_grad_partials", [&] { gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, at::cuda::getCurrentCUDAStream()>>>( @@ -302,12 +310,19 @@ torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_t 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 unsigned col_blocks = static_cast((n + threads - 1) / threads); + 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(col_blocks64); + h3_tiles::check_lists(at::cuda::getCurrentCUDAStream(), nullptr, 0, 0, nullptr, nullptr, 0, + seg_first_tile.data_ptr(), segments, partial.size(0)); H3_DISPATCH(dtype, "h3_gate_grad_fold", [&] { gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, at::cuda::getCurrentCUDAStream()>>>( diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 6dbdb31a4..25e907b70 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -505,7 +505,14 @@ std::vector h3_rmsnorm_backward_partials( const int threads = 256; const unsigned col_blocks = static_cast((n + threads - 1) / threads); const int64_t tiles = dw_begin.numel(); + TORCH_CHECK(2 * n <= std::numeric_limits::max() && (2 * n + threads - 1) / threads <= 65535, + "N=", n, " is too large for the kernels (int indices, grid.y limit of 65535)"); auto dw_partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + // The row and tile lists are caller-built: every row must be a row of x and every tile must lie + // inside its list before the kernels read through them. + h3_tiles::check_lists(stream, dw_rows.data_ptr(), dw_rows.numel(), rows, + dw_begin.data_ptr(), dw_end.data_ptr(), tiles, nullptr, 0, + 0); if (tiles > 0) { H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_partials", [&] { rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, @@ -523,6 +530,9 @@ std::vector h3_rmsnorm_backward_partials( check_tiles(*seg_rows, *seg_begin, *seg_end, x); const int64_t seg_tiles = seg_begin->numel(); const int64_t width = 2 * n; + h3_tiles::check_lists(stream, seg_rows->data_ptr(), seg_rows->numel(), rows, + seg_begin->data_ptr(), seg_end->data_ptr(), seg_tiles, + nullptr, 0, 0); auto seg_partial = torch::empty({seg_tiles, width}, x.options().dtype(at::kFloat)); const unsigned wide_blocks = static_cast((width + threads - 1) / threads); if (seg_tiles > 0) { @@ -572,14 +582,22 @@ std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, to seg_partial->dim() == 2 && seg_partial->is_contiguous() && seg_partial->scalar_type() == at::kFloat && seg_partial->size(1) == 2 * n, "seg_partial must be a contiguous float32 (tiles, 2N) CUDA tensor"); + TORCH_CHECK(seg_first_tile->numel() >= 1, "seg_first_tile must hold R + 1 entries"); const int64_t segments = seg_first_tile->numel() - 1; const int64_t width = 2 * n; + TORCH_CHECK(width <= std::numeric_limits::max() && (width + threads - 1) / threads <= 65535, + "N=", n, " is too large for the kernels (int indices, grid.y limit of 65535)"); auto table_grad = torch::empty({segments, width}, dw_partial.options()); const unsigned wide_blocks = static_cast((width + threads - 1) / threads); - fold_segments_kernel<<(segments), wide_blocks), threads, 0, stream>>>( - seg_partial->data_ptr(), seg_first_tile->data_ptr(), - table_grad.data_ptr(), static_cast(width)); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (segments > 0) { // a zero-sized grid is not a valid launch + h3_tiles::check_lists(stream, nullptr, 0, 0, nullptr, nullptr, 0, + seg_first_tile->data_ptr(), segments, seg_partial->size(0)); + fold_segments_kernel<<(segments), wide_blocks), threads, 0, + stream>>>(seg_partial->data_ptr(), + seg_first_tile->data_ptr(), + table_grad.data_ptr(), static_cast(width)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } return {dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; } diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index a959054e2..7cda74cc4 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -58,4 +58,45 @@ inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, C10_CUDA_KERNEL_LAUNCH_CHECK(); } +// Sequence-parallel entry points take explicit lists. `rows` (n_rows entries, each a row of an +// M-row activation, so in [0, row_limit)) and the tiles [tile_begin, tile_end) over it, and/or the +// segment boundaries seg_first_tile (segments + 1 entries, in [0, seg_limit], non-decreasing). A +// null pointer skips its part. Same device-side asserts, same single small launch. +static __global__ void check_lists_kernel(const int64_t* __restrict__ rows, int64_t n_rows, + int64_t row_limit, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, int64_t tiles, + const int64_t* __restrict__ seg_first_tile, + int64_t segments, int64_t seg_limit) { + const int64_t first_thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + if (rows != nullptr) { + for (int64_t i = first_thread; i < n_rows; i += stride) { + CUDA_KERNEL_ASSERT(rows[i] >= 0 && rows[i] < row_limit); + } + } + if (tile_begin != nullptr) { + for (int64_t i = first_thread; i < tiles; i += stride) { + CUDA_KERNEL_ASSERT(tile_begin[i] >= 0 && tile_begin[i] <= tile_end[i] && + tile_end[i] <= n_rows); + } + } + if (seg_first_tile != nullptr) { + for (int64_t i = first_thread; i <= segments; i += stride) { + CUDA_KERNEL_ASSERT(seg_first_tile[i] >= 0 && seg_first_tile[i] <= seg_limit); + CUDA_KERNEL_ASSERT(i == segments || seg_first_tile[i] <= seg_first_tile[i + 1]); + } + } +} + +inline void check_lists(cudaStream_t stream, const int64_t* rows, int64_t n_rows, int64_t row_limit, + const int64_t* tile_begin, const int64_t* tile_end, int64_t tiles, + const int64_t* seg_first_tile, int64_t segments, int64_t seg_limit) { + const int64_t items = std::max({n_rows, tiles, segments + 1}); + const unsigned blocks = static_cast(std::min((items + 255) / 256, 1024)); + check_lists_kernel<<>>(rows, n_rows, row_limit, tile_begin, tile_end, + tiles, seg_first_tile, segments, seg_limit); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + } // namespace h3_tiles diff --git a/tests/models/minimax_h3/test_h3_adaln_row_gather.py b/tests/models/minimax_h3/test_h3_adaln_row_gather.py index 26a9c7921..65a927577 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -393,7 +393,7 @@ def test_the_backward_stops_on_metadata_that_does_not_describe_the_tiles(self): assert done.returncode != 0, mutation assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) - def test_backward_with_no_table_rows_returns_an_empty_table_and_rejects_no_boundaries(self): + def test_backward_with_no_table_rows_or_no_boundaries_is_rejected(self): from rl_engine.backends.extension import _C _cuda_op() @@ -401,8 +401,8 @@ def test_backward_with_no_table_rows_returns_an_empty_table_and_rejects_no_bound zero = torch.zeros(0, dtype=torch.int64, device="cuda") pos = torch.arange(4, device="cuda") one = torch.zeros(1, dtype=torch.int64, device="cuda") # R = 0 - empty = _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, one, torch.float32) - assert empty.shape == (0, 48) + with pytest.raises(RuntimeError, match="must be positive"): + _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, one, torch.float32) with pytest.raises(RuntimeError, match="R \\+ 1 entries"): _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, zero, torch.float32) diff --git a/tests/models/minimax_h3/test_h3_native_validation.py b/tests/models/minimax_h3/test_h3_native_validation.py index bcb06f1f2..1c8040e5d 100644 --- a/tests/models/minimax_h3/test_h3_native_validation.py +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -286,7 +286,8 @@ def test_native_table_bounds_fail_in_isolated_process(native_extension, operatio import sys import textwrap - script = textwrap.dedent(""" + script = textwrap.dedent( + """ import sys import torch from rl_engine import _C @@ -318,7 +319,8 @@ def test_native_table_bounds_fail_in_isolated_process(native_extension, operatio sys.exit(0) raise raise AssertionError("invalid index reached native table lookup without an assertion") - """) + """ + ) result = subprocess.run( [sys.executable, "-c", script, operation, str(bad_index)], capture_output=True, @@ -352,7 +354,8 @@ def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metada import sys import textwrap - script = textwrap.dedent(""" + script = textwrap.dedent( + """ import sys import torch from rl_engine import _C @@ -372,9 +375,119 @@ def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metada sys.exit(0) raise raise AssertionError("invalid metadata was not rejected") - """) + """ + ) result = subprocess.run( [sys.executable, "-c", script, metadata, bound], capture_output=True, text=True, timeout=60 ) assert result.returncode == 0, result.stdout + result.stderr assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout + + +def test_rmsnorm_rejects_modulation_rows_that_overlap(native_extension, native_case): + """A table of rows needs a row stride of at least N; an overlapping view is not one.""" + + table = torch.arange(32, dtype=torch.float32, device="cuda") + overlapping = table.as_strided((2, 8), (4, 1)) + with pytest.raises(RuntimeError, match="row stride 4 is smaller than N=8"): + native_extension.h3_rmsnorm_forward( + native_case["x"], native_case["weight"], 1e-5, overlapping, overlapping, + native_case["index"], + ) # fmt: skip + + +def test_rmsnorm_reads_a_view_of_a_wider_table_by_its_row_pitch(native_extension, native_case): + """``table.chunk(2, dim=-1)`` has row stride 2N; the kernel must index rows by it.""" + + g = torch.Generator(device="cpu").manual_seed(0) + views = torch.randn(2, 16, generator=g).cuda().chunk(2, dim=-1) + copies = tuple(v.contiguous() for v in views) + assert views[0].stride(0) == 16 and copies[0].stride(0) == 8 + out = [ + native_extension.h3_rmsnorm_forward( + native_case["x"], native_case["weight"], 1e-5, shift, scale, native_case["index"] + )[0] + for shift, scale in (views, copies) + ] + assert torch.equal(out[0], out[1]) + + +_BACKWARD_CHILD = """ +import sys +import torch +from rl_engine.backends.extension import _C + +kind, mutation = sys.argv[1:] +x = torch.randn(4, 8, device="cuda") +weight = torch.ones(8, device="cuda") +gate = torch.ones(2, 8, device="cuda") +index = torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda") +i64 = dict(dtype=torch.int64, device="cuda") +sorted_pos = torch.tensor([0, 2, 1, 3], **i64) +begin, end = torch.tensor([0, 2], **i64), torch.tensor([2, 4], **i64) +first = torch.tensor([0, 1, 2], **i64) +if mutation == "tile_end_past_rows": + end = torch.tensor([2, 9], **i64) +elif mutation == "position_past_rows": + sorted_pos = torch.tensor([0, 2, 1, 4], **i64) +elif mutation == "segments_past_tiles": + first = torch.tensor([0, 1, 5], **i64) +if kind == "rmsnorm_backward": + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) + _C.h3_rmsnorm_backward(torch.ones_like(x), x, weight, rstd, shift, scale, index, sorted_pos, + begin, end, first) +elif kind == "gate_backward": + _C.h3_gate_residual_backward(torch.ones_like(x), x, gate, index, sorted_pos, begin, end, first) +elif kind == "gate_partials": # sequence parallel: explicit row list and tiles + _C.h3_gate_grad_partials(torch.ones_like(x), x, sorted_pos, begin, end) +elif kind == "gate_fold": + _C.h3_gate_grad_fold(torch.ones(2, 8, device="cuda"), first, torch.float32) +elif kind == "rmsnorm_partials": + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + _C.h3_rmsnorm_backward_partials(torch.ones_like(x), x, weight, rstd, None, None, None, + sorted_pos, begin, end, None, None, None) +elif kind == "rmsnorm_fold": + _C.h3_rmsnorm_fold_partials(torch.ones(2, 8, device="cuda"), weight, + torch.ones(2, 16, device="cuda"), first) +torch.cuda.synchronize() +""" + +_METADATA_CASES = { + "rmsnorm_backward": ("tile_end_past_rows", "position_past_rows", "segments_past_tiles"), + "gate_backward": ("tile_end_past_rows", "position_past_rows", "segments_past_tiles"), + "gate_partials": ("tile_end_past_rows", "position_past_rows"), + "gate_fold": ("segments_past_tiles",), + "rmsnorm_partials": ("tile_end_past_rows", "position_past_rows"), + "rmsnorm_fold": ("segments_past_tiles",), +} + + +@pytest.mark.parametrize("kind", sorted(_METADATA_CASES)) +def test_segment_metadata_that_does_not_describe_the_rows_ends_in_a_device_assert( + native_extension, kind +): + """Every entry point that reads caller-built tile lists checks them before it reads.""" + + import os + import subprocess + import sys + + if not hasattr(native_extension, "h3_gate_grad_partials"): + pytest.skip("rl_engine._C lacks the H3 SP entry points") + + def run(mutation): + return subprocess.run( + [sys.executable, "-c", _BACKWARD_CHILD, kind, mutation], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + + control = run("none") + assert control.returncode == 0, control.stderr + for mutation in _METADATA_CASES[kind]: + done = run(mutation) + assert done.returncode != 0, (kind, mutation) + assert "device-side assert triggered" in done.stderr, (kind, mutation, done.stderr) diff --git a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py index 4866e68f3..b3213fd1e 100644 --- a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py +++ b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py @@ -165,6 +165,79 @@ def rank(coll): ) assert torch.equal(dgate, ws1_gate[1]) + @pytest.mark.parametrize("sp, batch, seq", [(4, 2, 777), (3, 2, 1000), (8, 1, 40)]) + def test_padding_of_the_exchange_and_partial_gathers_is_never_read( + self, monkeypatch, sp, batch, seq + ): + """Pad every send block and every partial gather with NaN: WS1 equality still holds. + + ``recv_global`` is zero in the padding slots, which map to a real table row, so the + padding is only harmless because no tile list refers to it; a NaN would reach the sums + of any tile that did. + """ + + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3 import sp_norm_adaln as module + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + + hidden = 256 + c = sp_case(batch, seq, hidden, layout="interleaved", dtype=torch.float32, seed=sp + seq) + chunks = c["table"].chunk(6, dim=-1) + shift, scale, gate = chunks[3], chunks[4], chunks[2] + x2, g2, y2 = (c[k].view(-1, hidden) for k in ("residual", "grad", "y")) + idx, rows = c["index"], shift.shape[0] + tiles = _segment_tiles(idx.repeat(batch), rows) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5, shift, scale, idx) + ws1_norm = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd, shift, scale, idx, *tiles) + ws1_gate = _C.h3_gate_residual_backward(g2, y2, gate, idx, *tiles) + nan = float("nan") + gather_rows = module.gather_rows + + def exchange(self, collective, *tensors): + out = [] + for t in tensors: + if self.max_send == 0: + out.append(t) + continue + send = t.new_full((self.max_send, *t.shape[1:]), nan) # padding is NaN + send[: self.send_local.numel()] = t[self.send_local] + out.append(torch.cat([t, gather_rows(collective, send)])) + return out + + def gather_partials(self, collective, family, partial): + fam = self.families[family] + padded = partial.new_full((fam["width"], *partial.shape[1:]), nan) # padding is NaN + padded[: partial.shape[0]] = partial + return gather_rows(collective, padded).index_select(0, fam["gather_index"]).contiguous() + + monkeypatch.setattr(module._Plan, "exchange", exchange) + monkeypatch.setattr(module._Plan, "gather_partials", gather_partials) + + def rank(coll): + lay = sp_row_layout(seq, batch, sp, coll.rank) + part = slice(lay.lo, lay.hi) + local = { + k: c[k][:, part].reshape(-1, hidden).contiguous() for k in ("residual", "grad", "y") + } + rs = rstd.view(batch, seq)[:, part].reshape(-1).contiguous() + plan = module.sp_plan(lay, idx, rows) + norm = module.sp_rmsnorm_backward( + local["grad"], local["residual"], c["weight"], rs, shift, scale, plan, coll + ) + return norm, module.sp_gate_residual_backward( + local["grad"], local["y"], gate, plan, coll + ) + + for (dx, dw, dsh, dsc), (dy, dgate) in _ThreadGroup(sp).run(rank): + assert torch.isfinite(dx).all() and torch.isfinite(dy).all() + assert torch.equal(dw, ws1_norm[1]) + assert torch.equal(dsh, ws1_norm[2].to(dsh.dtype)) + assert torch.equal(dsc, ws1_norm[3].to(dsc.dtype)) + assert torch.equal(dgate, ws1_gate[1]) + def test_plain_norm_ranks_reproduce_ws1(self): _require_ext() from rl_engine import _C @@ -284,7 +357,9 @@ def _worker_never_reports(rank, world, init, queue, path, target, kwargs): @pytest.mark.parametrize("mode", ["clean_exit", "crash", "reported_error", "timeout"]) def test_run_world_cleans_up_failed_workers(monkeypatch, mode): import time + import torch.multiprocessing as mp + from rl_engine.validation.models import h3_ws2 workers = { From 9aa200d236bd7410342c021067575db81b1f1a44 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:46:07 +0000 Subject: [PATCH 16/19] docs(h3): document the det_linear partial and fold stubs used by the tensor-parallel path h3_det_linear_backward_input_partials returns the (C, T, K) FP32 per-chunk partials (64 output columns per chunk, ascending), h3_det_linear_fold_chunks left-folds them over C and casts once; both were bare signatures in _C.pyi. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/_C.pyi | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 5408f1557..c910add64 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -374,10 +374,23 @@ def h3_det_linear_backward_input( """ ... -def h3_det_linear_backward_input_partials( - grad: torch.Tensor, weight: torch.Tensor -) -> torch.Tensor: ... -def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: ... +def h3_det_linear_backward_input_partials(grad: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Return the ``(C, T, K)`` FP32 per-chunk partials of ``grad @ weight`` before the fold. + + Same operands as ``h3_det_linear_backward_input``. Chunk ``c`` of ``C = ceil(N / 64)`` sums the + output columns ``64c .. 64c + 63`` in ascending order; the tensor-parallel path reduces its + shard's chunks with ``h3_det_linear_fold_chunks`` after the shards are gathered. + """ + ... + +def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Left-fold ``partial (C, T, K)`` FP32 over ``C`` in ascending order, cast once to ``out_dtype``. + + Require a non-empty contiguous FP32 CUDA tensor and an FP32 or BF16 ``out_dtype``; return + ``(T, K)``. + """ + ... + def h3_det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True ) -> list[torch.Tensor]: @@ -425,6 +438,7 @@ def h3_adaln_row_gather_backward( inside ``sorted_pos``, segments that cut the tile list exactly and in order. """ ... + def h3_rmsnorm_forward( x: torch.Tensor, weight: torch.Tensor, @@ -469,6 +483,7 @@ def h3_rmsnorm_backward( modulation they must be omitted. """ ... + def h3_gate_residual_forward( residual: torch.Tensor, y: torch.Tensor, gate: torch.Tensor, index: torch.Tensor ) -> torch.Tensor: @@ -501,4 +516,3 @@ def h3_gate_residual_backward( checked on the device before use. """ ... - From 002fb393da77917cccb92cf70a2390a38019cea8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:52:23 +0000 Subject: [PATCH 17/19] fix(h3): check the modulation index in the sequence-parallel partials, floor the list check's launch, pin the no-H3-symbols build * `h3_rmsnorm_backward_partials` runs `h3_tiles::check_index_range` over the S index values before its dweight kernel, which reads them inside its row loop without an assert (an assert there costs about 11% of the backward). This entry point can be called without the dx entry point having run, so it cannot lean on dx's per-row assert. The native-validation child-process test calls it with an index past the table and a negative one (and a valid control) and requires the device-side assert. * `check_lists` launches at least one block, like `check_tile_metadata`. * `test_h3_extension_without_h3_symbols.py`: with an extension that lacks every H3 binding (what a ROCm, fast-math or relaxed-flag build produces), every CUDA op, including the sequence-parallel one, refuses to construct, so the registry falls back instead of failing with AttributeError; the SP availability probe is False when any single symbol it calls is missing. * `_C.pyi`: contracts for the sequence-parallel entry points. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../test_h3_extension_without_h3_symbols.py | 73 +++++++++++++++++++ .../minimax_h3/test_h3_native_validation.py | 11 +++ 2 files changed, 84 insertions(+) create mode 100644 tests/models/minimax_h3/test_h3_extension_without_h3_symbols.py diff --git a/tests/models/minimax_h3/test_h3_extension_without_h3_symbols.py b/tests/models/minimax_h3/test_h3_extension_without_h3_symbols.py new file mode 100644 index 000000000..7d5e465e3 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_extension_without_h3_symbols.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""A build that leaves the H3 sources out has no ``h3_*`` symbols; no CUDA op may then construct. + +``KERNEL_ALIGN_USE_FAST_MATH=1``, ROCm and relaxed ``NVCC_*_FLAGS`` builds omit the H3 kernels and +their bindings (build_tools/extensions.py). Every H3 CUDA op checks for its symbols when it is +constructed, so the registry falls back to the PyTorch reference instead of failing on the first +call with an AttributeError. +""" + +from __future__ import annotations + +import importlib + +import pytest + +PACKAGE = "rl_engine.backends.cuda.model_specific.minimax_h3" +# the modules that define the probes and own the ``_C`` they look in +PROBE_MODULES = ( + "timestep_sinusoid", + "det_linear", + "adaln_row_gather", + "rmsnorm", + "gate_residual", + "sp_norm_adaln", +) +# module, class, constructor arguments +OPS = ( + ("timestep_sinusoid", "H3TimestepSinusoidCudaOp", ()), + ("timestep_mlp", "H3TimestepMLPCudaOp", ()), + ("adaln_projection", "H3AdaLNProjectionCudaOp", ()), + ("adaln_row_gather", "H3AdaLNRowGatherCudaOp", ()), + ("adaln_modulation", "H3AdaLNModulationCudaOp", ()), + ("rmsnorm", "H3RMSNormCudaOp", ()), + ("gate_residual", "H3GateResidualCudaOp", ()), + ("final_adaln_out", "H3FinalAdaLNOutCudaOp", ()), + ("sp_norm_adaln", "H3SPNormAdaLNCudaOp", (None, 1, 1)), +) + + +class ExtensionWithoutH3: + """Stands in for a compiled ``rl_engine._C`` that has none of the H3 bindings.""" + + +@pytest.fixture +def extension_without_h3(monkeypatch): + for name in PROBE_MODULES: + module = importlib.import_module(f"{PACKAGE}.{name}") + monkeypatch.setattr(module, "_C", ExtensionWithoutH3()) + monkeypatch.setattr(module, "_EXT_AVAILABLE", True) + + +@pytest.mark.parametrize(("module", "op", "args"), OPS, ids=[o[1] for o in OPS]) +def test_every_cuda_op_refuses_to_construct(extension_without_h3, module, op, args): + cls = getattr(importlib.import_module(f"{PACKAGE}.{module}"), op) + with pytest.raises(RuntimeError): + cls(*args) + + +def test_the_sequence_parallel_op_checks_every_symbol_it_calls(monkeypatch): + sp = importlib.import_module(f"{PACKAGE}.sp_norm_adaln") + monkeypatch.setattr(sp, "_EXT_AVAILABLE", True) + + def extension(missing=None): + names = {n: staticmethod(lambda *a, **k: None) for n in sp._SYMBOLS if n != missing} + return type("Extension", (), names)() + + monkeypatch.setattr(sp, "_C", extension()) + assert sp.sp_norm_adaln_available() is True + for name in sp._SYMBOLS: # one missing symbol is enough to refuse: the op does not half-work + monkeypatch.setattr(sp, "_C", extension(missing=name)) + assert sp.sp_norm_adaln_available() is False, name diff --git a/tests/models/minimax_h3/test_h3_native_validation.py b/tests/models/minimax_h3/test_h3_native_validation.py index 1c8040e5d..657f0b4cb 100644 --- a/tests/models/minimax_h3/test_h3_native_validation.py +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -432,6 +432,10 @@ def test_rmsnorm_reads_a_view_of_a_wider_table_by_its_row_pitch(native_extension sorted_pos = torch.tensor([0, 2, 1, 4], **i64) elif mutation == "segments_past_tiles": first = torch.tensor([0, 1, 5], **i64) +elif mutation == "index_past_table": + index = torch.tensor([0, 1, 0, 7], **i64) +elif mutation == "index_negative": + index = torch.tensor([0, 1, 0, -1], **i64) if kind == "rmsnorm_backward": _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) @@ -450,6 +454,12 @@ def test_rmsnorm_reads_a_view_of_a_wider_table_by_its_row_pitch(native_extension elif kind == "rmsnorm_fold": _C.h3_rmsnorm_fold_partials(torch.ones(2, 8, device="cuda"), weight, torch.ones(2, 16, device="cuda"), first) +elif kind == "rmsnorm_partials_modulated": + # called without the dx entry point having run: the partials check the index values themselves + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) + _C.h3_rmsnorm_backward_partials(torch.ones_like(x), x, weight, rstd, shift, scale, index, + sorted_pos, begin, end, sorted_pos, begin, end) torch.cuda.synchronize() """ @@ -459,6 +469,7 @@ def test_rmsnorm_reads_a_view_of_a_wider_table_by_its_row_pitch(native_extension "gate_partials": ("tile_end_past_rows", "position_past_rows"), "gate_fold": ("segments_past_tiles",), "rmsnorm_partials": ("tile_end_past_rows", "position_past_rows"), + "rmsnorm_partials_modulated": ("index_past_table", "index_negative"), "rmsnorm_fold": ("segments_past_tiles",), } From 2768c0f17e378f3b3616cfc438cdfb79686b10e5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:19:01 +0000 Subject: [PATCH 18/19] test(h3): restore the rmsnorm native validation tests dropped from this branch The file is on test-h3 and on every branch below this one, but earlier merges of the tensor-parallel branch into this one left it out. Three of its tests (modulation row outside the table asserts in the kernel, a view of a wider table is read by its row pitch, the backward stops on tile metadata that does not describe the rows) exist nowhere else. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../test_h3_rmsnorm_native_validation.py | 282 ++++++++++++++++++ 1 file changed, 282 insertions(+) create mode 100644 tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py diff --git a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py new file mode 100644 index 000000000..e97d7f6c8 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Direct native calls reject malformed pointers and layouts before launching CUDA work.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +import torch + +pytestmark = [ + pytest.mark.cuda_only, + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), +] + +TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") + + +@pytest.fixture +def native(): + from rl_engine.backends.extension import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not all( + hasattr(_C, name) for name in ("h3_rmsnorm_forward", "h3_rmsnorm_backward") + ): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return _C + + +@pytest.fixture +def inputs(native): + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device=x.device) + _, rstd = native.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device=x.device).chunk(2, dim=1) + return { + "grad": torch.ones_like(x), + "x": x, + "weight": weight, + "rstd": rstd, + "shift": shift, + "scale": scale, + "index": torch.tensor([0, 1], dtype=torch.int64, device=x.device), + "sorted_pos": torch.tensor([0, 2, 1, 3], dtype=torch.int64, device=x.device), + "tile_begin": torch.tensor([0, 2], dtype=torch.int64, device=x.device), + "tile_end": torch.tensor([2, 4], dtype=torch.int64, device=x.device), + "seg_first_tile": torch.tensor([0, 1, 2], dtype=torch.int64, device=x.device), + } + + +def _call(native, inputs, entrypoint): + if entrypoint == "forward": + return native.h3_rmsnorm_forward( + inputs["x"], + inputs["weight"], + 1e-5, + inputs["shift"], + inputs["scale"], + inputs["index"], + ) + return native.h3_rmsnorm_backward(**inputs) + + +def _malformed(tensor, defect): + if defect == "cpu": + return tensor.cpu() + if defect == "dtype": + dtype = torch.float16 if tensor.dtype == torch.float32 else torch.int32 + return tensor.to(dtype) + if defect == "rank": + return tensor.unsqueeze(0) + if defect == "stride": + return torch.stack((tensor, tensor), dim=1)[:, 0] + if defect == "length": + return tensor[:-1] + raise ValueError(f"unknown defect: {defect}") + + +def test_valid_plain_and_modulated_backward(native, inputs): + plain = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + assert len(native.h3_rmsnorm_backward(**plain)) == 2 + dx, dweight, dshift, dscale = native.h3_rmsnorm_backward(**inputs) + assert dx.shape == inputs["x"].shape + assert dweight.shape == inputs["weight"].shape + assert dscale.shape == inputs["scale"].shape + torch.testing.assert_close(dshift, torch.full_like(inputs["shift"], 2.0)) + torch.cuda.synchronize(inputs["x"].device) + + +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_empty_row_index(native, inputs, entrypoint): + inputs["index"] = inputs["index"][:0] + with pytest.raises(RuntimeError, match="row index must be a non-empty"): + _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_zero_hidden_size(native, inputs, entrypoint): + inputs["x"] = inputs["x"][:, :0].contiguous() + inputs["weight"] = inputs["weight"][:0] + with pytest.raises(RuntimeError, match="at least one column"): + _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +@pytest.mark.parametrize("mode", ["plain", "modulated"]) +def test_rejects_invalid_rstd(native, inputs, defect, mode): + inputs["rstd"] = _malformed(inputs["rstd"], defect) + if mode == "plain": + inputs = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.parametrize("name", TILE_NAMES) +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +def test_rejects_invalid_segment_tiles(native, inputs, name, defect): + inputs[name] = _malformed(inputs[name], defect) + message = "tile_begin and tile_end" if defect == "length" and name.startswith("tile_") else name + with pytest.raises(RuntimeError, match=message): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.parametrize("name", TILE_NAMES) +def test_rejects_missing_segment_tiles(native, inputs, name): + inputs[name] = None + with pytest.raises(RuntimeError, match="needs the sorted segment tiles"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.parametrize("name", TILE_NAMES) +def test_plain_backward_rejects_segment_tiles(native, inputs, name): + plain = {key: inputs[key] for key in ("grad", "x", "weight", "rstd", name)} + with pytest.raises(RuntimeError, match="sorted segment tiles require modulation"): + native.h3_rmsnorm_backward(**plain) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize( + ("entrypoint", "name"), + [ + ("forward", "weight"), + ("forward", "index"), + ("backward", "weight"), + ("backward", "index"), + ("backward", "rstd"), + ("backward", "sorted_pos"), + ("backward", "tile_begin"), + ("backward", "tile_end"), + ("backward", "seg_first_tile"), + ], +) +def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name): + other_device = (inputs["x"].device.index + 1) % torch.cuda.device_count() + inputs[name] = inputs[name].to(f"cuda:{other_device}") + message = "row index" if name == "index" else name + with pytest.raises(RuntimeError, match=message): + _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +@pytest.mark.parametrize("bad_row", [-1, 2], ids=["negative", "past-table"]) +def test_modulation_row_outside_the_table_asserts_in_the_kernel(native, entrypoint, bad_row): + """check_range=False callers get a device-side assert, not a read past the shift/scale table.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + entrypoint, bad_row = sys.argv[1], int(sys.argv[2]) + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device="cuda") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) # a table of two rows + index = torch.tensor([0, bad_row], dtype=torch.int64, device="cuda") + if entrypoint == "forward": + _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + else: + _C.h3_rmsnorm_backward( + grad=torch.ones_like(x), x=x, weight=weight, rstd=rstd, shift=shift, scale=scale, + index=index, + sorted_pos=torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cuda"), + tile_begin=torch.tensor([0, 2], dtype=torch.int64, device="cuda"), + tile_end=torch.tensor([2, 4], dtype=torch.int64, device="cuda"), + seg_first_tile=torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"), + ) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(bad_row)], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + assert completed.returncode != 0, "the kernel accepted a modulation row outside the table" + assert ( + "device-side assert triggered" in completed.stderr + ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + + +def test_rejects_modulation_rows_that_overlap(native, inputs): + """A table of rows needs a row stride of at least N; an overlapping view is not one.""" + + table = torch.arange(32, dtype=torch.float32, device="cuda") + inputs["shift"] = inputs["scale"] = table.as_strided((2, 8), (4, 1)) + with pytest.raises(RuntimeError, match="row stride 4 is smaller than N=8"): + _call(native, inputs, "forward") + + +def test_a_view_of_a_wider_table_is_read_by_its_row_pitch(native, inputs): + """``table.chunk(2, dim=-1)`` has row stride 2N; the kernel must index rows by it.""" + + g = torch.Generator(device="cpu").manual_seed(0) + table = torch.randn(2, 16, generator=g).cuda() + views = table.chunk(2, dim=-1) + copies = tuple(v.contiguous() for v in views) + assert views[0].stride(0) == 16 and copies[0].stride(0) == 8 + out = [] + for shift, scale in (views, copies): + inputs["shift"], inputs["scale"] = shift, scale + out.append(_call(native, inputs, "forward")[0]) + assert torch.equal(out[0], out[1]) + + +def _backward_in_child(mutation: str) -> subprocess.CompletedProcess: + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + mutation = sys.argv[1] + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device="cuda") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) + args = dict( + grad=torch.ones_like(x), x=x, weight=weight, rstd=rstd, shift=shift, scale=scale, + index=torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda"), + sorted_pos=torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cuda"), + tile_begin=torch.tensor([0, 2], dtype=torch.int64, device="cuda"), + tile_end=torch.tensor([2, 4], dtype=torch.int64, device="cuda"), + seg_first_tile=torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"), + ) + if mutation == "tile_end_past_rows": + args["tile_end"] = torch.tensor([2, 9], dtype=torch.int64, device="cuda") + elif mutation == "position_past_rows": + args["sorted_pos"] = torch.tensor([0, 2, 1, 4], dtype=torch.int64, device="cuda") + elif mutation == "segments_past_tiles": + args["seg_first_tile"] = torch.tensor([0, 1, 5], dtype=torch.int64, device="cuda") + _C.h3_rmsnorm_backward(**args) + torch.cuda.synchronize() + """ + ) + return subprocess.run( + [sys.executable, "-c", probe, mutation], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + + +def test_the_backward_stops_on_tile_metadata_that_does_not_describe_the_rows(native): + control = _backward_in_child("none") + assert control.returncode == 0, control.stderr + for mutation in ("tile_end_past_rows", "position_past_rows", "segments_past_tiles"): + done = _backward_in_child(mutation) + assert done.returncode != 0, mutation + assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) From 7b8594136f50e9a86050d420a01f43f82f662b9f Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:33:07 +0000 Subject: [PATCH 19/19] fix(h3): skip the weight and bias gradients of a frozen tensor-parallel shard The projection backward folded dW and db for every call. A frozen shard now returns None for them and skips the folds, which are local to the rank. The input-gradient gather is a collective, so it stays unconditional. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/tp_adaln_projection.py | 13 ++++++-- tests/models/minimax_h3/test_h3_tp_adaln.py | 32 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py index 371f0ae74..28768e881 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py @@ -161,13 +161,22 @@ def backward(ctx, grad_table): f"got {tuple(grad_table.shape)}" ) grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() - dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) + # The weight and bias folds are local to this rank: skip them for a frozen shard. The input + # gradient below stays unconditional, because its gather is a collective every rank must + # reach, and ``temb`` is replicated. + _, need_weight, need_bias = ctx.needs_input_grad[:3] + dw = db = None + if need_weight or need_bias: + dw, *rest = det_linear_backward_weight( + grad_local, act, weight.dtype, with_bias=need_bias + ) + db = rest[0] if need_bias else None partial = det_linear_backward_input_partials(grad_local, weight) d_act = det_linear_fold_chunks(gather_rows(ctx.collective, partial), torch.float32) record_backward( "tp_adaln_3mod", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) - return silu_backward_fp32(d_act, temb), dw, db, None, None + return silu_backward_fp32(d_act, temb), dw if need_weight else None, db, None, None class H3TPAdaLNProjectionCudaOp: diff --git a/tests/models/minimax_h3/test_h3_tp_adaln.py b/tests/models/minimax_h3/test_h3_tp_adaln.py index 6d3737824..74281f750 100644 --- a/tests/models/minimax_h3/test_h3_tp_adaln.py +++ b/tests/models/minimax_h3/test_h3_tp_adaln.py @@ -168,6 +168,38 @@ def test_backward_rejects_a_table_gradient_of_the_wrong_width(self): _TPAdaLNProjection.backward(node, bad.to(table.dtype)) assert _TPAdaLNProjection.backward(node, grad.to(table.dtype))[0].shape == temb.shape + @pytest.mark.parametrize("frozen", [("weight",), ("bias",), ("weight", "bias")]) + def test_backward_skips_the_gradients_of_a_frozen_shard(self, frozen): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(2, 3 * 6 * 256, 136) + op = H3TPAdaLNProjectionCudaOp(one, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 1, 0) + leaves = { + "temb": temb.clone().requires_grad_(True), + "weight": w.clone().requires_grad_("weight" not in frozen), + "bias": b.clone().requires_grad_("bias" not in frozen), + } + table = op.forward_table(*leaves.values()) + got = table.grad_fn.apply(grad.to(table.dtype)) + reference = { + name: leaf.detach().clone().requires_grad_(True) for name, leaf in leaves.items() + } + for name in leaves: + reference[name].requires_grad_(name not in frozen) + reference["temb"].requires_grad_(True) + op.forward_table(*reference.values()).backward(grad.to(table.dtype)) + for name, value in zip(("temb", "weight", "bias"), got[:3], strict=True): + if name in frozen: + assert value is None + else: + assert torch.equal(value, reference[name].grad) + @pytest.mark.parametrize("world", [2, 4, 8]) @pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes")