From d275a06f96169f9d4b3b79c7084182b95668031d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:40:27 +0000 Subject: [PATCH 01/42] [WS1][CUDA][MiniMax-H3] adaln_row_gather: fused six-way AdaLN row gather RFC #420 step 3, row `adaln_row_gather` (stacked on adaln_projection_3mod). - Forward: one launch gathers all six (S, H) modulation tensors by timestep_index * 3 + token_tag with 16-byte copies; bitwise equal to diffusers' six index_select calls. Out-of-range tags or timesteps fail closed (RFC probes H2/H3), never clamped. - Backward: deterministic FP32 segmented sum (stable sort by row, fixed 256-position tiles, ordered fold, one cast; no atomics) instead of index_select's atomic BF16 scatter-add. The PyTorch reference gets an equally deterministic backward. - H3AdaLNModulationCudaOp: projection + gather as one autograd node, bitwise-equal forward, FP32 table gradient into the projection backward. - Chain replay gains the gather stage and a backward mode (separate, fused and diffusers chains vs an FP64 golden); the end-to-end test now covers the whole chain forward and backward. - Report/figure hooks, registry, gtest (reduction class: the VJP is a segmented sum), doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 243 ++++++++++++++++++ csrc/ops.cpp | 15 ++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-adaln-row-gather.md | 141 ++++++++++ rl_engine/_C.pyi | 15 ++ rl_engine/kernels/gtest/operator_inputs.py | 20 ++ rl_engine/kernels/gtest/operator_specs.py | 15 ++ .../kernels/ops/cuda/h3/adaln_modulation.py | 133 ++++++++++ .../kernels/ops/cuda/h3/adaln_row_gather.py | 113 ++++++++ .../ops/pytorch/h3/adaln_row_gather.py | 141 ++++++++++ rl_engine/kernels/registry.py | 13 + rl_engine/testing/h3_chain.py | 99 +++++++ rl_engine/testing/h3_provider.py | 9 + rl_engine/testing/h3_report.py | 95 ++++++- rl_engine/tests/test_dispatch.py | 7 +- scripts/h3_chain_replay.py | 24 ++ scripts/plot_h3_evidence.py | 99 ++++++- setup.py | 1 + tests/h3/test_h3_adaln_modulation.py | 117 +++++++++ tests/h3/test_h3_adaln_row_gather.py | 207 +++++++++++++++ tests/h3/test_h3_conditioning_e2e.py | 19 ++ tests/test_ws1_gtest_gpu.py | 1 + 23 files changed, 1520 insertions(+), 9 deletions(-) create mode 100644 csrc/cuda/h3/adaln_row_gather.cu create mode 100644 docs/operators/h3-adaln-row-gather.md create mode 100644 rl_engine/kernels/ops/cuda/h3/adaln_modulation.py create mode 100644 rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py create mode 100644 tests/h3/test_h3_adaln_modulation.py create mode 100644 tests/h3/test_h3_adaln_row_gather.py diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu new file mode 100644 index 000000000..192faa187 --- /dev/null +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -0,0 +1,243 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 AdaLN row gather (RFC #420 `adaln_row_gather`). +// +// rows: (3T, C * H) modulation rows, row r = timestep * 3 + modality and +// column block c = shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, +// gate_mlp. For every packed-sequence position s: +// +// r[s] = timestep_indices[s] * modality_num + token_tags[s] +// out[c, s, :] = rows[r[s], c * H : (c + 1) * H] +// +// Forward is a pure copy: bitwise equal to C index_select calls, one launch, +// independent of S, packing order and repeats. +// +// Backward d_rows[r, j] = sum_{s : r[s] = r} grad[c(j), s, h(j)] is the one +// cross-row reduction. Positions are visited in a fixed order: the caller +// sorts them by (r, s) with a stable sort and cuts every segment into tiles of +// at most kTile positions. Each tile is an ascending FP32 chain from 0; the +// tiles of a segment are then left-folded in ascending order and cast once. +// No atomics: the result is a function of (r[s], grad) only. + +#include +#include +#include +#include +#include + +#include + +namespace { + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_float(__half v) { return __half2float(v); } +__device__ __forceinline__ float to_float(double v) { return static_cast(v); } + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ float from_float(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); +} +template <> +__device__ __forceinline__ __half from_float<__half>(float v) { return __float2half(v); } +template <> +__device__ __forceinline__ double from_float(float v) { return v; } + +template +struct CudaType { + using type = T; +}; +template <> +struct CudaType { + using type = __nv_bfloat16; +}; +template <> +struct CudaType { + using type = __half; +}; + +template +__device__ __forceinline__ int64_t adaln_row(const idx_t* ti, const idx_t* tags, int64_t s, + int64_t modality_num) { + return static_cast(ti[s]) * modality_num + static_cast(tags[s]); +} + +// One block per (s, c); 16-byte copies when every row start is 16-byte aligned. +template +__global__ void adaln_row_gather_kernel(const T* __restrict__ rows, int64_t row_stride, + const idx_t* __restrict__ ti, + const idx_t* __restrict__ tags, T* __restrict__ out, + int64_t seq, int64_t hidden, int64_t chunks, + int64_t modality_num) { + for (int64_t s = blockIdx.x; s < seq; s += gridDim.x) { + const int64_t c = blockIdx.y; + const int64_t r = adaln_row(ti, tags, s, modality_num); + const T* src = rows + r * row_stride + c * hidden; + T* dst = out + (c * seq + s) * hidden; + if constexpr (kVector) { + constexpr int kVec = 16 / sizeof(T); + const uint4* src4 = reinterpret_cast(src); + uint4* dst4 = reinterpret_cast(dst); + for (int64_t i = threadIdx.x; i < hidden / kVec; i += blockDim.x) dst4[i] = src4[i]; + } else { + for (int64_t i = threadIdx.x; i < hidden; i += blockDim.x) dst[i] = src[i]; + } + } +} + +// partial[tile, j] = sum over sorted positions p in tile (ascending) of grad[c, s_p, h] +template +__global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, + const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int64_t seq, + int64_t hidden, int64_t width) { + const int64_t tile = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= width) return; + const int64_t c = j / hidden; + 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) { + acc += to_float(g[sorted_pos[p] * hidden]); + } + partial[tile * width + j] = acc; +} + +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) { + 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) { + acc += partial[tile * width + j]; + } + out[r * width + j] = from_float(acc); +} + +void check_index(const torch::Tensor& t, const char* name, int64_t seq, + const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device(), name, " must be on ", like.device()); + TORCH_CHECK(t.dim() == 1 && t.size(0) == seq, name, " must be (", seq, ",)"); + TORCH_CHECK(t.scalar_type() == at::kLong || t.scalar_type() == at::kInt, name, + " must be int64 or int32"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +// rows (R, C * H) with unit column stride -> out (C, S, H) contiguous. +torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor timestep_indices, + torch::Tensor token_tags, int64_t chunks, + int64_t modality_num) { + TORCH_CHECK(rows.is_cuda() && rows.dim() == 2 && rows.stride(1) == 1, + "rows must be a 2-D CUDA tensor with unit column stride"); + TORCH_CHECK(chunks > 0 && rows.size(1) % chunks == 0, "rows width ", rows.size(1), + " is not a multiple of ", chunks); + TORCH_CHECK(modality_num > 0 && rows.size(0) % modality_num == 0, "rows count ", + rows.size(0), " is not a multiple of ", modality_num); + const int64_t seq = timestep_indices.numel(); + TORCH_CHECK(seq > 0, "the packed sequence must not be empty"); + check_index(timestep_indices, "timestep_indices", seq, rows); + check_index(token_tags, "token_tags", seq, rows); + TORCH_CHECK(timestep_indices.scalar_type() == token_tags.scalar_type(), + "timestep_indices and token_tags must share an integer dtype"); + const int64_t hidden = rows.size(1) / chunks; + const c10::cuda::CUDAGuard device_guard(rows.device()); + auto out = torch::empty({chunks, seq, hidden}, rows.options()); + auto stream = at::cuda::getCurrentCUDAStream(); + const dim3 grid(static_cast(std::min(seq, 1 << 20)), + static_cast(chunks)); + const int threads = 128; + const int64_t elem = rows.element_size(); + const bool vector = (hidden * elem) % 16 == 0 && (rows.stride(0) * elem) % 16 == 0 && + reinterpret_cast(rows.data_ptr()) % 16 == 0; + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, rows.scalar_type(), "h3_adaln_row_gather_forward", [&] { + using T = typename CudaType::type; + const T* src = reinterpret_cast(rows.data_ptr()); + T* dst = reinterpret_cast(out.data_ptr()); + auto launch = [&](auto idx_tag) { + using idx_t = decltype(idx_tag); + const idx_t* ti = timestep_indices.data_ptr(); + const idx_t* tags = token_tags.data_ptr(); + if (vector) { + adaln_row_gather_kernel<<>>( + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num); + } else { + adaln_row_gather_kernel<<>>( + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num); + } + }; + if (timestep_indices.scalar_type() == at::kLong) { + launch(int64_t{0}); + } else { + launch(int32_t{0}); + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// grad (C, S, H) contiguous; sorted_pos (S,) positions sorted stably by row; +// tiles [tile_begin, tile_end) in sorted order; seg_first_tile (R + 1,). +torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sorted_pos, + torch::Tensor tile_begin, torch::Tensor tile_end, + torch::Tensor seg_first_tile, + c10::ScalarType out_dtype) { + 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() && + t->scalar_type() == at::kLong, + "tile metadata must be contiguous 1-D int64 CUDA tensors"); + } + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); + const int64_t chunks = grad.size(0); + const int64_t seq = grad.size(1); + const int64_t hidden = grad.size(2); + 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(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)); + auto partial = torch::empty({std::max(tiles, 1), width}, + grad.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((width + threads - 1) / threads); + if (tiles > 0) { + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, grad.scalar_type(), "h3_adaln_row_gather_partial", [&] { + using G = typename CudaType::type; + adaln_row_gather_partial_kernel + <<(tiles), col_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), + sorted_pos.data_ptr(), tile_begin.data_ptr(), + tile_end.data_ptr(), partial.data_ptr(), seq, hidden, width); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, out_dtype, "h3_adaln_row_gather_fold", [&] { + using O = typename CudaType::type; + 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); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index f245f8e15..27159d5e9 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -510,6 +510,13 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, c10::ScalarType w_dtype, bool with_bias); +torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor timestep_indices, + torch::Tensor token_tags, int64_t chunks, + int64_t modality_num); +torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sorted_pos, + torch::Tensor tile_begin, torch::Tensor tile_end, + torch::Tensor seg_first_tile, + c10::ScalarType out_dtype); #endif // PyBind11 Module Registration @@ -783,5 +790,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight, "Deterministic dW/dbias as ascending-row FP32 folds", py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true); + m.def("h3_adaln_row_gather_forward", &h3_adaln_row_gather_forward, + "Fused six-way AdaLN row gather by timestep_index * 3 + token_tag (pure copy)", + py::arg("rows"), py::arg("timestep_indices"), py::arg("token_tags"), + py::arg("chunks") = 6, py::arg("modality_num") = 3); + m.def("h3_adaln_row_gather_backward", &h3_adaln_row_gather_backward, + "Deterministic segmented sum (sorted tiles folded in order) for the row gather", + py::arg("grad"), py::arg("sorted_pos"), py::arg("tile_begin"), py::arg("tile_end"), + py::arg("seg_first_tile"), py::arg("out_dtype")); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index b5a3c718b..ff55306a7 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -33,6 +33,7 @@ nav: - operators/h3-timestep-sinusoid.md - operators/h3-timestep-mlp.md - operators/h3-adaln-projection.md + - operators/h3-adaln-row-gather.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index d713f52d1..33e405acd 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -35,4 +35,5 @@ Every operator page should include: - [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) - [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) - [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) +- [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md new file mode 100644 index 000000000..54636dd09 --- /dev/null +++ b/docs/operators/h3-adaln-row-gather.md @@ -0,0 +1,141 @@ +# MiniMax-H3 AdaLN Row Gather + +## Summary + +`adaln_row_gather` selects every packed-sequence position's six modulation vectors from +one block's AdaLN table (RFC #420, WS1 step 3): + +```text +adaln_indices = timestep_indices * 3 + token_tags (S,) 0 video, 1 text, 2 audio +shift_msa[adaln_indices], scale_msa[...], gate_msa[...], +shift_mlp[...], scale_mlp[...], gate_mlp[...] six (S, 5376) +``` + +`rows` is the `(3T, 6H)` view of the [AdaLN projection](h3-adaln-projection.md) table, and +its six column blocks are the six modulation tensors. `timestep_indices` and `token_tags` +are semantic inputs (RFC #420 §4): they are validated, never clamped. + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("adaln_row_gather", device="cuda") +outs = op(table.view(-1, 6 * 5376), timestep_indices, token_tags) # six (S, 5376) +outs = op.gather_chunks((shift_msa, ..., gate_mlp), timestep_indices, token_tags) # drop-in +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.kernels.ops.cuda.h3.adaln_row_gather.H3AdaLNRowGatherCudaOp` | `rl_engine._C.h3_adaln_row_gather_{forward,backward}` | Forward bitwise equal to `index_select`; deterministic backward | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.adaln_row_gather.NativeH3AdaLNRowGatherOp` | n/a | `index_select` forward; deterministic per-row FP32 backward | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +The raw diffusers path, including its atomic backward, is kept as +`rl_engine.testing.h3_provider.provider_adaln_row_gather` for evidence. + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `rows` | `(3T, 6H)` | bf16 / fp16 / fp32 | Unit column stride. Any row stride (16-byte copies when aligned) | +| `timestep_indices` | `(S,)`, `S >= 1` | int64 or int32 | In `[0, T)` | +| `token_tags` | `(S,)` | same as `timestep_indices` | In `{0, 1, 2}` | +| outputs | six `(S, H)` | `rows.dtype` | Contiguous; slices of one `(6, S, H)` buffer | + +These inputs fail closed: out-of-range tags (a fourth modality, probe H2) or timesteps +(an offset past `T`, probe H3), length or dtype mismatches, empty sequences, and tables +that are not 3 rows per timestep or 6 blocks wide. The range check is one host +read-back; pass `check_range=False` when the indices are already validated. + +## Numerics + +- **Forward** is a byte copy: one launch, one block per `(position, chunk)`, using + 16-byte vector copies. It is bitwise equal to diffusers' six `index_select` calls for + every packing, length and index dtype tested (S up to 131072). +- **Backward**: `d_rows[r] = sum over {s : r[s] = r} of grad[s]` is the op's only + reduction. Positions are sorted stably by `(row, position)` and cut into tiles of 256. + Each tile is an ascending FP32 chain from 0, and a row's tiles are left-folded in order + and rounded once. There are no atomics, so the result depends only on the indices and + the gradient. `index_select`'s own backward is an atomic scatter-add in the gradient + dtype (BF16), which is neither deterministic nor accurate. +- **Tolerance class.** The VJP is a segmented reduction, so gtest judges the op as + `reduction`. The forward is asserted bitwise in `tests/h3`. + +Measured on a B200 (torch 2.13.0+cu130): + +| Comparison (T = 3, S = 4097, H = 5376, BF16) | CUDA | provider (`index_select`) | +| --- | --- | --- | +| Forward vs `index_select` | bitwise equal | | +| `d_rows` repeat-bitwise | yes | no | +| `d_rows` correctly rounded from the FP64 sum | 99.99999% | 7.8% | +| `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | + +### Fused modulation: `H3AdaLNModulationCudaOp` + +`rl_engine.kernels.ops.cuda.h3.adaln_modulation.H3AdaLNModulationCudaOp(temb, weight, bias, +timestep_indices, token_tags)` runs [`adaln_projection_3mod`](h3-adaln-projection.md) and +this gather as one autograd node. Its forward uses the same kernels and is bitwise equal +to calling the two ops in sequence. + +The difference is in the backward. When the ops are called separately, autograd returns +the gather's FP32 segment sum to the projection in the table's BF16 dtype, which rounds it +once more. The fused node feeds the FP32 table gradient straight into the projection +backward. Use it wherever both ops run back to back, as they do in an H3 block. + +### Whole-chain backward + +`scripts/h3_chain_replay.py --backward` runs on the pinned weights for +T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. It checks six parameter gradients in each +of the 16 cases (the time embedder's `linear_{1,2}.{weight,bias}` and block 0's AdaLN +weight and bias) against an FP64 golden in which every declared cast is straight-through: + +| S >= 257 (72 gradients) | RL-Kernel, separate ops | RL-Kernel, fused | diffusers | +| --- | --- | --- | --- | +| repeat-bitwise | 72 / 72 | 72 / 72 | 0 / 72 | +| time embedder: max abs error / golden max | 5.3e-4 to 2.4e-3 | 6.5e-7 to 4.9e-6 | 5.6e-3 to 1.2e-1 | +| AdaLN weight/bias: correctly rounded BF16, worst case | 59.8% | 99.57% | 1.6% | + +At S = 3 all three chains are deterministic. Both RL-Kernel variants reach 1.5e-6 on the +time embedder, against 3.2e-3 for diffusers. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op adaln_row_gather +``` + +B200, T = 3, H = 5376, BF16. "write BW" counts the 6·S·H outputs; the table stays in L2. + +| S | CUDA forward | provider forward | CUDA backward | provider backward | +| --- | --- | --- | --- | --- | +| 4097 | 62 µs | 92 µs | 1.10 ms | 1.44 ms | +| 32768 | 315 µs | 587 µs | 2.10 ms | 8.66 ms | +| 131072 | 1.24 ms (6.8 TB/s write) | 2.37 ms | 6.84 ms | 34.9 ms | + +The backward's peak memory is about twice the provider's, because autograd stacks the +six output gradients and the per-tile FP32 partials are materialised. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_adaln_row_gather.py tests/h3/test_h3_adaln_modulation.py -v +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # whole chain, forward + backward +python scripts/check_operator.py --op adaln_row_gather --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op adaln_row_gather \ + --out docs/usage/evidence/h3-adaln-row-gather-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-adaln-row-gather-b200/report.json +python scripts/h3_chain_replay.py --timesteps 1,2,3,4 --seq-lens 3,257,4097,32768 \ + --backward --out docs/usage/evidence/h3-adaln-row-gather-b200/chain_replay.json +``` + +## Known Limitations + +- Like diffusers, the op materialises six `(S, H)` tensors, which is 8 GB at + S = 131072. Fusing the gather into the norm/modulate and gate-residual consumers + (`h3_rmsnorm`, `adaln_gate_residual`) avoids this. Those are separate rows. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index b8cf1d58a..144b06e27 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -341,3 +341,18 @@ def h3_det_linear_backward_input( def h3_det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True ) -> list[torch.Tensor]: ... +def h3_adaln_row_gather_forward( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + chunks: int = 6, + modality_num: int = 3, +) -> torch.Tensor: ... +def h3_adaln_row_gather_backward( + grad: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 255a48c0c..ef77fe7bc 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -49,6 +49,7 @@ def make_operator_inputs( "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, + "adaln_row_gather": _make_adaln_row_gather_inputs, } try: return builders[op_name](args, dtype, device) @@ -84,6 +85,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: f"x{H3_TIME_EMBED}", "adaln_projection_3mod": f"{_h3_num_timesteps(args)}x{H3_TIME_EMBED}" f"x{6 * 3 * _h3_hidden(args)}", + "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" + f"->{batch * seq}", } try: return names[op_name] @@ -411,6 +414,23 @@ def _make_adaln_projection_3mod_inputs( } +def _make_adaln_row_gather_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # T distinct timesteps (--batch) and a packed sequence of batch * seq rows + # mixing all three modalities. + batch, seq = _batch_seq(args) + num_timesteps = _h3_num_timesteps(args) + generator = _generator(args, device, offset=29) + packed = batch * seq + token_tags = torch.randint(0, 3, (packed,), generator=generator, device=device) + timestep_indices = torch.randint( + 0, num_timesteps, (packed,), generator=generator, device=device + ) + rows = _floating_tensor((3 * num_timesteps, 6 * _h3_hidden(args)), args, dtype, device, 0) + return {"rows": rows, "timestep_indices": timestep_indices, "token_tags": token_tags} + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 91d877b62..11ce047db 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -285,6 +285,21 @@ def _load_object(path: str) -> Any: }, grad_input_names=("temb", "weight", "bias"), ), + # The gather's forward is a copy (asserted bitwise in tests/h3), but its VJP + # sums every packed position of a table row, so it is judged as a reduction. + "adaln_row_gather": OperatorSpec( + name="adaln_row_gather", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.h3.adaln_row_gather.NativeH3AdaLNRowGatherOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.kernels.ops.pytorch.h3.adaln_row_gather.NativeH3AdaLNRowGatherOp" + ), + "cuda": "rl_engine.kernels.ops.cuda.h3.adaln_row_gather.H3AdaLNRowGatherCudaOp", + }, + grad_input_names=("rows",), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_modulation.py b/rl_engine/kernels/ops/cuda/h3/adaln_modulation.py new file mode 100644 index 000000000..ea13bd232 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/adaln_modulation.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused H3 AdaLN modulation: ``adaln_projection_3mod`` followed by ``adaln_row_gather``. + +The forward is exactly the two RFC #420 ops back to back (same kernels, same +bits). The difference is the backward. With two separate ops, autograd hands +the gather's table gradient to the projection in the table's dtype (BF16), +which rounds a long FP32 segment sum once more before it reaches the +projection weights and ``temb``. Here the segment sum stays FP32 into the +projection backward, so the parameter gradients carry no extra BF16 rounding. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import ( + _segment_tiles, + adaln_row_gather_available, +) +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.kernels.ops.pytorch.h3.adaln_projection import validate_h3_adaln_projection +from rl_engine.kernels.ops.pytorch.h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_adaln_row_gather_forward" +) +BACKWARD_IMPL = "fp32_segment_sum_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_adaln_row_gather_backward[fp32]" + "+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.kernels.ops.cuda.h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNModulationCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias, timestep_indices, token_tags, hidden): + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias) + rows = table.view(-1, H3_ADALN_CHUNKS * hidden) + packed = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(temb, act, weight, timestep_indices, token_tags) + ctx.table_shape = table.shape + return packed + + @staticmethod + def backward(ctx, grad_packed): + temb, act, weight, timestep_indices, token_tags = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad[:3] + num_timesteps, width = ctx.table_shape + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_timesteps * H3_MODALITY_NUM) + d_rows = _C.h3_adaln_row_gather_backward(grad_packed.contiguous(), *tiles, torch.float32) + d_table = d_rows.view(num_timesteps, width) + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + if need_temb: + d_temb = silu_backward_fp32( + det_linear_backward_input(d_table, weight, torch.float32), temb + ) + record_backward( + "adaln_modulation_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return ( + d_temb, + dw if need_weight else None, + db if need_bias else None, + None, + None, + None, + ) + + +class H3AdaLNModulationCudaOp: + """``adaln_projection_3mod`` + ``adaln_row_gather`` with an FP32 table gradient. + + Returns the six ``(S, H)`` modulation tensors for every packed position, + bitwise equal to running the two ops separately. + """ + + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and adaln_row_gather_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_adaln_row_gather_* ops") + + def __call__(self, temb, weight, bias, timestep_indices, token_tags, *, check_range=True): + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNModulationCudaOp needs CUDA tensors") + _validate_indices(timestep_indices, token_tags, temb.shape[0], check_range) + packed = _H3AdaLNModulationCuda.apply( + temb, + weight, + bias, + timestep_indices.contiguous(), + token_tags.contiguous(), + hidden, + ) + return tuple(packed.unbind(0)) + + +def _validate_indices(timestep_indices, token_tags, num_timesteps, check_range): + # Reuse the gather's checks against a stand-in of the right row count. + stand_in = torch.empty( + (num_timesteps * H3_MODALITY_NUM, H3_ADALN_CHUNKS), device=timestep_indices.device + ) + validate_h3_adaln_row_gather(stand_in, timestep_indices, token_tags, check_range=check_range) diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py b/rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py new file mode 100644 index 000000000..3569d5010 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 AdaLN row gather (RFC #420 ``adaln_row_gather``).""" + +from __future__ import annotations + +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.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.kernels.ops.pytorch.h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = "rl_engine._C.h3_adaln_row_gather_forward" +BACKWARD_IMPL = "stable_sorted_segment_tiles_fp32_ordered_fold" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_adaln_row_gather_backward" +# Positions per backward tile; part of the fixed summation order. +BACKWARD_TILE = 256 + + +def adaln_row_gather_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_adaln_row_gather_forward") + and hasattr(_C, "h3_adaln_row_gather_backward") + ) + + +def _segment_tiles(index: torch.Tensor, num_rows: int): + """Sorted positions and fixed-size tiles per row segment (all integer ops).""" + + sorted_pos = torch.argsort(index, stable=True) + counts = torch.bincount(index, minlength=num_rows) + seg_start = torch.cumsum(counts, 0) - counts + tiles_per_seg = (counts + BACKWARD_TILE - 1) // BACKWARD_TILE + seg_first_tile = torch.zeros(num_rows + 1, dtype=torch.int64, device=index.device) + seg_first_tile[1:] = torch.cumsum(tiles_per_seg, 0) + tile_seg = torch.repeat_interleave(torch.arange(num_rows, device=index.device), tiles_per_seg) + tile_local = torch.arange(tile_seg.numel(), device=index.device) - seg_first_tile[tile_seg] + tile_begin = seg_start[tile_seg] + tile_local * BACKWARD_TILE + tile_end = torch.minimum(tile_begin + BACKWARD_TILE, seg_start[tile_seg] + counts[tile_seg]) + return sorted_pos, tile_begin.contiguous(), tile_end.contiguous(), seg_first_tile + + +class _H3AdaLNRowGatherCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, rows, timestep_indices, token_tags): + out = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(timestep_indices, token_tags) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return out + + @staticmethod + def backward(ctx, grad_out): + timestep_indices, token_tags = ctx.saved_tensors + num_rows, rows_dtype = ctx.rows_meta + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_rows) + grad_rows = _C.h3_adaln_row_gather_backward(grad_out.contiguous(), *tiles, rows_dtype) + record_backward( + "adaln_row_gather", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad_rows, None, None + + +class H3AdaLNRowGatherCudaOp: + """CUDA candidate: one launch gathers all six modulation tensors. + + The forward is a byte copy (bitwise equal to six ``index_select`` calls). + The backward replaces ``index_select``'s atomic BF16 scatter-add with a + deterministic FP32 segmented sum in stable sorted order, cast once. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not adaln_row_gather_available(): + raise RuntimeError( + "rl_engine._C lacks h3_adaln_row_gather_*; rebuild with " + "csrc/cuda/h3/adaln_row_gather.cu" + ) + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + if not rows.is_cuda: + raise ValueError("H3AdaLNRowGatherCudaOp needs CUDA tensors") + if rows.stride(1) != 1: + rows = rows.contiguous() + packed = _H3AdaLNRowGatherCuda.apply( + rows, timestep_indices.contiguous(), token_tags.contiguous() + ) + return tuple(packed.unbind(0)) + + def gather_chunks(self, chunks, timestep_indices, token_tags, **kwargs): + """Drop-in for diffusers' six ``(3T, H)`` tensors: concatenate, then gather.""" + + if len(chunks) != H3_ADALN_CHUNKS: + raise ValueError(f"expected {H3_ADALN_CHUNKS} modulation tensors, got {len(chunks)}") + return self.forward(torch.cat(list(chunks), dim=1), timestep_indices, token_tags, **kwargs) + + def forward_fp32(self, rows, timestep_indices, token_tags): + return tuple(out.float() for out in self.forward(rows, timestep_indices, token_tags)) diff --git a/rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py b/rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py new file mode 100644 index 000000000..93341d602 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 AdaLN row gather (RFC #420 row ``adaln_row_gather``). + +H3 computes one row index per packed-sequence position and uses it to pick +every position's six modulation vectors:: + + adaln_indices = timestep_indices * 3 + token_tags (S,) + shift_msa[adaln_indices], scale_msa[adaln_indices], ... six (S, H) + +``rows`` is the ``(3T, 6H)`` view of one block's AdaLN projection +(``table.view(-1, 6H)``); its six column blocks are the six modulation +tensors. ``timestep_indices`` and ``token_tags`` are semantic inputs +(RFC #420 §4): out-of-range values are an error, never clamped. +""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_INDEX_DTYPES = (torch.int64, torch.int32) + + +def validate_h3_adaln_row_gather( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + *, + check_range: bool = True, +) -> int: + """Return the hidden size; raise on anything H3 does not define.""" + + if not isinstance(rows, torch.Tensor) or rows.dim() != 2: + raise ValueError("rows must be a 2-D (3T, 6H) tensor") + if not rows.is_floating_point(): + raise TypeError(f"rows must be floating point, got {rows.dtype}") + num_rows, width = rows.shape + if num_rows == 0 or num_rows % H3_MODALITY_NUM != 0: + raise ValueError(f"rows must hold 3 rows per timestep, got {num_rows}") + if width == 0 or width % H3_ADALN_CHUNKS != 0: + raise ValueError(f"rows width must be 6 * H, got {width}") + for name, index in (("timestep_indices", timestep_indices), ("token_tags", token_tags)): + if not isinstance(index, torch.Tensor) or index.dim() != 1: + raise ValueError(f"{name} must be a 1-D (S,) tensor") + if index.dtype not in _INDEX_DTYPES: + raise TypeError(f"{name} must be int64 or int32, got {index.dtype}") + if index.device != rows.device: + raise ValueError(f"{name} is on {index.device}, rows are on {rows.device}") + if timestep_indices.shape != token_tags.shape: + raise ValueError( + f"timestep_indices {tuple(timestep_indices.shape)} and token_tags " + f"{tuple(token_tags.shape)} must have the same length" + ) + if timestep_indices.numel() == 0: + raise ValueError("the packed sequence must not be empty") + if timestep_indices.dtype != token_tags.dtype: + raise TypeError("timestep_indices and token_tags must share an integer dtype") + if check_range: + num_timesteps = num_rows // H3_MODALITY_NUM + bad = ( + (timestep_indices < 0) + | (timestep_indices >= num_timesteps) + | (token_tags < 0) + | (token_tags >= H3_MODALITY_NUM) + ) + if bool(bad.any()): # one host sync + raise IndexError( + f"timestep_indices must lie in [0, {num_timesteps}) and token_tags in " + f"[0, {H3_MODALITY_NUM}) (0 video, 1 text, 2 audio)" + ) + return width // H3_ADALN_CHUNKS + + +def h3_adaln_indices(timestep_indices: torch.Tensor, token_tags: torch.Tensor) -> torch.Tensor: + return timestep_indices.long() * H3_MODALITY_NUM + token_tags.long() + + +class _DeterministicRowGather(torch.autograd.Function): + """``index_select`` forward; FP32 per-row segment sums for the backward. + + ``index_select``'s own backward is an atomic scatter-add in the grad dtype + (BF16 for H3), which is neither deterministic nor accurate. There are + only 3T table rows, so the reference sums each row's positions with one + FP32 ``torch.sum`` per row (a fixed reduction for a fixed packing) and + rounds once. + """ + + @staticmethod + def forward(ctx, rows, index): + ctx.save_for_backward(index) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return torch.stack( + [chunk.index_select(0, index) for chunk in rows.chunk(H3_ADALN_CHUNKS, dim=-1)] + ) + + @staticmethod + def backward(ctx, grad): + (index,) = ctx.saved_tensors + num_rows, dtype = ctx.rows_meta + grad32 = grad.float() # (6, S, H) + out = grad32.new_zeros((num_rows, grad32.shape[0], grad32.shape[2])) + for row in range(num_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + out[row] = grad32.index_select(1, positions).sum(dim=1) + return out.reshape(num_rows, -1).to(dtype), None + + +class NativeH3AdaLNRowGatherOp: + """PyTorch reference: ``index_select`` on each of the six column blocks. + + The forward is byte-identical to diffusers; the backward is deterministic + (see ``_DeterministicRowGather``). The raw diffusers path, including its + atomic backward, is ``rl_engine.testing.h3_provider.provider_adaln_row_gather``. + ``forward_fp32`` returns the same rows in FP32 with an FP64 gradient. + + The forward is a copy, but the VJP is a segmented reduction over the + packed sequence, so the op is judged with the ``reduction`` tolerance + class; forward bitwise equality is asserted separately. + """ + + op_class = "reduction" + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple(_DeterministicRowGather.apply(rows, index).unbind(0)) + + def forward_fp32(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple( + chunk.index_select(0, index).float() + for chunk in rows.double().chunk(H3_ADALN_CHUNKS, dim=-1) + ) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 22bbab259..ba6c57f05 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -215,6 +215,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_ADALN_PROJECTION = ( "rl_engine.kernels.ops.cuda.h3.adaln_projection.H3AdaLNProjectionCudaOp" ) + PYTORCH_H3_ADALN_ROW_GATHER = ( + "rl_engine.kernels.ops.pytorch.h3.adaln_row_gather.NativeH3AdaLNRowGatherOp" + ) + CUDA_H3_ADALN_ROW_GATHER = ( + "rl_engine.kernels.ops.cuda.h3.adaln_row_gather.H3AdaLNRowGatherCudaOp" + ) def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -637,6 +643,10 @@ def __init__(self): OpBackend.CUDA_H3_ADALN_PROJECTION, OpBackend.PYTORCH_H3_ADALN_PROJECTION, ], + "adaln_row_gather": [ + OpBackend.CUDA_H3_ADALN_ROW_GATHER, + OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, + ], }, "rocm": { "logp": [ @@ -680,6 +690,7 @@ def __init__(self): "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -727,6 +738,7 @@ def __init__(self): "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -755,6 +767,7 @@ def __init__(self): "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 8671ef141..06bfe252b 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -32,6 +32,7 @@ from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps from rl_engine.testing.h3_provider import ( provider_adaln_modulation, + provider_adaln_row_gather, provider_time_embedder, provider_time_proj, ) @@ -94,8 +95,35 @@ def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: golden_atol=5e-2, # reduction / bfloat16 golden_rtol=2e-2, ), + Stage( + name="adaln_row_gather", + op_type="adaln_row_gather", + candidate=lambda op, ctx, up: op.gather_chunks( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + provider=lambda ctx, up: provider_adaln_row_gather( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + golden=lambda op, ctx, up: op.forward_fp32( + torch.cat(list(up), dim=1), ctx["timestep_indices"], ctx["token_tags"] + ), + provider_bitwise_isolated=True, # a copy + golden_atol=5e-2, # carries the projection's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ), ] +# Parameters whose gradients the backward replay reports, in chain order. +GRAD_LEAVES = ( + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +) +BACKWARD_MODES = ("candidate", "candidate_fused", "provider") + def _flat(value: Any) -> list[torch.Tensor]: if isinstance(value, torch.Tensor): @@ -181,6 +209,77 @@ def run_case( return report +def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torch.Tensor]: + """Parameter gradients of the full chain for one execution mode. + + ``candidate`` runs the four RL-Kernel ops separately, ``candidate_fused`` + runs projection + gather as ``H3AdaLNModulationCudaOp``, ``provider`` the + diffusers path, ``golden`` the FP64 references. + """ + + leaves = { + name: ctx["weights"][name].detach().clone().requires_grad_(True) for name in GRAD_LEAVES + } + run_ctx = {**ctx, "weights": {**ctx["weights"], **leaves}} + value = None + if mode == "candidate_fused": + from rl_engine.kernels.ops.cuda.h3.adaln_modulation import H3AdaLNModulationCudaOp + + for stage in STAGES[:2]: + value = stage.candidate(registry.get_op(stage.op_type, device="cuda"), run_ctx, value) + value = H3AdaLNModulationCudaOp()( + value, *adaln_params(run_ctx), ctx["timestep_indices"], ctx["token_tags"] + ) + else: + for stage in STAGES: + if mode == "provider": + value = stage.provider(run_ctx, value) + elif mode == "candidate": + op = registry.get_op(stage.op_type, device="cuda") + value = stage.candidate(op, run_ctx, value) + 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)]) + return [leaves[name].grad for name in GRAD_LEAVES] + + +def run_backward_case( + registry: KernelRegistry, weights, *, num_timesteps: int, seq_len: int, seed: int = 0 +) -> dict[str, Any]: + """Determinism and accuracy of the chain's parameter gradients, per mode.""" + + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + generator = torch.Generator(device="cuda").manual_seed(seed + 1) + hidden = weights[GRAD_LEAVES[4]].shape[0] // 18 + upstream = [ + torch.randn(seq_len, hidden, device="cuda", generator=generator).bfloat16() + for _ in range(6) + ] + golden = chain_grads("golden", registry, ctx, upstream) + runs = { + mode: [chain_grads(mode, registry, ctx, upstream) for _ in range(2)] + for mode in BACKWARD_MODES + } + report: dict[str, Any] = {"num_timesteps": num_timesteps, "seq_len": seq_len, "leaves": {}} + for index, name in enumerate(GRAD_LEAVES): + gold = golden[index].float() + entry: dict[str, Any] = {"golden_absmax": float(gold.abs().max())} + for mode, (first, second) in runs.items(): + grad = first[index] + err = (grad.float() - gold).abs().max() + entry[mode] = { + "repeat_bitwise_equal": bool(torch.equal(first[index], second[index])), + "max_abs_vs_golden": float(err), + "max_abs_vs_golden_over_absmax": float(err / gold.abs().max().clamp_min(1e-30)), + "correctly_rounded_fraction": float( + (grad == golden[index].to(grad.dtype)).float().mean() + ), + } + report["leaves"][name] = entry + return report + + def git_state() -> dict[str, Any]: def run(*args: str) -> str: return subprocess.check_output(["git", *args], cwd=REPO_ROOT, text=True).strip() diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index f64f7174f..6cedcdf60 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -56,3 +56,12 @@ def provider_adaln_modulation( temb = F.linear(F.silu(temb).to(weight.dtype), weight, bias) temb = temb.view(-1, 6 * hidden_size) return temb.chunk(6, dim=-1) + + +def provider_adaln_row_gather( + modulation: tuple[torch.Tensor, ...], timestep_indices: torch.Tensor, token_tags: torch.Tensor +) -> tuple[torch.Tensor, ...]: + """``adaln_indices = timestep_indices * 3 + token_tags`` then six ``index_select``.""" + + adaln_indices = timestep_indices * 3 + token_tags + return tuple(tensor.index_select(0, adaln_indices) for tensor in modulation) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 3f84eb520..e0a5939b5 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -16,9 +16,10 @@ import torch from rl_engine.kernels.registry import KernelRegistry -from rl_engine.testing.h3_cases import h3_timesteps +from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps from rl_engine.testing.h3_provider import ( provider_adaln_modulation, + provider_adaln_row_gather, provider_time_embedder, provider_time_proj, ) @@ -258,13 +259,105 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, } +# --------------------------------------------------------------------------- # +# adaln_row_gather +# --------------------------------------------------------------------------- # + +GATHER_SEQ_LENS = (4097, 32768, 131072) + + +def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_row_gather", device="cuda") + num_timesteps, hidden = 3, 5376 + rows = torch.randn(3 * num_timesteps, 6 * hidden, device="cuda").bfloat16() + chunks = rows.chunk(6, dim=-1) + cases = [] + for seq in GATHER_SEQ_LENS: + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + grads = [torch.randn(seq, hidden, device="cuda").bfloat16() for _ in range(6)] + + def backward(fn, ti=ti, tags=tags, grads=grads): + leaf = rows.detach().requires_grad_(True) + torch.autograd.backward(list(fn(leaf, ti, tags)), grads) + + cases.append( + { + "op": "adaln_row_gather", + "case": f"S={seq}", + "backend": type(op).__name__, + # bytes written (six (S, H) BF16 outputs); the 3T-row table stays in L2 + "bytes": 6 * seq * hidden * 2, + "candidate": lambda ti=ti, tags=tags: op.forward(rows, ti, tags, check_range=False), + "provider": lambda ti=ti, tags=tags: provider_adaln_row_gather(chunks, ti, tags), + "candidate_backward": lambda b=backward: b( + lambda r, ti, tags: op.forward(r, ti, tags, check_range=False) + ), + "provider_backward": lambda b=backward: b( + lambda r, ti, tags: provider_adaln_row_gather(r.chunk(6, dim=-1), ti, tags) + ), + } + ) + return cases + + +def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: + from rl_engine.testing.h3_chain import run_backward_case + from rl_engine.testing.h3_weights import load_h3_conditioning_weights + + op = registry.get_op("adaln_row_gather", device="cuda") + golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) + rows = torch.randn(9, 6 * 5376, device="cuda").bfloat16() + forward_bitwise = {} + for seq in (1, 257, 4097, 32768): + ti, tags = h3_packed_layout(seq, 3, seed=seq) + ours = op(rows, ti, tags) + theirs = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + forward_bitwise[str(seq)] = all(torch.equal(a, b) for a, b in zip(ours, theirs)) + + # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = [torch.randn(4097, 5376, device="cuda").bfloat16() for _ in range(6)] + + def grad_of(fn): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(fn(leaf)), grads) + return leaf.grad + + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward(list(golden.forward_fp32(ref, ti, tags)), [g.float() for g in grads]) + gold = ref.grad.bfloat16() + backward = {} + for name, fn in ( + ("cuda", lambda leaf: op(leaf, ti, tags)), + ("provider", lambda leaf: provider_adaln_row_gather(leaf.chunk(6, dim=-1), ti, tags)), + ): + first, second = grad_of(fn), grad_of(fn) + backward[name] = { + "repeat_bitwise_equal": bool(torch.equal(first, second)), + "correctly_rounded_fraction": float((first == gold).float().mean()), + } + + weights = load_h3_conditioning_weights("cuda") + chain = [ + run_backward_case(registry, weights, num_timesteps=t, seq_len=s) + for t, s in ((1, 257), (3, 257), (1, 4097), (3, 4097), (4, 32768)) + ] + return { + "forward_bitwise_vs_index_select": forward_bitwise, + "op_backward": backward, + "chain_backward": chain, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, "adaln_projection_3mod": _projection_perf, + "adaln_row_gather": _gather_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, "timestep_mlp_fp32": _mlp_accuracy, "adaln_projection_3mod": _projection_accuracy, + "adaln_row_gather": _gather_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index c20e157e4..c36ade2f2 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -226,7 +226,12 @@ def fake_load_backend(backend): ] -H3_OPS = ("timestep_sinusoid_h3", "timestep_mlp_fp32", "adaln_projection_3mod") +H3_OPS = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", +) @pytest.mark.parametrize("op_name", H3_OPS) diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py index e262ca053..52a1cddde 100644 --- a/scripts/h3_chain_replay.py +++ b/scripts/h3_chain_replay.py @@ -40,6 +40,12 @@ def main() -> None: parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") parser.add_argument("--seed", type=int, default=0) parser.add_argument("--out", type=Path, default=None) + parser.add_argument( + "--backward", + action="store_true", + help="also replay the chain backward (needs every stage): parameter-gradient " + "determinism and accuracy for the separate-op, fused and diffusers chains", + ) args = parser.parse_args() if not torch.cuda.is_available(): @@ -76,6 +82,23 @@ def main() -> None: ) print(f"T={num_timesteps} S={seq_len}: {summary}; first_drift={case['first_drift']}") + backward_cases = [] + if args.backward: + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_backward_case( + registry, weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=args.seed + ) + backward_cases.append(case) + parts = [] + for mode in h3_chain.BACKWARD_MODES: + entries = case["leaves"].values() + det = all(e[mode]["repeat_bitwise_equal"] for e in entries) + worst = max(e[mode]["max_abs_vs_golden_over_absmax"] for e in entries) + parts.append(f"{mode}: {'det' if det else 'NONDET'} {worst:.1e}") + summary = " | ".join(parts) + print(f"backward T={num_timesteps} S={seq_len}: {summary}") + evidence = { "kind": "h3_conditioning_chain_replay", "rfc": manifest["rfc"], @@ -86,6 +109,7 @@ def main() -> None: "environment": h3_chain.environment(), "stages": [s.name for s in stages], "cases": cases, + "backward_cases": backward_cases, } if args.out is not None: args.out.parent.mkdir(parents=True, exist_ok=True) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 6385b4aa4..f2039f0aa 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -46,14 +46,24 @@ def _style(ax, title: str, ylabel: str) -> None: ax.spines["bottom"].set_color(GRID) -def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt: str) -> None: - """Grouped bars with a 2px-style gap and value labels on top.""" +def _bars(ax, labels: list[str], series: list[tuple], fmt: str, log: bool = False) -> None: + """Grouped bars with a gap and value labels on top; a 4th tuple item is a hatch.""" n = len(series) width = 0.8 / n - for i, (name, color, values) in enumerate(series): + for i, (name, color, values, *hatch) in enumerate(series): xs = [x + (i - (n - 1) / 2) * width for x in range(len(labels))] - bars = ax.bar(xs, values, width * 0.92, color=color, label=name, zorder=2) + bars = ax.bar( + xs, + values, + width * 0.92, + color=color if not hatch else SURFACE, + edgecolor=color, + hatch=hatch[0] if hatch else None, + linewidth=1.2, + label=name, + zorder=2, + ) for bar, value in zip(bars, values): ax.annotate( fmt.format(value), @@ -65,9 +75,13 @@ def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt color=INK, ) ax.set_xticks(range(len(labels)), labels) - top = max(v for _, _, values in series for v in values) - ax.set_ylim(0, top * 1.35) - ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=n) + top = max(v for _, _, values, *_ in series for v in values) + if log: + ax.set_yscale("log") + ax.set_ylim(min(v for _, _, values, *_ in series for v in values) / 2, top * 8) + else: + ax.set_ylim(0, top * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=min(n, 4)) # Exact zeros cannot sit on a log axis; draw them on this floor and label them. @@ -261,10 +275,81 @@ def plot_projection(report: dict[str, Any]): return fig +def plot_gather(report: dict[str, Any]): + fig, (left, right) = _figure(report, "adaln_row_gather (T = 3, H = 5376, BF16)") + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("index_select fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "index_select bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + chain = report["accuracy"]["chain_backward"] + _style( + right, + "Whole-chain grads of the FP32 time embedder vs FP64", + "max abs error / golden max", + ) + names = { + "candidate": ("RL-Kernel, separate ops", THIRD, "o"), + "candidate_fused": ("RL-Kernel, fused modulation", RL_KERNEL, "D"), + "provider": ("diffusers", PROVIDER, "s"), + } + labels = [f"T={c['num_timesteps']}\nS={c['seq_len']}" for c in chain] + for j, (mode, (name, color, marker)) in enumerate(names.items()): + leaves = [e for c in chain for e in c["leaves"].values()] + det = sum(e[mode]["repeat_bitwise_equal"] for e in leaves) + adaln = min( + e[mode]["correctly_rounded_fraction"] + for c in chain + for n, e in c["leaves"].items() + if n.startswith("transformer_blocks") + ) + label = ( + f"{name}\n repeat-bitwise {det}/{len(leaves)} · " + f"AdaLN BF16 grads ≥{100 * adaln:.1f}% correctly rounded" + ) + for i, case in enumerate(chain): + values = [ + e[mode]["max_abs_vs_golden_over_absmax"] + for n, e in case["leaves"].items() + if n.startswith("time_embedder") + ] + right.scatter( + [i + (j - 1) * 0.22] * len(values), + values, + s=30, + color=color, + marker=marker, + alpha=0.85, + linewidths=0, + zorder=3, + label=label if i == 0 else None, + ) + right.set_yscale("log") + right.set_ylim(1e-8, 1e4) + right.set_xticks(range(len(labels)), labels, fontsize=7) + right.legend(frameon=False, fontsize=7, labelcolor=INK, loc="upper left") + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, "adaln_projection_3mod": plot_projection, + "adaln_row_gather": plot_gather, } diff --git a/setup.py b/setup.py index b388bdd29..fcefa9016 100644 --- a/setup.py +++ b/setup.py @@ -159,6 +159,7 @@ def get_extensions(): # MiniMax-H3 (RFC #420) conditioning path; plain-arch, runs on SM90 and SM100. cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") cuda_sources.append("csrc/cuda/h3/det_linear.cu") + cuda_sources.append("csrc/cuda/h3/adaln_row_gather.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/test_h3_adaln_modulation.py b/tests/h3/test_h3_adaln_modulation.py new file mode 100644 index 000000000..97b4927d0 --- /dev/null +++ b/tests/h3/test_h3_adaln_modulation.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused AdaLN modulation (projection + row gather) with an FP32 table gradient. + +The forward must be bitwise equal to the two RFC #420 ops run separately; the +backward must be deterministic and must not re-round the table gradient to +BF16 (so ``d_temb`` reaches FP64-golden accuracy). +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.testing.h3_cases import h3_packed_layout + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _ops(): + from rl_engine.kernels.ops.cuda.h3.adaln_modulation import H3AdaLNModulationCudaOp + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks the H3 kernels") + return H3AdaLNModulationCudaOp(), H3AdaLNProjectionCudaOp(), H3AdaLNRowGatherCudaOp() + + +def _params(hidden, dim, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = 18 * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).bfloat16().cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).bfloat16().cuda() + temb = (torch.randn(3, dim, generator=g) * 2).cuda() + return temb, weight, bias + + +def _grads(seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + +def _leaf_grads(fn, temb, weight, bias, grads): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(fn(*leaves)), grads) + return [leaf.grad for leaf in leaves] + + +@requires_cuda +class TestFusedModulation: + def test_forward_bitwise_equal_to_separate_ops(self): + fused, proj, gather = _ops() + temb, weight, bias = _params(64, 2688) + ti, tags = h3_packed_layout(1000, 3) + ours = fused(temb, weight, bias, ti, tags) + separate = gather.gather_chunks(proj(temb, weight, bias), ti, tags) + assert all(torch.equal(a, b) for a, b in zip(ours, separate)) + + def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): + fused, proj, gather = _ops() + weight = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.weight"].cuda() + bias = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.bias"].cuda() + temb = torch.randn(3, 2688, device="cuda") * 2 + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = _grads(4097, 5376, seed=1) + ours = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + separate = _leaf_grads( + lambda t, w, b: gather.gather_chunks(proj(t, w, b), ti, tags), + temb, + weight, + bias, + grads, + ) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + act = act + (act.to(torch.bfloat16).double() - act).detach() # identity-VJP cast + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376) + index = ti * 3 + tags + torch.autograd.backward( + [chunk.index_select(0, index) for chunk in table.chunk(6, dim=-1)], + [g.double() for g in grads], + ) + + def rel(grad, golden): + return ((grad.double() - golden).abs().max() / golden.abs().max()).item() + + # d_temb: no BF16 rounding anywhere on its path, so it reaches FP32 accuracy. + assert rel(ours[0], ref[0].grad) < 1e-5 + assert rel(ours[0], ref[0].grad) < rel(separate[0], ref[0].grad) / 100 + # dW, db: only their own final BF16 rounding remains. + for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad)): + assert (grad == golden.to(torch.bfloat16)).float().mean() > 0.9999 + + def test_backward_repeat_bitwise(self): + fused, _, _ = _ops() + temb, weight, bias = _params(64, 2688, seed=3) + ti, tags = h3_packed_layout(3000, 3, seed=3) + grads = _grads(3000, 64, seed=3) + first = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + again = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + assert all(torch.equal(a, b) for a, b in zip(first, again)) + + def test_rejects_bad_indices(self): + fused, _, _ = _ops() + temb, weight, bias = _params(8, 16) + ti, tags = h3_packed_layout(10, 3) + with pytest.raises(IndexError): + fused(temb, weight, bias, ti + 3, tags) + with pytest.raises(TypeError): + fused(temb.bfloat16(), weight, bias, ti, tags) diff --git a/tests/h3/test_h3_adaln_row_gather.py b/tests/h3/test_h3_adaln_row_gather.py new file mode 100644 index 000000000..0c58c7e49 --- /dev/null +++ b/tests/h3/test_h3_adaln_row_gather.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_row_gather``: ``timestep_index * 3 + token_tag`` lookup of six tensors. + +* forward is a copy: bitwise equal to diffusers' six ``index_select`` calls, + for every packing, length and index dtype; +* semantic indices fail closed (out-of-range tags or timesteps, RFC probes + H2/H3), and flipping a tag or offsetting a timestep picks a different row; +* backward is a deterministic FP32 segmented sum: repeat-bitwise and + correctly rounded, unlike the atomic BF16 ``index_select`` backward. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.h3.adaln_row_gather import NativeH3AdaLNRowGatherOp +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_adaln_row_gather + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks h3_adaln_row_gather_*") + return H3AdaLNRowGatherCudaOp() + + +def _rows(num_timesteps, hidden=5376, dtype=torch.bfloat16, seed=0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(3 * num_timesteps, 6 * hidden, generator=g).to(dtype).to(device) + + +class TestReference: + def test_matches_diffusers_index_select(self): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = h3_packed_layout(11, 2, device="cpu") + index = ti * 3 + tags + chunks = rows.chunk(6, dim=-1) + for ours, chunk in zip(NativeH3AdaLNRowGatherOp().forward(rows, ti, tags), chunks): + assert torch.equal(ours, chunk.index_select(0, index)) + + @pytest.mark.parametrize( + "mutate, error", + [ + (lambda ti, tags: (ti, tags.clone().fill_(3)), IndexError), # H2: no 4th modality + (lambda ti, tags: (ti, tags - 1), IndexError), + (lambda ti, tags: (ti + 2, tags), IndexError), # H3: timestep offset past T + (lambda ti, tags: (ti, tags[:-1]), ValueError), + (lambda ti, tags: (ti[:0], tags[:0]), ValueError), + (lambda ti, tags: (ti.float(), tags), TypeError), + (lambda ti, tags: (ti.int(), tags), TypeError), # mixed index dtypes + ], + ) + def test_rejects_invalid_indices(self, mutate, error): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = mutate(*h3_packed_layout(9, 2, device="cpu")) + with pytest.raises(error): + NativeH3AdaLNRowGatherOp().forward(rows, ti, tags) + + def test_rejects_bad_rows(self): + ti, tags = h3_packed_layout(4, 1, device="cpu") + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(4, 24), ti, tags) # not 3 per t + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(3, 25), ti, tags) # not 6 * H + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("seq", [1, 2, 3, 257, 4097, 32768]) + @pytest.mark.parametrize("num_timesteps", [1, 3]) + def test_bitwise_equal_to_index_select(self, seq, num_timesteps): + rows = _rows(num_timesteps, seed=seq) + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + ours = _cuda_op()(rows, ti, tags) + ref = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + assert len(ours) == 6 + for a, b in zip(ours, ref): + assert a.shape == (seq, 5376) and a.is_contiguous() and torch.equal(a, b) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [24, 5, 5376]) + def test_dtypes_and_unaligned_hidden(self, dtype, hidden): + rows = _rows(2, hidden=hidden, dtype=dtype) + ti, tags = h3_packed_layout(33, 2) + for a, b in zip( + _cuda_op()(rows, ti.int(), tags.int()), NativeH3AdaLNRowGatherOp()(rows, ti, tags) + ): + assert torch.equal(a, b) + + def test_strided_rows_view(self): + table = _rows(4, hidden=16) + rows = table[::2] # row stride 2 * 6H, still 3 rows per timestep after slicing 12 -> 6 + ti, tags = h3_packed_layout(20, 2) + for a, b in zip(_cuda_op()(rows, ti, tags), NativeH3AdaLNRowGatherOp()(rows, ti, tags)): + assert torch.equal(a, b) + + def test_tag_flip_and_timestep_offset_select_other_rows(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(30, 2) + base = _cuda_op()(rows, ti, tags) + flipped = _cuda_op()(rows, ti, (tags + 1) % 3) # H2 + shifted = _cuda_op()(rows, 1 - ti, tags) # H3 within range + assert not torch.equal(base[0], flipped[0]) + assert not torch.equal(base[0], shifted[0]) + # Every output row is exactly the addressed table row, nothing else. + index = ti * 3 + tags + for c, out in enumerate(base): + assert torch.equal(out, rows[index, c * 16 : (c + 1) * 16]) + + def test_rejects_out_of_range_on_device(self): + rows = _rows(1, hidden=16) + ti, tags = h3_packed_layout(8, 1) + with pytest.raises(IndexError): + _cuda_op()(rows, ti, tags + 3) + + def test_gather_chunks_drop_in(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(10, 2) + chunks = rows.chunk(6, dim=-1) + for a, b in zip(_cuda_op().gather_chunks(chunks, ti, tags), _cuda_op()(rows, ti, tags)): + assert torch.equal(a, b) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + def _cuda_grad(self, rows, ti, tags, grads): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(_cuda_op()(leaf, ti, tags)), grads) + return leaf.grad + + def test_correctly_rounded_fp32_segment_sum(self): + rows = _rows(3, seed=1) + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = self._grads(4097, 5376) + ours = self._cuda_grad(rows, ti, tags, grads) + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward( + list(NativeH3AdaLNRowGatherOp().forward_fp32(ref, ti, tags)), [g.float() for g in grads] + ) + assert ours.dtype == torch.bfloat16 + # tolerance_contract.json gradient_accuracy / elementwise / bfloat16. + torch.testing.assert_close(ours.double(), ref.grad, atol=2e-2, rtol=1.6e-2) + assert (ours == ref.grad.bfloat16()).float().mean() > 0.9999 + + def test_repeat_bitwise(self): + rows = _rows(2, hidden=512, seed=2) + ti, tags = h3_packed_layout(5000, 2, seed=2) + grads = self._grads(5000, 512, seed=2) + first = self._cuda_grad(rows, ti, tags, grads) + for _ in range(3): + assert torch.equal(self._cuda_grad(rows, ti, tags, grads), first) + + def test_unreferenced_rows_get_zero_grad(self): + rows = _rows(3, hidden=16) + ti = torch.zeros(7, dtype=torch.long, device="cuda") + tags = torch.zeros(7, dtype=torch.long, device="cuda") # only row 0 is used + grad = self._cuda_grad(rows, ti, tags, self._grads(7, 16)) + assert torch.count_nonzero(grad[1:]) == 0 + assert torch.count_nonzero(grad[0]) > 0 + + def test_segments_longer_than_one_tile(self): + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import BACKWARD_TILE + + seq = 3 * BACKWARD_TILE + 17 + rows = _rows(1, hidden=8, dtype=torch.float32) + ti = torch.zeros(seq, dtype=torch.long, device="cuda") + tags = torch.zeros(seq, dtype=torch.long, device="cuda") + grads = [torch.ones(seq, 8, device="cuda") for _ in range(6)] + grad = self._cuda_grad(rows, ti, tags, grads) + assert torch.equal(grad[0], torch.full((48,), float(seq), device="cuda")) + + def test_reference_backward_is_deterministic_and_accurate(self): + rows = _rows(2, hidden=256, seed=4) + ti, tags = h3_packed_layout(3000, 2, seed=4) + grads = self._grads(3000, 256, seed=4) + + def ref_grad(): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(NativeH3AdaLNRowGatherOp()(leaf, ti, tags)), grads) + return leaf.grad + + first = ref_grad() + assert torch.equal(ref_grad(), first) + torch.testing.assert_close( + first.float(), self._cuda_grad(rows, ti, tags, grads).float(), atol=1e-1, rtol=2e-2 + ) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_row_gather", device="cuda") + assert type(op).__name__ == "H3AdaLNRowGatherCudaOp" diff --git a/tests/h3/test_h3_conditioning_e2e.py b/tests/h3/test_h3_conditioning_e2e.py index edd75a3b2..cdf830f5f 100644 --- a/tests/h3/test_h3_conditioning_e2e.py +++ b/tests/h3/test_h3_conditioning_e2e.py @@ -53,3 +53,22 @@ def test_chain_forward(registry, chain_weights, num_timesteps, seq_len): assert entry["isolated_vs_provider"]["bitwise_equal"], entry["stage"] drift = report["first_isolated_drift"] assert drift is None or not stages[drift].provider_bitwise_isolated + + +@pytest.mark.parametrize("num_timesteps, seq_len", [(1, 257), (3, 4097)]) +def test_chain_backward(registry, chain_weights, num_timesteps, seq_len): + """Parameter gradients of the whole chain: deterministic, and FP32-accurate when fused.""" + + report = h3_chain.run_backward_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + for name, entry in report["leaves"].items(): + for mode in ("candidate", "candidate_fused"): + assert entry[mode]["repeat_bitwise_equal"], (mode, name) + fused = entry["candidate_fused"] + if name.startswith("time_embedder"): + # FP32 parameters: no BF16 rounding anywhere on the fused path. + assert fused["max_abs_vs_golden_over_absmax"] < 1e-5, (name, fused) + else: + # BF16 parameters: only their own final rounding remains. + assert fused["correctly_rounded_fraction"] > 0.99, (name, fused) diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index fdef637f7..455a88638 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -52,6 +52,7 @@ def test_all_ws1_single_ops_are_registered(): "timestep_sinusoid_h3", "timestep_mlp_fp32", "adaln_projection_3mod", + "adaln_row_gather", } <= names From 1f7253c439301cf6105ac4983647a34f1fcc97c6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:45:10 +0000 Subject: [PATCH 02/42] [WS1][MiniMax-H3] adaln_row_gather: B200 evidence report, chain replay and figure Produced by scripts/h3_evidence.py and scripts/h3_chain_replay.py --backward at fa551c2 on a clean tree; 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-adaln-row-gather.md | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 54636dd09..1fb114792 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -95,7 +95,7 @@ weight and bias) against an FP64 golden in which every declared cast is straight | S >= 257 (72 gradients) | RL-Kernel, separate ops | RL-Kernel, fused | diffusers | | --- | --- | --- | --- | | repeat-bitwise | 72 / 72 | 72 / 72 | 0 / 72 | -| time embedder: max abs error / golden max | 5.3e-4 to 2.4e-3 | 6.5e-7 to 4.9e-6 | 5.6e-3 to 1.2e-1 | +| time embedder: max abs error / golden max | 5.3e-4 to 2.4e-3 | 6.5e-7 to 4.9e-6 | 3.8e-3 to 1.3e-1, varying run to run | | AdaLN weight/bias: correctly rounded BF16, worst case | 59.8% | 99.57% | 1.6% | At S = 3 all three chains are deterministic. Both RL-Kernel variants reach 1.5e-6 on the @@ -111,13 +111,26 @@ B200, T = 3, H = 5376, BF16. "write BW" counts the 6·S·H outputs; the table st | S | CUDA forward | provider forward | CUDA backward | provider backward | | --- | --- | --- | --- | --- | -| 4097 | 62 µs | 92 µs | 1.10 ms | 1.44 ms | -| 32768 | 315 µs | 587 µs | 2.10 ms | 8.66 ms | -| 131072 | 1.24 ms (6.8 TB/s write) | 2.37 ms | 6.84 ms | 34.9 ms | +| 4097 | 0.06 ms | 0.09 ms | 1.15 ms | 1.47 ms | +| 32768 | 0.32 ms | 0.59 ms | 2.18 ms | 8.72 ms | +| 131072 | 1.23 ms (6.8 TB/s write) | 2.37 ms | 6.86 ms | 34.9 ms | The backward's peak memory is about twice the provider's, because autograd stacks the six output gradients and the per-tile FP32 partials are materialised. +## Evidence + +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../usage/evidence/h3-adaln-row-gather-b200/figure.png) + +There are two data files: + +- [`report.json`](../usage/evidence/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`](../usage/evidence/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}. + +Both were written from a clean tree at commit `fa551c2`. + ## Tests ```bash From 024ee0813511ab2e430e0b4b28848b42a85c82f2 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 02:37:29 +0000 Subject: [PATCH 03/42] [WS1][CUDA][MiniMax-H3] h3_rmsnorm: RMSNorm with fused AdaLN modulation RFC #420 step 4, row `h3_rmsnorm` (stacked on adaln_row_gather). - CUDA kernel replaying PyTorch's own RMSNorm reduction order (vectorized_layer_norm_kernel: (32, 4) block per row, 4-element vectors, shuffle-down and cross-warp trees, rsqrtf, w * (rstd * x)): bitwise equal to nn.RMSNorm on all four pinned H3 norm weights and batch/position invariant. - Fused AdaLN modulation n * (1 + scale[i]) + shift[i] with the rows gathered in-kernel from strided table views and each step rounded where the eager expression rounds: bitwise equal to diffusers for the block (adaln_indices) and norm_out (timestep_indices), without materialising the (S, H) gathered tensors; 2.5x faster forward at S = 131072. - Deterministic backward: row-local dx, dweight over fixed 256-row tiles, dshift/dscale as sorted segment sums (diffusers: non-deterministic BF16 atomics, ~15x larger error). - Manifest pins the block/refiner-final/norm_out norm weights and norm_out.linear (all in shard 1); chain gains a per-mode history and the block norm1 + modulation stage (covered by the end-to-end test). - Registry, gtest, report/figure hooks, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/rmsnorm_modulate.cu | 403 ++++++++++++++++++++ csrc/ops.cpp | 21 + docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-rmsnorm.md | 107 ++++++ rl_engine/_C.pyi | 21 + rl_engine/kernels/gtest/operator_inputs.py | 15 + rl_engine/kernels/gtest/operator_specs.py | 11 + rl_engine/kernels/ops/cuda/h3/rmsnorm.py | 103 +++++ rl_engine/kernels/ops/pytorch/h3/rmsnorm.py | 143 +++++++ rl_engine/kernels/registry.py | 6 + rl_engine/testing/h3_chain.py | 70 +++- rl_engine/testing/h3_manifest.json | 12 +- rl_engine/testing/h3_provider.py | 16 + rl_engine/testing/h3_report.py | 120 ++++++ rl_engine/tests/test_dispatch.py | 1 + scripts/plot_h3_evidence.py | 65 ++++ setup.py | 1 + tests/h3/test_h3_rmsnorm.py | 203 ++++++++++ tests/test_ws1_gtest_gpu.py | 1 + 20 files changed, 1312 insertions(+), 9 deletions(-) create mode 100644 csrc/cuda/h3/rmsnorm_modulate.cu create mode 100644 docs/operators/h3-rmsnorm.md create mode 100644 rl_engine/kernels/ops/cuda/h3/rmsnorm.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/rmsnorm.py create mode 100644 tests/h3/test_h3_rmsnorm.py diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu new file mode 100644 index 000000000..5212324b1 --- /dev/null +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -0,0 +1,403 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 RMSNorm and fused AdaLN modulation (RFC #420 `h3_rmsnorm`). +// +// n[r] = cast(w * (rstd[r] * x[r])) rstd = rsqrtf(sum(x^2) / N + eps) +// out[r] = cast(cast(n[r] * cast(1 + scale[i])) + shift[i]) i = row_index[r % S] +// +// Forward contract h3-rmsnorm-v1: the statistics replay PyTorch's +// vectorized_layer_norm_kernel (torch cf30153): one +// (32, 4) block per row, 4-element vectors, thread t sums vectors t, t+128, ... +// in order, a shuffle-down tree (16..1), a cross-warp tree, sum / N, rsqrtf, +// then w * (rstd * x) and one cast. nn.RMSNorm therefore matches bitwise, and +// each modulation step rounds to the tensor dtype exactly where the eager +// `n * (1.0 + scale) + shift` expression does. Rows are independent: the result +// does not depend on batch size, position or the other rows. +// +// Backward (FP32, one cast per output, no atomics): +// d_n = g * cast(1 + scale) (g without modulation) +// 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) + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +constexpr int kVec = 4; // PyTorch's layer-norm vec_size +constexpr int kWarps = 4; // num_threads() / warp_size = 128 / 32 +constexpr int kRowTile = 256; // dweight row tile + +__device__ __forceinline__ float to_f(float v) { return v; } +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ float from_f(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } + +// Round to T and back: the value an eager T-dtype op would produce. +template +__device__ __forceinline__ float round_t(float v) { return to_f(from_f(v)); } + +// PyTorch's rms statistics: sequential per thread, shuffle-down, cross-warp tree. +// Returns sum(v^2) / N on every thread. +template +__device__ __forceinline__ float row_mean_square(const T* __restrict__ row, int n, float* buf) { + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + float s = 0.0f; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const float v = to_f(row[i * kVec + k]); + s = s + v * v; + } + } + for (int off = 16; off > 0; off >>= 1) s = s + __shfl_down_sync(0xffffffffu, s, off); + for (int off = blockDim.y / 2; off > 0; off /= 2) { + if (threadIdx.x == 0 && threadIdx.y >= off && threadIdx.y < 2 * off) buf[threadIdx.y - off] = s; + __syncthreads(); + if (threadIdx.x == 0 && threadIdx.y < off) s = s + buf[threadIdx.y]; + __syncthreads(); + } + if (thrx == 0) buf[0] = s / float(n); + __syncthreads(); + const float out = buf[0]; + __syncthreads(); + return out; +} + +// Same tree for an arbitrary per-element FP32 term; returns the row sum. +__device__ __forceinline__ float row_sum_tree(float s, float* buf) { + for (int off = 16; off > 0; off >>= 1) s = s + __shfl_down_sync(0xffffffffu, s, off); + for (int off = blockDim.y / 2; off > 0; off /= 2) { + if (threadIdx.x == 0 && threadIdx.y >= off && threadIdx.y < 2 * off) buf[threadIdx.y - off] = s; + __syncthreads(); + if (threadIdx.x == 0 && threadIdx.y < off) s = s + buf[threadIdx.y]; + __syncthreads(); + } + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + if (thrx == 0) buf[0] = s; + __syncthreads(); + const float out = buf[0]; + __syncthreads(); + return out; +} + +struct Modulation { + const void* shift; // (R, N) rows with stride `row_stride` (elements), or nullptr + const void* scale; + 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] +}; + +template +__global__ void __launch_bounds__(32 * kWarps) + rmsnorm_modulate_fwd_kernel(const T* __restrict__ x, const T* __restrict__ w, T* __restrict__ y, + float* __restrict__ rstd_out, int n, float eps, Modulation mod) { + __shared__ float buf[kWarps]; + const int64_t r = blockIdx.x; + const T* xr = x + r * n; + const float rstd = rsqrtf(row_mean_square(xr, n, buf) + eps); + const T* shift = nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + shift = static_cast(mod.shift) + i * mod.row_stride; + 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; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + float v = round_t(__fmul_rn(to_f(w[j]), __fmul_rn(rstd, to_f(xr[j])))); + if (shift != nullptr) { + // Separately rounded ops (no FMA contraction), as the eager expression. + const float t1 = round_t(__fadd_rn(1.0f, to_f(scale[j]))); + v = round_t(__fadd_rn(round_t(__fmul_rn(v, t1)), to_f(shift[j]))); + } + y[r * n + j] = from_f(v); + } + } + if (thrx == 0) rstd_out[r] = rstd; +} + +template +__device__ __forceinline__ float d_norm_out(const T* g_row, const T* scale, int j) { + const float g = to_f(g_row[j]); + return scale == nullptr ? g : g * round_t(1.0f + to_f(scale[j])); +} + +template +__global__ void __launch_bounds__(32 * kWarps) + rmsnorm_modulate_dx_kernel(const T* __restrict__ g, const T* __restrict__ x, + const T* __restrict__ w, const float* __restrict__ rstd, + T* __restrict__ dx, int n, Modulation mod) { + __shared__ float buf[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 int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + float s = 0.0f; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + s = fmaf(to_f(w[j]) * d_norm_out(gr, scale, j), to_f(xr[j]), s); + } + } + const float dot = row_sum_tree(s, buf); + const float rs = rstd[r]; + const float coef = rs * rs * rs * dot / float(n); + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + dx[r * n + j] = from_f(rs * to_f(w[j]) * d_norm_out(gr, scale, j) - to_f(xr[j]) * coef); + } + } +} + +// partial[tile, j] = sum over rows of the tile (ascending) of d_n * x * rstd +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) { + 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); + 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; + 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; +} + +template +__global__ void fold_tiles_kernel(const float* __restrict__ partial, T* __restrict__ out, + int64_t tiles, int n) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t t = 0; t < tiles; ++t) acc += partial[t * n + j]; + out[j] = from_f(acc); +} + +// Table gradient tiles: positions sorted by (index, position); columns j < n +// accumulate g (shift), columns j >= n accumulate g * n_out (scale). +template +__global__ void modulation_grad_partial_kernel( + const T* __restrict__ g, const T* __restrict__ x, const T* __restrict__ w, + const float* __restrict__ rstd, const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int n) { + const int64_t tile = blockIdx.x; + const int j2 = blockIdx.y * blockDim.x + threadIdx.x; + if (j2 >= 2 * n) return; + const bool is_scale = j2 >= n; + const int j = is_scale ? j2 - n : j2; + const float wj = to_f(w[j]); + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t r = sorted_pos[p]; + const float gv = to_f(g[r * n + j]); + if (is_scale) { + const float nv = round_t(wj * (rstd[r] * to_f(x[r * n + j]))); + acc = fmaf(gv, nv, acc); + } else { + acc += gv; + } + } + partial[tile * 2 * n + j2] = acc; +} + +__global__ void fold_segments_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + float* __restrict__ out, int width) { + const int64_t seg = blockIdx.x; + const int j = blockIdx.y * blockDim.x + threadIdx.x; + if (j >= width) return; + float acc = 0.0f; + for (int64_t t = seg_first_tile[seg]; t < seg_first_tile[seg + 1]; ++t) acc += partial[t * width + j]; + out[seg * width + j] = acc; +} + +void check_rows(const torch::Tensor& t, const char* name, const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device(), name, " must be on ", like.device()); + TORCH_CHECK(t.dim() == 2 && t.is_contiguous(), name, " must be a contiguous 2-D tensor"); + TORCH_CHECK(t.scalar_type() == like.scalar_type(), name, " must have dtype ", like.scalar_type()); +} + +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}; + if (!index.has_value()) { + TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); + return mod; + } + TORCH_CHECK(shift.has_value() && scale.has_value(), "modulation needs both shift and scale"); + const auto& sh = *shift; + const auto& sc = *scale; + const auto& ix = *index; + for (const auto* t : {&sh, &sc}) { + TORCH_CHECK(t->is_cuda() && t->device() == x.device() && t->scalar_type() == x.scalar_type(), + "shift/scale must be CUDA tensors with x's dtype"); + TORCH_CHECK(t->dim() == 2 && t->size(1) == n && t->stride(1) == 1, + "shift/scale must be (R, N) with unit column stride"); + } + TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), + "shift and scale must be views of the same table layout"); + TORCH_CHECK(ix.is_cuda() && ix.scalar_type() == at::kLong && ix.dim() == 1 && ix.is_contiguous(), + "row index must be a contiguous int64 CUDA tensor"); + TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", + ix.size(0), ")"); + mod.shift = sh.data_ptr(); + mod.scale = sc.data_ptr(); + mod.row_stride = sh.stride(0); + mod.index = ix.data_ptr(); + mod.seq = ix.size(0); + return mod; +} + +void check_xw(const torch::Tensor& x, const torch::Tensor& w) { + TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.is_contiguous(), "x must be a contiguous (M, N) CUDA tensor"); + TORCH_CHECK(x.size(0) > 0, "x must have at least one row"); + TORCH_CHECK(w.is_cuda() && w.dim() == 1 && w.size(0) == x.size(1) && w.is_contiguous() && + w.scalar_type() == x.scalar_type(), + "weight must be a contiguous (N,) CUDA tensor with x's dtype"); + TORCH_CHECK(x.size(1) % kVec == 0, "N=", x.size(1), " must be a multiple of ", kVec, + " (PyTorch's vectorized RMSNorm path)"); +} + +// float32 / float16 / bfloat16 only (float64 has no PyTorch-matching vector path here). +#define H3_DISPATCH(TYPE, NAME, ...) \ + [&] { \ + switch (TYPE) { \ + case at::kFloat: { using T = float; __VA_ARGS__(); break; } \ + case at::kHalf: { using T = __half; __VA_ARGS__(); break; } \ + case at::kBFloat16: { using T = __nv_bfloat16; __VA_ARGS__(); break; } \ + default: TORCH_CHECK(false, NAME, ": unsupported dtype ", TYPE); \ + } \ + }() + +} // namespace + +// x (M, N); optional shift/scale (R, N) strided views and index (S,), M % S == 0. +std::vector h3_rmsnorm_forward(torch::Tensor x, torch::Tensor weight, double eps, + c10::optional shift, + c10::optional scale, + c10::optional index) { + check_xw(x, weight); + TORCH_CHECK(x.scalar_type() != at::kDouble, "float64 is not supported"); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + const c10::cuda::CUDAGuard guard(x.device()); + auto y = torch::empty_like(x); + auto rstd = torch::empty({x.size(0)}, x.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_forward", [&] { + rmsnorm_modulate_fwd_kernel<<(x.size(0)), dim3(32, kWarps), 0, stream>>>( + reinterpret_cast(x.data_ptr()), reinterpret_cast(weight.data_ptr()), + reinterpret_cast(y.data_ptr()), rstd.data_ptr(), static_cast(n), + static_cast(eps), mod); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {y, rstd}; +} + +// Returns {dx, dweight} and, with modulation, {d_shift, d_scale} as FP32 (R, N). +std::vector h3_rmsnorm_backward( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, c10::optional sorted_pos, + c10::optional tile_begin, c10::optional tile_end, + c10::optional seg_first_tile) { + 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.scalar_type() == at::kFloat && rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics"); + const int64_t rows = x.size(0); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + const c10::cuda::CUDAGuard guard(x.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto dx = torch::empty_like(x); + const int64_t tiles = (rows + kRowTile - 1) / kRowTile; + auto partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + auto dweight = torch::empty_like(weight); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward", [&] { + const T* g = reinterpret_cast(grad.data_ptr()); + const T* xp = reinterpret_cast(x.data_ptr()); + const T* wp = reinterpret_cast(weight.data_ptr()); + rmsnorm_modulate_dx_kernel<<(rows), dim3(32, kWarps), 0, stream>>>( + g, xp, wp, rstd.data_ptr(), reinterpret_cast(dx.data_ptr()), + static_cast(n), mod); + rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, 0, + stream>>>(g, xp, rstd.data_ptr(), + partial.data_ptr(), rows, + static_cast(n), mod); + fold_tiles_kernel<<>>( + partial.data_ptr(), reinterpret_cast(dweight.data_ptr()), tiles, + static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (mod.index == nullptr) return {dx, dweight}; + + TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && + seg_first_tile.has_value(), + "modulated backward needs the sorted segment tiles"); + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t seg_tiles = tile_begin->numel(); + const int64_t width = 2 * n; + auto seg_partial = torch::empty({std::max(seg_tiles, 1), width}, x.options().dtype(at::kFloat)); + auto table_grad = torch::empty({segments, 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_grad", [&] { + 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(), + sorted_pos->data_ptr(), tile_begin->data_ptr(), + tile_end->data_ptr(), seg_partial.data_ptr(), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + 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 {dx, dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 27159d5e9..f2288efe4 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -517,6 +517,16 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor torch::Tensor tile_begin, torch::Tensor tile_end, torch::Tensor seg_first_tile, c10::ScalarType out_dtype); +std::vector h3_rmsnorm_forward(torch::Tensor x, torch::Tensor weight, double eps, + c10::optional shift, + c10::optional scale, + c10::optional index); +std::vector h3_rmsnorm_backward( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, c10::optional sorted_pos, + c10::optional tile_begin, c10::optional tile_end, + c10::optional seg_first_tile); #endif // PyBind11 Module Registration @@ -798,5 +808,16 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Deterministic segmented sum (sorted tiles folded in order) for the row gather", py::arg("grad"), py::arg("sorted_pos"), py::arg("tile_begin"), py::arg("tile_end"), py::arg("seg_first_tile"), py::arg("out_dtype")); + m.def("h3_rmsnorm_forward", &h3_rmsnorm_forward, + "RMSNorm replaying PyTorch's reduction order, with optional fused AdaLN modulation", + py::arg("x"), py::arg("weight"), py::arg("eps"), py::arg("shift") = py::none(), + py::arg("scale") = py::none(), py::arg("index") = py::none()); + m.def("h3_rmsnorm_backward", &h3_rmsnorm_backward, + "Deterministic RMSNorm(+modulation) backward: row-local dx, tiled dweight, sorted table grads", + py::arg("grad"), py::arg("x"), py::arg("weight"), py::arg("rstd"), + py::arg("shift") = py::none(), py::arg("scale") = py::none(), + py::arg("index") = py::none(), py::arg("sorted_pos") = py::none(), + py::arg("tile_begin") = py::none(), py::arg("tile_end") = py::none(), + py::arg("seg_first_tile") = py::none()); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index ff55306a7..af8c1863f 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -34,6 +34,7 @@ nav: - operators/h3-timestep-mlp.md - operators/h3-adaln-projection.md - operators/h3-adaln-row-gather.md + - operators/h3-rmsnorm.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 33e405acd..b687a65b5 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -36,4 +36,5 @@ Every operator page should include: - [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) - [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) - [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) +- [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md new file mode 100644 index 000000000..3bd0b2f6e --- /dev/null +++ b/docs/operators/h3-rmsnorm.md @@ -0,0 +1,107 @@ +# MiniMax-H3 RMSNorm and AdaLN Modulation + +## Summary + +`h3_rmsnorm` covers every RMSNorm in MiniMax-H3 (RFC #420, WS1 step 4). That is the block +`norm1`/`norm2`, the token-refiner norms, the refiner `final_norm` and `norm_out.norm`. All of +them are `nn.RMSNorm(5376, eps=1e-5)` with a BF16 affine weight. In a transformer block and in +`norm_out`, the normalised rows are immediately modulated by per-row AdaLN parameters, and the +op fuses that step: + +```text +n = rms_norm(x, weight, eps) +out = n * (1.0 + scale[index]) + shift[index] +``` + +`index` is `adaln_indices` in a block (rows of the [projection](h3-adaln-projection.md) table) +and `timestep_indices` in `norm_out`. `shift`/`scale` are `(R, H)` row views of the AdaLN table, +gathered inside the kernel, so the `(S, H)` tensors that +[`adaln_row_gather`](h3-adaln-row-gather.md) materialises are never needed. + +## Entry Point + +```python +op = kernel_registry.get_op("h3_rmsnorm", device="cuda") +n = op(x, weight) # plain RMSNorm +out = op.forward_modulated(x, weight, shift_msa, scale_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+, validated on SM100) | `rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp` | `rl_engine._C.h3_rmsnorm_{forward,backward}` | Bitwise equal to `nn.RMSNorm` and to the diffusers modulation | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp` | n/a | `forward*`: provider path; `forward*_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `x` | `(..., S, N)` | bf16 / fp16 / fp32 | `N % 4 == 0`; H3: `N = 5376` | +| `weight` | `(N,)` | `x.dtype` | | +| `shift`, `scale` | `(R, N)` row views, same stride | `x.dtype` | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; one entry per position, shared across the batch | + +Out-of-range indices, mismatched dtypes or shapes, a non-positive eps, and `N % 4 != 0` all fail +closed. + +## Numerics + +- **Statistics (contract `h3-rmsnorm-v1`).** The kernel replays PyTorch's own + `vectorized_layer_norm_kernel` (torch `cf30153`): + - one `(32, 4)` block per row, reading 4-element vectors; + - thread `t` sums vectors `t, t+128, …` in order; + - a shuffle-down tree (16 → 1), then a cross-warp tree; + - `rsqrtf(Σx²/N + eps)`, then `w * (rstd * x)` with one cast. + + The plain norm is therefore **bitwise equal to `nn.RMSNorm`**. Rows are independent of batch + size and position. +- **Modulation.** `1 + scale`, `n * (…)` and `+ shift` are each rounded to the tensor dtype, + exactly where the eager expression rounds. The fused output is **bitwise equal to diffusers' + `norm(x) * (1.0 + scale.index_select(0, i)) + shift.index_select(0, i)`**. +- **Backward.** Everything is in FP32 with one cast per output and no atomics: + - `dx` is row-local, using the same reduction tree; + - `dweight` is folded over fixed 256-row tiles in ascending order; + - `dshift`/`dscale` are segmented sums over positions sorted stably by table row, in the + same scheme as [`adaln_row_gather`](h3-adaln-row-gather.md). + + Diffusers' `index_select` backward accumulates `dshift`/`dscale` with BF16 atomics, so it is + non-deterministic and about 15× less accurate. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op h3_rmsnorm +``` + +B200, block `norm1` + MSA modulation, B = 1, H = 5376, BF16: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.06 ms | 0.11 ms | 1.00 ms | 0.95 ms | +| 32768 | 0.31 ms | 0.75 ms | 2.50 ms | 5.20 ms | +| 131072 | 1.16 ms | 2.90 ms | 7.48 ms | 20.6 ms | + +The forward is a single pass that never materialises the gathered rows. At small S the backward +is dominated by the fixed cost of the stable sort and tile setup. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_rmsnorm.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end, incl. block norm1 + modulation +python scripts/check_operator.py --op h3_rmsnorm --candidate cuda --device cuda \ + --dtype bf16 --batch 2 --seq 257 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op h3_rmsnorm --out docs/usage/evidence/h3-rmsnorm-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-rmsnorm-b200/report.json +``` + +## Known Limitations + +- Bitwise parity with `nn.RMSNorm` is tied to PyTorch's vectorized layer-norm path (`N % 4 == 0`, + aligned tensors). Other shapes are rejected rather than approximated. +- The token-refiner blocks' own norm weights live in shard 14, which is not pinned. They are the + same module (`nn.RMSNorm(5376, eps=1e-5)`) as the four pinned norms tested here. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 144b06e27..01e6e5348 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -356,3 +356,24 @@ def h3_adaln_row_gather_backward( seg_first_tile: torch.Tensor, out_dtype: torch.dtype, ) -> torch.Tensor: ... +def h3_rmsnorm_forward( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_backward( + 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, + sorted_pos: torch.Tensor | None = None, + tile_begin: torch.Tensor | None = None, + tile_end: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index ef77fe7bc..02f7f2e37 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -50,6 +50,7 @@ def make_operator_inputs( "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, "adaln_row_gather": _make_adaln_row_gather_inputs, + "h3_rmsnorm": _make_h3_rmsnorm_inputs, } try: return builders[op_name](args, dtype, device) @@ -87,6 +88,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: f"x{6 * 3 * _h3_hidden(args)}", "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" f"->{batch * seq}", + "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -431,6 +433,19 @@ def _make_adaln_row_gather_inputs( return {"rows": rows, "timestep_indices": timestep_indices, "token_tags": token_tags} +def _make_h3_rmsnorm_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + weight = _floating_tensor((hidden,), args, torch.float32, device, 1).abs() + 0.5 + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "weight": weight.to(dtype), + "eps": 1e-5, + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 11ce047db..197db8874 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -300,6 +300,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("rows",), ), + "h3_rmsnorm": OperatorSpec( + name="h3_rmsnorm", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp", + "cuda": "rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp", + }, + grad_input_names=("x", "weight"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/rmsnorm.py b/rl_engine/kernels/ops/cuda/h3/rmsnorm.py new file mode 100644 index 000000000..53cc8204a --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/rmsnorm.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 RMSNorm with fused AdaLN modulation (RFC #420 ``h3_rmsnorm``).""" + +from __future__ import annotations + +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 _segment_tiles +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +KERNEL_ID = "rl_engine._C.h3_rmsnorm_forward" +BACKWARD_IMPL = "row_local_dx_tiled_dweight_sorted_table_grads" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_rmsnorm_backward" + + +def h3_rmsnorm_available() -> bool: + return bool( + _EXT_AVAILABLE and hasattr(_C, "h3_rmsnorm_forward") and hasattr(_C, "h3_rmsnorm_backward") + ) + + +class _H3RMSNormCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, index, eps): + hidden = x.shape[-1] + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward(x2, weight.contiguous(), float(eps), shift, scale, index) + ctx.save_for_backward(x2, weight, rstd, shift, scale, index) + ctx.x_shape = x.shape + ctx.modulated = index is not None + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + if not ctx.modulated: + dx, dweight = _C.h3_rmsnorm_backward(g2, x2, weight.contiguous(), rstd) + d_shift = d_scale = None + else: + positions = index.repeat(x2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, shift.shape[0]) + dx, dweight, d_shift, d_scale = _C.h3_rmsnorm_backward( + g2, x2, weight.contiguous(), rstd, shift, scale, index, *tiles + ) + d_shift, d_scale = d_shift.to(shift.dtype), d_scale.to(scale.dtype) + record_backward( + "h3_rmsnorm", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None + + +class H3RMSNormCudaOp: + """CUDA candidate. + + The statistics replay PyTorch's own RMSNorm reduction order, so the plain + norm is bitwise equal to ``nn.RMSNorm``; the modulated form gathers its + ``shift``/``scale`` rows in the kernel (no ``(S, H)`` intermediates) and + rounds each step where the eager expression does, so it is bitwise equal to + diffusers' ``n * (1.0 + scale[i]) + shift[i]`` as well. Rows are + independent of batch size and position. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_rmsnorm_available(): + raise RuntimeError( + "rl_engine._C lacks h3_rmsnorm_*; rebuild with csrc/cuda/h3/rmsnorm_modulate.cu" + ) + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + return _H3RMSNormCuda.apply(x, weight, None, None, None, eps) + + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index, check_range=check_range) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + return _H3RMSNormCuda.apply(x, weight, shift, scale, index.contiguous(), eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + return self.forward(x, weight, eps).float() diff --git a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py new file mode 100644 index 000000000..6ab850a07 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 RMSNorm and its fused AdaLN modulation (RFC #420 row ``h3_rmsnorm``). + +Every H3 RMSNorm is ``nn.RMSNorm(5376, eps=1e-5)`` with an affine weight in the +block dtype: the block ``norm1``/``norm2``, the token-refiner norms and the +final ``norm_out.norm``. Inside a transformer block (and in ``norm_out``) the +normalised rows are immediately modulated by per-row AdaLN parameters:: + + n = rms_norm(x, weight, eps) + out = n * (1.0 + scale[index]) + shift[index] # eager, rounded after every op + +``index`` is ``adaln_indices`` in a block and ``timestep_indices`` in +``norm_out``; ``shift``/``scale`` are ``(R, H)`` rows (views of the AdaLN +table). ``x`` is ``(..., S, H)`` and ``index`` has one entry per position ``S``. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +H3_NORM_EPS = 1e-5 +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> int: + for name, tensor in (("x", x), ("weight", weight)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype not in _DTYPES: + raise TypeError(f"{name} must be fp32, bf16 or fp16, got {tensor.dtype}") + if weight.dtype != x.dtype: + raise TypeError(f"weight dtype {weight.dtype} must match x dtype {x.dtype}") + if weight.device != x.device: + raise ValueError(f"weight is on {weight.device}, x is on {x.device}") + if x.dim() < 2 or x.numel() == 0: + raise ValueError(f"x must be a non-empty (..., N) tensor, got {tuple(x.shape)}") + if weight.shape != (x.shape[-1],): + raise ValueError(f"weight must be ({x.shape[-1]},), got {tuple(weight.shape)}") + if not eps > 0: + raise ValueError(f"eps must be positive, got {eps}") + return x.shape[-1] + + +def validate_h3_modulation( + x: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + hidden = x.shape[-1] + for name, tensor in (("shift", shift), ("scale", scale)): + if tensor.device != x.device or (same_dtype and tensor.dtype != x.dtype): + raise TypeError(f"{name} must match x's dtype and device") + if tensor.dim() != 2 or tensor.shape[1] != hidden: + raise ValueError(f"{name} must be (R, {hidden}), got {tuple(tensor.shape)}") + if shift.shape != scale.shape: + raise ValueError("shift and scale must have the same shape") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != x.device: + raise TypeError("index must be a 1-D int64 tensor on x's device") + if x.dim() < 2 or x.shape[-2] != index.shape[0]: + raise ValueError( + f"x must be (..., S, {hidden}) with S = len(index) = {index.shape[0]}, " + f"got {tuple(x.shape)}" + ) + if check_range and bool(((index < 0) | (index >= shift.shape[0])).any()): + raise IndexError(f"index must lie in [0, {shift.shape[0]})") + + +class _DeterministicModulation(torch.autograd.Function): + """Eager ``n * (1.0 + scale[i]) + shift[i]``; the ``shift``/``scale`` + gradients as FP32 per-row sums instead of ``index_select``'s atomic + scatter-add (``d_n`` is the eager VJP).""" + + @staticmethod + def forward(ctx, n, shift, scale, index): + t1 = 1.0 + scale.index_select(0, index) + ctx.save_for_backward(n, t1, index) + ctx.rows = shift.shape[0] + return n * t1 + shift.index_select(0, index) + + @staticmethod + def backward(ctx, grad): + n, t1, index = ctx.saved_tensors + hidden = n.shape[-1] + g = grad.float().reshape(-1, index.shape[0], hidden) + per_shift = g.sum(dim=0) + per_scale = (g * n.float().reshape(-1, index.shape[0], hidden)).sum(dim=0) + d_shift = per_shift.new_zeros((ctx.rows, hidden)) + d_scale = per_scale.new_zeros((ctx.rows, hidden)) + for row in range(ctx.rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_shift[row] = per_shift.index_select(0, positions).sum(dim=0) + d_scale[row] = per_scale.index_select(0, positions).sum(dim=0) + return grad * t1, d_shift.to(t1.dtype), d_scale.to(t1.dtype), None + + +class NativeH3RMSNormOp: + """PyTorch reference. + + ``forward`` / ``forward_modulated`` are the provider path (``F.rms_norm`` and + the eager modulation expression, with a deterministic shift/scale gradient; + the raw diffusers path is ``h3_provider.provider_norm_modulate``). ``forward_fp32`` / + ``forward_modulated_fp32`` are the golden: the same math in FP64 with no + intermediate rounding, returned in FP32. + """ + + op_class = "reduction" + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + hidden = validate_h3_rmsnorm(x, weight, eps) + return F.rms_norm(x, (hidden,), weight, eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + rstd = torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) + return (x64 * rstd * weight.double()).float() + + def forward_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + validate_h3_modulation(x, shift, scale, index) + n = self.forward(x, weight, eps) + return _DeterministicModulation.apply(n, shift, scale, index) + + def forward_modulated_fp32(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + # The golden may be fed FP32 golden modulation rows for a BF16 x. + validate_h3_modulation(x, shift, scale, index, same_dtype=False) + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) * weight.double() + out = n * (1.0 + scale.double().index_select(0, index)) + shift.double().index_select( + 0, index + ) + return out.float() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index ba6c57f05..8fa1d1453 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -221,6 +221,8 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_ADALN_ROW_GATHER = ( "rl_engine.kernels.ops.cuda.h3.adaln_row_gather.H3AdaLNRowGatherCudaOp" ) + PYTORCH_H3_RMSNORM = "rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp" + CUDA_H3_RMSNORM = "rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -647,6 +649,7 @@ def __init__(self): OpBackend.CUDA_H3_ADALN_ROW_GATHER, OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, ], + "h3_rmsnorm": [OpBackend.CUDA_H3_RMSNORM, OpBackend.PYTORCH_H3_RMSNORM], }, "rocm": { "logp": [ @@ -691,6 +694,7 @@ def __init__(self): "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -739,6 +743,7 @@ def __init__(self): "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -768,6 +773,7 @@ def __init__(self): "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 06bfe252b..fe010c7e3 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -33,11 +33,13 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_norm_modulate, provider_time_embedder, provider_time_proj, ) REPO_ROOT = Path(__file__).resolve().parents[2] +HIDDEN = 5376 @dataclass(frozen=True) @@ -113,6 +115,45 @@ def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: ), ] + +def adaln_indices(ctx: dict[str, Any]) -> torch.Tensor: + return ctx["timestep_indices"] * 3 + ctx["token_tags"] + + +def _norm1_modulated(ctx: dict[str, Any], modulation, call): + shift_msa, scale_msa = modulation[0], modulation[1] + weight = ctx["weights"]["transformer_blocks.0.norm1.weight"] + return call(ctx["hidden"], weight, shift_msa, scale_msa, adaln_indices(ctx)) + + +STAGES.append( + Stage( + name="h3_rmsnorm", + op_type="h3_rmsnorm", + # Block norm1 followed by the MSA shift/scale of the projection stage. + candidate=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated + ), + provider=lambda ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], provider_norm_modulate + ), + golden=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated_fp32 + ), + provider_bitwise_isolated=True, # replays nn.RMSNorm's reduction order + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + +# The backward replay covers the conditioning chain up to the row gather. +CONDITIONING_STAGES = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", +) + # Parameters whose gradients the backward replay reports, in chain order. GRAD_LEAVES = ( "time_embedder.linear_1.weight", @@ -151,6 +192,13 @@ def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dic "timestep_indices": timestep_indices, "token_tags": token_tags, "weights": weights, + # Stand-in for the packed hidden states (the patch/text projections are + # other RFC rows): BF16 (1, S, H), seeded. + "hidden": torch.randn( + (1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 2) + ) + .to(torch.bfloat16) + .cuda(), } @@ -179,16 +227,24 @@ def run_case( } chained = provider = golden = None first_drift = first_isolated = None + # Each mode sees its own earlier outputs (a stage may read any earlier stage). + history = {"candidate": {}, "provider": {}, "golden": {}} + ctx_c = {**ctx, "history": history["candidate"]} + ctx_p = {**ctx, "history": history["provider"]} + ctx_g = {**ctx, "history": history["golden"]} for stage in stages: op = registry.get_op(stage.op_type, device="cuda") gold = golden_op(registry, stage.op_type) with torch.no_grad(): - isolated = stage.candidate(op, ctx, provider) + isolated = stage.candidate(op, ctx_p, provider) chained_input = chained - chained = stage.candidate(op, ctx, chained_input) - repeat = stage.candidate(op, ctx, chained_input) - golden = stage.golden(gold, ctx, golden) - provider = stage.provider(ctx, provider) + chained = stage.candidate(op, ctx_c, chained_input) + repeat = stage.candidate(op, ctx_c, chained_input) + golden = stage.golden(gold, ctx_g, golden) + provider = stage.provider(ctx_p, provider) + history["candidate"][stage.name] = chained + history["provider"][stage.name] = provider + history["golden"][stage.name] = golden entry = { "stage": stage.name, "backend": type(op).__name__, @@ -220,7 +276,7 @@ def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torc leaves = { name: ctx["weights"][name].detach().clone().requires_grad_(True) for name in GRAD_LEAVES } - run_ctx = {**ctx, "weights": {**ctx["weights"], **leaves}} + run_ctx = {**ctx, "weights": {**ctx["weights"], **leaves}, "history": {}} value = None if mode == "candidate_fused": from rl_engine.kernels.ops.cuda.h3.adaln_modulation import H3AdaLNModulationCudaOp @@ -231,7 +287,7 @@ def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torc value, *adaln_params(run_ctx), ctx["timestep_indices"], ctx["token_tags"] ) else: - for stage in STAGES: + for stage in (s for s in STAGES if s.name in CONDITIONING_STAGES): if mode == "provider": value = stage.provider(run_ctx, value) elif mode == "candidate": diff --git a/rl_engine/testing/h3_manifest.json b/rl_engine/testing/h3_manifest.json index 52b1989b1..176ddcf04 100644 --- a/rl_engine/testing/h3_manifest.json +++ b/rl_engine/testing/h3_manifest.json @@ -35,7 +35,13 @@ "time_embedder.linear_2.weight": {"dtype": "float32", "shape": [2688, 5376]}, "time_embedder.linear_2.bias": {"dtype": "float32", "shape": [2688]}, "transformer_blocks.0.adaln_proj.linear.weight": {"dtype": "bfloat16", "shape": [96768, 2688]}, - "transformer_blocks.0.adaln_proj.linear.bias": {"dtype": "bfloat16", "shape": [96768]} + "transformer_blocks.0.adaln_proj.linear.bias": {"dtype": "bfloat16", "shape": [96768]}, + "transformer_blocks.0.norm1.weight": {"dtype": "bfloat16", "shape": [5376]}, + "transformer_blocks.0.norm2.weight": {"dtype": "bfloat16", "shape": [5376]}, + "token_refiner.final_norm.weight": {"dtype": "bfloat16", "shape": [5376]}, + "norm_out.norm.weight": {"dtype": "bfloat16", "shape": [5376]}, + "norm_out.linear.weight": {"dtype": "bfloat16", "shape": [10752, 2688]}, + "norm_out.linear.bias": {"dtype": "bfloat16", "shape": [10752]} }, "reference_implementation": { "repo": "huggingface/diffusers", @@ -53,7 +59,9 @@ "precision": { "timestep_sinusoid": "float32", "timestep_mlp": "float32 weights, float32 activations", - "adaln_projection": "SiLU in float32, cast to bfloat16, bfloat16 weight/bias/output" + "adaln_projection": "SiLU in float32, cast to bfloat16, bfloat16 weight/bias/output", + "rmsnorm": "nn.RMSNorm(5376, eps=1e-5), bfloat16 affine weight, statistics in float32", + "final_modulation": "norm_out.linear(silu(temb).bfloat16) -> (shift, scale), indexed by timestep_indices" } } } diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index 6cedcdf60..05c7ebec9 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -65,3 +65,19 @@ def provider_adaln_row_gather( adaln_indices = timestep_indices * 3 + token_tags return tuple(tensor.index_select(0, adaln_indices) for tensor in modulation) + + +def provider_norm_modulate( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``norm(x) * (1.0 + scale[i]) + shift[i]`` with ``index_select``; block and norm_out.""" + + norm_hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return norm_hidden_states * (1.0 + scale.index_select(0, indices)) + shift.index_select( + 0, indices + ) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index e0a5939b5..dfecec419 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -20,6 +20,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_norm_modulate, provider_time_embedder, provider_time_proj, ) @@ -349,15 +350,134 @@ def grad_of(fn): } +# --------------------------------------------------------------------------- # +# h3_rmsnorm (block norm1 + MSA modulation) +# --------------------------------------------------------------------------- # + +NORM_SEQ_LENS = (4097, 32768, 131072) + + +def _norm_inputs(seq: int, seed: int = 0): + weight = h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0].bfloat16() + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).bfloat16() + shift, scale = table.view(9, 6, 5376)[:, 0], table.view(9, 6, 5376)[:, 1] + ti, tags = h3_packed_layout(seq, 3, seed=seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + return x, weight, shift, scale, ti * 3 + tags + + +def _norm_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("h3_rmsnorm", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, weight, shift, scale, index = _norm_inputs(seq) + grad = torch.randn_like(x) + + def backward(fn, x=x, weight=weight, shift=shift, scale=scale, grad=grad): + leaves = [t.detach().requires_grad_(True) for t in (x, weight, shift, scale)] + fn(*leaves).backward(grad) + + cases.append( + { + "op": "h3_rmsnorm", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 2 * seq * 5376 * 2, # read x, write y + "candidate": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + op.forward_modulated(x, w, sh, sc, i, check_range=False) + ), + "provider": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + provider_norm_modulate(x, w, sh, sc, i) + ), + "candidate_backward": lambda b=backward, i=index: b( + lambda x_, w_, sh_, sc_: op.forward_modulated( + x_, w_, sh_, sc_, i, check_range=False + ) + ), + "provider_backward": lambda b=backward, i=index: b( + lambda x_, w_, sh_, sc_: provider_norm_modulate(x_, w_, sh_, sc_, i) + ), + } + ) + return cases + + +def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("h3_rmsnorm", device="cuda") + weights = load_h3_conditioning_weights("cuda") + plain = {} + for name in ( + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", + ): + x = torch.randn(2, 777, 5376, device="cuda").bfloat16() + ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) + plain[name] = bool(torch.equal(op(x, weights[name]), ref)) + x, weight, shift, scale, index = _norm_inputs(4097, seed=1) + modulated = bool( + torch.equal( + op.forward_modulated(x, weight, shift, scale, index), + provider_norm_modulate(x, weight, shift, scale, index), + ) + ) + invariant = bool( + torch.equal( + op.forward_modulated(x[:, 100:140], weight, shift, scale, index[100:140])[0], + op.forward_modulated(x, weight, shift, scale, index)[0, 100:140], + ) + ) + + # Backward against FP64, for the CUDA op and the diffusers expression. + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, weight, shift, scale)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden(x_, w_, sh_, sc_): + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * w_ + return n * (1 + sc_.index_select(0, index)) + sh_.index_select(0, index) + + ref = grads(golden, torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op.forward_modulated(*t, index)), + ("provider", lambda *t: provider_norm_modulate(*t, index)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref) + }, + } + return { + "plain_bitwise_vs_nn_rmsnorm": plain, + "modulated_bitwise_vs_diffusers": modulated, + "rows_batch_invariant": invariant, + "backward": backward, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, "adaln_projection_3mod": _projection_perf, "adaln_row_gather": _gather_perf, + "h3_rmsnorm": _norm_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, "timestep_mlp_fp32": _mlp_accuracy, "adaln_projection_3mod": _projection_accuracy, "adaln_row_gather": _gather_accuracy, + "h3_rmsnorm": _norm_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index c36ade2f2..49d502f12 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -231,6 +231,7 @@ def fake_load_backend(backend): "timestep_mlp_fp32", "adaln_projection_3mod", "adaln_row_gather", + "h3_rmsnorm", ) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index f2039f0aa..da8bc6774 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -345,11 +345,76 @@ def plot_gather(report: dict[str, Any]): return fig +def plot_rmsnorm(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "h3_rmsnorm (block norm1 + MSA modulation, H = 5376, BF16)" + ) + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("diffusers fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "diffusers bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + acc = report["accuracy"] + bwd = acc["backward"] + keys = ["dx", "dweight", "dshift", "dscale"] + fwd_ok = acc["modulated_bitwise_vs_diffusers"] and all( + acc["plain_bitwise_vs_nn_rmsnorm"].values() + ) + _style(right, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") + + def repeat(mode: str) -> str: + return "yes" if bwd[mode]["repeat_bitwise_equal"] else "no" + + _bars( + right, + keys, + [ + ( + f"RL-Kernel CUDA (repeat-bitwise: {repeat('cuda')})", + RL_KERNEL, + [bwd["cuda"]["rel_error"][k] for k in keys], + ), + ( + f"diffusers (repeat-bitwise: {repeat('provider')})", + PROVIDER, + [bwd["provider"]["rel_error"][k] for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + right.text( + 0.98, + 0.80, + f"forward bitwise equal to diffusers: {'yes' if fwd_ok else 'NO'}", + transform=right.transAxes, + ha="right", + fontsize=8, + color=INK, + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, "adaln_projection_3mod": plot_projection, "adaln_row_gather": plot_gather, + "h3_rmsnorm": plot_rmsnorm, } diff --git a/setup.py b/setup.py index fcefa9016..fb671d0f0 100644 --- a/setup.py +++ b/setup.py @@ -160,6 +160,7 @@ def get_extensions(): cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") cuda_sources.append("csrc/cuda/h3/det_linear.cu") cuda_sources.append("csrc/cuda/h3/adaln_row_gather.cu") + cuda_sources.append("csrc/cuda/h3/rmsnorm_modulate.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/test_h3_rmsnorm.py b/tests/h3/test_h3_rmsnorm.py new file mode 100644 index 000000000..f0d013bfb --- /dev/null +++ b/tests/h3/test_h3_rmsnorm.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``h3_rmsnorm``: RMSNorm (eps 1e-5) and its fused AdaLN modulation. + +* forward is bitwise equal to ``nn.RMSNorm`` and to diffusers' + ``norm(x) * (1.0 + scale[i]) + shift[i]`` (block: ``adaln_indices``; + ``norm_out``: ``timestep_indices``) on the pinned norm weights; +* rows are independent of batch size and position; +* backward is deterministic and checked against an FP64 golden; +* malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import NativeH3RMSNormOp +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_norm_modulate + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NORMS = ( + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", +) + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.rmsnorm import H3RMSNormCudaOp, h3_rmsnorm_available + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3RMSNormCudaOp() + + +def _x(shape, dtype=torch.bfloat16, seed=0, scale=2.0): + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(shape, generator=g) * scale).to(dtype).cuda() + + +def _table(rows, hidden, chunks, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(rows, chunks * hidden, generator=g) * 0.5).to(dtype).cuda() + return table.chunk(chunks, dim=-1) # strided (R, H) views, like the AdaLN outputs + + +def _fp64_modulated(x, weight, shift, scale, index, eps=1e-5): + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * weight.double() + return n * (1 + scale.double().index_select(0, index)) + shift.double().index_select(0, index) + + +class TestReference: + def test_golden_close_to_provider(self): + x, w = _x((3, 64), torch.float32).cpu(), torch.rand(64) + 0.5 + op = NativeH3RMSNormOp() + torch.testing.assert_close(op.forward(x, w), op.forward_fp32(x, w), atol=1e-6, rtol=1e-6) + + def test_rejects_bad_inputs(self): + op = NativeH3RMSNormOp() + x, w = torch.randn(2, 8), torch.ones(8) + with pytest.raises(TypeError): + op.forward(x, w.double()) + with pytest.raises(TypeError): + op.forward(x, w.bfloat16()) + with pytest.raises(ValueError): + op.forward(x, torch.ones(7)) + with pytest.raises(ValueError): + op.forward(x, w, eps=0.0) + shift, scale = torch.zeros(3, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op.forward_modulated(x, w, shift, scale, torch.tensor([0, 3])) + with pytest.raises(ValueError): # one index per position S + op.forward_modulated(x, w, shift, scale, torch.tensor([0])) + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("name", NORMS) + def test_bitwise_equal_to_nn_rmsnorm_on_pinned_weights(self, h3_weights_cpu, name): + weight = h3_weights_cpu[name].cuda() + module = torch.nn.RMSNorm(5376, eps=1e-5).cuda().bfloat16() + module.weight.data.copy_(weight) + x = _x((2, 777, 5376), seed=1) + assert torch.equal(_cuda_op()(x, weight), module(x)) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [12, 256, 4096, 5376]) + def test_bitwise_equal_to_f_rms_norm(self, dtype, hidden): + x = _x((5, 33, hidden), dtype, seed=hidden, scale=30.0) + w = (torch.rand(hidden, device="cuda") * 2).to(dtype) + ref = torch.nn.functional.rms_norm(x, (hidden,), w, 1e-5) + assert torch.equal(_cuda_op()(x, w), ref) + + def test_block_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["transformer_blocks.0.norm1.weight"].cuda() + shift, scale, *_ = _table(9, 5376, 6, seed=2) # (3T, H) rows, T = 3 + ti, tags = h3_packed_layout(4097, 3, seed=2) + index = ti * 3 + tags + x = _x((2, 4097, 5376), seed=2) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, index)) + + def test_norm_out_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["norm_out.norm.weight"].cuda() + shift, scale = _table(3, 5376, 2, seed=3) # (T, H) rows indexed by timestep + ti, _ = h3_packed_layout(1000, 3, seed=3) + x = _x((1, 1000, 5376), seed=3) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, ti) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, ti)) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_modulation_bitwise_in_every_dtype(self, dtype): + """FP32 must not contract the separately rounded eager ops into an FMA.""" + + w = (torch.rand(5376, device="cuda") + 0.5).to(dtype) + shift, scale, *_ = _table(9, 5376, 6, dtype=dtype, seed=9) + ti, tags = h3_packed_layout(333, 3, seed=9) + index = ti * 3 + tags + x = _x((2, 333, 5376), dtype, seed=9) + ours = _cuda_op().forward_modulated(x, w, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, w, shift, scale, index)) + + def test_rows_are_batch_and_position_invariant(self): + w = (torch.rand(5376, device="cuda") + 0.5).bfloat16() + shift, scale, *_ = _table(9, 5376, 6, seed=4) + ti, tags = h3_packed_layout(300, 3, seed=4) + index = ti * 3 + tags + x = _x((3, 300, 5376), seed=4) + full = _cuda_op().forward_modulated(x, w, shift, scale, index) + part = _cuda_op().forward_modulated(x[2:3, 50:90], w, shift, scale, index[50:90]) + assert torch.equal(part[0], full[2, 50:90]) + + def test_rejects_unsupported(self): + w = torch.ones(6, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="multiple of 4"): + _cuda_op()(_x((2, 6)), w) + with pytest.raises(ValueError): + _cuda_op()(_x((2, 8)).cpu(), torch.ones(8).bfloat16()) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, fn, tensors, grad): + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad) + return [leaf.grad for leaf in leaves] + + def test_modulated_backward_against_fp64(self): + g = torch.Generator(device="cpu").manual_seed(5) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.5).bfloat16().cuda() + ti, tags = h3_packed_layout(2049, 3, seed=5) + index = ti * 3 + tags + x = _x((2, 2049, 5376), seed=5) + grad = _x((2, 2049, 5376), seed=6, scale=1.0) + op = _cuda_op() + + def ours(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return op.forward_modulated(x_, w_, shift, scale, index) + + def golden(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return _fp64_modulated(x_, w_, shift, scale, index) + + got = self._grads(ours, (x, w, table), grad) + ref = self._grads(golden, (x.double(), w.double(), table.double()), grad.double()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + rel = ((g_.double() - r).abs().max() / r.abs().max()).item() + # Two BF16 roundings remain: the eager VJP uses round(1 + scale), and + # every gradient is rounded once to BF16 (2^-9 relative each). + assert rel < 1e-2, (name, rel) + again = self._grads(ours, (x, w, table), grad) + assert all(torch.equal(a, b) for a, b in zip(got, again)) + + def test_plain_backward_against_fp64_and_row_local_dx(self): + w = (torch.rand(5376, generator=torch.Generator().manual_seed(7)) + 0.5).bfloat16().cuda() + x = _x((4, 64, 5376), seed=7) + grad = _x((4, 64, 5376), seed=8, scale=1.0) + op = _cuda_op() + got = self._grads(lambda a, b: op(a, b), (x, w), grad) + + def golden(a, b): + return a * torch.rsqrt(a.square().mean(-1, keepdim=True) + 1e-5) * b + + ref = self._grads(golden, (x.double(), w.double()), grad.double()) + for g, r in zip(got, ref): + assert ((g.double() - r).abs().max() / r.abs().max()).item() < 5e-3 + part = self._grads(lambda a, b: op(a, b), (x[1:2, 10:20], w), grad[1:2, 10:20]) + assert torch.equal(part[0][0], got[0][1, 10:20]) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + assert ( + type(KernelRegistry().get_op("h3_rmsnorm", device="cuda")).__name__ == "H3RMSNormCudaOp" + ) diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 455a88638..2b165c088 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -53,6 +53,7 @@ def test_all_ws1_single_ops_are_registered(): "timestep_mlp_fp32", "adaln_projection_3mod", "adaln_row_gather", + "h3_rmsnorm", } <= names From 729978430a95ee8a82d7bdf07147a729ebca88c6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:13:55 +0000 Subject: [PATCH 04/42] [WS1][MiniMax-H3] h3_rmsnorm: B200 evidence report and figure Produced by scripts/h3_evidence.py at f36aa60 on a clean tree; 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-rmsnorm.md | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 3bd0b2f6e..957670552 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -67,7 +67,7 @@ closed. same scheme as [`adaln_row_gather`](h3-adaln-row-gather.md). Diffusers' `index_select` backward accumulates `dshift`/`dscale` with BF16 atomics, so it is - non-deterministic and about 15× less accurate. + non-deterministic and about 15–20× less accurate (its error varies from run to run). ## Performance Notes @@ -79,13 +79,24 @@ B200, block `norm1` + MSA modulation, B = 1, H = 5376, BF16: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.06 ms | 0.11 ms | 1.00 ms | 0.95 ms | +| 4097 | 0.06 ms | 0.11 ms | 1.08 ms | 0.95 ms | | 32768 | 0.31 ms | 0.75 ms | 2.50 ms | 5.20 ms | -| 131072 | 1.16 ms | 2.90 ms | 7.48 ms | 20.6 ms | +| 131072 | 1.17 ms | 2.90 ms | 7.40 ms | 20.6 ms | The forward is a single pass that never materialises the gathered rows. At small S the backward is dominated by the fixed cost of the stable sort and tile setup. +## Evidence + +![h3_rmsnorm on B200: latency and backward accuracy](../usage/evidence/h3-rmsnorm-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-rmsnorm-b200/report.json), written by +`scripts/h3_evidence.py` from a clean tree at commit `f36aa60`. It also records: + +- bitwise equality with `nn.RMSNorm` for all four pinned norm weights; +- bitwise equality with diffusers for the modulation; +- row invariance. + ## Tests ```bash From 2bbef2263b4300fef343726a12e85b7571179d5c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:54:31 +0000 Subject: [PATCH 05/42] fix(h3): validate native gather indices and replay stage dependencies Check timestep and modality bounds on the device before computing gather addresses, including unchecked calls and overflowing int64 indices. Reject non-prefix replay stage selections and partial backward chains before loading weights. Validation: rebuilt the CUDA extension on B200; tests/h3: 175 passed, including 12 isolated native bounds probes and 16 CLI regressions. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 20 +++++-- scripts/h3_chain_replay.py | 26 +++++++-- tests/h3/test_h3_adaln_row_gather.py | 55 ++++++++++++++++++ tests/h3/test_h3_chain_replay_cli.py | 85 ++++++++++++++++++++++++++++ 4 files changed, 174 insertions(+), 12 deletions(-) create mode 100644 tests/h3/test_h3_chain_replay_cli.py diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 192faa187..4b3a168a0 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -63,8 +63,13 @@ struct CudaType { template __device__ __forceinline__ int64_t adaln_row(const idx_t* ti, const idx_t* tags, int64_t s, - int64_t modality_num) { - return static_cast(ti[s]) * modality_num + static_cast(tags[s]); + int64_t modality_num, int64_t num_timesteps) { + const int64_t timestep = static_cast(ti[s]); + const int64_t tag = static_cast(tags[s]); + // Check each semantic index before multiplication, including values that overflow int64. + CUDA_KERNEL_ASSERT(timestep >= 0 && timestep < num_timesteps); + CUDA_KERNEL_ASSERT(tag >= 0 && tag < modality_num); + return timestep * modality_num + tag; } // One block per (s, c); 16-byte copies when every row start is 16-byte aligned. @@ -73,10 +78,10 @@ __global__ void adaln_row_gather_kernel(const T* __restrict__ rows, int64_t row_ const idx_t* __restrict__ ti, const idx_t* __restrict__ tags, T* __restrict__ out, int64_t seq, int64_t hidden, int64_t chunks, - int64_t modality_num) { + int64_t modality_num, int64_t num_timesteps) { for (int64_t s = blockIdx.x; s < seq; s += gridDim.x) { const int64_t c = blockIdx.y; - const int64_t r = adaln_row(ti, tags, s, modality_num); + const int64_t r = adaln_row(ti, tags, s, modality_num, num_timesteps); const T* src = rows + r * row_stride + c * hidden; T* dst = out + (c * seq + s) * hidden; if constexpr (kVector) { @@ -146,6 +151,7 @@ torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor time " is not a multiple of ", chunks); TORCH_CHECK(modality_num > 0 && rows.size(0) % modality_num == 0, "rows count ", rows.size(0), " is not a multiple of ", modality_num); + const int64_t num_timesteps = rows.size(0) / modality_num; const int64_t seq = timestep_indices.numel(); TORCH_CHECK(seq > 0, "the packed sequence must not be empty"); check_index(timestep_indices, "timestep_indices", seq, rows); @@ -174,10 +180,12 @@ torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor time const idx_t* tags = token_tags.data_ptr(); if (vector) { adaln_row_gather_kernel<<>>( - src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num); + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num, + num_timesteps); } else { adaln_row_gather_kernel<<>>( - src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num); + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num, + num_timesteps); } }; if (timestep_indices.scalar_type() == at::kLong) { diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py index 52a1cddde..d5fc6a94d 100644 --- a/scripts/h3_chain_replay.py +++ b/scripts/h3_chain_replay.py @@ -30,12 +30,20 @@ def _ints(text: str) -> list[int]: + """Parse a comma-separated list of integer case sizes.""" + return [int(v) for v in text.split(",")] def main() -> None: + """Validate the chain prefix, replay cases, and optionally save evidence.""" + parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--stages", default=",".join(s.name for s in h3_chain.STAGES)) + parser.add_argument( + "--stages", + default=",".join(s.name for s in h3_chain.STAGES), + help="comma list of stages forming a nonempty ordered prefix of the chain", + ) parser.add_argument("--timesteps", type=_ints, default=[1, 2, 4], help="comma list of T") parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") parser.add_argument("--seed", type=int, default=0) @@ -48,13 +56,19 @@ def main() -> None: ) args = parser.parse_args() - if not torch.cuda.is_available(): - raise SystemExit("the chain replay needs a CUDA device") wanted = args.stages.split(",") - unknown = sorted(set(wanted) - {s.name for s in h3_chain.STAGES}) + stage_names = [s.name for s in h3_chain.STAGES] + unknown = sorted(set(wanted) - set(stage_names)) if unknown: - raise SystemExit(f"unknown stages: {unknown}") - stages = [s for s in h3_chain.STAGES if s.name in wanted] + parser.error(f"unknown stages: {unknown}") + if wanted != stage_names[: len(wanted)]: + parser.error("--stages must form a nonempty ordered prefix of the chain") + stages = h3_chain.STAGES[: len(wanted)] + if args.backward and len(stages) != len(h3_chain.STAGES): + parser.error("--backward requires every stage of the chain") + + if not torch.cuda.is_available(): + raise SystemExit("the chain replay needs a CUDA device") torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False diff --git a/tests/h3/test_h3_adaln_row_gather.py b/tests/h3/test_h3_adaln_row_gather.py index 0c58c7e49..8f6558034 100644 --- a/tests/h3/test_h3_adaln_row_gather.py +++ b/tests/h3/test_h3_adaln_row_gather.py @@ -13,6 +13,12 @@ from __future__ import annotations +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest import torch @@ -123,6 +129,55 @@ def test_rejects_out_of_range_on_device(self): with pytest.raises(IndexError): _cuda_op()(rows, ti, tags + 3) + @pytest.mark.parametrize("entrypoint", ["native", "unchecked_wrapper"]) + @pytest.mark.parametrize( + "timestep, tag, index_dtype", + [ + pytest.param(-1, 0, "int32", id="negative-timestep"), + pytest.param(2, 0, "int64", id="timestep-past-table"), + pytest.param(0, 3, "int32", id="tag-past-modality-valid-row"), + pytest.param(1, -1, "int64", id="negative-tag-valid-row"), + # Multiplication by 3 wraps this value to row 2 in signed int64. + pytest.param((2**64 + 2) // 3, 0, "int64", id="timestep-overflow-valid-row"), + pytest.param(0, 2**63 - 1, "int64", id="int64-max-tag"), + ], + ) + def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): + """Reject semantic index violations without poisoning pytest's CUDA context.""" + + _cuda_op() + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.kernels.ops.base import _C + from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import H3AdaLNRowGatherCudaOp + + entrypoint, timestep, tag, index_dtype = sys.argv[1:] + rows = torch.zeros((6, 6), dtype=torch.float32, device="cuda") + dtype = getattr(torch, index_dtype) + ti = torch.tensor([int(timestep)], dtype=dtype, device="cuda") + tags = torch.tensor([int(tag)], dtype=dtype, device="cuda") + if entrypoint == "native": + _C.h3_adaln_row_gather_forward(rows, ti, tags, 6, 3) + else: + H3AdaLNRowGatherCudaOp().forward(rows, ti, tags, check_range=False) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(timestep), str(tag), index_dtype], + cwd=Path(__file__).resolve().parents[2], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + assert completed.returncode != 0, "native gather accepted invalid semantic indices" + assert "device-side assert triggered" in completed.stderr, ( + f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + ) + def test_gather_chunks_drop_in(self): rows = _rows(2, hidden=16) ti, tags = h3_packed_layout(10, 2) diff --git a/tests/h3/test_h3_chain_replay_cli.py b/tests/h3/test_h3_chain_replay_cli.py new file mode 100644 index 000000000..32f695af4 --- /dev/null +++ b/tests/h3/test_h3_chain_replay_cli.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Reject incomplete chain dependencies before checking CUDA or loading weights.""" + +from __future__ import annotations + +import sys + +import pytest + +from scripts import h3_chain_replay + +STAGE_NAMES = [stage.name for stage in h3_chain_replay.h3_chain.STAGES] + + +@pytest.fixture +def no_cuda_check(monkeypatch): + """Fail if invalid CLI arguments reach a CUDA availability check.""" + + def unexpected_check(): + """Report an environment check performed before argument rejection.""" + + pytest.fail("invalid arguments must be rejected before checking CUDA") + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", unexpected_check) + + +@pytest.mark.parametrize( + "stages, message", + [ + ("", "unknown stages"), + ("unknown", "unknown stages"), + (STAGE_NAMES[1], "ordered prefix"), + (STAGE_NAMES[-1], "ordered prefix"), + (",".join(STAGE_NAMES[:1] + STAGE_NAMES[2:]), "ordered prefix"), + (",".join(reversed(STAGE_NAMES)), "ordered prefix"), + (",".join([STAGE_NAMES[0], STAGE_NAMES[0]]), "ordered prefix"), + ], +) +def test_reject_invalid_stages(monkeypatch, capsys, no_cuda_check, stages, message): + """Reject unknown stages and selections that break chain dependencies.""" + + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", "--stages", stages]) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert message in capsys.readouterr().err + + +@pytest.mark.parametrize("count", range(1, len(STAGE_NAMES))) +def test_backward_requires_all_stages(monkeypatch, capsys, no_cuda_check, count): + """Require the full chain for parameter-gradient replay.""" + + monkeypatch.setattr( + sys, + "argv", + ["h3_chain_replay.py", "--stages", ",".join(STAGE_NAMES[:count]), "--backward"], + ) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert "--backward requires every stage" in capsys.readouterr().err + + +@pytest.mark.parametrize("count", range(1, len(STAGE_NAMES) + 1)) +def test_accept_ordered_prefix(monkeypatch, count): + """Allow every nonempty prefix to advance to the device requirement.""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr( + sys, "argv", ["h3_chain_replay.py", "--stages", ",".join(STAGE_NAMES[:count])] + ) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() + + +@pytest.mark.parametrize("options", [[], ["--backward"]]) +def test_accept_all_stages_by_default(monkeypatch, options): + """Keep the complete chain as the default for forward and backward replay.""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", *options]) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() From d113891c6b67a7262f08adc1837550a47e5c7e3d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:56:45 +0000 Subject: [PATCH 06/42] ci: enable CodeRabbit automatic reviews for test-h3 Include the RFC 420 target branch in CodeRabbit automatic reviews so pushes to the H3 stack are reviewed instead of silently skipped. Validation: configuration parses and validates against the published CodeRabbit v2 schema; the additional branch pattern matches only test-h3. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .coderabbit.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .coderabbit.yaml diff --git a/.coderabbit.yaml b/.coderabbit.yaml new file mode 100644 index 000000000..cc811057e --- /dev/null +++ b/.coderabbit.yaml @@ -0,0 +1,5 @@ +# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json +reviews: + auto_review: + base_branches: + - "^test-h3$" From 5c79216e34368257f14014bcac133be11fb4c216 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:57:41 +0000 Subject: [PATCH 07/42] ci: inherit existing CodeRabbit review settings Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .coderabbit.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/.coderabbit.yaml b/.coderabbit.yaml index cc811057e..ebea58ae2 100644 --- a/.coderabbit.yaml +++ b/.coderabbit.yaml @@ -1,4 +1,5 @@ # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json +inheritance: true reviews: auto_review: base_branches: From d2a6546cbb424f8201428ed9183fd513bc33f9f5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:21:02 +0000 Subject: [PATCH 08/42] fix(h3): validate RMSNorm inputs and isolate backward benchmarks Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- benchmarks/benchmark_h3_conditioning.py | 13 +- csrc/cuda/h3/rmsnorm_modulate.cu | 43 +++- rl_engine/kernels/ops/pytorch/h3/rmsnorm.py | 6 +- rl_engine/testing/h3_report.py | 120 +++++++--- tests/h3/test_h3_benchmark_timing.py | 211 ++++++++++++++++++ tests/h3/test_h3_rmsnorm.py | 17 +- tests/h3/test_h3_rmsnorm_native_validation.py | 158 +++++++++++++ 7 files changed, 519 insertions(+), 49 deletions(-) create mode 100644 tests/h3/test_h3_benchmark_timing.py create mode 100644 tests/h3/test_h3_rmsnorm_native_validation.py diff --git a/benchmarks/benchmark_h3_conditioning.py b/benchmarks/benchmark_h3_conditioning.py index de1aaceb7..5a39fc175 100644 --- a/benchmarks/benchmark_h3_conditioning.py +++ b/benchmarks/benchmark_h3_conditioning.py @@ -4,7 +4,8 @@ """Benchmark the MiniMax-H3 conditioning ops (RFC #420) against the provider path. CUDA-event medians, bandwidth and peak memory per case, plus the backend -the registry dispatched. Cases live in ``rl_engine/testing/h3_report.py``. +the registry dispatched. Timings alternate candidate/provider execution order; +backward timings exclude forward setup. Cases live in ``rl_engine/testing/h3_report.py``. python benchmarks/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 python benchmarks/benchmark_h3_conditioning.py --op all --json out.json @@ -40,6 +41,7 @@ def main() -> None: registry = KernelRegistry() env = environment() print(f"device={env['gpu']} torch={env['torch']} cuda={env['cuda']}") + print("timing order alternates each iteration; backward timings exclude forward setup") results = [] for name in list(PERF_CASES) if args.op == "all" else [args.op]: for case in PERF_CASES[name](registry): @@ -51,7 +53,14 @@ def main() -> None: for key in TIMED_KEYS if f"{key}_us" in row ] - print(f"{name} {row['case']} [{row['backend']}]: " + "; ".join(parts)) + order = row["execution_order"] + phases = [" -> ".join(order[f"iteration_{i}"]) for i in (0, 1)] + print( + f"{name} {row['case']} [{row['backend']}]: " + + "; ".join(parts) + + "; alternating order: " + + " / ".join(phases) + ) if args.json is not None: args.json.write_text(json.dumps({"environment": env, "results": results}, indent=2) + "\n") diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 5212324b1..4eef07be1 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -255,6 +255,12 @@ void check_rows(const torch::Tensor& t, const char* name, const torch::Tensor& l TORCH_CHECK(t.scalar_type() == like.scalar_type(), name, " must have dtype ", like.scalar_type()); } +void check_index_tensor(const torch::Tensor& t, const char* name, const torch::Tensor& x) { + TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.scalar_type() == at::kLong && + t.dim() == 1 && t.is_contiguous(), + name, " must be a contiguous 1-D int64 CUDA tensor on x's device"); +} + Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, @@ -276,8 +282,8 @@ Modulation make_modulation(const c10::optional& shift, } TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), "shift and scale must be views of the same table layout"); - TORCH_CHECK(ix.is_cuda() && ix.scalar_type() == at::kLong && ix.dim() == 1 && ix.is_contiguous(), - "row index must be a contiguous int64 CUDA tensor"); + check_index_tensor(ix, "row index", x); + TORCH_CHECK(ix.numel() > 0, "row index must be nonempty"); TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", ix.size(0), ")"); mod.shift = sh.data_ptr(); @@ -291,9 +297,11 @@ Modulation make_modulation(const c10::optional& shift, void check_xw(const torch::Tensor& x, const torch::Tensor& w) { TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.is_contiguous(), "x must be a contiguous (M, N) CUDA tensor"); TORCH_CHECK(x.size(0) > 0, "x must have at least one row"); - TORCH_CHECK(w.is_cuda() && w.dim() == 1 && w.size(0) == x.size(1) && w.is_contiguous() && + TORCH_CHECK(w.is_cuda() && w.device() == x.device() && w.dim() == 1 && + w.size(0) == x.size(1) && w.is_contiguous() && w.scalar_type() == x.scalar_type(), - "weight must be a contiguous (N,) CUDA tensor with x's dtype"); + "weight must be a contiguous (N,) CUDA tensor with x's dtype and device"); + TORCH_CHECK(x.size(1) > 0, "x must have at least one column"); TORCH_CHECK(x.size(1) % kVec == 0, "N=", x.size(1), " must be a multiple of ", kVec, " (PyTorch's vectorized RMSNorm path)"); } @@ -344,11 +352,31 @@ std::vector h3_rmsnorm_backward( 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.scalar_type() == at::kFloat && rstd.numel() == x.size(0), - "rstd must be the forward's float32 (M,) statistics"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && + rstd.scalar_type() == at::kFloat && rstd.dim() == 1 && + rstd.is_contiguous() && rstd.size(0) == x.size(0), + "rstd must be contiguous float32 (M,) statistics on x's device"); const int64_t rows = x.size(0); const int64_t n = x.size(1); const Modulation mod = make_modulation(shift, scale, index, x, n); + if (mod.index != nullptr) { + TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && + seg_first_tile.has_value(), + "modulated backward needs the sorted segment tiles"); + check_index_tensor(*sorted_pos, "sorted_pos", x); + check_index_tensor(*tile_begin, "tile_begin", x); + check_index_tensor(*tile_end, "tile_end", x); + check_index_tensor(*seg_first_tile, "seg_first_tile", x); + TORCH_CHECK(sorted_pos->size(0) == rows, "sorted_pos must have M entries"); + TORCH_CHECK(seg_first_tile->size(0) == shift->size(0) + 1, + "seg_first_tile must have R + 1 entries"); + TORCH_CHECK(tile_begin->size(0) == tile_end->size(0), + "tile_begin and tile_end must have the same number of entries"); + } else { + TORCH_CHECK(!sorted_pos.has_value() && !tile_begin.has_value() && !tile_end.has_value() && + !seg_first_tile.has_value(), + "sorted segment tiles require modulation"); + } const c10::cuda::CUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); auto dx = torch::empty_like(x); @@ -375,9 +403,6 @@ std::vector h3_rmsnorm_backward( C10_CUDA_KERNEL_LAUNCH_CHECK(); if (mod.index == nullptr) return {dx, dweight}; - TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && - seg_first_tile.has_value(), - "modulated backward needs the sorted segment tiles"); const int64_t segments = seg_first_tile->numel() - 1; const int64_t seg_tiles = tile_begin->numel(); const int64_t width = 2 * n; diff --git a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py index 6ab850a07..aeee51754 100644 --- a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py +++ b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py @@ -126,8 +126,10 @@ def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: rstd = torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) return (x64 * rstd * weight.double()).float() - def forward_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): - validate_h3_modulation(x, shift, scale, index) + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ): + validate_h3_modulation(x, shift, scale, index, check_range=check_range) n = self.forward(x, weight, eps) return _DeterministicModulation.apply(n, shift, scale, index) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index dfecec419..c8dfcba78 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -36,41 +36,83 @@ ) -def time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 200) -> float: - """Median CUDA-event time of ``fn`` in microseconds (includes the Python wrapper).""" +def _prepared_call( + fn: Callable[..., Any], setup: Callable[[], Any] | None = None +) -> Callable[[], Any]: + if setup is None: + return fn + prepared = setup() + return lambda: fn(prepared) + + +def _sample_us(fn: Callable[..., Any], setup: Callable[[], Any] | None = None) -> float: + run = _prepared_call(fn, setup) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1e3 + + +def time_us( + fn: Callable[..., Any], + warmup: int = 20, + iters: int = 200, + *, + setup: Callable[[], Any] | None = None, +) -> float: + """Median CUDA-event microseconds; optional setup runs outside the timed region. + + When supplied, ``setup`` runs before every call and its return value is passed + to ``fn``. Backward measurements use it to build a fresh forward graph. + """ for _ in range(warmup): - fn() + _prepared_call(fn, setup)() torch.cuda.synchronize() - samples = [] - for _ in range(iters): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - fn() - end.record() - end.synchronize() - samples.append(start.elapsed_time(end) * 1e3) - return statistics.median(samples) - - -def peak_mib(fn: Callable[[], Any]) -> float: + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) + + +def peak_mib(fn: Callable[..., Any], *, setup: Callable[[], Any] | None = None) -> float: + """Incremental peak CUDA memory of the call, excluding optional setup.""" + + run = _prepared_call(fn, setup) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() base = torch.cuda.memory_allocated() - fn() + run() torch.cuda.synchronize() return (torch.cuda.max_memory_allocated() - base) / 2**20 def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + """Interleave timed keys, reversing their execution order on alternate iterations.""" + row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} - for key in TIMED_KEYS: - if key in case: - us = time_us(case[key], warmup, iters) - row[f"{key}_us"] = us - row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 - row[f"{key}_peak_mib"] = peak_mib(case[key]) + keys = [key for key in TIMED_KEYS if key in case] + orders = (keys, list(reversed(keys))) + row["execution_order"] = { + "policy": "alternating", + "iteration_0": orders[0], + "iteration_1": orders[1], + } + if any(key.endswith("_backward") for key in keys): + row["backward_timing_scope"] = "backward_only" + for iteration in range(warmup): + for key in orders[iteration % 2]: + _prepared_call(case[key], case.get(f"{key}_setup"))() + torch.cuda.synchronize() + samples = {key: [] for key in keys} + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) + for key in keys: + us = statistics.median(samples[key]) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) return row @@ -277,9 +319,12 @@ def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) grads = [torch.randn(seq, hidden, device="cuda").bfloat16() for _ in range(6)] - def backward(fn, ti=ti, tags=tags, grads=grads): + def prepare_backward(fn, ti=ti, tags=tags): leaf = rows.detach().requires_grad_(True) - torch.autograd.backward(list(fn(leaf, ti, tags)), grads) + return list(fn(leaf, ti, tags)) + + def backward(outputs, grads=grads): + torch.autograd.backward(outputs, grads) cases.append( { @@ -290,11 +335,15 @@ def backward(fn, ti=ti, tags=tags, grads=grads): "bytes": 6 * seq * hidden * 2, "candidate": lambda ti=ti, tags=tags: op.forward(rows, ti, tags, check_range=False), "provider": lambda ti=ti, tags=tags: provider_adaln_row_gather(chunks, ti, tags), - "candidate_backward": lambda b=backward: b( - lambda r, ti, tags: op.forward(r, ti, tags, check_range=False) + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: op.forward(r, indices, row_tags, check_range=False) ), - "provider_backward": lambda b=backward: b( - lambda r, ti, tags: provider_adaln_row_gather(r.chunk(6, dim=-1), ti, tags) + "provider_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: provider_adaln_row_gather( + r.chunk(6, dim=-1), indices, row_tags + ) ), } ) @@ -374,9 +423,12 @@ def _norm_perf(registry: KernelRegistry) -> list[dict[str, Any]]: x, weight, shift, scale, index = _norm_inputs(seq) grad = torch.randn_like(x) - def backward(fn, x=x, weight=weight, shift=shift, scale=scale, grad=grad): + def prepare_backward(fn, x=x, weight=weight, shift=shift, scale=scale): leaves = [t.detach().requires_grad_(True) for t in (x, weight, shift, scale)] - fn(*leaves).backward(grad) + return fn(*leaves) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -390,12 +442,14 @@ def backward(fn, x=x, weight=weight, shift=shift, scale=scale, grad=grad): "provider": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( provider_norm_modulate(x, w, sh, sc, i) ), - "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 x_, w_, sh_, sc_: op.forward_modulated( x_, w_, sh_, sc_, i, check_range=False ) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda x_, w_, sh_, sc_: provider_norm_modulate(x_, w_, sh_, sc_, i) ), } diff --git a/tests/h3/test_h3_benchmark_timing.py b/tests/h3/test_h3_benchmark_timing.py new file mode 100644 index 000000000..e6517bd70 --- /dev/null +++ b/tests/h3/test_h3_benchmark_timing.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU regressions for benchmark ordering and backward timing boundaries.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.testing import h3_report + + +@pytest.fixture +def cuda_clock(monkeypatch): + clock = SimpleNamespace(time=0.0, active=False, log=[], allocated=0, peak=0, events=0) + + class Event: + def __init__(self, *, enable_timing): + assert enable_timing + self.start = clock.events % 2 == 0 + clock.events += 1 + + def record(self): + clock.active = self.start + clock.log.append("start" if self.start else "end") + self.time = clock.time + + def synchronize(self): + assert not clock.active + + def elapsed_time(self, end): + return end.time - self.time + + def reset_peak(): + clock.log.append("reset_peak") + clock.peak = clock.allocated + + def allocated(): + clock.log.append("memory_baseline") + return clock.allocated + + monkeypatch.setattr(torch.cuda, "Event", Event) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: clock.log.append("synchronize")) + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", reset_peak) + monkeypatch.setattr(torch.cuda, "memory_allocated", allocated) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: clock.peak) + return clock + + +def test_setup_is_outside_backward_events_and_memory_baseline(cuda_clock): + states = [] + + def setup(): + assert not cuda_clock.active + cuda_clock.log.append("setup") + cuda_clock.time += 50.0 + cuda_clock.allocated += 64 * 2**20 + return object() + + def backward(state): + states.append(state) + cuda_clock.log.append("backward") + cuda_clock.time += 2.0 + cuda_clock.allocated += 3 * 2**20 + cuda_clock.peak = cuda_clock.allocated + + assert h3_report.time_us(backward, warmup=1, iters=2, setup=setup) == 2000.0 + assert len({id(state) for state in states}) == 3 + assert cuda_clock.log.count("start") == 2 + for index, entry in enumerate(cuda_clock.log): + if entry == "start": + assert cuda_clock.log[index - 1] == "setup" + assert cuda_clock.log[index + 1 : index + 3] == ["backward", "end"] + + cuda_clock.log.clear() + assert h3_report.peak_mib(backward, setup=setup) == 3.0 + assert cuda_clock.log == [ + "setup", + "synchronize", + "reset_peak", + "memory_baseline", + "backward", + "synchronize", + ] + + +@pytest.mark.parametrize("keys", [("candidate", "provider"), h3_report.TIMED_KEYS]) +def test_measure_interleaves_and_records_both_orders(cuda_clock, keys): + calls = [] + case = {"op": "test", "case": "tiny", "backend": "cpu", "bytes": 4096} + for position, key in enumerate(keys, start=1): + + def run(key=key, duration=position): + calls.append(key) + cuda_clock.time += duration + + case[key] = run + + row = h3_report.measure(case, warmup=2, iters=4) + first, second = list(keys), list(reversed(keys)) + assert calls == (first + second) * 3 + first # warmup, samples, then memory calls + assert row["execution_order"] == { + "policy": "alternating", + "iteration_0": first, + "iteration_1": second, + } + for position, key in enumerate(keys, start=1): + assert row[f"{key}_us"] == position * 1000.0 + assert row[f"{key}_gbps"] == pytest.approx(4096 / (position * 1e-3) / 1e9) + assert row[f"{key}_peak_mib"] == 0.0 + assert not any(key.endswith("_setup") for key in row) + + +@pytest.mark.parametrize("operator", ["norm", "gather"]) +def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( + monkeypatch, cuda_clock, operator +): + prepared_leaves = [] + previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + + def track_backward(outputs): + tensors = outputs if isinstance(outputs, tuple) else (outputs,) + + def backward_hook(grad): + cuda_clock.time += 2.0 / len(tensors) + return grad + + for tensor in tensors: + tensor.register_hook(backward_hook) + return outputs + + def track(fn, leaves, *args): + if leaves[0].requires_grad: + assert not cuda_clock.active + prepared_leaves.append(leaves) + cuda_clock.time += 1.0 + outputs = fn(*leaves, *args) + return track_backward(outputs) if leaves[0].requires_grad else outputs + + if operator == "norm": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 4), + torch.randn(3, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_norm_modulate + + def forward(x, weight, shift, scale, index, **kwargs): + return track(provider, [x, weight, shift, scale], index) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_norm_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) + op = SimpleNamespace(forward_modulated=forward) + build = h3_report._norm_perf + else: + randn = torch.randn + + def cpu_randn(*args, **kwargs): + kwargs["device"] = "cpu" + return randn(*args, **kwargs) + + provider = h3_report.provider_adaln_row_gather + + def forward(rows, ti, tags, **kwargs): + def gather(leaf, ti, tags): + return provider(leaf.chunk(6, dim=-1), ti, tags) + + return track(gather, [rows], ti, tags) + + def provider_forward(chunks, ti, tags): + if chunks[0].requires_grad: + assert not cuda_clock.active + prepared_leaves.append([chunks[0]._base]) + cuda_clock.time += 1.0 + outputs = provider(chunks, ti, tags) + return track_backward(outputs) if chunks[0].requires_grad else outputs + + monkeypatch.setattr(torch, "randn", cpu_randn) + monkeypatch.setattr(h3_report, "GATHER_SEQ_LENS", (2, 3)) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: (torch.arange(seq) % num, torch.arange(seq) % 3), + ) + monkeypatch.setattr(h3_report, "provider_adaln_row_gather", provider_forward) + op = SimpleNamespace(forward=forward) + build = h3_report._gather_perf + + try: + registry = SimpleNamespace(get_op=lambda *args, **kwargs: op) + for case in build(registry): + row = h3_report.measure(case, warmup=1, iters=2) + assert row["backward_timing_scope"] == "backward_only" + assert row["candidate_backward_us"] == pytest.approx(2000.0) + assert row["provider_backward_us"] == pytest.approx(2000.0) + finally: + torch.set_num_threads(previous_threads) + + assert prepared_leaves + assert len({id(leaves[0]) for leaves in prepared_leaves}) == len(prepared_leaves) + assert all(leaf.grad is not None for leaves in prepared_leaves for leaf in leaves) diff --git a/tests/h3/test_h3_rmsnorm.py b/tests/h3/test_h3_rmsnorm.py index f0d013bfb..53e46b16b 100644 --- a/tests/h3/test_h3_rmsnorm.py +++ b/tests/h3/test_h3_rmsnorm.py @@ -37,9 +37,9 @@ def _cuda_op(): return H3RMSNormCudaOp() -def _x(shape, dtype=torch.bfloat16, seed=0, scale=2.0): +def _x(shape, dtype=torch.bfloat16, seed=0, scale=2.0, device="cuda"): g = torch.Generator(device="cpu").manual_seed(seed) - return (torch.randn(shape, generator=g) * scale).to(dtype).cuda() + return (torch.randn(shape, generator=g) * scale).to(device=device, dtype=dtype) def _table(rows, hidden, chunks, dtype=torch.bfloat16, seed=0): @@ -56,10 +56,21 @@ def _fp64_modulated(x, weight, shift, scale, index, eps=1e-5): class TestReference: def test_golden_close_to_provider(self): - x, w = _x((3, 64), torch.float32).cpu(), torch.rand(64) + 0.5 + x, w = _x((3, 64), torch.float32, device="cpu"), torch.rand(64) + 0.5 op = NativeH3RMSNormOp() torch.testing.assert_close(op.forward(x, w), op.forward_fp32(x, w), atol=1e-6, rtol=1e-6) + def test_modulated_with_prevalidated_indices(self): + x, w = _x((2, 8), torch.float32, device="cpu"), torch.ones(8) + shift, scale = torch.randn(3, 8), torch.randn(3, 8) + index = torch.tensor([2, 0]) + op = NativeH3RMSNormOp() + expected = op.forward(x, w) * (1 + scale[index]) + shift[index] + torch.testing.assert_close( + op.forward_modulated(x, w, shift, scale, index, check_range=False), + expected, + ) + def test_rejects_bad_inputs(self): op = NativeH3RMSNormOp() x, w = torch.randn(2, 8), torch.ones(8) diff --git a/tests/h3/test_h3_rmsnorm_native_validation.py b/tests/h3/test_h3_rmsnorm_native_validation.py new file mode 100644 index 000000000..a0d6f4df7 --- /dev/null +++ b/tests/h3/test_h3_rmsnorm_native_validation.py @@ -0,0 +1,158 @@ +# 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 nonempty"): + _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 df519d8e23786a66c0f7afb50fc2afb453e7a984 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:24:11 +0000 Subject: [PATCH 09/42] docs(h3): refresh B200 evidence with backward-only alternating samples Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 18 +++++++++++------- docs/operators/h3-rmsnorm.md | 12 +++++++----- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 1fb114792..f60b2371f 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -70,7 +70,7 @@ Measured on a B200 (torch 2.13.0+cu130): | --- | --- | --- | | Forward vs `index_select` | bitwise equal | | | `d_rows` repeat-bitwise | yes | no | -| `d_rows` correctly rounded from the FP64 sum | 99.99999% | 7.8% | +| `d_rows` correctly rounded from the FP64 sum | 99.99965% | 7.8% | | `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | ### Fused modulation: `H3AdaLNModulationCudaOp` @@ -111,12 +111,14 @@ B200, T = 3, H = 5376, BF16. "write BW" counts the 6·S·H outputs; the table st | S | CUDA forward | provider forward | CUDA backward | provider backward | | --- | --- | --- | --- | --- | -| 4097 | 0.06 ms | 0.09 ms | 1.15 ms | 1.47 ms | -| 32768 | 0.32 ms | 0.59 ms | 2.18 ms | 8.72 ms | -| 131072 | 1.23 ms (6.8 TB/s write) | 2.37 ms | 6.86 ms | 34.9 ms | +| 4097 | 0.07 ms | 0.11 ms | 0.88 ms | 1.28 ms | +| 32768 | 0.33 ms | 0.62 ms | 1.75 ms | 8.04 ms | +| 131072 | 1.17 ms | 2.42 ms | 5.62 ms | 32.45 ms | -The backward's peak memory is about twice the provider's, because autograd stacks the -six output gradients and the per-tile FP32 partials are materialised. +Backward timings and peak memory exclude leaf creation and the forward pass. The CUDA +backward stacks the six output gradients and materialises per-tile FP32 partials, adding +8130 MiB at S = 131072; the provider adds under 2 MiB. Candidate/provider execution order +alternates each iteration and is recorded in the report. ## Evidence @@ -129,7 +131,9 @@ There are two data files: - [`chain_replay.json`](../usage/evidence/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}. -Both were written from a clean tree at commit `fa551c2`. +`report.json` was regenerated from a clean tree at commit `80e4609` with backward-only +timings and alternating execution order. `chain_replay.json` was written from a clean tree +at commit `fa551c2`. ## Tests diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 957670552..3c69706e7 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -79,19 +79,21 @@ B200, block `norm1` + MSA modulation, B = 1, H = 5376, BF16: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.06 ms | 0.11 ms | 1.08 ms | 0.95 ms | -| 32768 | 0.31 ms | 0.75 ms | 2.50 ms | 5.20 ms | -| 131072 | 1.17 ms | 2.90 ms | 7.40 ms | 20.6 ms | +| 4097 | 0.08 ms | 0.12 ms | 1.23 ms | 0.85 ms | +| 32768 | 0.35 ms | 0.76 ms | 2.18 ms | 4.40 ms | +| 131072 | 1.21 ms | 2.90 ms | 6.17 ms | 17.55 ms | The forward is a single pass that never materialises the gathered rows. At small S the backward -is dominated by the fixed cost of the stable sort and tile setup. +is dominated by the fixed cost of the stable sort and tile setup. Backward timings and peak +memory exclude leaf creation and the forward pass. Candidate/provider execution order alternates +each iteration and is recorded in the report. ## Evidence ![h3_rmsnorm on B200: latency and backward accuracy](../usage/evidence/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](../usage/evidence/h3-rmsnorm-b200/report.json), written by -`scripts/h3_evidence.py` from a clean tree at commit `f36aa60`. It also records: +`scripts/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; From e880a103c8dc8e36eff468d22c811a2843a09ae7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:33:12 +0000 Subject: [PATCH 10/42] test(h3): assert adaln_row_gather batch invariance per logical row A position's six gathered rows are bitwise the same whichever other positions share the call (single positions, spans, the tail, a shuffled repacking), and a table row's gradient from the full packing equals the one from a packing of only its own positions, across tile boundaries. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/h3/test_h3_adaln_row_gather.py | 38 +++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/tests/h3/test_h3_adaln_row_gather.py b/tests/h3/test_h3_adaln_row_gather.py index 8f6558034..4bbe0dc10 100644 --- a/tests/h3/test_h3_adaln_row_gather.py +++ b/tests/h3/test_h3_adaln_row_gather.py @@ -8,7 +8,10 @@ * semantic indices fail closed (out-of-range tags or timesteps, RFC probes H2/H3), and flipping a tag or offsetting a timestep picks a different row; * backward is a deterministic FP32 segmented sum: repeat-bitwise and - correctly rounded, unlike the atomic BF16 ``index_select`` backward. + correctly rounded, unlike the atomic BF16 ``index_select`` backward; +* batch invariance: a position's output does not depend on which other + positions share the call, and a row's gradient does not depend on the + positions that reference other rows. """ from __future__ import annotations @@ -219,6 +222,39 @@ def test_repeat_bitwise(self): for _ in range(3): assert torch.equal(self._cuda_grad(rows, ti, tags, grads), first) + def test_forward_rows_invariant_to_batch_size_and_position(self): + rows = _rows(3, hidden=512, seed=5) + ti, tags = h3_packed_layout(4097, 3, seed=5) + full = _cuda_op()(rows, ti, tags) + perm = torch.randperm(4097, generator=torch.Generator().manual_seed(5)).cuda() + for pick in ( + torch.tensor([0]), + torch.tensor([2048]), + torch.arange(100, 357), + torch.arange(4000, 4097), + perm, + ): + pick = pick.cuda() + part = _cuda_op()(rows, ti[pick], tags[pick]) + for a, b in zip(part, full): + assert torch.equal(a, b[pick]) + + def test_backward_row_gradient_independent_of_other_rows_tokens(self): + # Each row's segment is tiled on its own, so a row's gradient from the + # full packing equals the one from a packing of only its own positions + # (same order), even though ~455 positions per row cross tile boundaries. + rows = _rows(3, hidden=512, seed=6) + ti, tags = h3_packed_layout(4097, 3, seed=6) + grads = self._grads(4097, 512, seed=6) + full = self._cuda_grad(rows, ti, tags, grads) + flat = ti * 3 + tags + for r in range(rows.shape[0]): + own = (flat == r).nonzero().squeeze(1) + assert own.numel() > 256 + alone = self._cuda_grad(rows, ti[own], tags[own], [g[own] for g in grads]) + assert torch.equal(alone[r], full[r]) + assert torch.count_nonzero(alone[torch.arange(rows.shape[0], device="cuda") != r]) == 0 + def test_unreferenced_rows_get_zero_grad(self): rows = _rows(3, hidden=16) ti = torch.zeros(7, dtype=torch.long, device="cuda") From 93c692f5479a6f2b54944ad108a2f29b8262af4a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:44:55 +0000 Subject: [PATCH 11/42] fix(h3): store norm(x) and 1 + scale in BF16 in the modulated RMSNorm golden forward_modulated_fp32 computed both in FP64, but the model stores them in x's dtype. 1 + scale is shared by every position of a row, so against the old golden the CUDA d_weight missed the BF16 reduction tolerance from S=1024 and d_scale at S=4097. With both declared straight-through every gradient passes at S=257/1024/4097; the new test fails on the old golden. No evidence changes: the report uses its own FP64 reference. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-rmsnorm.md | 2 +- rl_engine/kernels/ops/pytorch/h3/rmsnorm.py | 16 +++++++--- tests/h3/test_h3_rmsnorm.py | 35 +++++++++++++++++++++ 3 files changed, 48 insertions(+), 5 deletions(-) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 3c69706e7..21c6b5c6f 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -31,7 +31,7 @@ out = op.forward_modulated(x, weight, shift_msa, scale_msa, adaln_indices) | Backend | Wrapper | Native symbols | Status | | --- | --- | --- | --- | | CUDA (SM80+, validated on SM100) | `rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp` | `rl_engine._C.h3_rmsnorm_{forward,backward}` | Bitwise equal to `nn.RMSNorm` and to the diffusers modulation | -| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp` | n/a | `forward*`: provider path; `forward*_fp32`: FP64 golden | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp` | n/a | `forward*`: provider path; `forward*_fp32`: FP64 golden (the modulated one stores `norm(x)` and `1 + scale` in BF16, straight-through, as the model does) | | ROCm | n/a | n/a | Falls back to the PyTorch reference | ## Tensor Contract diff --git a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py index aeee51754..1727fd1d0 100644 --- a/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py +++ b/rl_engine/kernels/ops/pytorch/h3/rmsnorm.py @@ -107,8 +107,9 @@ class NativeH3RMSNormOp: ``forward`` / ``forward_modulated`` are the provider path (``F.rms_norm`` and the eager modulation expression, with a deterministic shift/scale gradient; the raw diffusers path is ``h3_provider.provider_norm_modulate``). ``forward_fp32`` / - ``forward_modulated_fp32`` are the golden: the same math in FP64 with no - intermediate rounding, returned in FP32. + ``forward_modulated_fp32`` are the golden: the same math in FP64, returned in + FP32. The modulated golden rounds, straight-through, only the two values the + model stores in ``x``'s dtype: ``norm(x)`` and ``1 + scale``. """ op_class = "reduction" @@ -139,7 +140,14 @@ def forward_modulated_fp32(self, x, weight, shift, scale, index, eps: float = H3 validate_h3_rmsnorm(x, weight, eps) x64 = x.double() n = x64 * torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) * weight.double() - out = n * (1.0 + scale.double().index_select(0, index)) + shift.double().index_select( - 0, index + # The norm is an x.dtype module and diffusers forms 1 + scale in x.dtype. + # 1 + scale is shared by every position of a row, so leaving it unrounded + # skews d_norm_weight systematically with S; n's rounding enters d_scale + # summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.double().index_select(0, index) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() ) + out = n * one_plus_scale + shift.double().index_select(0, index) return out.float() diff --git a/tests/h3/test_h3_rmsnorm.py b/tests/h3/test_h3_rmsnorm.py index 53e46b16b..bf655e423 100644 --- a/tests/h3/test_h3_rmsnorm.py +++ b/tests/h3/test_h3_rmsnorm.py @@ -189,6 +189,41 @@ def golden(x_, w_, t_): again = self._grads(ours, (x, w, table), grad) assert all(torch.equal(a, b) for a, b in zip(got, again)) + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_modulated_gradients_meet_contract_against_golden(self, seq): + # Against forward_modulated_fp32 before it stored norm(x) and 1 + scale + # in BF16, d_weight missed from S=1024 and d_scale at S=4097. + from rl_engine.kernels.gtest.tolerance import load_contract, resolve_tolerance + + spec = resolve_tolerance( + load_contract(), + judgment="gradient_accuracy", + op_class="reduction", + dtype=torch.bfloat16, + ) + g = torch.Generator(device="cpu").manual_seed(seq) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.1).bfloat16().cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seq) + index = ti * 3 + tags + x = _x((1, seq, 5376), seed=seq) + grad = _x((1, seq, 5376), seed=seq + 1, scale=1.0) + op, native = _cuda_op(), NativeH3RMSNormOp() + + def run(fn, upstream): + def call(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return fn(x_, w_, shift, scale, index) + + return self._grads(call, (x, w, table), upstream) + + got = run(op.forward_modulated, grad) + ref = run(native.forward_modulated_fp32, grad.float()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + torch.testing.assert_close( + g_.float(), r.float(), atol=spec.atol, rtol=spec.rtol, msg=lambda m: f"{name}: {m}" + ) + def test_plain_backward_against_fp64_and_row_local_dx(self): w = (torch.rand(5376, generator=torch.Generator().manual_seed(7)) + 0.5).bfloat16().cuda() x = _x((4, 64, 5376), seed=7) From 9bf060591972ac62235c3829eb899b4261222ffa Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:16:45 +0000 Subject: [PATCH 12/42] docs(h3): adaln_row_gather vs existing implementations (B200, 6c900ae) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 1fb114792..a40b1a03e 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,6 +131,29 @@ There are two data files: Both were written from a clean tree at commit `fa551c2`. +## Existing implementations (RFC #420 reuse rule) + +![adaln_row_gather vs existing implementations](../usage/evidence/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 | +|---|---|---|---| +| diffusers six index_select calls | **no** (param/table grads not repeatable) | 0.0e+00 / 6.0e-02 / 1381 µs | 0.0e+00 / 2.1e-01 / 9753 µs | +| rl-kernel H3AdaLNRowGatherCudaOp | yes | 0.0e+00 / 2.5e-03 / 1106 µs | 0.0e+00 / 3.0e-03 / 2776 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +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`](../usage/evidence/h3-prior-art-b200/adaln_row_gather.json) +was written from a clean tree at `6c900ae` by + +```bash +python scripts/h3_prior_art.py --op adaln_row_gather --out docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 656b64ba6aad69bff68f15ca9fa440b941f77764 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:30:12 +0000 Subject: [PATCH 13/42] docs(h3): norm_modulate vs existing implementations (B200, ee83dec) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-rmsnorm.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 21c6b5c6f..fe58b4b8b 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -99,6 +99,34 @@ The data is in [`report.json`](../usage/evidence/h3-rmsnorm-b200/report.json), w - bitwise equality with diffusers for the modulation; - row invariance. +## Existing implementations (RFC #420 reuse rule) + +![norm_modulate vs existing implementations](../usage/evidence/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 | +|---|---|---|---| +| diffusers composition (F.rms_norm + index_select modulation) | **no** (param/table grads not repeatable) | 8.6e-03 / 4.5e-02 / 872 µs | 7.5e-03 / 1.9e-01 / 4434 µs | +| torch F.rms_norm (no modulation) | yes | 2.1e-03 / 2.1e-03 / 331 µs | 2.0e-03 / 2.8e-03 / 850 µs | +| TE 2.20.2 RMSNorm (no modulation) | yes | 2.1e-03 / 2.1e-03 / 471 µs | 2.0e-03 / 2.8e-03 / 1597 µs | +| Liger 0.8.4 modulated RMSNorm + index_select | **no** (param/table grads not repeatable) | 6.6e-03 / 4.7e-02 / 810 µs | 6.9e-03 / 2.0e-01 / 3911 µs | +| Liger 0.8.4 modulated RMSNorm + rl-kernel row gather | yes | 6.6e-03 / 6.5e-03 / 1824 µs | 6.9e-03 / 8.0e-03 / 4763 µs | +| SGLang 0.5.21 fused_norm_scale_shift (forward only) | yes | 4.4e-03 / — / 85 µs (fwd only) | 4.1e-03 / — / 436 µs (fwd only) | +| rl-kernel H3RMSNormCudaOp.forward_modulated | yes | 8.6e-03 / 4.2e-03 / 1093 µs | 7.5e-03 / 5.0e-03 / 2653 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +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`](../usage/evidence/h3-prior-art-b200/norm_modulate.json) +was written from a clean tree at `ee83dec` by + +```bash +python scripts/h3_prior_art.py --op norm_modulate --out docs/usage/evidence/h3-prior-art-b200/norm_modulate.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/norm_modulate.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 78b8942a5c45d603d663f97323fe587b9a5c8978 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:39:10 +0000 Subject: [PATCH 14/42] docs(h3): point row-gather validation to migrated tests Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 2813e8934..e8a5adf5c 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -62,7 +62,7 @@ read-back; pass `check_range=False` when the indices are already validated. the gradient. `index_select`'s own backward is an atomic scatter-add in the gradient dtype (BF16), which is neither deterministic nor accurate. - **Tolerance class.** The VJP is a segmented reduction, so gtest judges the op as - `reduction`. The forward is asserted bitwise in `tests/h3`. + `reduction`. The forward is asserted bitwise in `tests/models/minimax_h3`. Measured on a B200 (torch 2.13.0+cu130): From b52916ac9517ba5156071b3dc4c783c54573f48c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:14:32 +0000 Subject: [PATCH 15/42] test(h3): skip the CUDA chain backward e2e when the H3 kernels are not built test_chain_backward runs H3AdaLNModulationCudaOp directly, which a KERNEL_ALIGN_USE_FAST_MATH=1 build does not provide; use the shared require_h3_kernels fixture like test_chain_forward. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_conditioning_e2e.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_conditioning_e2e.py b/tests/models/minimax_h3/test_h3_conditioning_e2e.py index 3372ff62b..07cd49776 100644 --- a/tests/models/minimax_h3/test_h3_conditioning_e2e.py +++ b/tests/models/minimax_h3/test_h3_conditioning_e2e.py @@ -65,7 +65,7 @@ def test_chain_forward(require_h3_kernels, registry, chain_weights, num_timestep @pytest.mark.parametrize("num_timesteps, seq_len", [(1, 257), (3, 4097)]) -def test_chain_backward(registry, chain_weights, num_timesteps, seq_len): +def test_chain_backward(require_h3_kernels, registry, chain_weights, num_timesteps, seq_len): """Parameter gradients of the whole chain: deterministic, and FP32-accurate when fused.""" report = h3_chain.run_backward_case( From 62f8009e7ae70381d99fa722c7ddbbe1719ae4dc Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:23:35 +0000 Subject: [PATCH 16/42] test(h3): fix the flaky FP32 table-gradient check in adaln_modulation test_backward_keeps_table_gradient_fp32 drew temb from an unseeded torch.randn and failed about 1 run in 30 with a dW correctly-rounded fraction near 0.9997. The kernel is deterministic (150 repeats bitwise equal); the golden was inconsistent: it rounded an FP64 SiLU to BF16, while the kernel and diffusers round the FP32 SiLU. Near a BF16 tie the two pick different neighbours for 1-2 of the 8064 activations, which shifts a whole dW column (10 of 300 draws fell below 0.9999). Round the golden's activation from F.silu(temb.float()) and seed temb. Over 300 temb draws the fraction is then 0.999981-0.999982 and the d_temb error at most 2.0e-6 (limits 0.9999 and 1e-5); 30 repeated runs pass. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_adaln_modulation.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py index cd6647739..066445fc0 100644 --- a/tests/models/minimax_h3/test_h3_adaln_modulation.py +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -70,7 +70,8 @@ def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): fused, proj, gather = _ops() weight = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.weight"].cuda() bias = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.bias"].cuda() - temb = torch.randn(3, 2688, device="cuda") * 2 + g = torch.Generator(device="cuda").manual_seed(0) + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 ti, tags = h3_packed_layout(4097, 3, seed=1) grads = _grads(4097, 5376, seed=1) ours = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) @@ -84,7 +85,11 @@ def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] act = ref[0] * torch.sigmoid(ref[0]) - act = act + (act.to(torch.bfloat16).double() - act).detach() # identity-VJP cast + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column (about 1 temb draw in 30 dropped below 0.9999 that way). + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # identity-VJP cast table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376) index = ti * 3 + tags torch.autograd.backward( From 7aeb6ba16b34c9b14de3cf5e2cc6c420b84ec505 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:25:43 +0000 Subject: [PATCH 17/42] evidence(h3): adaln_row_gather direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 46347118. 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. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index e8a5adf5c..1d0387183 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -154,6 +154,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(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) + +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 + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png +``` + ## Tests ```bash From cad31b95545347344cf11a6fef6faf949cd72b17 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:29:03 +0000 Subject: [PATCH 18/42] evidence(h3): norm_modulate direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 228e50b9. 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. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-rmsnorm.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 13b585606..7a16f4a57 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -127,6 +127,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(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) + +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 + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json \ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png +``` + ## Tests ```bash From ecb30e25a04177c7f42f6e2fe08a9f4dee4a5216 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:20:24 +0000 Subject: [PATCH 19/42] evidence(h3): regenerate adaln_row_gather report and chain replay after the golden fix Written by tools/validation/models/h3_evidence.py and h3_chain_replay.py (--timesteps 1,2,3,4 --seq-lens 3,257,4097,32768 --backward) from a clean tree at d3b37fe1 on an idle B200. Every CUDA field and pass/fail flag matches the previous files; only timings and the diffusers backward errors (not repeatable run to run) changed. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 1d0387183..c870b6983 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -129,7 +129,10 @@ There are two data files: - [`chain_replay.json`](../../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}. -Both were written from a clean tree at commit `fa551c2`. +Both were written from a clean tree at commit `d3b37fe1` on an otherwise idle B200, regenerated +after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU (as the model +does). Every CUDA field and every pass/fail flag is unchanged; only timings and the diffusers +backward errors, which vary from run to run, moved. ## Existing implementations (RFC #420 reuse rule) From a1e37ef3894a5dc627d70a85c99e79d89ceefea8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:24:25 +0000 Subject: [PATCH 20/42] evidence(h3): regenerate the row-gather report on the newer runner after the golden fix This branch's adaln_row_gather report.json uses the newer evidence runner (backward-only timings, alternating order). Regenerated by tools/validation/models/h3_evidence.py from a clean tree at 14448146 on an idle B200; chain_replay.json comes from the #486 regeneration at d3b37fe1. Every CUDA pass/fail flag is unchanged; the op-backward correctly rounded fraction (unseeded draw) is 99.99931% vs 99.99965% before, and the doc quotes the new value. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 3cfd6adf6..4ffbb0b90 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -70,7 +70,7 @@ Measured on a B200 (torch 2.13.0+cu130): | --- | --- | --- | | Forward vs `index_select` | bitwise equal | | | `d_rows` repeat-bitwise | yes | no | -| `d_rows` correctly rounded from the FP64 sum | 99.99965% | 7.8% | +| `d_rows` correctly rounded from the FP64 sum | 99.99931% | 7.9% | | `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | ### Fused modulation: `H3AdaLNModulationCudaOp` @@ -131,9 +131,12 @@ There are two data files: - [`chain_replay.json`](../../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 `80e4609` with backward-only +`report.json` was regenerated from a clean tree at commit `14448146` with backward-only timings and alternating execution order. `chain_replay.json` was written from a clean tree -at commit `fa551c2`. +at commit `d3b37fe1`. 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. The op-backward check draws unseeded rows and gradients, so its correctly +rounded fraction moves slightly between runs (99.99965% in the previous report). ## Existing implementations (RFC #420 reuse rule) From a6c85451abb1c818b244a3a9dde16f8d283cdcc3 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:26:07 +0000 Subject: [PATCH 21/42] fix(h3): seed the row-gather op-backward check in the evidence report _gather_accuracy drew its table rows and upstream gradients from the global CUDA generator, so the CUDA op-backward correctly rounded fraction moved between runs (99.99965% and 99.99931% in two reports). Draw both from a seeded generator; three runs now give the same CUDA fraction. The diffusers fraction still varies run to run (atomic scatter-add). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 8b17494a2..505a49e56 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -351,7 +351,9 @@ def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("adaln_row_gather", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) - rows = torch.randn(9, 6 * 5376, device="cuda").bfloat16() + # Seeded so the op-backward fractions below are reproducible from this commit. + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(9, 6 * 5376, device="cuda", generator=gen).bfloat16() forward_bitwise = {} for seq in (1, 257, 4097, 32768): ti, tags = h3_packed_layout(seq, 3, seed=seq) @@ -361,7 +363,7 @@ def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. ti, tags = h3_packed_layout(4097, 3, seed=1) - grads = [torch.randn(4097, 5376, device="cuda").bfloat16() for _ in range(6)] + grads = [torch.randn(4097, 5376, device="cuda", generator=gen).bfloat16() for _ in range(6)] def grad_of(fn): leaf = rows.detach().clone().requires_grad_(True) From cbbfaedc8702bb70a6b74d18c8b1c7b0c5c48346 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:29:33 +0000 Subject: [PATCH 22/42] evidence(h3): regenerate the row-gather report with the seeded op-backward check Written by tools/validation/models/h3_evidence.py from a clean tree at 40f4400c on an idle B200. With rows and gradients seeded, the CUDA op-backward correctly rounded fraction is 99.99931% (99.99999% before, which came from the default CUDA generator's state after the runner's earlier draws). Every other CUDA field and pass/fail flag is unchanged. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index c870b6983..9fa8a690e 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -70,7 +70,7 @@ Measured on a B200 (torch 2.13.0+cu130): | --- | --- | --- | | Forward vs `index_select` | bitwise equal | | | `d_rows` repeat-bitwise | yes | no | -| `d_rows` correctly rounded from the FP64 sum | 99.99999% | 7.8% | +| `d_rows` correctly rounded from the FP64 sum | 99.99931% | 7.8% | | `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | ### Fused modulation: `H3AdaLNModulationCudaOp` @@ -129,10 +129,14 @@ There are two data files: - [`chain_replay.json`](../../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}. -Both were written from a clean tree at commit `d3b37fe1` on an otherwise idle B200, regenerated +`report.json` was rewritten from a clean tree at commit `40f4400c`, after the op-backward check +was seeded (it previously took its rows and gradients from the process's default CUDA +generator, so its fraction depended on how many draws the runner made before it; 99.99999% in +the previous report). `chain_replay.json` was written from a clean tree at commit `d3b37fe1`. +Both ran on an otherwise idle B200 and were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU (as the model -does). Every CUDA field and every pass/fail flag is unchanged; only timings and the diffusers -backward errors, which vary from run to run, moved. +does). Apart from the seeded op-backward fraction, every CUDA field and every pass/fail flag +is unchanged; only timings and the diffusers backward errors, which vary from run to run, moved. ## Existing implementations (RFC #420 reuse rule) From fe0350c1cf8015aa4b6b05b6c865adfed9a492af Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:30:04 +0000 Subject: [PATCH 23/42] evidence(h3): regenerate this branch's row-gather report with the seeded check Written by tools/validation/models/h3_evidence.py from a clean tree at f020f7d5 (same code as this merge's parent tree) on an idle B200. The seeded CUDA op-backward fraction is 99.99931%, matching #486's seeded run; every other CUDA field and pass/fail flag is unchanged. The doc no longer says the fraction varies between runs: the old draw was deterministic per runner version, not per run. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 4ffbb0b90..88fee5f3d 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,12 +131,13 @@ There are two data files: - [`chain_replay.json`](../../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 `14448146` with backward-only -timings and alternating execution order. `chain_replay.json` was written from a clean tree -at commit `d3b37fe1`. 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. The op-backward check draws unseeded rows and gradients, so its correctly -rounded fraction moves slightly between runs (99.99965% in the previous report). +`report.json` was regenerated from a clean tree at commit `f020f7d5` 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 +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. ## Existing implementations (RFC #420 reuse rule) From 2553881ba979522645a592706a77dc374c870a79 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:57:53 +0000 Subject: [PATCH 24/42] fix(h3): seed the row-gather timing inputs in the evidence report _gather_perf drew its table rows and upstream gradients from the process's default CUDA generator. Seed them. Only timing inputs change; no reported accuracy field depends on them. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 4c5eaba33..347f3f856 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -317,12 +317,15 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: op = registry.get_op("adaln_row_gather", device="cuda") num_timesteps, hidden = 3, 5376 - rows = torch.randn(3 * num_timesteps, 6 * hidden, device="cuda").bfloat16() + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(3 * num_timesteps, 6 * hidden, device="cuda", generator=gen).bfloat16() chunks = rows.chunk(6, dim=-1) cases = [] for seq in GATHER_SEQ_LENS: ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) - grads = [torch.randn(seq, hidden, device="cuda").bfloat16() for _ in range(6)] + grads = [ + torch.randn(seq, hidden, device="cuda", generator=gen).bfloat16() for _ in range(6) + ] def backward(fn, ti=ti, tags=tags, grads=grads): leaf = rows.detach().requires_grad_(True) From 95186cb701161fab13b592c3b528d42dcb2ad274 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 20:00:41 +0000 Subject: [PATCH 25/42] fix(h3): seed the plain-RMSNorm inputs of the norm evidence check _norm_accuracy drew the inputs of its four plain-norm bitwise checks from the process's default CUDA generator. Draw them from one seeded generator. All four checks stay bitwise equal to F.rms_norm, and the rest of the accuracy section matches the committed report (the CUDA dscale error differs only in the 16th digit, from the FP64 reference's atomic index_select backward), so it needs no regeneration. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index d42942492..1102aa396 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -490,6 +490,7 @@ def backward(output, grad=grad): def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("h3_rmsnorm", device="cuda") weights = load_h3_conditioning_weights("cuda") + gen = torch.Generator(device="cuda").manual_seed(0) plain = {} for name in ( "transformer_blocks.0.norm1.weight", @@ -497,7 +498,7 @@ def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: "token_refiner.final_norm.weight", "norm_out.norm.weight", ): - x = torch.randn(2, 777, 5376, device="cuda").bfloat16() + x = torch.randn(2, 777, 5376, device="cuda", generator=gen).bfloat16() ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) plain[name] = bool(torch.equal(op(x, weights[name]), ref)) x, weight, shift, scale, index = _norm_inputs(4097, seed=1) From 8725b3d25e2b55d7304cfc8758827f48b13c5a82 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 20:11:43 +0000 Subject: [PATCH 26/42] test(h3): let the CPU timing shim accept seeded CUDA generators The gather case of test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events redirects torch.randn to the CPU. _gather_perf now draws from a seeded CUDA generator, which cannot produce a CPU tensor, so the shim swaps it for a CPU generator with the same seed (one per CUDA generator, so the draw sequence continues). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_benchmark_timing.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/models/minimax_h3/test_h3_benchmark_timing.py b/tests/models/minimax_h3/test_h3_benchmark_timing.py index 7dba40be6..95af1ff18 100644 --- a/tests/models/minimax_h3/test_h3_benchmark_timing.py +++ b/tests/models/minimax_h3/test_h3_benchmark_timing.py @@ -164,9 +164,17 @@ def forward(x, weight, shift, scale, index, **kwargs): build = h3_report._norm_perf else: randn = torch.randn + cpu_generators = {} def cpu_randn(*args, **kwargs): kwargs["device"] = "cpu" + # A seeded CUDA generator cannot draw a CPU tensor: continue the same seed on + # one CPU generator per CUDA generator. + gen = kwargs.get("generator") + if gen is not None and gen.device.type != "cpu": + if id(gen) not in cpu_generators: + cpu_generators[id(gen)] = torch.Generator().manual_seed(gen.initial_seed()) + kwargs["generator"] = cpu_generators[id(gen)] return randn(*args, **kwargs) provider = h3_report.provider_adaln_row_gather From e26e369bda5a5ba8d9c2366fb6237fea3d14f9a0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:14 +0000 Subject: [PATCH 27/42] 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-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 4 files changed, 25 insertions(+), 25 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a76cc0417..a9c8832e2 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) | 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 | |---|---|---|---|---| @@ -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 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 @@ -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_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 9fa8a690e..a6cd10d0d 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -120,13 +120,13 @@ six output gradients and the per-tile FP32 partials are materialised. ## 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/pr486/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/pr486/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 rewritten from a clean tree at commit `40f4400c`, after the op-backward check @@ -140,7 +140,7 @@ is unchanged; only timings and the diffusers backward errors, which vary from ru ## 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 | |---|---|---|---| @@ -151,7 +151,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 @@ -168,10 +168,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-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 From 2985fd20e3ce9d9d0bd567716350bbfb38f2e3bd Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:16 +0000 Subject: [PATCH 28/42] 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-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 5 files changed, 31 insertions(+), 31 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a76cc0417..a9c8832e2 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) | 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 | |---|---|---|---|---| @@ -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 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 @@ -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_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-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 7a16f4a57..60f111647 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -90,9 +90,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; @@ -101,7 +101,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 | |---|---|---|---| @@ -117,7 +117,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 @@ -134,10 +134,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 From 1a8e7b26e67861866ce4cfacf35b043e3dbef760 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:56 +0000 Subject: [PATCH 29/42] 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-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 6 +++--- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a9c8832e2..3f79732cd 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 @@ -172,7 +172,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 a6cd10d0d..a6bec87f8 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -129,10 +129,10 @@ 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 rewritten from a clean tree at commit `40f4400c`, after the op-backward check +`report.json` was rewritten from a clean tree at commit `a6c85451`, after the op-backward check was seeded (it previously took its rows and gradients from the process's default CUDA generator, so its fraction depended on how many draws the runner made before it; 99.99999% in -the previous report). `chain_replay.json` was written from a clean tree at commit `d3b37fe1`. +the previous report). `chain_replay.json` was written from a clean tree at commit `61d23ade`. Both ran on an otherwise idle B200 and were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU (as the model does). Apart from the seeded op-backward fraction, every CUDA field and every pass/fail flag @@ -171,7 +171,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-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 \ From 26a7cab3cad361ef47b50059cd10c4abe5f77289 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:57 +0000 Subject: [PATCH 30/42] 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-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a9c8832e2..3f79732cd 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 @@ -172,7 +172,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-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 60f111647..b64c5e44e 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -93,7 +93,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; @@ -137,7 +137,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 \ From 343fa40bd64386ad0bdc11ead4891377666ab4a5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:43:51 +0000 Subject: [PATCH 31/42] fix(h3): use one timing harness across the stacked PRs The stacked PRs carried three versions of measure(). #487-#494 have the most recent one (alternating order, setup outside the timed region, backward-only scope); take that block and its test here so the whole stack agrees and merging the PRs in order does not conflict. Also move the noqa: E402 of prepare_h3_weights.py to the first line of the import. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 91 ++++++++++++++--------- tests/models/minimax_h3/test_h3_report.py | 2 +- 2 files changed, 56 insertions(+), 37 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 347f3f856..e15f837a8 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -35,37 +35,59 @@ ) -def time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 200) -> float: - """Median CUDA-event time of ``fn`` in microseconds (includes the Python wrapper).""" +def _prepared_call( + fn: Callable[..., Any], setup: Callable[[], Any] | None = None +) -> Callable[[], Any]: + if setup is None: + return fn + prepared = setup() + return lambda: fn(prepared) + + +def _sample_us(fn: Callable[..., Any], setup: Callable[[], Any] | None = None) -> float: + run = _prepared_call(fn, setup) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1e3 + + +def time_us( + fn: Callable[..., Any], + warmup: int = 20, + iters: int = 200, + *, + setup: Callable[[], Any] | None = None, +) -> float: + """Median CUDA-event microseconds; optional setup runs outside the timed region. + + When supplied, ``setup`` runs before every call and its return value is passed + to ``fn``. Backward measurements use it to build a fresh forward graph. + """ for _ in range(warmup): - fn() + _prepared_call(fn, setup)() torch.cuda.synchronize() - samples = [] - for _ in range(iters): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - fn() - end.record() - end.synchronize() - samples.append(start.elapsed_time(end) * 1e3) - return statistics.median(samples) + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) -def peak_mib(fn: Callable[[], Any]) -> float: - """Synchronize one CUDA call and report peak allocated memory above baseline in MiB.""" +def peak_mib(fn: Callable[..., Any], *, setup: Callable[[], Any] | None = None) -> float: + """Incremental peak CUDA memory of the call, excluding optional setup.""" + run = _prepared_call(fn, setup) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() base = torch.cuda.memory_allocated() - fn() + run() torch.cuda.synchronize() return (torch.cuda.max_memory_allocated() - base) / 2**20 def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: - """Time case callables in rotating order after warmup and return their evidence. + """Interleave timed keys, reversing their execution order on alternate iterations. Record raw CUDA-event samples, executed orders, median microseconds, decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. @@ -75,32 +97,29 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st keys = [key for key in TIMED_KEYS if key in case] if not keys: return row - for _ in range(warmup): - for key in keys: - case[key]() + orders = (keys, list(reversed(keys))) + row["execution_order"] = { + "policy": "alternating", + "iteration_0": orders[0], + "iteration_1": orders[1], + } + if any(key.endswith("_backward") for key in keys): + row["backward_timing_scope"] = "backward_only" + for iteration in range(warmup): + for key in orders[iteration % 2]: + _prepared_call(case[key], case.get(f"{key}_setup"))() torch.cuda.synchronize() samples = {key: [] for key in keys} - orders = [] - for index in range(iters): - # Rotate the first implementation so drift does not favor one backend. - offset = index % len(keys) - order = keys[offset:] + keys[:offset] - orders.append(order) - for key in order: - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - case[key]() - end.record() - end.synchronize() - samples[key].append(start.elapsed_time(end) * 1e3) + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) row["timing_samples_us"] = samples - row["timing_order"] = orders + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] for key in keys: us = statistics.median(samples[key]) row[f"{key}_us"] = us row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 - row[f"{key}_peak_mib"] = peak_mib(case[key]) + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) return row diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py index 98e6ab714..b79973eb1 100644 --- a/tests/models/minimax_h3/test_h3_report.py +++ b/tests/models/minimax_h3/test_h3_report.py @@ -53,7 +53,7 @@ def provider(): monkeypatch.setattr(h3_report.torch.cuda, "Event", Event) monkeypatch.setattr(h3_report.torch.cuda, "synchronize", lambda: None) - monkeypatch.setattr(h3_report, "peak_mib", lambda fn: 0.0) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) report = h3_report.measure( {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 ) From 9ed5d9c1296828b76b99914f9f0580f280d121c7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:41:03 +0000 Subject: [PATCH 32/42] fix(h3): reject indices on another device than temb, pre-SM80 devices, and an elided gather example - H3AdaLNModulationCudaOp validated the indices against a stand-in allocated on the indices' own device, so the device check was always true; the stand-in now sits on temb's device and a mismatch raises ValueError before the native call (also with check_range=False). - H3AdaLNModulationCudaOp refuses a pre-SM80 device in __init__ like the projection op, so the registry falls back to the PyTorch reference. - The row-gather doc passes all six tensors to gather_chunks instead of an Ellipsis. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 4 +++- .../minimax_h3/adaln_modulation.py | 18 ++++++++++-------- .../minimax_h3/test_h3_adaln_modulation.py | 17 +++++++++++++++++ 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index a6bec87f8..8d0eb0584 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -22,7 +22,9 @@ from rl_engine.runtime.registry import kernel_registry op = kernel_registry.get_op("adaln_row_gather", device="cuda") outs = op(table.view(-1, 6 * 5376), timestep_indices, token_tags) # six (S, 5376) -outs = op.gather_chunks((shift_msa, ..., gate_mlp), timestep_indices, token_tags) # drop-in +outs = op.gather_chunks( + (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp), timestep_indices, token_tags +) # drop-in, six (T*3, 5376) tensors ``` ## Backends diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py index 763dba546..9e7cbaf70 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py @@ -16,8 +16,6 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( _segment_tiles, adaln_row_gather_available, @@ -26,9 +24,12 @@ det_linear_available, det_linear_backward_input, det_linear_backward_weight, + det_linear_device_supported, det_linear_forward, silu_backward_fp32, ) +from rl_engine.backends.extension import _C +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 validate_h3_adaln_projection from rl_engine.reference.minimax_h3.adaln_row_gather import ( @@ -108,12 +109,14 @@ class H3AdaLNModulationCudaOp: def __init__(self) -> None: if not (det_linear_available() and adaln_row_gather_available()): raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_adaln_row_gather_* ops") + if not det_linear_device_supported(): + raise RuntimeError("the BF16 h3 det_linear forward needs SM80 or newer") def __call__(self, temb, weight, bias, timestep_indices, token_tags, *, check_range=True): hidden = validate_h3_adaln_projection(temb, weight, bias) if not temb.is_cuda: raise ValueError("H3AdaLNModulationCudaOp needs CUDA tensors") - _validate_indices(timestep_indices, token_tags, temb.shape[0], check_range) + _validate_indices(timestep_indices, token_tags, temb.shape[0], temb.device, check_range) packed = _H3AdaLNModulationCuda.apply( temb, weight, @@ -125,9 +128,8 @@ def __call__(self, temb, weight, bias, timestep_indices, token_tags, *, check_ra return tuple(packed.unbind(0)) -def _validate_indices(timestep_indices, token_tags, num_timesteps, check_range): - # Reuse the gather's checks against a stand-in of the right row count. - stand_in = torch.empty( - (num_timesteps * H3_MODALITY_NUM, H3_ADALN_CHUNKS), device=timestep_indices.device - ) +def _validate_indices(timestep_indices, token_tags, num_timesteps, device, check_range): + # Reuse the gather's checks against a stand-in of the right row count on temb's device, so + # indices on another device are rejected here. + stand_in = torch.empty((num_timesteps * H3_MODALITY_NUM, H3_ADALN_CHUNKS), device=device) validate_h3_adaln_row_gather(stand_in, timestep_indices, token_tags, check_range=check_range) diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py index 066445fc0..dc27820dd 100644 --- a/tests/models/minimax_h3/test_h3_adaln_modulation.py +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -56,8 +56,25 @@ def _leaf_grads(fn, temb, weight, bias, grads): return [leaf.grad for leaf in leaves] +def test_pre_sm80_device_is_refused_at_construction_so_the_registry_can_fall_back(monkeypatch): + from rl_engine.backends.cuda.model_specific.minimax_h3 import adaln_modulation as module + + monkeypatch.setattr(module, "det_linear_available", lambda: True) + monkeypatch.setattr(module, "adaln_row_gather_available", lambda: True) + monkeypatch.setattr(module, "det_linear_device_supported", lambda: False) + with pytest.raises(RuntimeError, match="SM80"): + module.H3AdaLNModulationCudaOp() + + @requires_cuda class TestFusedModulation: + def test_indices_on_another_device_than_temb_are_rejected(self): + fused, _, _ = _ops() + temb, weight, bias = _params(64, 2688) + ti, tags = h3_packed_layout(1000, 3) + with pytest.raises(ValueError, match="rows are on"): + fused(temb, weight, bias, ti.cpu(), tags.cpu(), check_range=False) + def test_forward_bitwise_equal_to_separate_ops(self): fused, proj, gather = _ops() temb, weight, bias = _params(64, 2688) From 0f67234e5f3dffa6820c8b31801bbe096cce31c0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:06:55 +0000 Subject: [PATCH 33/42] fix(h3): bound-check the RMSNorm modulation row and the row-gather backward metadata - rmsnorm_modulate.cu: Modulation carries the table row count; the forward, dx and dweight kernels assert the selected row is in [0, R) before computing a shift/scale offset, as the row-gather kernel does for check_range=False callers. - adaln_row_gather.cu: h3_adaln_row_gather_backward requires the tile metadata on grad's device and a non-empty seg_first_tile (an empty one made num_rows -1). - docs/operators/h3-rmsnorm.md names the commit that wrote report.json in the current history (d2a6546, recorded as 80e4609 in the report and the figure footer) and the prior-art commit. - Tests: device-side assert for rows -1 and R in forward and backward (subprocess), metadata checks for the gather backward. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 7 +-- csrc/cuda/h3/rmsnorm_modulate.cu | 18 +++++-- docs/operators/h3-rmsnorm.md | 5 +- .../minimax_h3/test_h3_adaln_row_gather.py | 47 ++++++++++++++++- .../test_h3_rmsnorm_native_validation.py | 51 +++++++++++++++++++ 5 files changed, 117 insertions(+), 11 deletions(-) diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..838142b1e 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -207,10 +207,11 @@ 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() && - t->scalar_type() == at::kLong, - "tile metadata must be contiguous 1-D int64 CUDA tensors"); + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->dim() == 1 && + t->is_contiguous() && t->scalar_type() == at::kLong, + "tile metadata must be contiguous 1-D int64 tensors on grad's device"); } + TORCH_CHECK(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"); const int64_t chunks = grad.size(0); const int64_t seq = grad.size(1); diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 4eef07be1..0165a1bb1 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -105,8 +105,17 @@ 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 rows; // R: table rows; an index outside [0, R) is a caller error }; +// Table row of x row `r`. check_range=False callers skip the Python range check, so the native side +// asserts the index before it is used as an offset (as the row-gather kernel does). +__device__ __forceinline__ int64_t modulation_row(const Modulation& mod, int64_t r) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.rows); + return i; +} + template __global__ void __launch_bounds__(32 * kWarps) rmsnorm_modulate_fwd_kernel(const T* __restrict__ x, const T* __restrict__ w, T* __restrict__ y, @@ -118,7 +127,7 @@ __global__ void __launch_bounds__(32 * kWarps) const T* shift = nullptr; const T* scale = nullptr; if (mod.index != nullptr) { - const int64_t i = mod.index[r % mod.seq]; + const int64_t i = modulation_row(mod, r); shift = static_cast(mod.shift) + i * mod.row_stride; scale = static_cast(mod.scale) + i * mod.row_stride; } @@ -156,7 +165,7 @@ __global__ void __launch_bounds__(32 * kWarps) 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 + ? static_cast(mod.scale) + modulation_row(mod, r) * mod.row_stride : nullptr; const int numx = blockDim.x * blockDim.y; const int thrx = threadIdx.x + threadIdx.y * blockDim.x; @@ -193,7 +202,7 @@ __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* 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 + modulation_row(mod, r) * mod.row_stride : nullptr; acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); } @@ -265,7 +274,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; @@ -291,6 +300,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.rows = sh.size(0); return mod; } diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index b64c5e44e..50fd4bf17 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -93,7 +93,8 @@ 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 `d2a6546`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546` (`80e4609` in the report and the figure footer, which predate +the history rewrite of 2026-10-23). It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -118,7 +119,7 @@ forward + backward time on an otherwise idle B200. Batch invariance is bitwise a 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`](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 +was written from a clean tree at `f6ad7a23` (`ee83dec` in the JSON, same rewrite) by ```bash python tools/validation/models/h3_prior_art.py --op norm_modulate --out reports/experiments/h3-prior-art-b200/norm_modulate.json 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 2b7df3803..1a5d548ab 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -149,7 +149,8 @@ def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): """Reject semantic index violations without poisoning pytest's CUDA context.""" _cuda_op() - probe = textwrap.dedent(""" + probe = textwrap.dedent( + """ import sys import torch from rl_engine.backends.extension import _C @@ -167,7 +168,8 @@ def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): else: H3AdaLNRowGatherCudaOp().forward(rows, ti, tags, check_range=False) torch.cuda.synchronize() - """) + """ + ) completed = subprocess.run( [sys.executable, "-c", probe, entrypoint, str(timestep), str(tag), index_dtype], cwd=Path(__file__).resolve().parents[3], @@ -296,3 +298,44 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("adaln_row_gather", device="cuda") assert type(op).__name__ == "H3AdaLNRowGatherCudaOp" + + +def _backward_args(device="cuda"): + kw = {"device": device, "dtype": torch.int64} + return { + "grad": torch.ones((6, 4, 8), device="cuda"), + "sorted_pos": torch.tensor([0, 2, 1, 3], **kw), + "tile_begin": torch.tensor([0, 2], **kw), + "tile_end": torch.tensor([2, 4], **kw), + "seg_first_tile": torch.tensor([0, 1, 2], **kw), + "out_dtype": torch.float32, + } + + +@requires_cuda +class TestNativeBackwardMetadata: + def test_a_valid_call_runs(self): + _cuda_op() + from rl_engine.backends.extension import _C + + assert _C.h3_adaln_row_gather_backward(**_backward_args()).shape == (2, 48) + + def test_empty_segment_table_is_rejected(self): + _cuda_op() + from rl_engine.backends.extension import _C + + args = _backward_args() + args["seg_first_tile"] = args["seg_first_tile"][:0] + with pytest.raises(RuntimeError, match="seg_first_tile"): + _C.h3_adaln_row_gather_backward(**args) + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") + @pytest.mark.parametrize("name", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) + def test_metadata_on_another_device_than_grad_is_rejected(self, name): + _cuda_op() + from rl_engine.backends.extension import _C + + args = _backward_args() + args[name] = args[name].to("cuda:1") + with pytest.raises(RuntimeError, match="grad's device"): + _C.h3_adaln_row_gather_backward(**args) diff --git a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py index ff94c7924..a2b032344 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py @@ -5,6 +5,12 @@ from __future__ import annotations +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest import torch @@ -156,3 +162,48 @@ 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("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}" From f2f51f3b28719097e5e00d1db6fd6e8d6c32dd46 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:14:55 +0000 Subject: [PATCH 34/42] fix(h3): state the position contract of the gather backward and the no-activation projection - _repeat_positions() builds the table row of every flattened position and raises when the row count is not a whole number of repeats of the index (the native forward already requires it); the backwards of the stacked ops use it instead of index.repeat(rows // len(index)), which truncates silently. - The fused modulation passes ACT_NONE to det_linear_forward explicitly, like the projection op. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/adaln_modulation.py | 3 ++- .../minimax_h3/adaln_row_gather.py | 17 ++++++++++++++++- .../minimax_h3/test_h3_adaln_row_gather.py | 17 +++++++++++++++-- 3 files changed, 33 insertions(+), 4 deletions(-) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py index 9e7cbaf70..a856ce616 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py @@ -21,6 +21,7 @@ adaln_row_gather_available, ) from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + ACT_NONE, det_linear_available, det_linear_backward_input, det_linear_backward_weight, @@ -55,7 +56,7 @@ class _H3AdaLNModulationCuda(torch.autograd.Function): def forward(ctx, temb, weight, bias, timestep_indices, token_tags, hidden): temb = temb.contiguous() act = F.silu(temb).to(weight.dtype) - (table,) = det_linear_forward(act, weight, bias) + (table,) = det_linear_forward(act, weight, bias, activation=ACT_NONE) rows = table.view(-1, H3_ADALN_CHUNKS * hidden) packed = _C.h3_adaln_row_gather_forward( rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py index 92835f0cc..2298ad294 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py @@ -7,8 +7,8 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.extension import _C, _EXT_AVAILABLE +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_row_gather import ( h3_adaln_indices, @@ -46,6 +46,21 @@ def _segment_tiles(index: torch.Tensor, num_rows: int): return sorted_pos, tile_begin.contiguous(), tile_end.contiguous(), seg_first_tile +def _repeat_positions(index: torch.Tensor, rows: int) -> torch.Tensor: + """Table row of every flattened position: ``index`` repeated ``rows // len(index)`` times. + + The native forward only accepts a row count that is a multiple of ``len(index)``; the backward + states the same contract instead of letting floor division build a shorter position list. + """ + + if index.numel() == 0 or rows % index.shape[0] != 0: + raise ValueError( + f"{rows} flattened rows are not a whole number of repeats of an index of " + f"length {index.shape[0]}" + ) + return index.repeat(rows // index.shape[0]) + + class _H3AdaLNRowGatherCuda(torch.autograd.Function): @staticmethod def forward(ctx, rows, timestep_indices, token_tags): 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 2b7df3803..a8de11866 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -83,6 +83,17 @@ def test_rejects_bad_rows(self): NativeH3AdaLNRowGatherOp().forward(torch.zeros(3, 25), ti, tags) # not 6 * H +def test_backward_positions_are_a_whole_number_of_index_repeats(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _repeat_positions + + index = torch.tensor([2, 0, 1]) + assert _repeat_positions(index, 6).tolist() == [2, 0, 1, 2, 0, 1] + with pytest.raises(ValueError, match="whole number of repeats"): + _repeat_positions(index, 7) + with pytest.raises(ValueError, match="whole number of repeats"): + _repeat_positions(index[:0], 4) + + @requires_cuda class TestCudaForward: @pytest.mark.parametrize("seq", [1, 2, 3, 257, 4097, 32768]) @@ -149,7 +160,8 @@ def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): """Reject semantic index violations without poisoning pytest's CUDA context.""" _cuda_op() - probe = textwrap.dedent(""" + probe = textwrap.dedent( + """ import sys import torch from rl_engine.backends.extension import _C @@ -167,7 +179,8 @@ def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): else: H3AdaLNRowGatherCudaOp().forward(rows, ti, tags, check_range=False) torch.cuda.synchronize() - """) + """ + ) completed = subprocess.run( [sys.executable, "-c", probe, entrypoint, str(timestep), str(tag), index_dtype], cwd=Path(__file__).resolve().parents[3], From fb72001cdded4d7f17056d4b582c5c456e65ac22 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:19:52 +0000 Subject: [PATCH 35/42] fix(h3): check the position repeat count in the modulated RMSNorm backward Uses _repeat_positions() instead of index.repeat(rows // len(index)), so a row count that is not a whole number of index repeats raises instead of tiling a shorter position list. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../backends/cuda/model_specific/minimax_h3/rmsnorm.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py index a6ed3a28b..5188f3f80 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py @@ -7,9 +7,12 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _repeat_positions, + _segment_tiles, +) from rl_engine.backends.extension import _C, _EXT_AVAILABLE -from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.rmsnorm import ( H3_NORM_EPS, validate_h3_modulation, @@ -46,7 +49,7 @@ def backward(ctx, grad): dx, dweight = _C.h3_rmsnorm_backward(g2, x2, weight.contiguous(), rstd) d_shift = d_scale = None else: - positions = index.repeat(x2.shape[0] // index.shape[0]) + positions = _repeat_positions(index, x2.shape[0]) tiles = _segment_tiles(positions, shift.shape[0]) dx, dweight, d_shift, d_scale = _C.h3_rmsnorm_backward( g2, x2, weight.contiguous(), rstd, shift, scale, index, *tiles From c416a10168892bbd7dd61bacd0d7ede80bd77832 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:45:06 +0000 Subject: [PATCH 36/42] fix(h3): check the row-gather tile metadata in one pre-pass, time the gather backward without the forward, scan rows once adaln_row_gather.cu: a small pre-pass kernel (check_tile_metadata in the new csrc/cuda/h3/tile_metadata.cuh, which the other segment-tile kernels of the stack reuse) checks the tile metadata before anything reads through it: every sorted position is a position of the sequence, every tile lies inside sorted_pos, and the segment boundaries are non-decreasing and inside the tile list (including the case of no tiles with a non-trivial seg_first_tile, where the partial buffer holds one uninitialized row). A stale or short list ends in a device-side assert instead of an out-of-bounds read. The checks are not made inside the partial and fold kernels: an assert in or before their loops changes how the compiler schedules the loads and slowed the native backward by 6-12% (and, inside the loop, by several times); the pre-pass costs about 4 us per call (native backward, idle kernel timing, min of 400 on an idle B200: 105.4 -> 109.2 us at S = 4097, 662.9 -> 665.8 us at S = 32768, 2578.0 -> 2581.3 us at S = 131072). A seg_first_tile with no boundary at all is a TORCH_CHECK, a table with no rows returns an empty gradient instead of launching a zero-sized grid, and the forward and backward check the grid.y limit (65535) for chunks and for the column blocks with a message that names the shape. Tests run corrupted metadata in a child process and expect the device assert, plus the host checks. h3_report: the row-gather backward cases build the forward graph in candidate_backward_setup / provider_backward_setup, so the timed region and the peak memory hold only autograd.backward, which is what backward_timing_scope = backward_only says. A test pins that every perf case with a timed backward has its setup. The numbers in docs/operators/h3-adaln-row-gather.md and the B200/H100 reports predate this and still include the forward pass. The reference backward (_DeterministicRowGather) takes the positions of all 3T rows from one stable sort and one host read (positions_by_row) instead of one torch.nonzero scan and sync per row; the sums are over the same ascending positions, and a test compares the gradient bitwise with the loop it replaces. adaln_modulation: det_linear_backward_weight skips the bias fold when only the weight needs a gradient (a test freezes each input in turn and compares the others bitwise). chain_grads zips with strict=True. The gather_chunks output shape in the doc is (S, 5376), not (3T, 5376). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 17 ++- csrc/cuda/h3/tile_metadata.cuh | 61 ++++++++++ docs/operators/h3-adaln-row-gather.md | 2 +- .../minimax_h3/adaln_modulation.py | 4 +- .../reference/minimax_h3/adaln_row_gather.py | 26 +++- rl_engine/validation/models/h3_chain.py | 4 +- rl_engine/validation/models/h3_report.py | 19 ++- .../minimax_h3/test_h3_adaln_modulation.py | 19 +++ .../minimax_h3/test_h3_adaln_row_gather.py | 112 ++++++++++++++++++ .../minimax_h3/test_h3_report_timing.py | 21 ++++ 10 files changed, 269 insertions(+), 16 deletions(-) create mode 100644 csrc/cuda/h3/tile_metadata.cuh diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..f8c2d3e31 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -28,6 +28,8 @@ #include +#include "tile_metadata.cuh" + namespace { __device__ __forceinline__ float to_float(float v) { return v; } @@ -109,8 +111,10 @@ __global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, const int64_t c = j / hidden; const int64_t h = j - c * hidden; const g_t* g = grad + c * seq * hidden + h; + const int64_t begin = tile_begin[tile]; + const int64_t end = tile_end[tile]; float acc = 0.0f; - for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + for (int64_t p = begin; p < end; ++p) { acc += to_float(g[sorted_pos[p] * hidden]); } partial[tile * width + j] = acc; @@ -151,6 +155,7 @@ torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor time " is not a multiple of ", chunks); TORCH_CHECK(modality_num > 0 && rows.size(0) % modality_num == 0, "rows count ", rows.size(0), " is not a multiple of ", modality_num); + TORCH_CHECK(chunks <= 65535, "chunks=", chunks, " exceeds the grid.y limit of 65535"); const int64_t num_timesteps = rows.size(0) / modality_num; const int64_t seq = timestep_indices.numel(); TORCH_CHECK(seq > 0, "the packed sequence must not be empty"); @@ -216,16 +221,24 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor const int64_t seq = grad.size(1); const int64_t hidden = grad.size(2); const int64_t width = chunks * hidden; + TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); const int64_t num_rows = seg_first_tile.numel() - 1; const int64_t tiles = tile_begin.numel(); TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries"); const c10::cuda::CUDAGuard device_guard(grad.device()); auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype)); + if (num_rows == 0) return out; // a zero-sized grid is not a valid launch auto partial = torch::empty({std::max(tiles, 1), width}, grad.options().dtype(at::kFloat)); auto stream = at::cuda::getCurrentCUDAStream(); const int threads = 256; - const unsigned col_blocks = static_cast((width + threads - 1) / threads); + const int64_t col_blocks64 = (width + threads - 1) / threads; + TORCH_CHECK(col_blocks64 <= 65535, "width ", width, " needs ", col_blocks64, + " column blocks, more than the grid.y limit of 65535"); + const unsigned col_blocks = static_cast(col_blocks64); + h3_tiles::check_tile_metadata(stream, sorted_pos.data_ptr(), seq, + tile_begin.data_ptr(), tile_end.data_ptr(), tiles, + seg_first_tile.data_ptr(), num_rows); if (tiles > 0) { AT_DISPATCH_FLOATING_TYPES_AND2( at::kHalf, at::kBFloat16, grad.scalar_type(), "h3_adaln_row_gather_partial", [&] { diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh new file mode 100644 index 000000000..a959054e2 --- /dev/null +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Segment-tile metadata shared by the H3 backward kernels that sum the rows of a gradient into the +// rows of a small table (adaln_row_gather, rmsnorm modulation, gate residual). +// +// The caller sorts the S positions by table row (`sorted_pos`), cuts every row's run into tiles +// `[tile_begin, tile_end)` over `sorted_pos`, and gives the first tile of each row in +// `seg_first_tile` (R + 1 entries). The partial kernels read gradients through `sorted_pos`, and the +// fold kernels read the partial rows of a segment's tiles, so a stale or short list would read out +// of bounds. +// +// check_tile_metadata() verifies the list on the device before any of that happens. It is a separate +// pass on purpose: an assert in or before the loops of the partial and fold kernels changes how the +// compiler schedules their loads and slowed them by 6-12% (and by several times inside the loop), +// while this pass costs one small launch (about 4 us). +#pragma once + +#include +#include + +#include +#include + +namespace h3_tiles { + +static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ sorted_pos, + int64_t positions, + 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) { + const int64_t first_thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = first_thread; i < positions; i += stride) { + CUDA_KERNEL_ASSERT(sorted_pos[i] >= 0 && sorted_pos[i] < positions); + } + 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] <= positions); + } + for (int64_t i = first_thread; i <= segments; i += stride) { + CUDA_KERNEL_ASSERT(seg_first_tile[i] >= 0 && seg_first_tile[i] <= tiles); + CUDA_KERNEL_ASSERT(i == segments || seg_first_tile[i] <= seg_first_tile[i + 1]); + } +} + +// `positions` entries in sorted_pos (each in [0, positions)), `tiles` tiles inside sorted_pos, and +// `segments` row segments (seg_first_tile has segments + 1 entries). +inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, int64_t positions, + const int64_t* tile_begin, const int64_t* tile_end, int64_t tiles, + const int64_t* seg_first_tile, int64_t segments) { + const int64_t items = std::max({positions, tiles, segments + 1}); + const unsigned blocks = static_cast(std::min((items + 255) / 256, 1024)); + check_tile_metadata_kernel<<>>(sorted_pos, positions, tile_begin, + tile_end, tiles, seg_first_tile, segments); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +} // namespace h3_tiles diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 8d0eb0584..2ce2990c1 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -24,7 +24,7 @@ op = kernel_registry.get_op("adaln_row_gather", device="cuda") outs = op(table.view(-1, 6 * 5376), timestep_indices, token_tags) # six (S, 5376) outs = op.gather_chunks( (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp), timestep_indices, token_tags -) # drop-in, six (T*3, 5376) tensors +) # drop-in: six (3T, 5376) tables in, six (S, 5376) tensors out (S = len(timestep_indices)) ``` ## Backends diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py index a856ce616..4c42eb606 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py @@ -76,7 +76,9 @@ def backward(ctx, grad_packed): d_table = d_rows.view(num_timesteps, width) d_temb = dw = db = None if need_weight or need_bias: - dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + # The bias fold is skipped when only the weight needs a gradient. + dw, *rest = det_linear_backward_weight(d_table, act, weight.dtype, with_bias=need_bias) + db = rest[0] if need_bias else None if need_temb: d_temb = silu_backward_fp32( det_linear_backward_input(d_table, weight, torch.float32), temb diff --git a/rl_engine/reference/minimax_h3/adaln_row_gather.py b/rl_engine/reference/minimax_h3/adaln_row_gather.py index 2cd713d45..937177f1b 100644 --- a/rl_engine/reference/minimax_h3/adaln_row_gather.py +++ b/rl_engine/reference/minimax_h3/adaln_row_gather.py @@ -78,14 +78,31 @@ def h3_adaln_indices(timestep_indices: torch.Tensor, token_tags: torch.Tensor) - return timestep_indices.long() * H3_MODALITY_NUM + token_tags.long() +def positions_by_row(index: torch.Tensor, num_rows: int) -> list[torch.Tensor]: + """Ascending positions of every table row, with one host read for all rows. + + Equal to ``torch.nonzero(index == row).flatten()`` for each ``row``: a stable sort keeps the + positions of a row in ascending order, and the per-row counts split the sorted list. It + replaces one device scan and one host sync per row, which grew with the 3T table rows. + """ + + order = torch.argsort(index, stable=True) + counts = torch.bincount(index, minlength=num_rows).tolist() + positions, start = [], 0 + for count in counts: + positions.append(order[start : start + count]) + start += count + return positions + + class _DeterministicRowGather(torch.autograd.Function): """``index_select`` forward; FP32 per-row segment sums for the backward. ``index_select``'s own backward is an atomic scatter-add in the grad dtype (BF16 for H3), which is neither deterministic nor accurate. There are - only 3T table rows, so the reference sums each row's positions with one - FP32 ``torch.sum`` per row (a fixed reduction for a fixed packing) and - rounds once. + only 3T table rows (T is the handful of distinct timesteps), so the reference sums each row's + positions with one FP32 ``torch.sum`` per row (a fixed reduction for a fixed packing) and + rounds once. The positions of all rows come from one sort and one host read. """ @staticmethod @@ -102,8 +119,7 @@ def backward(ctx, grad): num_rows, dtype = ctx.rows_meta grad32 = grad.float() # (6, S, H) out = grad32.new_zeros((num_rows, grad32.shape[0], grad32.shape[2])) - for row in range(num_rows): - positions = torch.nonzero(index == row).flatten() + for row, positions in enumerate(positions_by_row(index, num_rows)): if positions.numel(): out[row] = grad32.index_select(1, positions).sum(dim=1) return out.reshape(num_rows, -1).to(dtype), None diff --git a/rl_engine/validation/models/h3_chain.py b/rl_engine/validation/models/h3_chain.py index 8927c18dc..7b53eb56f 100644 --- a/rl_engine/validation/models/h3_chain.py +++ b/rl_engine/validation/models/h3_chain.py @@ -258,7 +258,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/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 6d4d82a59..8f8ba5501 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -354,9 +354,12 @@ def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: torch.randn(seq, hidden, device="cuda", generator=gen).bfloat16() for _ in range(6) ] - def backward(fn, ti=ti, tags=tags, grads=grads): + def prepare_backward(fn, ti=ti, tags=tags): leaf = rows.detach().requires_grad_(True) - torch.autograd.backward(list(fn(leaf, ti, tags)), grads) + return list(fn(leaf, ti, tags)) + + def backward(outputs, grads=grads): + torch.autograd.backward(outputs, grads) cases.append( { @@ -367,11 +370,15 @@ def backward(fn, ti=ti, tags=tags, grads=grads): "bytes": 6 * seq * hidden * 2, "candidate": lambda ti=ti, tags=tags: op.forward(rows, ti, tags, check_range=False), "provider": lambda ti=ti, tags=tags: provider_adaln_row_gather(chunks, ti, tags), - "candidate_backward": lambda b=backward: b( - lambda r, ti, tags: op.forward(r, ti, tags, check_range=False) + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: op.forward(r, indices, row_tags, check_range=False) ), - "provider_backward": lambda b=backward: b( - lambda r, ti, tags: provider_adaln_row_gather(r.chunk(6, dim=-1), ti, tags) + "provider_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: provider_adaln_row_gather( + r.chunk(6, dim=-1), indices, row_tags + ) ), } ) diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py index dc27820dd..9adc6b4aa 100644 --- a/tests/models/minimax_h3/test_h3_adaln_modulation.py +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -124,6 +124,25 @@ def rel(grad, golden): for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad)): assert (grad == golden.to(torch.bfloat16)).float().mean() > 0.9999 + @pytest.mark.parametrize("frozen", ["bias", "weight", "temb", "bias_and_temb"]) + def test_frozen_inputs_get_no_gradient_and_the_others_are_unchanged(self, frozen): + fused, _, _ = _ops() + temb, weight, bias = _params(64, 2688, seed=5) + ti, tags = h3_packed_layout(1500, 3, seed=5) + grads = _grads(1500, 64, seed=5) + full = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + + leaves = [t.detach().clone() for t in (temb, weight, bias)] + names = ("temb", "weight", "bias") + for leaf, name in zip(leaves, names): + leaf.requires_grad_(name not in frozen) + torch.autograd.backward(list(fused(*leaves, ti, tags)), grads) + for leaf, name, expected in zip(leaves, names, full): + if name in frozen: + assert leaf.grad is None + else: + assert torch.equal(leaf.grad, expected) + def test_backward_repeat_bitwise(self): fused, _, _ = _ops() temb, weight, bias = _params(64, 2688, seed=3) 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 a8de11866..ff9662f75 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -94,6 +94,43 @@ def test_backward_positions_are_a_whole_number_of_index_repeats(): _repeat_positions(index[:0], 4) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_positions_by_row_equal_the_per_row_scan(device): + from rl_engine.reference.minimax_h3.adaln_row_gather import positions_by_row + + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("needs a CUDA device") + g = torch.Generator(device="cpu").manual_seed(0) + index = torch.randint(0, 7, (500,), generator=g).to(device) # rows 7..11 are never referenced + positions = positions_by_row(index, 12) + assert len(positions) == 12 + for row, found in enumerate(positions): + assert torch.equal(found, torch.nonzero(index == row).flatten()) + + +@requires_cuda +def test_reference_backward_is_bitwise_the_per_row_scan_it_replaced(): + from rl_engine.reference.minimax_h3.adaln_row_gather import _DeterministicRowGather + + rows = _rows(3, hidden=64, seed=7) + ti, tags = h3_packed_layout(777, 3, seed=7) + index = ti.long() * 3 + tags.long() + gen = torch.Generator(device="cuda").manual_seed(7) + grad = torch.randn(6, 777, 64, device="cuda", generator=gen).bfloat16() + + leaf = rows.detach().clone().requires_grad_(True) + outputs = _DeterministicRowGather.apply(leaf, index) + outputs.backward(grad) + + grad32 = grad.float() + expected = grad32.new_zeros((rows.shape[0], 6, 64)) + for row in range(rows.shape[0]): # the loop of the earlier implementation + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + expected[row] = grad32.index_select(1, positions).sum(dim=1) + assert torch.equal(leaf.grad, expected.reshape(rows.shape[0], -1).to(rows.dtype)) + + @requires_cuda class TestCudaForward: @pytest.mark.parametrize("seq", [1, 2, 3, 257, 4097, 32768]) @@ -303,6 +340,81 @@ def ref_grad(): first.float(), self._cuda_grad(rows, ti, tags, grads).float(), atol=1e-1, rtol=2e-2 ) + @staticmethod + def _backward_probe(mutation: str): + """Run the native backward on corrupted tile metadata in a child process.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + + mutation = sys.argv[1] + seq, hidden = 40, 8 + index = (torch.arange(seq, device="cuda") % 3) + grad = torch.ones(6, seq, hidden, device="cuda") + sorted_pos, begin, end, first = _segment_tiles(index, 3) + if mutation == "tile_end_past_seq": + end = end.clone(); end[-1] = seq + 5 + elif mutation == "position_past_seq": + sorted_pos = sorted_pos.clone(); sorted_pos[0] = seq + elif mutation == "segments_past_tiles": + first = first.clone(); first[-1] = begin.numel() + 3 + elif mutation == "no_tiles": + begin = begin[:0]; end = end[:0] + _C.h3_adaln_row_gather_backward(grad, sorted_pos, begin, end, first, torch.float32) + 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_metadata_that_does_not_describe_the_tiles(self): + _cuda_op() + control = self._backward_probe("none") + assert control.returncode == 0, control.stderr + for mutation in ( + "tile_end_past_seq", + "position_past_seq", + "segments_past_tiles", + "no_tiles", + ): + done = self._backward_probe(mutation) + 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): + from rl_engine.backends.extension import _C + + _cuda_op() + grad = torch.ones(6, 4, 8, device="cuda") + 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="R \\+ 1 entries"): + _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, zero, torch.float32) + + def test_forward_rejects_more_chunks_than_one_grid_dimension_holds(self): + from rl_engine.backends.extension import _C + + _cuda_op() + rows = torch.zeros(3, 65536, device="cuda") + ti = torch.zeros(1, dtype=torch.int64, device="cuda") + with pytest.raises(RuntimeError, match="grid.y"): + _C.h3_adaln_row_gather_forward(rows, ti, ti, 65536, 3) + def test_registry_dispatches_cuda(self): from rl_engine.runtime.registry import KernelRegistry diff --git a/tests/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py index 853680a67..576a96f1e 100644 --- a/tests/models/minimax_h3/test_h3_report_timing.py +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -52,3 +52,24 @@ def test_the_reported_medians_come_from_the_stored_samples(): assert row[f"{key}_peak_mib"] >= 0 # the larger matmul is slower than the smaller one in the same run assert row["provider_us"] > row["candidate_us"] + + +def test_backward_timings_of_every_perf_case_exclude_the_forward_pass(monkeypatch): + """A case with a timed backward must prepare the graph outside the CUDA-event window.""" + + from rl_engine.runtime.registry import KernelRegistry + from rl_engine.validation.models import h3_report + + # The structure is what is checked; the cases need not be as large as the evidence runs. + for name in ("GATHER_SEQ_LENS", "NORM_SEQ_LENS"): + if hasattr(h3_report, name): + monkeypatch.setattr(h3_report, name, (257,)) + registry = KernelRegistry() + checked = 0 + for op, build in h3_report.PERF_CASES.items(): + for case in build(registry): + for side in ("candidate", "provider"): + if f"{side}_backward" in case: + checked += 1 + assert f"{side}_backward_setup" in case, f"{op} {case['case']} {side}" + assert checked From 3e4c84915a3042647b2f7fadc7636184ffba843b Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:53:08 +0000 Subject: [PATCH 37/42] fix(h3): check the rmsnorm table layout and tile metadata, scan rows once, and align the reference validators rmsnorm_modulate.cu: make_modulation rejects a shift/scale table whose row stride is smaller than N (an expanded or overlapping view is not a table of rows; a view of a wider table, such as table.chunk(2, dim=-1) with row stride 2N, is still read by its row pitch and a test pins that it is bitwise the copy). The modulated backward runs the shared tile-metadata pre-pass (tile_metadata.cuh, #486) before the segment kernels: sorted positions, tiles and segment boundaries that do not describe the rows end in a device-side assert instead of an out-of-bounds read, and the fold launch is skipped for zero segments. N is checked against the int and grid.y limits of the segment kernels. The column indices of the segment kernels were already int64 in the tile term, so nothing changes there. Tests run corrupted metadata in a child process. rmsnorm.py (CUDA op): a table that is not made of dense rows (stride(0) < H) is copied before the call, like mixed layouts; a wider-table view goes through as it is. Reference: the modulation backward takes the positions of all rows from positions_by_row (one sort and one host read instead of two torch.nonzero scans per row; a test compares the gradient bitwise with the loop it replaces). forward_modulated_fp32 takes check_range like forward_modulated. A device mismatch of shift/scale is a ValueError, as in the other H3 validators; a dtype mismatch stays a TypeError. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/rmsnorm_modulate.cu | 25 ++++++- .../cuda/model_specific/minimax_h3/rmsnorm.py | 11 ++- rl_engine/reference/minimax_h3/rmsnorm.py | 17 +++-- tests/models/minimax_h3/test_h3_rmsnorm.py | 50 +++++++++++++ .../test_h3_rmsnorm_native_validation.py | 73 +++++++++++++++++++ 5 files changed, 165 insertions(+), 11 deletions(-) diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 0165a1bb1..3f352f417 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -29,8 +29,11 @@ #include #include +#include #include +#include "tile_metadata.cuh" + namespace { constexpr int kVec = 4; // PyTorch's layer-norm vec_size @@ -291,6 +294,10 @@ Modulation make_modulation(const c10::optional& shift, } TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), "shift and scale must be views of the same table layout"); + // Rows are read as dense N-element spans at `row_stride`: a view of a wider table is fine, an + // expanded or overlapping one (row_stride < N) is not a table of rows. + TORCH_CHECK(sh.size(0) <= 1 || sh.stride(0) >= n, "shift/scale row stride ", sh.stride(0), + " is smaller than N=", n); check_index_tensor(ix, "row index", x); TORCH_CHECK(ix.numel() > 0, "row index must be nonempty"); TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", @@ -416,9 +423,16 @@ std::vector h3_rmsnorm_backward( const int64_t segments = seg_first_tile->numel() - 1; const int64_t seg_tiles = tile_begin->numel(); const int64_t width = 2 * n; + TORCH_CHECK(width <= std::numeric_limits::max(), "N=", n, " is too large for the kernels"); + TORCH_CHECK((width + threads - 1) / threads <= 65535, "N=", n, + " needs more column blocks than the grid.y limit of 65535"); auto seg_partial = torch::empty({std::max(seg_tiles, 1), width}, x.options().dtype(at::kFloat)); auto table_grad = torch::empty({segments, width}, x.options().dtype(at::kFloat)); const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + // The tile list is caller-built metadata: check it before the kernels read through it. + h3_tiles::check_tile_metadata(stream, sorted_pos->data_ptr(), rows, + tile_begin->data_ptr(), tile_end->data_ptr(), + seg_tiles, seg_first_tile->data_ptr(), segments); if (seg_tiles > 0) { H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_modulation_grad", [&] { modulation_grad_partial_kernel @@ -430,9 +444,12 @@ std::vector h3_rmsnorm_backward( }); C10_CUDA_KERNEL_LAUNCH_CHECK(); } - 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 + 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 {dx, dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; } diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py index 5188f3f80..61e4a17b6 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py @@ -98,7 +98,16 @@ def forward_modulated( validate_h3_modulation(x, shift, scale, index, check_range=check_range) if not x.is_cuda: raise ValueError("H3RMSNormCudaOp needs CUDA tensors") - if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + # A view of a wider table (``table.chunk(2, dim=-1)``: unit column stride, row stride 2H) + # goes to the kernel as it is, which indexes rows by the stride. Anything that is not a + # table of dense rows (mixed layouts, or rows that overlap because the stride is below + # H) is copied first. + if ( + shift.stride(1) != 1 + or scale.stride(1) != 1 + or shift.stride(0) != scale.stride(0) + or (shift.shape[0] > 1 and shift.stride(0) < shift.shape[1]) + ): shift, scale = shift.contiguous(), scale.contiguous() return _H3RMSNormCuda.apply(x, weight, shift, scale, index.contiguous(), eps) diff --git a/rl_engine/reference/minimax_h3/rmsnorm.py b/rl_engine/reference/minimax_h3/rmsnorm.py index 1727fd1d0..ddbdaf6bf 100644 --- a/rl_engine/reference/minimax_h3/rmsnorm.py +++ b/rl_engine/reference/minimax_h3/rmsnorm.py @@ -21,6 +21,8 @@ import torch import torch.nn.functional as F +from rl_engine.reference.minimax_h3.adaln_row_gather import positions_by_row + H3_NORM_EPS = 1e-5 _DTYPES = (torch.float32, torch.bfloat16, torch.float16) @@ -55,8 +57,10 @@ def validate_h3_modulation( ) -> None: hidden = x.shape[-1] for name, tensor in (("shift", shift), ("scale", scale)): - if tensor.device != x.device or (same_dtype and tensor.dtype != x.dtype): - raise TypeError(f"{name} must match x's dtype and device") + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if same_dtype and tensor.dtype != x.dtype: + raise TypeError(f"{name} must match x's dtype, got {tensor.dtype} for {x.dtype}") if tensor.dim() != 2 or tensor.shape[1] != hidden: raise ValueError(f"{name} must be (R, {hidden}), got {tuple(tensor.shape)}") if shift.shape != scale.shape: @@ -93,8 +97,7 @@ def backward(ctx, grad): per_scale = (g * n.float().reshape(-1, index.shape[0], hidden)).sum(dim=0) d_shift = per_shift.new_zeros((ctx.rows, hidden)) d_scale = per_scale.new_zeros((ctx.rows, hidden)) - for row in range(ctx.rows): - positions = torch.nonzero(index == row).flatten() + for row, positions in enumerate(positions_by_row(index, ctx.rows)): if positions.numel(): d_shift[row] = per_shift.index_select(0, positions).sum(dim=0) d_scale[row] = per_scale.index_select(0, positions).sum(dim=0) @@ -134,9 +137,11 @@ def forward_modulated( n = self.forward(x, weight, eps) return _DeterministicModulation.apply(n, shift, scale, index) - def forward_modulated_fp32(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + def forward_modulated_fp32( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ): # The golden may be fed FP32 golden modulation rows for a BF16 x. - validate_h3_modulation(x, shift, scale, index, same_dtype=False) + validate_h3_modulation(x, shift, scale, index, check_range=check_range, same_dtype=False) validate_h3_rmsnorm(x, weight, eps) x64 = x.double() n = x64 * torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) * weight.double() diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py index 826b75700..6d91f60bb 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm.py @@ -74,6 +74,27 @@ def test_modulated_with_prevalidated_indices(self): expected, ) + def test_modulation_backward_is_bitwise_the_per_row_scan_it_replaced(self): + from rl_engine.reference.minimax_h3.rmsnorm import _DeterministicModulation + + g = torch.Generator(device="cpu").manual_seed(11) + n = torch.randn(300, 16, generator=g) + shift, scale = torch.randn(7, 16, generator=g), torch.randn(7, 16, generator=g) + index = torch.randint(0, 5, (300,), generator=g) # rows 5 and 6 are never referenced + grad = torch.randn(300, 16, generator=g) + leaves = [t.clone().requires_grad_(True) for t in (n, shift, scale)] + _DeterministicModulation.apply(*leaves, index).backward(grad) + + per_shift = grad.reshape(-1, 300, 16).sum(dim=0) + per_scale = (grad * n).reshape(-1, 300, 16).sum(dim=0) + d_shift, d_scale = torch.zeros(7, 16), torch.zeros(7, 16) + for row in range(7): # the loop of the earlier implementation + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_shift[row] = per_shift.index_select(0, positions).sum(dim=0) + d_scale[row] = per_scale.index_select(0, positions).sum(dim=0) + assert torch.equal(leaves[1].grad, d_shift) and torch.equal(leaves[2].grad, d_scale) + def test_rejects_bad_inputs(self): op = NativeH3RMSNormOp() x, w = torch.randn(2, 8), torch.ones(8) @@ -88,6 +109,13 @@ def test_rejects_bad_inputs(self): shift, scale = torch.zeros(3, 8), torch.zeros(3, 8) with pytest.raises(IndexError): op.forward_modulated(x, w, shift, scale, torch.tensor([0, 3])) + with pytest.raises(IndexError): # the golden takes check_range like the provider path + op.forward_modulated_fp32(x, w, shift, scale, torch.tensor([0, 3])) + op.forward_modulated_fp32(x, w, shift, scale, torch.tensor([0, 2]), check_range=False) + with pytest.raises(ValueError, match="is on meta"): # a device mismatch is a ValueError + op.forward_modulated(x, w, shift.to("meta"), scale, torch.tensor([0, 1])) + with pytest.raises(TypeError, match="dtype"): # a dtype mismatch stays a TypeError + op.forward_modulated(x, w, shift.double(), scale, torch.tensor([0, 1])) with pytest.raises(ValueError): # one index per position S op.forward_modulated(x, w, shift, scale, torch.tensor([0])) @@ -139,6 +167,28 @@ def test_modulation_bitwise_in_every_dtype(self, dtype): ours = _cuda_op().forward_modulated(x, w, shift, scale, index) assert torch.equal(ours, provider_norm_modulate(x, w, shift, scale, index)) + def test_overlapping_table_rows_are_copied_and_wider_table_views_are_not(self): + w = (torch.rand(5376, device="cuda") + 0.5).bfloat16() + table = _x((9, 2 * 5376), seed=21, scale=0.5) + wide_shift, wide_scale = table.chunk(2, dim=-1) # row stride 2H, handed over as views + ti, tags = h3_packed_layout(100, 3, seed=21) + index = ti * 3 + tags + x = _x((1, 100, 5376), seed=21) + op = _cuda_op() + expected = op.forward_modulated( + x, w, wide_shift.contiguous(), wide_scale.contiguous(), index + ) + assert torch.equal(op.forward_modulated(x, w, wide_shift, wide_scale, index), expected) + # An expanded table (row stride 0) is a table of equal rows: copied, not rejected. + row = _x((1, 5376), seed=22, scale=0.5) + expanded = row.expand(9, 5376) + assert expanded.stride(0) == 0 + dense = row.repeat(9, 1) + assert torch.equal( + op.forward_modulated(x, w, expanded, expanded, index), + op.forward_modulated(x, w, dense, dense, index), + ) + def test_rows_are_batch_and_position_invariant(self): w = (torch.rand(5376, device="cuda") + 0.5).bfloat16() shift, scale, *_ = _table(9, 5376, 6, seed=4) diff --git a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py index a2b032344..dda99680f 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py @@ -207,3 +207,76 @@ def test_modulation_row_outside_the_table_asserts_in_the_kernel(native, entrypoi 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 86d6f11c42af16532bd7644721133f0c3bf2595c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:58:48 +0000 Subject: [PATCH 38/42] docs(h3): row-gather backward timings measured without the forward pass The B200 report and figure are regenerated from a clean tree at c416a101 on an otherwise idle GPU (utilization 0% before and after) with the backward cases built by candidate_backward_setup / provider_backward_setup. The backward columns of the latency table are now backward-only: 1.01 vs 1.30 ms at S = 4097, 1.76 vs 8.00 ms at S = 32768, 5.59 vs 32.2 ms at S = 131072 (CUDA vs index_select); the peak-memory sentence describes the backward-only increment. The H100 report and chain_replay.json are not regenerated and still carry the earlier timing. Links point at the pr_assets commit that holds the new files. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-row-gather.md | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 2ce2990c1..93c66a4a5 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -113,25 +113,33 @@ B200, T = 3, H = 5376, BF16. "write BW" counts the 6·S·H outputs; the table st | S | CUDA forward | provider forward | CUDA backward | provider backward | | --- | --- | --- | --- | --- | -| 4097 | 0.06 ms | 0.09 ms | 1.15 ms | 1.47 ms | -| 32768 | 0.32 ms | 0.59 ms | 2.18 ms | 8.72 ms | -| 131072 | 1.23 ms (6.8 TB/s write) | 2.37 ms | 6.86 ms | 34.9 ms | +| 4097 | 0.07 ms | 0.11 ms | 1.01 ms | 1.30 ms | +| 32768 | 0.34 ms | 0.63 ms | 1.76 ms | 8.00 ms | +| 131072 | 1.24 ms (6.8 TB/s write) | 2.40 ms | 5.59 ms | 32.2 ms | -The backward's peak memory is about twice the provider's, because autograd stacks the -six output gradients and the per-tile FP32 partials are materialised. +The backward columns time `autograd.backward` only: the forward pass and the graph construction run +outside the timed region (`candidate_backward_setup` / `provider_backward_setup`), and the execution +order alternates and is recorded. + +Backward-only peak memory (increment over the forward graph): 255, 2034 and 8130 MiB for the CUDA op +at S = 4097, 32768 and 131072, against about 1 MiB for the provider. Autograd stacks the six +output gradients into one `(6, S, H)` tensor for the CUDA op and the per-tile FP32 partials are +materialised, while the provider's atomic scatter-add writes straight into the `(3T, 6H)` table +gradient. The peak is the price of the deterministic FP32 segment sum. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/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/999c0c75078605229ad74dc600dc963d89c7ed9e/pr486/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/999c0c75078605229ad74dc600dc963d89c7ed9e/pr486/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`](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 rewritten from a clean tree at commit `a6c85451`, after the op-backward check +`report.json` was rewritten from a clean tree at commit `c416a101` on an otherwise idle B200 (backward +timings exclude the forward pass; the previous revision came from `a6c85451`), after the op-backward check was seeded (it previously took its rows and gradients from the process's default CUDA generator, so its fraction depended on how many draws the runner made before it; 99.99999% in the previous report). `chain_replay.json` was written from a clean tree at commit `61d23ade`. From 65e86208d96b1d8a89fdc0d0c67c0076461f027f Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:21:18 +0000 Subject: [PATCH 39/42] fix(h3): check that the gather segments cut the tile list exactly, and show the device asserts survive a release build * The tile-metadata pre-pass also requires `seg_first_tile[0] == 0` and `seg_first_tile[R] == tiles`. The old bound (`<= tiles`) kept every read inside `partial`, but a tile that no segment owns, or a first segment that starts after tile 0, would drop contributions from a row's sum without an error. Two more corrupted-metadata cases run in the child-process test. The launch uses at least one block, so an empty list is a no-op instead of an invalid configuration. * `CUDA_KERNEL_ASSERT` is not removed by NDEBUG on CUDA (c10 expands it to `if (!cond) __assert_fail(...)`; only Android, Apple and FreeBSD compile it out). The header says so, and `test_h3_assert_in_release_build.py` compiles a one-kernel program with `-O3 -DNDEBUG` against the torch headers in use and requires the "device-side assert triggered" stop (and a clean run for a valid value) so the claim is checked for the torch this runs against. * `test_backward_writes_the_requested_dtype_whatever_the_gradient_dtype` pins that the gather backward allocates the table gradient in `out_dtype` for FP32, BF16 and FP16 gradients (exact integer data, torch.equal against the FP64 sum). * `_C.pyi`: contracts for `h3_adaln_row_gather_forward` / `_backward`; the modulation tests zip strictly (Ruff B905). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/tile_metadata.cuh | 13 ++- rl_engine/_C.pyi | 24 ++++- .../minimax_h3/test_h3_adaln_modulation.py | 10 +- .../minimax_h3/test_h3_adaln_row_gather.py | 33 +++++++ .../test_h3_assert_in_release_build.py | 91 +++++++++++++++++++ 5 files changed, 163 insertions(+), 8 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_assert_in_release_build.py diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index a959054e2..96318225d 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -14,6 +14,12 @@ // pass on purpose: an assert in or before the loops of the partial and fold kernels changes how the // compiler schedules their loads and slowed them by 6-12% (and by several times inside the loop), // while this pass costs one small launch (about 4 us). +// +// The checks are CUDA_KERNEL_ASSERTs. Unlike assert(), that macro is not removed by NDEBUG on CUDA +// (c10 expands it to a plain `if (!cond) __assert_fail(...)`; only Android, Apple and FreeBSD +// compile it out, and ROCm aborts), so they stay in the release build this extension ships as +// well; tests/models/minimax_h3/test_h3_assert_in_release_build.py compiles with -DNDEBUG -O3 to +// keep that true for the torch it runs against. #pragma once #include @@ -43,6 +49,10 @@ static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ so for (int64_t i = first_thread; i <= segments; i += stride) { CUDA_KERNEL_ASSERT(seg_first_tile[i] >= 0 && seg_first_tile[i] <= tiles); CUDA_KERNEL_ASSERT(i == segments || seg_first_tile[i] <= seg_first_tile[i + 1]); + // The segments cut the tile list exactly. A tile that no segment owns would be missing from + // every row's sum, silently. + CUDA_KERNEL_ASSERT(i != 0 || seg_first_tile[0] == 0); + CUDA_KERNEL_ASSERT(i != segments || seg_first_tile[segments] == tiles); } } @@ -51,7 +61,8 @@ static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ so inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, int64_t positions, const int64_t* tile_begin, const int64_t* tile_end, int64_t tiles, const int64_t* seg_first_tile, int64_t segments) { - const int64_t items = std::max({positions, tiles, segments + 1}); + // At least one block, so an empty list (segments + 1 == 0) is a no-op and not an invalid launch. + const int64_t items = std::max({positions, tiles, segments + 1, int64_t{1}}); const unsigned blocks = static_cast(std::min((items + 255) / 256, 1024)); check_tile_metadata_kernel<<>>(sorted_pos, positions, tile_begin, tile_end, tiles, seg_first_tile, segments); diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 35fee7775..cef79ab51 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -392,7 +392,18 @@ def h3_adaln_row_gather_forward( token_tags: torch.Tensor, chunks: int = 6, modality_num: int = 3, -) -> torch.Tensor: ... +) -> torch.Tensor: + """Gather AdaLN table rows into packed per-position chunks; return ``(chunks, S, H)``. + + ``rows`` is a 2-D CUDA ``(R, chunks * H)`` table with unit column stride and ``R`` a multiple of + ``modality_num``; ``timestep_indices`` and ``token_tags`` are ``(S,)`` contiguous integer + tensors (int64 or int32, one dtype) on its device, ``S > 0``. Position ``s`` reads table row + ``timestep_indices[s] * modality_num + token_tags[s]`` and receives its ``chunks`` slices of + width ``H`` in ``rows``' dtype (a pure copy). Out-of-range indices stop the kernel with a + device-side assert. + """ + ... + def h3_adaln_row_gather_backward( grad: torch.Tensor, sorted_pos: torch.Tensor, @@ -400,4 +411,13 @@ def h3_adaln_row_gather_backward( tile_end: torch.Tensor, seg_first_tile: torch.Tensor, out_dtype: torch.dtype, -) -> torch.Tensor: ... +) -> torch.Tensor: + """Sum a ``(chunks, S, H)`` gradient into table rows; return ``(R, chunks * H)`` in ``out_dtype``. + + ``sorted_pos`` holds the ``S`` positions sorted stably by table row, ``[tile_begin, tile_end)`` + cut that order into fixed-size tiles, and ``seg_first_tile`` (``R + 1`` entries) gives the + first tile of each row (see ``_segment_tiles``). Tiles are summed in FP32 and folded in order, + with no atomics. The metadata is checked on the device first: positions inside ``[0, S)``, tiles + inside ``sorted_pos``, segments that cut the tile list exactly and in order. + """ + ... diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py index 9adc6b4aa..cc3424a04 100644 --- a/tests/models/minimax_h3/test_h3_adaln_modulation.py +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -81,7 +81,7 @@ def test_forward_bitwise_equal_to_separate_ops(self): ti, tags = h3_packed_layout(1000, 3) ours = fused(temb, weight, bias, ti, tags) separate = gather.gather_chunks(proj(temb, weight, bias), ti, tags) - assert all(torch.equal(a, b) for a, b in zip(ours, separate)) + assert all(torch.equal(a, b) for a, b in zip(ours, separate, strict=True)) def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): fused, proj, gather = _ops() @@ -121,7 +121,7 @@ def rel(grad, golden): assert rel(ours[0], ref[0].grad) < 1e-5 assert rel(ours[0], ref[0].grad) < rel(separate[0], ref[0].grad) / 100 # dW, db: only their own final BF16 rounding remains. - for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad)): + for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad), strict=True): assert (grad == golden.to(torch.bfloat16)).float().mean() > 0.9999 @pytest.mark.parametrize("frozen", ["bias", "weight", "temb", "bias_and_temb"]) @@ -134,10 +134,10 @@ def test_frozen_inputs_get_no_gradient_and_the_others_are_unchanged(self, frozen leaves = [t.detach().clone() for t in (temb, weight, bias)] names = ("temb", "weight", "bias") - for leaf, name in zip(leaves, names): + for leaf, name in zip(leaves, names, strict=True): leaf.requires_grad_(name not in frozen) torch.autograd.backward(list(fused(*leaves, ti, tags)), grads) - for leaf, name, expected in zip(leaves, names, full): + for leaf, name, expected in zip(leaves, names, full, strict=True): if name in frozen: assert leaf.grad is None else: @@ -150,7 +150,7 @@ def test_backward_repeat_bitwise(self): grads = _grads(3000, 64, seed=3) first = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) again = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) - assert all(torch.equal(a, b) for a, b in zip(first, again)) + assert all(torch.equal(a, b) for a, b in zip(first, again, strict=True)) def test_rejects_bad_indices(self): fused, _, _ = _ops() 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 ff9662f75..220311269 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -364,6 +364,10 @@ def _backward_probe(mutation: str): sorted_pos = sorted_pos.clone(); sorted_pos[0] = seq elif mutation == "segments_past_tiles": first = first.clone(); first[-1] = begin.numel() + 3 + elif mutation == "first_segment_not_at_tile_zero": + first = first.clone(); first[0] = 1 + elif mutation == "last_tile_owned_by_no_segment": + first = first.clone(); first[-1] = first[-1] - 1 elif mutation == "no_tiles": begin = begin[:0]; end = end[:0] _C.h3_adaln_row_gather_backward(grad, sorted_pos, begin, end, first, torch.float32) @@ -387,12 +391,41 @@ def test_the_backward_stops_on_metadata_that_does_not_describe_the_tiles(self): "tile_end_past_seq", "position_past_seq", "segments_past_tiles", + "first_segment_not_at_tile_zero", + "last_tile_owned_by_no_segment", "no_tiles", ): done = self._backward_probe(mutation) assert done.returncode != 0, mutation assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) + @pytest.mark.parametrize("grad_dtype", [torch.float32, torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) + def test_backward_writes_the_requested_dtype_whatever_the_gradient_dtype( + self, grad_dtype, out_dtype + ): + """The table gradient has ``out_dtype``, not the dtype of the incoming gradient.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + from rl_engine.backends.extension import _C + + _cuda_op() + seq, hidden = 40, 8 + index = torch.arange(seq, device="cuda") % 3 + generator = torch.Generator(device="cpu").manual_seed(0) + # Small integers: every sum is exact in all three dtypes, so the result can be compared + # with torch.equal. + grad = torch.randint(-3, 4, (6, seq, hidden), generator=generator).to("cuda", grad_dtype) + tiles = _segment_tiles(index, 3) + got = _C.h3_adaln_row_gather_backward(grad, *tiles, out_dtype) + assert got.dtype == out_dtype and got.shape == (3, 6 * hidden) + expected = torch.stack( + [grad.double()[:, index == r, :].sum(1).reshape(-1) for r in range(3)] + ).to(out_dtype) + assert torch.equal(got, expected) + def test_backward_with_no_table_rows_returns_an_empty_table_and_rejects_no_boundaries(self): from rl_engine.backends.extension import _C diff --git a/tests/models/minimax_h3/test_h3_assert_in_release_build.py b/tests/models/minimax_h3/test_h3_assert_in_release_build.py new file mode 100644 index 000000000..57043d30f --- /dev/null +++ b/tests/models/minimax_h3/test_h3_assert_in_release_build.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The H3 kernels' CUDA_KERNEL_ASSERT bounds checks survive a release (-DNDEBUG) build. + +The native entry points check caller-built metadata and index values with ``CUDA_KERNEL_ASSERT`` +on the device (a host-side check would add a synchronisation to every call). Unlike ``assert()``, +c10 keeps that macro under ``NDEBUG`` on CUDA, so a release build still stops on a bad value +instead of reading out of bounds. This compiles a one-kernel program with the release flags and +the torch headers this repository builds against, and checks that the assertion fires. The tests +that run the built extension on bad metadata and bad indices live next to each op +(``test_native_bounds_assertions``, ``test_the_backward_stops_on_metadata...``). +""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import textwrap +from pathlib import Path + +import pytest +import torch +from torch.utils import cpp_extension + +SOURCE = textwrap.dedent( + """ + #include + #include + #include + + __global__ void check(const long* v, int n) { + const int i = threadIdx.x; + if (i < n) CUDA_KERNEL_ASSERT(v[i] >= 0 && v[i] < 100); + } + + int main(int argc, char** argv) { + long host[4] = {0, 1, 2, 3}; + if (argc > 1) host[2] = std::atol(argv[1]); + long* device = nullptr; + cudaMalloc(&device, sizeof(host)); + cudaMemcpy(device, host, sizeof(host), cudaMemcpyHostToDevice); + check<<<1, 4>>>(device, 4); + const cudaError_t status = cudaDeviceSynchronize(); + std::printf("%s\\n", cudaGetErrorString(status)); + return status == cudaSuccess ? 0 : 2; + } + """ +) + + +def _nvcc() -> str | None: + for root in (os.environ.get("CUDA_HOME"), "/usr/local/cuda"): + if root and (Path(root) / "bin" / "nvcc").exists(): + return str(Path(root) / "bin" / "nvcc") + return shutil.which("nvcc") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.skipif(_nvcc() is None, reason="needs nvcc") +def test_cuda_kernel_assert_stays_active_with_ndebug(tmp_path): + major, minor = torch.cuda.get_device_capability() + arch = f"{major}{minor}" + (tmp_path / "check.cu").write_text(SOURCE) + build = subprocess.run( + [ + _nvcc(), + "-O3", + "-DNDEBUG", + "-std=c++17", + f"-gencode=arch=compute_{arch},code=sm_{arch}", + *(f"-I{path}" for path in cpp_extension.include_paths()), + str(tmp_path / "check.cu"), + "-o", + str(tmp_path / "check"), + ], + capture_output=True, + text=True, + timeout=300, + ) + assert build.returncode == 0, build.stderr[-1500:] + + ok = subprocess.run([str(tmp_path / "check"), "2"], capture_output=True, text=True, timeout=60) + assert ok.returncode == 0, ok.stdout + ok.stderr # the control: a valid value passes + for bad in ("-1", "100", "123456"): + done = subprocess.run( + [str(tmp_path / "check"), bad], capture_output=True, text=True, timeout=60 + ) + assert done.returncode != 0, bad + assert "device-side assert triggered" in done.stdout, (bad, done.stdout, done.stderr) From 560111898846b358238e1dbcbadc950178c73d3c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:33:52 +0000 Subject: [PATCH 40/42] fix(h3): bound the sizes the rmsnorm kernels take as int, check the modulation index once before the dweight loop * `check_xw`, called by every rmsnorm entry point, rejects N or M above INT_MAX (the kernels take N as `int` and launch one block per row), and the modulated backward bounds the table rows and tiles it launches as blocks. A test allocates a (1, 2**31) FP16 row and checks forward and backward. * The per-row `CUDA_KERNEL_ASSERT` on the modulation index inside `rmsnorm_dweight_partial_kernel` cost about 11% of the modulated backward (A/B, interleaved min of 300, B200: 1770 -> 1586 us at S = 32768, 6346 -> 5694 us at S = 131072, 378 -> 371 us at S = 4097 once removed). The loop now reads the index unchecked and the entry point runs `h3_tiles::check_index_range` over the S index values first, a 4 us launch like the tile-metadata pre-pass. The forward and dx kernels keep their once-per-row assert. A child-process test requires the device-side assert for negative, past-the-table and far indices in both forward and backward, with a valid-index control. * The backward header states the exact rounding chain of dx (coef = ((rstd*rstd)*rstd)*dot/N, one cast, FMA contraction allowed with fixed flags) and that the backward is measured against FP64, not claimed bitwise to autograd; the empty `Modulation` initialiser says why `seq` is 1. * `_C.pyi`: contracts for `h3_rmsnorm_forward` / `h3_rmsnorm_backward`. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/rmsnorm_modulate.cu | 36 ++++++++-- csrc/cuda/h3/tile_metadata.cuh | 20 ++++++ tests/models/minimax_h3/test_h3_rmsnorm.py | 79 ++++++++++++++++++++++ 3 files changed, 131 insertions(+), 4 deletions(-) diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 3f352f417..6d0146408 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -18,6 +18,12 @@ // Backward (FP32, one cast per output, no atomics): // d_n = g * cast(1 + scale) (g without modulation) // dx = rstd * w * d_n - x * rstd^3 * sum(w * d_n * x) / N (row-local) +// evaluated as coef = ((rstd * rstd) * rstd) * dot / N and +// dx = cast(rstd * w * d_n - x * coef), with dot a fixed block tree over fmaf terms. +// nvcc may contract the last line to an FMA; the build flags are fixed, so the result +// is the same on every run. That is the contract of the backward: it is measured +// against an FP64 golden (the *_against_fp64 tests), and not claimed bitwise equal to +// autograd's RMSNorm backward, which forms rstd^3 differently. // 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) @@ -112,13 +118,21 @@ struct Modulation { }; // Table row of x row `r`. check_range=False callers skip the Python range check, so the native side -// asserts the index before it is used as an offset (as the row-gather kernel does). +// asserts the index before it is used as an offset (as the row-gather kernel does). Kernels that read +// it once per row use this; the dweight partials read it inside their row loop, where an assert +// costs about 11% of the backward, so they use the unchecked read and the entry points run +// h3_tiles::check_index_range over the S index values before launching them. __device__ __forceinline__ int64_t modulation_row(const Modulation& mod, int64_t r) { const int64_t i = mod.index[r % mod.seq]; CUDA_KERNEL_ASSERT(i >= 0 && i < mod.rows); return i; } +__device__ __forceinline__ int64_t modulation_row_checked_by_entry_point(const Modulation& mod, + int64_t r) { + return mod.index[r % mod.seq]; +} + template __global__ void __launch_bounds__(32 * kWarps) rmsnorm_modulate_fwd_kernel(const T* __restrict__ x, const T* __restrict__ w, T* __restrict__ y, @@ -204,9 +218,10 @@ __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) + - modulation_row(mod, r) * mod.row_stride - : nullptr; + const T* scale = mod.index != nullptr + ? static_cast(mod.scale) + + modulation_row_checked_by_entry_point(mod, r) * mod.row_stride + : nullptr; 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; @@ -277,6 +292,8 @@ Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, int64_t n) { + // No modulation: kernels test `index == nullptr`. seq = 1 (not 0) keeps `r % seq` defined for + // any code that computes it before looking at the index. 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"); @@ -321,6 +338,11 @@ void check_xw(const torch::Tensor& x, const torch::Tensor& w) { TORCH_CHECK(x.size(1) > 0, "x must have at least one column"); TORCH_CHECK(x.size(1) % kVec == 0, "N=", x.size(1), " must be a multiple of ", kVec, " (PyTorch's vectorized RMSNorm path)"); + // The kernels take N as `int` and launch one block per row on grid.x. + TORCH_CHECK(x.size(1) <= std::numeric_limits::max(), "N=", x.size(1), + " is too large for the kernels (at most ", std::numeric_limits::max(), ")"); + TORCH_CHECK(x.size(0) <= std::numeric_limits::max(), "M=", x.size(0), + " rows do not fit one grid.x launch (at most ", std::numeric_limits::max(), ")"); } // float32 / float16 / bfloat16 only (float64 has no PyTorch-matching vector path here). @@ -396,6 +418,8 @@ std::vector h3_rmsnorm_backward( } const c10::cuda::CUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); + // The dweight partials read the index inside their row loop: check its values first. + if (mod.index != nullptr) h3_tiles::check_index_range(stream, mod.index, mod.seq, mod.rows); auto dx = torch::empty_like(x); const int64_t tiles = (rows + kRowTile - 1) / kRowTile; auto partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); @@ -424,6 +448,10 @@ std::vector h3_rmsnorm_backward( const int64_t seg_tiles = tile_begin->numel(); const int64_t width = 2 * n; TORCH_CHECK(width <= std::numeric_limits::max(), "N=", n, " is too large for the kernels"); + TORCH_CHECK(segments <= std::numeric_limits::max() && + seg_tiles <= std::numeric_limits::max(), + "table rows (", segments, ") and tiles (", seg_tiles, + ") must each fit one grid.x launch"); TORCH_CHECK((width + threads - 1) / threads <= 65535, "N=", n, " needs more column blocks than the grid.y limit of 65535"); auto seg_partial = torch::empty({std::max(seg_tiles, 1), width}, x.options().dtype(at::kFloat)); diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index 96318225d..09c6ad4be 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -69,4 +69,24 @@ inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, C10_CUDA_KERNEL_LAUNCH_CHECK(); } +// Table-row indices read inside a kernel's loop: `n` entries, each a row of a table with `limit` +// rows. The assert would cost the loop its scheduling (the same 6-12% as above, measured on the +// rmsnorm dweight partials), so the entry point checks the S index values once, here, before it +// launches the kernels that read them. +static __global__ void check_index_range_kernel(const int64_t* __restrict__ index, int64_t n, + int64_t limit) { + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n; i += stride) { + CUDA_KERNEL_ASSERT(index[i] >= 0 && index[i] < limit); + } +} + +inline void check_index_range(cudaStream_t stream, const int64_t* index, int64_t n, + int64_t limit) { + const unsigned blocks = + static_cast(std::max(1, std::min((n + 255) / 256, 1024))); + check_index_range_kernel<<>>(index, n, limit); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + } // namespace h3_tiles diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py index 6d91f60bb..be899893b 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm.py @@ -13,6 +13,12 @@ from __future__ import annotations +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest import torch @@ -206,6 +212,22 @@ def test_rejects_unsupported(self): with pytest.raises(ValueError): _cuda_op()(_x((2, 8)).cpu(), torch.ones(8).bfloat16()) + def test_rejects_a_row_wider_than_the_int_the_kernels_count_with(self): + """N goes to the kernels as ``int``: forward and backward both refuse it up front.""" + + from rl_engine.backends.extension import _C + + _cuda_op() + n = 2**31 # a multiple of 4 above INT_MAX + if torch.cuda.mem_get_info()[0] < 10 * 2**30: + pytest.skip("needs about 8 GiB of free device memory for the two operands") + x = torch.empty(1, n, device="cuda", dtype=torch.float16) + w = torch.empty(n, device="cuda", dtype=torch.float16) + with pytest.raises(RuntimeError, match="too large"): + _C.h3_rmsnorm_forward(x, w, 1e-6) + with pytest.raises(RuntimeError, match="too large"): + _C.h3_rmsnorm_backward(x, x, w, torch.empty(1, device="cuda")) + @requires_cuda class TestCudaBackward: @@ -300,3 +322,60 @@ def test_registry_dispatches_cuda(self): assert ( type(KernelRegistry().get_op("h3_rmsnorm", device="cuda")).__name__ == "H3RMSNormCudaOp" ) + + +def _index_probe(entry: str, bad: int): + """Run one native entry point with a table-row index outside the table, in a child process.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + entry, bad = sys.argv[1], int(sys.argv[2]) + rows, hidden, table_rows = 600, 64, 4 + x = torch.randn(rows, hidden, device="cuda").bfloat16() + w = torch.ones(hidden, device="cuda").bfloat16() + table = torch.randn(table_rows, 2 * hidden, device="cuda").bfloat16() + shift, scale = table.chunk(2, dim=-1) + index = torch.arange(rows, device="cuda") % table_rows + index[rows - 3] = bad # the last tile of rows, past the first dx blocks + if entry == "forward": + _C.h3_rmsnorm_forward(x, w, 1e-5, shift, scale, index) + else: + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + valid = index.clamp(0, table_rows - 1) + y, rstd = _C.h3_rmsnorm_forward(x, w, 1e-5, shift, scale, valid) + tiles = _segment_tiles(valid, table_rows) + _C.h3_rmsnorm_backward(torch.ones_like(x), x, w, rstd, shift, scale, index, *tiles) + torch.cuda.synchronize() + """ + ) + return subprocess.run( + [sys.executable, "-c", probe, entry, str(bad)], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=120, + ) + + +@requires_cuda +@pytest.mark.parametrize("entry", ["forward", "backward"]) +@pytest.mark.parametrize("bad", [-1, 4, 10**9], ids=["negative", "past-table", "far"]) +def test_a_modulation_index_outside_the_table_stops_the_kernels(entry, bad): + """check_range=False and direct calls get a device-side assert, not an out-of-bounds read. + + The dweight partials read the index inside their row loop, where an assert costs about 11% of + the backward; the entry point checks the index values once before launching them. + """ + + _cuda_op() + assert _index_probe(entry, 1).returncode == 0 # the control: the same call with a valid index + done = _index_probe(entry, bad) + assert done.returncode != 0, entry + assert "device-side assert triggered" in done.stderr, done.stderr From e366ecef628030ab33a4cdbcb1a7491e2b958370 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:00:14 +0000 Subject: [PATCH 41/42] fix(h3): require the segment tiles to partition the positions in order check_tile_metadata bounded every tile and cut the tile list exactly by segment, but a gap between tiles (positions dropped from every row's sum) or an overlap (positions counted twice) passed all of it. Each tile must now start where the one before it ends, the first at 0 and the last at the number of positions, and a tile list that is empty over a non-empty sequence is refused on the host. Mutation tests cover a gap, an overlap, a first tile not at 0, a last tile that stops short, and the empty list with and without matching segment bounds. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/tile_metadata.cuh | 12 ++++++++++- .../minimax_h3/test_h3_adaln_row_gather.py | 21 ++++++++++++++++++- 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index 96318225d..ca11e47e6 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -10,7 +10,9 @@ // fold kernels read the partial rows of a segment's tiles, so a stale or short list would read out // of bounds. // -// check_tile_metadata() verifies the list on the device before any of that happens. It is a separate +// check_tile_metadata() verifies the list on the device before any of that happens: indices in +// range, the tiles an exact in-order partition of the positions, and the segments an exact cut of +// the tile list. It is a separate // pass on purpose: an assert in or before the loops of the partial and fold kernels changes how the // compiler schedules their loads and slowed them by 6-12% (and by several times inside the loop), // while this pass costs one small launch (about 4 us). @@ -45,6 +47,11 @@ static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ so 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] <= positions); + // The tiles partition [0, positions) in order: each starts where the one before it ends, the + // first at 0 and the last at `positions`. A gap would drop positions from every row's sum and + // an overlap would count them twice, both without leaving the bounds checked above. + CUDA_KERNEL_ASSERT(tile_begin[i] == (i == 0 ? int64_t{0} : tile_end[i - 1])); + CUDA_KERNEL_ASSERT(i != tiles - 1 || tile_end[i] == positions); } for (int64_t i = first_thread; i <= segments; i += stride) { CUDA_KERNEL_ASSERT(seg_first_tile[i] >= 0 && seg_first_tile[i] <= tiles); @@ -61,6 +68,9 @@ static __global__ void check_tile_metadata_kernel(const int64_t* __restrict__ so inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, int64_t positions, const int64_t* tile_begin, const int64_t* tile_end, int64_t tiles, const int64_t* seg_first_tile, int64_t segments) { + // No tiles over a non-empty sequence would leave every row's sum at zero. + TORCH_CHECK(tiles > 0 || positions == 0, "the tile list is empty but there are ", positions, + " positions to sum"); // At least one block, so an empty list (segments + 1 == 0) is a no-op and not an invalid launch. const int64_t items = std::max({positions, tiles, segments + 1, int64_t{1}}); const unsigned blocks = static_cast(std::min((items + 255) / 256, 1024)); 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 220311269..892e86c85 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -370,6 +370,16 @@ def _backward_probe(mutation: str): first = first.clone(); first[-1] = first[-1] - 1 elif mutation == "no_tiles": begin = begin[:0]; end = end[:0] + elif mutation == "no_tiles_and_empty_segments": + begin = begin[:0]; end = end[:0]; first = torch.zeros_like(first) + elif mutation == "gap_between_tiles": + begin = begin.clone(); begin[1] += 1 + elif mutation == "overlapping_tiles": + begin = begin.clone(); begin[1] -= 1 + elif mutation == "first_tile_not_at_zero": + begin = begin.clone(); begin[0] += 1 + elif mutation == "last_tile_stops_short": + end = end.clone(); end[-1] -= 1 _C.h3_adaln_row_gather_backward(grad, sorted_pos, begin, end, first, torch.float32) torch.cuda.synchronize() """ @@ -393,11 +403,20 @@ def test_the_backward_stops_on_metadata_that_does_not_describe_the_tiles(self): "segments_past_tiles", "first_segment_not_at_tile_zero", "last_tile_owned_by_no_segment", - "no_tiles", + "gap_between_tiles", + "overlapping_tiles", + "first_tile_not_at_zero", + "last_tile_stops_short", ): done = self._backward_probe(mutation) assert done.returncode != 0, mutation assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) + # An empty tile list over a non-empty sequence is refused on the host, before any launch, + # whether or not the segment boundaries agree with it. + for mutation in ("no_tiles", "no_tiles_and_empty_segments"): + done = self._backward_probe(mutation) + assert done.returncode != 0, mutation + assert "tile list is empty" in done.stderr, (mutation, done.stderr) @pytest.mark.parametrize("grad_dtype", [torch.float32, torch.bfloat16, torch.float16]) @pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) From 2164211d8b3ec7dee0f4245cbcbdf578683a0dd0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:00:59 +0000 Subject: [PATCH 42/42] fix(h3): reject a missing shift, scale or index in the modulated RMSNorm validate_h3_modulation read index.dtype first, so forward_modulated(..., index=None) failed with an AttributeError. All three tensors are required for the modulated form; a missing one now raises a TypeError that names it. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/reference/minimax_h3/rmsnorm.py | 5 +++++ tests/models/minimax_h3/test_h3_rmsnorm.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/rl_engine/reference/minimax_h3/rmsnorm.py b/rl_engine/reference/minimax_h3/rmsnorm.py index ddbdaf6bf..dda2e1352 100644 --- a/rl_engine/reference/minimax_h3/rmsnorm.py +++ b/rl_engine/reference/minimax_h3/rmsnorm.py @@ -56,6 +56,11 @@ def validate_h3_modulation( same_dtype: bool = True, ) -> None: hidden = x.shape[-1] + # The modulated form always gathers rows: a missing table or index is an error, not the plain + # norm (which is ``forward``, without shift and scale). + for name, tensor in (("shift", shift), ("scale", scale), ("index", index)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor).__name__}") for name, tensor in (("shift", shift), ("scale", scale)): if tensor.device != x.device: raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py index be899893b..309775562 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm.py @@ -379,3 +379,21 @@ def test_a_modulation_index_outside_the_table_stops_the_kernels(entry, bad): done = _index_probe(entry, bad) assert done.returncode != 0, entry assert "device-side assert triggered" in done.stderr, done.stderr + + +@pytest.mark.parametrize("missing", ["shift", "scale", "index"]) +def test_the_modulated_form_rejects_a_missing_table_or_index(missing): + """``index=None`` must not fall back to the plain norm (or fail on a ``None`` attribute).""" + + x = torch.randn(4, 8, dtype=torch.bfloat16) + weight = torch.ones(8, dtype=torch.bfloat16) + args = { + "shift": torch.randn(2, 8, dtype=torch.bfloat16), + "scale": torch.randn(2, 8, dtype=torch.bfloat16), + "index": torch.tensor([0, 1, 1, 0]), + } + args[missing] = None + with pytest.raises(TypeError, match=f"{missing} must be a torch.Tensor"): + NativeH3RMSNormOp().forward_modulated( + x, weight, args["shift"], args["scale"], args["index"] + )