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/24] [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/24] [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 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 03/24] 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 04/24] 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 05/24] 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 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 06/24] 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 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 07/24] 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 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 08/24] 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 09/24] 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 10/24] 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 11/24] 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 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 12/24] 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 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 13/24] 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 14/24] 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 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 15/24] 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 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 16/24] 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 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 17/24] 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 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 18/24] 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 19/24] 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 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 20/24] 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 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 21/24] 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 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 22/24] 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 23/24] 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 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 24/24] 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])