From e7bc53aa89f3c4b2c3d6045c0b2ccfec8c091391 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 02:51:41 +0000 Subject: [PATCH 01/14] [WS1][CUDA][MiniMax-H3] adaln_gate_residual: gated residual with in-kernel gate gather RFC #420 step 6, row `adaln_gate_residual` (stacked on h3_rmsnorm). - residual + gate[adaln_index] * sublayer_output in one pass, the gate row read from the AdaLN table view in the kernel; the two roundings sit where the eager expression rounds (__fmul_rn/__fadd_rn, no FMA contraction), so the forward is bitwise equal to diffusers in bf16/fp16/fp32. 1.7x faster forward at S = 131072. - Backward: d_residual = grad; d_sublayer bitwise equal to the eager VJP (exact product, one rounding); d_gate a deterministic FP32 segment sum (diffusers: non-deterministic BF16 atomics, 13-49x larger error, varying run to run). The PyTorch reference gets an equally deterministic gate gradient. - Chain stage residual + gate_msa * (stand-in attention output), covered by the end-to-end test; registry, gtest, report/figure hooks, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 242 ++++++++++++++++++ csrc/ops.cpp | 15 ++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-adaln-gate-residual.md | 91 +++++++ rl_engine/_C.pyi | 13 + rl_engine/kernels/gtest/operator_inputs.py | 18 ++ rl_engine/kernels/gtest/operator_specs.py | 11 + .../kernels/ops/cuda/h3/gate_residual.py | 83 ++++++ .../kernels/ops/pytorch/h3/gate_residual.py | 97 +++++++ rl_engine/kernels/registry.py | 11 + rl_engine/testing/h3_chain.py | 38 +++ rl_engine/testing/h3_provider.py | 8 + rl_engine/testing/h3_report.py | 107 ++++++++ rl_engine/tests/test_dispatch.py | 1 + scripts/plot_h3_evidence.py | 72 ++++++ setup.py | 1 + tests/h3/test_h3_adaln_gate_residual.py | 115 +++++++++ tests/test_ws1_gtest_gpu.py | 1 + 19 files changed, 926 insertions(+) create mode 100644 csrc/cuda/h3/gate_residual.cu create mode 100644 docs/operators/h3-adaln-gate-residual.md create mode 100644 rl_engine/kernels/ops/cuda/h3/gate_residual.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/gate_residual.py create mode 100644 tests/h3/test_h3_adaln_gate_residual.py diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu new file mode 100644 index 000000000..18ede9e72 --- /dev/null +++ b/csrc/cuda/h3/gate_residual.cu @@ -0,0 +1,242 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 gated residual (RFC #420 `adaln_gate_residual`). +// +// out[r, j] = cast(residual[r, j] + cast(gate[i, j] * y[r, j])) i = index[r % S] +// +// The gate row is read from a strided view of the AdaLN table (gate_msa or +// gate_mlp) inside the kernel, and the two roundings sit exactly where the +// eager `residual + gate.index_select(0, i) * y` rounds (__fmul_rn/__fadd_rn +// keep FP32 from contracting them into one FMA), so the result is +// bitwise equal to diffusers. Elementwise: every output depends only on its +// own inputs. +// +// Backward: d_residual = grad; d_y = cast(grad * gate) (the product of two +// 16-bit values is exact in FP32, so this equals the eager VJP bitwise); +// d_gate[i] = sum over positions mapped to i of grad * y, an FP32 segmented +// sum over positions sorted stably by row, in fixed tiles folded in order. + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +__device__ __forceinline__ float to_f(float v) { return v; } +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ float from_f(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } + +template +__device__ __forceinline__ float round_t(float v) { return to_f(from_f(v)); } + +struct Gate { + const void* rows; // (R, N) view, row stride `stride` + int64_t stride; + const int64_t* index; // (S,) + int64_t seq; +}; + +// One block per row (grid-strided); 16-byte vectors when the row allows it. +// `kVecElems` elements are loaded together; each is computed independently. +template +__global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __restrict__ b, + T* __restrict__ out, int64_t rows, int64_t n, + Gate gate, bool vectorized) { + constexpr int kVecElems = 16 / sizeof(T); + for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { + const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const T* ar = a + r * n; + const T* br = kDy ? nullptr : b + r * n; + T* orow = out + r * n; + if (vectorized) { + for (int64_t v = threadIdx.x; v < n / kVecElems; v += blockDim.x) { + const uint4 av = reinterpret_cast(ar)[v]; + const uint4 gv = reinterpret_cast(g)[v]; + uint4 bv = av; + if (!kDy) bv = reinterpret_cast(br)[v]; + const T* ae = reinterpret_cast(&av); + const T* ge = reinterpret_cast(&gv); + const T* be = reinterpret_cast(&bv); + uint4 ov; + T* oe = reinterpret_cast(&ov); +#pragma unroll + for (int k = 0; k < kVecElems; ++k) { + if (kDy) { + oe[k] = from_f(to_f(ae[k]) * to_f(ge[k])); // grad * gate + } else { + const float p = round_t(__fmul_rn(to_f(ge[k]), to_f(be[k]))); // gate * y + oe[k] = from_f(__fadd_rn(to_f(ae[k]), p)); // residual + p + } + } + reinterpret_cast(orow)[v] = ov; + } + } else { + for (int64_t j = threadIdx.x; j < n; j += blockDim.x) { + if (kDy) { + orow[j] = from_f(to_f(ar[j]) * to_f(g[j])); + } else { + const float p = round_t(__fmul_rn(to_f(g[j]), to_f(br[j]))); + orow[j] = from_f(__fadd_rn(to_f(ar[j]), p)); + } + } + } + } +} + +// partial[tile, j] = sum over sorted positions of the tile (ascending) of grad * y +template +__global__ void gate_grad_partial_kernel(const T* __restrict__ grad, const T* __restrict__ y, + const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int64_t n) { + const int64_t tile = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t e = sorted_pos[p] * n + j; + acc = fmaf(to_f(grad[e]), to_f(y[e]), acc); + } + partial[tile * n + j] = acc; +} + +template +__global__ void gate_grad_fold_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + T* __restrict__ out, int64_t n) { + const int64_t seg = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t t = seg_first_tile[seg]; t < seg_first_tile[seg + 1]; ++t) acc += partial[t * n + j]; + out[seg * n + j] = from_f(acc); +} + +#define H3_DISPATCH(TYPE, NAME, ...) \ + [&] { \ + switch (TYPE) { \ + case at::kFloat: { using T = float; __VA_ARGS__(); break; } \ + case at::kHalf: { using T = __half; __VA_ARGS__(); break; } \ + case at::kBFloat16: { using T = __nv_bfloat16; __VA_ARGS__(); break; } \ + default: TORCH_CHECK(false, NAME, ": unsupported dtype ", TYPE); \ + } \ + }() + +Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torch::Tensor& x) { + TORCH_CHECK(gate.is_cuda() && gate.device() == x.device() && gate.scalar_type() == x.scalar_type(), + "gate must be a CUDA tensor with the activations' dtype"); + TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, + "gate must be (R, N) with unit column stride"); + TORCH_CHECK(index.is_cuda() && index.scalar_type() == at::kLong && index.dim() == 1 && + index.is_contiguous() && index.numel() > 0, + "index must be a non-empty contiguous int64 CUDA tensor"); + TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; +} + +void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device() && t.scalar_type() == like.scalar_type(), + name, " must match the residual's device and dtype"); + TORCH_CHECK(t.sizes() == like.sizes() && t.is_contiguous(), name, + " must be contiguous with the residual's shape"); +} + +unsigned row_blocks(int64_t rows) { + return static_cast(std::min(rows, 1 << 30)); +} + +bool can_vectorize(const torch::Tensor& gate, int64_t n, std::initializer_list acts) { + const int64_t vec = 16 / gate.element_size(); + if (n % vec != 0 || (gate.stride(0) * gate.element_size()) % 16 != 0) return false; + if (reinterpret_cast(gate.data_ptr()) % 16 != 0) return false; + for (const auto* t : acts) { + if (reinterpret_cast(t->data_ptr()) % 16 != 0) return false; + } + return true; +} + +} // namespace + +// residual, y: (M, N); gate: (R, N) view; index: (S,), M % S == 0. +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index) { + TORCH_CHECK(residual.is_cuda() && residual.dim() == 2 && residual.is_contiguous() && + residual.numel() > 0, + "residual must be a non-empty contiguous (M, N) CUDA tensor"); + check_act(y, "y", residual); + const Gate g = make_gate(gate, index, residual); + const c10::cuda::CUDAGuard guard(residual.device()); + auto out = torch::empty_like(residual); + auto stream = at::cuda::getCurrentCUDAStream(); + const bool vec = can_vectorize(gate, residual.size(1), {&residual, &y, &out}); + H3_DISPATCH(residual.scalar_type(), "h3_gate_residual_forward", [&] { + gate_residual_rows_kernel<<>>( + reinterpret_cast(residual.data_ptr()), reinterpret_cast(y.data_ptr()), + reinterpret_cast(out.data_ptr()), residual.size(0), residual.size(1), g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// Returns {d_y, d_gate (R, N) in the gate dtype}. d_residual is the incoming grad. +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + torch::Tensor tile_begin, + torch::Tensor tile_end, + torch::Tensor seg_first_tile) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + check_act(y, "y", grad); + const Gate g = make_gate(gate, index, grad); + for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kLong && t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors"); + } + const int64_t rows = grad.size(0); + const int64_t n = grad.size(1); + const int64_t segments = seg_first_tile.numel() - 1; + TORCH_CHECK(segments == gate.size(0), "segments must equal the gate rows"); + const int64_t tiles = tile_begin.numel(); + const c10::cuda::CUDAGuard guard(grad.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto dy = torch::empty_like(y); + auto partial = torch::empty({std::max(tiles, 1), n}, grad.options().dtype(at::kFloat)); + auto dgate = torch::empty({segments, n}, grad.options()); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward", [&] { + const T* gp = reinterpret_cast(grad.data_ptr()); + const T* yp = reinterpret_cast(y.data_ptr()); + const bool vec = can_vectorize(gate, n, {&grad, &dy}); + gate_residual_rows_kernel<<>>( + gp, nullptr, reinterpret_cast(dy.data_ptr()), rows, n, g, vec); + if (tiles > 0) { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, stream>>>( + gp, yp, sorted_pos.data_ptr(), tile_begin.data_ptr(), + tile_end.data_ptr(), partial.data_ptr(), n); + } + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(dgate.data_ptr()), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dy, dgate}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index f2288efe4..08085e721 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -527,6 +527,14 @@ std::vector h3_rmsnorm_backward( c10::optional index, c10::optional sorted_pos, c10::optional tile_begin, c10::optional tile_end, c10::optional seg_first_tile); +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index); +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + torch::Tensor tile_begin, + torch::Tensor tile_end, + torch::Tensor seg_first_tile); #endif // PyBind11 Module Registration @@ -819,5 +827,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("index") = py::none(), py::arg("sorted_pos") = py::none(), py::arg("tile_begin") = py::none(), py::arg("tile_end") = py::none(), py::arg("seg_first_tile") = py::none()); + m.def("h3_gate_residual_forward", &h3_gate_residual_forward, + "residual + gate[index] * y with the gate row gathered in-kernel (eager rounding order)", + py::arg("residual"), py::arg("y"), py::arg("gate"), py::arg("index")); + m.def("h3_gate_residual_backward", &h3_gate_residual_backward, + "Gated-residual backward: exact dy, deterministic sorted segment sum for dgate", + py::arg("grad"), py::arg("y"), py::arg("gate"), py::arg("index"), py::arg("sorted_pos"), + py::arg("tile_begin"), py::arg("tile_end"), py::arg("seg_first_tile")); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index af8c1863f..f6db1bc36 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -35,6 +35,7 @@ nav: - operators/h3-adaln-projection.md - operators/h3-adaln-row-gather.md - operators/h3-rmsnorm.md + - operators/h3-adaln-gate-residual.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index b687a65b5..7042d4561 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -37,4 +37,5 @@ Every operator page should include: - [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) - [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) - [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) +- [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md new file mode 100644 index 000000000..7a3010185 --- /dev/null +++ b/docs/operators/h3-adaln-gate-residual.md @@ -0,0 +1,91 @@ +# MiniMax-H3 AdaLN Gated Residual + +## Summary + +`adaln_gate_residual` adds each sublayer's output back to the residual stream through the +per-row AdaLN gate (RFC #420, WS1 step 6). In every H3 block this happens twice, after attention +with `gate_msa` and after the FFN with `gate_mlp`: + +```text +hidden = residual + gate[adaln_indices] * sublayer_output # RFC #420 §4 order +``` + +`gate` is an `(R, H)` row view of the [AdaLN projection](h3-adaln-projection.md) table. The op +gathers the gate row inside the kernel, so the gathered `(S, H)` gate is never materialised. + +## Entry Point + +```python +op = kernel_registry.get_op("adaln_gate_residual", device="cuda") +hidden = op(residual, attn_output, gate_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA | `rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp` | `rl_engine._C.h3_gate_residual_{forward,backward}` | Forward and `d_sublayer` bitwise equal to diffusers | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp` | n/a | Eager forward with a deterministic gate gradient; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `residual`, `sublayer_output` | `(..., S, H)` | bf16 / fp16 / fp32 | Same shape and dtype | +| `gate` | `(R, H)` row view | same | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; shared across the batch | + +## Numerics + +- **Forward.** Each element computes `p = round(gate * y)` and then `out = round(residual + p)`. + `__fmul_rn`/`__fadd_rn` keep FP32 from contracting the two operations into one FMA. The + roundings fall exactly where the eager expression rounds, so the forward is **bitwise equal to + diffusers** in BF16, FP16 and FP32. Rows are independent. +- **Backward.** + - `d_residual` is the incoming gradient. + - `d_sublayer = round(grad * gate)`. The product of two 16-bit values is exact in FP32, so this + is bitwise equal to the eager VJP. + - `d_gate` is an FP32 segmented sum of `grad * sublayer_output` over positions sorted stably + by table row, using the same tiles as [`adaln_row_gather`](h3-adaln-row-gather.md), rounded + once. + + Diffusers' `d_gate` uses `index_select`'s BF16 atomic scatter-add instead. It is + non-deterministic and about 35× further from the FP64 golden. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op adaln_gate_residual +``` + +B200, B = 1, H = 5376, BF16: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.05 ms | 0.06 ms | 1.21 ms | 0.68 ms | +| 32768 | 0.24 ms | 0.41 ms | 1.40 ms | 2.71 ms | +| 131072 | 0.92 ms | 1.59 ms | 3.31 ms | 10.6 ms | + +The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one +block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and +tile setup. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_adaln_gate_residual.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end, incl. the gated residual +python scripts/check_operator.py --op adaln_gate_residual --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op adaln_gate_residual \ + --out docs/usage/evidence/h3-adaln-gate-residual-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-adaln-gate-residual-b200/report.json +``` + +## Known Limitations + +- Until the attention and FFN rows land, the end-to-end chain feeds a seeded stand-in for the + sublayer output. The op itself is exercised on random and pinned-table inputs. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 01e6e5348..43fb55ff5 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -377,3 +377,16 @@ def h3_rmsnorm_backward( tile_end: torch.Tensor | None = None, seg_first_tile: torch.Tensor | None = None, ) -> list[torch.Tensor]: ... +def h3_gate_residual_forward( + residual: torch.Tensor, y: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... +def h3_gate_residual_backward( + grad: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 02f7f2e37..6afafc3ae 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -51,6 +51,7 @@ def make_operator_inputs( "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, "adaln_row_gather": _make_adaln_row_gather_inputs, "h3_rmsnorm": _make_h3_rmsnorm_inputs, + "adaln_gate_residual": _make_adaln_gate_residual_inputs, } try: return builders[op_name](args, dtype, device) @@ -89,6 +90,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" f"->{batch * seq}", "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", + "adaln_gate_residual": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -446,6 +448,22 @@ def _make_h3_rmsnorm_inputs( } +def _make_adaln_gate_residual_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # gate rows: 3 modality rows for each of --num-timesteps (default --batch) timesteps. + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + rows = 3 * _h3_num_timesteps(args) + generator = _generator(args, device, offset=31) + return { + "residual": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "y": _floating_tensor((batch, seq, hidden), args, dtype, device, 1), + "gate": (_floating_tensor((rows, hidden), args, torch.float32, device, 2) * 0.5).to(dtype), + "index": torch.randint(0, rows, (seq,), generator=generator, device=device), + } + + 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 197db8874..5254aba2e 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -311,6 +311,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("x", "weight"), ), + "adaln_gate_residual": OperatorSpec( + name="adaln_gate_residual", + op_class="reduction", # forward is elementwise; the gate VJP sums positions + gold_path="rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp", + "cuda": "rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp", + }, + grad_input_names=("residual", "y", "gate"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/gate_residual.py b/rl_engine/kernels/ops/cuda/h3/gate_residual.py new file mode 100644 index 000000000..2a0fc9e1d --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/gate_residual.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 gated residual (RFC #420 ``adaln_gate_residual``).""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles +from rl_engine.kernels.ops.pytorch.h3.gate_residual import validate_h3_gate_residual + +KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" +BACKWARD_IMPL = "exact_dy_sorted_segment_dgate" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_gate_residual_backward" + + +def h3_gate_residual_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_gate_residual_forward") + and hasattr(_C, "h3_gate_residual_backward") + ) + + +class _H3GateResidualCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, index): + hidden = residual.shape[-1] + res2 = residual.contiguous().view(-1, hidden) + y2 = y.contiguous().view(-1, hidden) + out = _C.h3_gate_residual_forward(res2, y2, gate, index) + ctx.save_for_backward(y2, gate, index) + ctx.shape = residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + positions = index.repeat(y2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, gate.shape[0]) + dy, dgate = _C.h3_gate_residual_backward(g2, y2, gate, index, *tiles) + record_backward( + "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad, dy.view(ctx.shape), dgate, None + + +class H3GateResidualCudaOp: + """CUDA candidate: ``residual + gate[index] * y`` in one elementwise pass. + + The gate row is gathered in the kernel from the AdaLN table view, and the + two roundings match the eager expression, so the output is bitwise equal + to diffusers. ``dy`` is bitwise equal to the eager VJP; ``dgate`` is a + deterministic FP32 segment sum instead of ``index_select``'s BF16 atomics. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_gate_residual_available(): + raise RuntimeError( + "rl_engine._C lacks h3_gate_residual_*; rebuild with csrc/cuda/h3/gate_residual.cu" + ) + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + if not residual.is_cuda: + raise ValueError("H3GateResidualCudaOp needs CUDA tensors") + if gate.stride(1) != 1: + gate = gate.contiguous() + return _H3GateResidualCuda.apply(residual, y, gate, index.contiguous()) + + def forward_fp32(self, residual, y, gate, index): + return self.forward(residual, y, gate, index).float() diff --git a/rl_engine/kernels/ops/pytorch/h3/gate_residual.py b/rl_engine/kernels/ops/pytorch/h3/gate_residual.py new file mode 100644 index 000000000..d2be659d5 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/gate_residual.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 gated residual (RFC #420 row ``adaln_gate_residual``). + +After each sublayer of a block, diffusers adds the sublayer output back to +the residual stream through the per-row AdaLN gate:: + + hidden = residual + gate.index_select(0, adaln_indices) * sublayer_output + +``gate`` is ``gate_msa`` (after attention) or ``gate_mlp`` (after the FFN), +an ``(R, H)`` row view of the AdaLN table; ``residual`` and the sublayer +output are ``(..., S, H)``. The order is fixed by RFC #420 §4: +``residual + gate[row] * sublayer_output``. +""" + +from __future__ import annotations + +import torch + +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_gate_residual( + residual: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + if residual.dtype not in _DTYPES: + raise TypeError(f"residual must be fp32, bf16 or fp16, got {residual.dtype}") + if y.shape != residual.shape or y.dtype != residual.dtype or y.device != residual.device: + raise ValueError("y must match the residual's shape, dtype and device") + if residual.dim() < 2 or residual.numel() == 0: + raise ValueError( + f"residual must be a non-empty (..., S, H) tensor, got {tuple(residual.shape)}" + ) + hidden = residual.shape[-1] + if gate.dim() != 2 or gate.shape[1] != hidden or gate.device != residual.device: + raise ValueError(f"gate must be (R, {hidden}) on the residual's device") + if same_dtype and gate.dtype != residual.dtype: + raise TypeError(f"gate dtype {gate.dtype} must match the residual dtype {residual.dtype}") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != residual.device: + raise TypeError("index must be a 1-D int64 tensor on the residual's device") + if residual.shape[-2] != index.shape[0]: + raise ValueError(f"residual must be (..., S, H) with S = len(index) = {index.shape[0]}") + if check_range and bool(((index < 0) | (index >= gate.shape[0])).any()): + raise IndexError(f"index must lie in [0, {gate.shape[0]})") + + +class _DeterministicGateResidual(torch.autograd.Function): + """Eager forward; the gate gradient as FP32 per-row sums instead of + ``index_select``'s BF16 atomic scatter-add (``d_residual``/``d_y`` are the + eager VJPs).""" + + @staticmethod + def forward(ctx, residual, y, gate, index): + gathered = gate.index_select(0, index) + ctx.save_for_backward(y, gathered, index) + ctx.gate_rows = gate.shape[0] + return residual + gathered * y + + @staticmethod + def backward(ctx, grad): + y, gathered, index = ctx.saved_tensors + contrib = (grad.float() * y.float()).reshape(-1, index.shape[0], y.shape[-1]) + per_position = contrib.sum(dim=0) # (S, H): fixed reduction over the batch + d_gate = per_position.new_zeros((ctx.gate_rows, y.shape[-1])) + for row in range(ctx.gate_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_gate[row] = per_position.index_select(0, positions).sum(dim=0) + return grad, grad * gathered, d_gate.to(gathered.dtype), None + + +class NativeH3GateResidualOp: + """PyTorch reference: ``forward`` is the eager provider expression (with a + deterministic gate gradient), ``forward_fp32`` the FP64 golden returned in + FP32. The raw diffusers path, atomic backward included, is + ``rl_engine.testing.h3_provider.provider_gate_residual``.""" + + op_class = "reduction" # forward is elementwise; the gate VJP sums positions + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + return _DeterministicGateResidual.apply(residual, y, gate, index) + + def forward_fp32(self, residual, y, gate, index): + validate_h3_gate_residual(residual, y, gate, index, same_dtype=False) + out = residual.double() + gate.double().index_select(0, index) * y.double() + return out.float() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 8fa1d1453..bcfc897b8 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -223,6 +223,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ) PYTORCH_H3_RMSNORM = "rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp" CUDA_H3_RMSNORM = "rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp" + PYTORCH_H3_GATE_RESIDUAL = ( + "rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp" + ) + CUDA_H3_GATE_RESIDUAL = "rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -650,6 +654,10 @@ def __init__(self): OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, ], "h3_rmsnorm": [OpBackend.CUDA_H3_RMSNORM, OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [ + OpBackend.CUDA_H3_GATE_RESIDUAL, + OpBackend.PYTORCH_H3_GATE_RESIDUAL, + ], }, "rocm": { "logp": [ @@ -695,6 +703,7 @@ def __init__(self): "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -744,6 +753,7 @@ def __init__(self): "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -774,6 +784,7 @@ def __init__(self): "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], }, # 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 fe010c7e3..7b13fc60f 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -33,6 +33,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_gate_residual, provider_norm_modulate, provider_time_embedder, provider_time_proj, @@ -146,6 +147,35 @@ def _norm1_modulated(ctx: dict[str, Any], modulation, call): ) ) +STAGES.append( + Stage( + name="adaln_gate_residual", + op_type="adaln_gate_residual", + # residual + gate_msa[row] * attention_output (the stand-in sublayer output). + candidate=lambda op, ctx, _up: op( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider=lambda ctx, _up: provider_gate_residual( + ctx["hidden"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ctx["sublayer"], + ), + golden=lambda op, ctx, _up: op.forward_fp32( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider_bitwise_isolated=True, # eager rounding order, elementwise + golden_atol=5e-2, # carries the gate's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ) +) + # The backward replay covers the conditioning chain up to the row gather. CONDITIONING_STAGES = ( "timestep_sinusoid_h3", @@ -199,6 +229,14 @@ def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dic ) .to(torch.bfloat16) .cuda(), + # Stand-in for the attention output the gate multiplies (the attention + # rows belong to other contributors): BF16 (1, S, H), seeded. + "sublayer": ( + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 3)) + * 3.0 + ) + .to(torch.bfloat16) + .cuda(), } diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index 05c7ebec9..8c084afc7 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -81,3 +81,11 @@ def provider_norm_modulate( return norm_hidden_states * (1.0 + scale.index_select(0, indices)) + shift.index_select( 0, indices ) + + +def provider_gate_residual( + residual: torch.Tensor, gate: torch.Tensor, indices: torch.Tensor, sublayer_output: torch.Tensor +) -> torch.Tensor: + """``residual + gate.index_select(0, adaln_indices) * sublayer_output`` (block).""" + + return residual + gate.index_select(0, indices) * sublayer_output diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index dfecec419..021dab4b0 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -20,6 +20,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_gate_residual, provider_norm_modulate, provider_time_embedder, provider_time_proj, @@ -467,12 +468,117 @@ def golden(x_, w_, sh_, sc_): } +# --------------------------------------------------------------------------- # +# adaln_gate_residual (gate_msa after attention) +# --------------------------------------------------------------------------- # + + +def _gate_inputs(seq: int, seed: int = 0, dtype=torch.bfloat16): + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).to(dtype) + residual = torch.randn(1, seq, 5376, device="cuda", generator=g).to(dtype) + y = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 3).to(dtype) + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate_view(table): + return table.view(table.shape[0], 6, -1)[:, 2] + + +def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_gate_residual", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + table, residual, y, index = _gate_inputs(seq) + gate = _gate_view(table) + grad = torch.randn_like(residual) + + def backward(fn, residual=residual, y=y, table=table, grad=grad): + leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] + fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + + cases.append( + { + "op": "adaln_gate_residual", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 3 * seq * 5376 * 2, # read residual and y, write out + "candidate": lambda r=residual, y=y, g=gate, i=index: op.forward( + r, y, g, i, check_range=False + ), + "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( + r, g, i, y + ), + "candidate_backward": lambda b=backward, i=index: b( + lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) + ), + "provider_backward": lambda b=backward, i=index: b( + lambda r, y_, g: provider_gate_residual(r, g, i, y_) + ), + } + ) + return cases + + +def _gate_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("adaln_gate_residual", device="cuda") + forward_bitwise = {} + for dtype in (torch.bfloat16, torch.float16, torch.float32): + table, residual, y, index = _gate_inputs(777, seed=1, dtype=dtype) + gate = _gate_view(table) + forward_bitwise[str(dtype).removeprefix("torch.")] = bool( + torch.equal( + op(residual, y, gate, index), provider_gate_residual(residual, gate, index, y) + ) + ) + table, residual, y, index = _gate_inputs(4097, seed=2) + invariant = bool( + torch.equal( + op(residual[:, 10:60], y[:, 10:60], _gate_view(table), index[10:60])[0], + op(residual, y, _gate_view(table), index)[0, 10:60], + ) + ) + grad = torch.randn( + residual.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(residual.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ref = grads(lambda r, y_, g: provider_gate_residual(r, g, index, y_), torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda r, y_, g: op(r, y_, g, index)), + ("provider", lambda r, y_, g: provider_gate_residual(r, g, index, y_)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("d_residual", "d_sublayer", "d_gate"), first, ref) + }, + } + return { + "forward_bitwise_vs_diffusers": forward_bitwise, + "rows_batch_invariant": invariant, + "backward": backward, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, "adaln_projection_3mod": _projection_perf, "adaln_row_gather": _gather_perf, "h3_rmsnorm": _norm_perf, + "adaln_gate_residual": _gate_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, @@ -480,4 +586,5 @@ def golden(x_, w_, sh_, sc_): "adaln_projection_3mod": _projection_accuracy, "adaln_row_gather": _gather_accuracy, "h3_rmsnorm": _norm_accuracy, + "adaln_gate_residual": _gate_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index 49d502f12..ee8e2c0ff 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -232,6 +232,7 @@ def fake_load_backend(backend): "adaln_projection_3mod", "adaln_row_gather", "h3_rmsnorm", + "adaln_gate_residual", ) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index da8bc6774..2c62208b7 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -409,12 +409,84 @@ def repeat(mode: str) -> str: return fig +def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: + _style(ax, "Forward and backward time (log, lower is better)", "ms") + _bars( + ax, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + (f"{provider} fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + f"{provider} bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + + +def _backward_errors(ax, backward: dict[str, Any], keys: list[str], note: str) -> None: + _style(ax, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") + + def label(name: str, mode: str) -> str: + return ( + f"{name} (repeat-bitwise: {'yes' if backward[mode]['repeat_bitwise_equal'] else 'no'})" + ) + + floor = 1e-9 + _bars( + ax, + keys, + [ + ( + label("RL-Kernel CUDA", "cuda"), + RL_KERNEL, + [max(backward["cuda"]["rel_error"][k], floor) for k in keys], + ), + ( + label("diffusers", "provider"), + PROVIDER, + [max(backward["provider"]["rel_error"][k], floor) for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + for text in ax.texts: + if text.get_text() == f"{floor:.1e}": + text.set_text("exact") + ax.set_ylim(floor / 3, 1e4) + ax.text(0.98, 0.80, note, transform=ax.transAxes, ha="right", fontsize=8, color=INK) + + +def plot_gate_residual(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "adaln_gate_residual (residual + gate_msa[row] * y, H = 5376)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + ok = all(acc["forward_bitwise_vs_diffusers"].values()) + _backward_errors( + right, + acc["backward"], + ["d_residual", "d_sublayer", "d_gate"], + f"forward bitwise equal to diffusers (bf16/fp16/fp32): {'yes' if ok else 'NO'}", + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, "adaln_projection_3mod": plot_projection, "adaln_row_gather": plot_gather, "h3_rmsnorm": plot_rmsnorm, + "adaln_gate_residual": plot_gate_residual, } diff --git a/setup.py b/setup.py index fb671d0f0..3f38e1fbe 100644 --- a/setup.py +++ b/setup.py @@ -161,6 +161,7 @@ def get_extensions(): cuda_sources.append("csrc/cuda/h3/det_linear.cu") cuda_sources.append("csrc/cuda/h3/adaln_row_gather.cu") cuda_sources.append("csrc/cuda/h3/rmsnorm_modulate.cu") + cuda_sources.append("csrc/cuda/h3/gate_residual.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_gate_residual.py b/tests/h3/test_h3_adaln_gate_residual.py new file mode 100644 index 000000000..d67210257 --- /dev/null +++ b/tests/h3/test_h3_adaln_gate_residual.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_gate_residual``: ``residual + gate[row] * sublayer_output``. + +* forward and ``d_sublayer`` are bitwise equal to diffusers' eager + expression in every dtype (the gate row is gathered in the kernel); +* ``d_gate`` is a deterministic FP32 segment sum, checked against FP64; +* rows are elementwise-independent; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.h3.gate_residual import NativeH3GateResidualOp +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_gate_residual + +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.gate_residual import ( + H3GateResidualCudaOp, + h3_gate_residual_available, + ) + + if not h3_gate_residual_available(): + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + return H3GateResidualCudaOp() + + +def _case(seq, hidden=5376, dtype=torch.bfloat16, seed=0, batch=2): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(9, 6 * hidden, generator=g) * 0.5).to(dtype).cuda() + residual = torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda() + y = (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate(table, chunk=2): + return table.view(table.shape[0], 6, -1)[:, chunk] # gate_msa (2) / gate_mlp (5) view + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3GateResidualOp() + res, y, gate = torch.zeros(4, 8), torch.zeros(4, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op(res, y, gate, torch.tensor([0, 1, 2, 3])) + with pytest.raises(ValueError): + op(res, y[:, :4], gate, torch.tensor([0, 1, 2, 0])) + with pytest.raises(TypeError): + op(res, y, gate.bfloat16(), torch.tensor([0, 1, 2, 0])) + with pytest.raises(ValueError): # S mismatch + op(res, y, gate, torch.tensor([0, 1])) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("seq, hidden", [(4097, 5376), (33, 12), (17, 13)]) + def test_forward_bitwise_equal_to_diffusers(self, dtype, seq, hidden): + table, residual, y, index = _case(seq, hidden, dtype, seed=seq) + for chunk in (2, 5): + ours = _cuda_op()(residual, y, _gate(table, chunk), index) + assert torch.equal( + ours, provider_gate_residual(residual, _gate(table, chunk), index, y) + ) + + def test_rows_are_batch_and_position_invariant(self): + table, residual, y, index = _case(500, seed=1) + full = _cuda_op()(residual, y, _gate(table), index) + part = _cuda_op()(residual[1:2, 100:150], y[1:2, 100:150], _gate(table), index[100:150]) + assert torch.equal(part[0], full[1, 100:150]) + + def test_backward(self): + table, residual, y, index = _case(4097, seed=2) + grad = torch.randn_like(residual) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ours = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + again = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + theirs = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy)) + golden = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy), torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + assert torch.equal(ours[0], theirs[0]) # d_residual + assert torch.equal(ours[1], theirs[1]) # d_sublayer: exact product, one rounding + rel = ((ours[2].double() - golden[2]).abs().max() / golden[2].abs().max()).item() + assert rel < 5e-3 # only the BF16 rounding of the table gradient + assert torch.count_nonzero(ours[2].view(9, 6, -1)[:, [0, 1, 3, 4, 5]]) == 0 + + def test_rejects_cpu_and_bad_index(self): + table, residual, y, index = _case(10, hidden=16) + with pytest.raises(IndexError): + _cuda_op()(residual, y, _gate(table), index + 9) + with pytest.raises(ValueError): + _cuda_op()(residual.cpu(), y.cpu(), _gate(table).cpu(), index.cpu()) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_gate_residual", device="cuda") + assert type(op).__name__ == "H3GateResidualCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 2b165c088..32c6322c3 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -54,6 +54,7 @@ def test_all_ws1_single_ops_are_registered(): "adaln_projection_3mod", "adaln_row_gather", "h3_rmsnorm", + "adaln_gate_residual", } <= names From cd688b354823109b19b1abd24d8fdfeb5b7a7fc5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:17:29 +0000 Subject: [PATCH 02/14] [WS1][MiniMax-H3] adaln_gate_residual: B200 evidence report and figure Produced by scripts/h3_evidence.py at a26b41a 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-gate-residual.md | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 7a3010185..f55412449 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -51,7 +51,7 @@ hidden = op(residual, attn_output, gate_msa, adaln_indices) once. Diffusers' `d_gate` uses `index_select`'s BF16 atomic scatter-add instead. It is - non-deterministic and about 35× further from the FP64 golden. + non-deterministic and 13–49× further from the FP64 golden (its error varies from run to run). ## Performance Notes @@ -63,14 +63,22 @@ B200, B = 1, H = 5376, BF16: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.05 ms | 0.06 ms | 1.21 ms | 0.68 ms | -| 32768 | 0.24 ms | 0.41 ms | 1.40 ms | 2.71 ms | -| 131072 | 0.92 ms | 1.59 ms | 3.31 ms | 10.6 ms | +| 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | +| 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | +| 131072 | 0.92 ms | 1.59 ms | 3.29 ms | 10.6 ms | The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and tile setup. +## Evidence + +![adaln_gate_residual on B200: latency and backward accuracy](../usage/evidence/h3-adaln-gate-residual-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-adaln-gate-residual-b200/report.json), +written by `scripts/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. + ## Tests ```bash From 7090f00fc986130391ebe91541f404586313df91 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:28:04 +0000 Subject: [PATCH 03/14] fix(h3): address native validation and evidence review comments Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 11 +- csrc/cuda/h3/rmsnorm_modulate.cu | 18 ++- docs/operators/h3-adaln-gate-residual.md | 2 + docs/operators/h3-rmsnorm.md | 2 + rl_engine/testing/h3_report.py | 61 +++++--- scripts/plot_h3_evidence.py | 14 ++ tests/h3/test_h3_native_validation.py | 124 ++++++++++++++++ tests/h3/test_h3_report.py | 177 +++++++++++++++++++++++ 8 files changed, 380 insertions(+), 29 deletions(-) create mode 100644 tests/h3/test_h3_native_validation.py create mode 100644 tests/h3/test_h3_report.py diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 18ede9e72..c2bd7921a 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -144,9 +144,9 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc "gate must be a CUDA tensor with the activations' dtype"); TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, "gate must be (R, N) with unit column stride"); - TORCH_CHECK(index.is_cuda() && index.scalar_type() == at::kLong && index.dim() == 1 && - index.is_contiguous() && index.numel() > 0, - "index must be a non-empty contiguous int64 CUDA tensor"); + TORCH_CHECK(index.is_cuda() && index.device() == x.device() && index.scalar_type() == at::kLong && + index.dim() == 1 && index.is_contiguous() && index.numel() > 0, + "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; } @@ -207,8 +207,9 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: check_act(y, "y", grad); const Gate g = make_gate(gate, index, grad); for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { - TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kLong && t->dim() == 1 && t->is_contiguous(), - "tile metadata must be contiguous int64 CUDA tensors"); + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 tensors on ", grad.device()); } const int64_t rows = grad.size(0); const int64_t n = grad.size(1); diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 5212324b1..2775572e6 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -276,8 +276,9 @@ Modulation make_modulation(const c10::optional& shift, } TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), "shift and scale must be views of the same table layout"); - TORCH_CHECK(ix.is_cuda() && ix.scalar_type() == at::kLong && ix.dim() == 1 && ix.is_contiguous(), - "row index must be a contiguous int64 CUDA tensor"); + TORCH_CHECK(ix.is_cuda() && ix.device() == x.device() && ix.scalar_type() == at::kLong && + ix.dim() == 1 && ix.is_contiguous() && ix.numel() > 0, + "row index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", ix.size(0), ")"); mod.shift = sh.data_ptr(); @@ -349,6 +350,16 @@ std::vector h3_rmsnorm_backward( const int64_t rows = x.size(0); const int64_t n = x.size(1); const Modulation mod = make_modulation(shift, scale, index, x, n); + if (mod.index != nullptr) { + TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && + seg_first_tile.has_value(), + "modulated backward needs the sorted segment tiles"); + for (const auto* t : {&*sorted_pos, &*tile_begin, &*tile_end, &*seg_first_tile}) { + TORCH_CHECK(t->is_cuda() && t->device() == x.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 tensors on ", x.device()); + } + } const c10::cuda::CUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); auto dx = torch::empty_like(x); @@ -375,9 +386,6 @@ std::vector h3_rmsnorm_backward( C10_CUDA_KERNEL_LAUNCH_CHECK(); if (mod.index == nullptr) return {dx, dweight}; - TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && - seg_first_tile.has_value(), - "modulated backward needs the sorted segment tiles"); const int64_t segments = seg_first_tile->numel() - 1; const int64_t seg_tiles = tile_begin->numel(); const int64_t width = 2 * n; diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index f55412449..2ef31ba54 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -16,6 +16,8 @@ gathers the gate row inside the kernel, so the gathered `(S, H)` gate is never m ## Entry Point ```python +from rl_engine.kernels.registry import kernel_registry + op = kernel_registry.get_op("adaln_gate_residual", device="cuda") hidden = op(residual, attn_output, gate_msa, adaln_indices) ``` diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 957670552..faaed48e2 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -21,6 +21,8 @@ gathered inside the kernel, so the `(S, H)` tensors that ## Entry Point ```python +from rl_engine.kernels.registry import kernel_registry + op = kernel_registry.get_op("h3_rmsnorm", device="cuda") n = op(x, weight) # plain RMSNorm out = op.forward_modulated(x, weight, shift_msa, scale_msa, adaln_indices) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 021dab4b0..489404bb6 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -25,7 +25,11 @@ provider_time_embedder, provider_time_proj, ) -from rl_engine.testing.h3_weights import h3_weights_dir, load_h3_conditioning_weights +from rl_engine.testing.h3_weights import ( + WEIGHTS_ENV, + h3_weights_dir, + load_h3_conditioning_weights, +) # Keys of a perf case that hold a timed callable, in display order. TIMED_KEYS = ( @@ -247,7 +251,7 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, invariant = all( all( torch.equal(s, f[3 * i : 3 * i + 3]) - for s, f in zip(op(temb[i : i + 1], weight, bias), full) + for s, f in zip(op(temb[i : i + 1], weight, bias), full, strict=True) ) for i in range(9) ) @@ -304,7 +308,6 @@ def backward(fn, ti=ti, tags=tags, grads=grads): 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]) @@ -314,7 +317,9 @@ def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: 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)) + forward_bitwise[str(seq)] = all( + torch.equal(a, b) for a, b in zip(ours, theirs, strict=True) + ) # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. ti, tags = h3_packed_layout(4097, 3, seed=1) @@ -339,16 +344,22 @@ def grad_of(fn): "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 { + result: dict[str, Any] = { "forward_bitwise_vs_index_select": forward_bitwise, "op_backward": backward, - "chain_backward": chain, + "chain_backward": [], } + if h3_weights_dir() is None: + result["chain_backward_skipped"] = ( + f"{WEIGHTS_ENV} not set; run scripts/prepare_h3_weights.py for chain measurements" + ) + else: + weights = load_h3_conditioning_weights("cuda") + result["chain_backward"] = [ + 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 result # --------------------------------------------------------------------------- # @@ -406,14 +417,19 @@ def backward(fn, x=x, weight=weight, shift=shift, scale=scale, grad=grad): def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("h3_rmsnorm", device="cuda") - weights = load_h3_conditioning_weights("cuda") - plain = {} - for name in ( + weight_source = "pinned_checkpoint" if h3_weights_dir() is not None else "synthetic" + names = [ "transformer_blocks.0.norm1.weight", "transformer_blocks.0.norm2.weight", "token_refiner.final_norm.weight", "norm_out.norm.weight", - ): + ] + weights = { + name: weight.bfloat16() + for name, weight in zip(names, h3_params(names, [(5376,)] * len(names)), strict=True) + } + plain = {} + for name in names: x = torch.randn(2, 777, 5376, device="cuda").bfloat16() ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) plain[name] = bool(torch.equal(op(x, weights[name]), ref)) @@ -454,13 +470,16 @@ def golden(x_, w_, sh_, sc_): ): first, second = grads(fn), grads(fn) backward[name] = { - "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), "rel_error": { key: float((g.double() - r).abs().max() / r.abs().max()) - for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref) + for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref, strict=True) }, } return { + "weight_source": weight_source, "plain_bitwise_vs_nn_rmsnorm": plain, "modulated_bitwise_vs_diffusers": modulated, "rows_batch_invariant": invariant, @@ -559,10 +578,14 @@ def grads(fn, dtype=None): ): first, second = grads(fn), grads(fn) backward[name] = { - "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), "rel_error": { key: float((g.double() - r).abs().max() / r.abs().max()) - for key, g, r in zip(("d_residual", "d_sublayer", "d_gate"), first, ref) + for key, g, r in zip( + ("d_residual", "d_sublayer", "d_gate"), first, ref, strict=True + ) }, } return { diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 15c800b5f..d8bd5ea43 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -302,6 +302,20 @@ def plot_gather(report: dict[str, Any]): "Whole-chain grads of the FP32 time embedder vs FP64", "max abs error / golden max", ) + if not chain: + right.text( + 0.5, + 0.5, + report["accuracy"].get("chain_backward_skipped", "No chain measurements available"), + transform=right.transAxes, + ha="center", + va="center", + color=INK, + fontsize=9, + wrap=True, + ) + right.set_axis_off() + return fig names = { "candidate": ("RL-Kernel, separate ops", THIRD, "o"), "candidate_fused": ("RL-Kernel, fused modulation", RL_KERNEL, "D"), diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py new file mode 100644 index 000000000..8dd919028 --- /dev/null +++ b/tests/h3/test_h3_native_validation.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Exercise H3 validation at the native entrypoints, bypassing Python guards.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@pytest.fixture +def native_extension(): + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + names = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward", + "h3_gate_residual_forward", + "h3_gate_residual_backward", + ) + if not _EXT_AVAILABLE or not all(hasattr(_C, name) for name in names): + pytest.skip("rl_engine._C lacks native H3 entrypoints") + return _C + + +@pytest.fixture +def native_case(): + device = "cuda:0" + x = torch.arange(16, device=device, dtype=torch.float32).view(2, 8) / 10 + return { + "x": x, + "weight": torch.ones(8, device=device), + "shift": torch.zeros(2, 8, device=device), + "scale": torch.zeros(2, 8, device=device), + "gate": torch.ones(2, 8, device=device), + "grad": torch.ones_like(x), + "rstd": torch.rsqrt(x.square().mean(-1) + 1e-5), + "index": torch.arange(2, device=device, dtype=torch.int64), + "sorted_pos": torch.arange(2, device=device, dtype=torch.int64), + "tile_begin": torch.arange(2, device=device, dtype=torch.int64), + "tile_end": torch.arange(1, 3, device=device, dtype=torch.int64), + "seg_first_tile": torch.arange(3, device=device, dtype=torch.int64), + } + + +def _native_call(extension, operation, boundary, case): + x, index = case["x"], case["index"] + tiles = tuple(case[name] for name in ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile")) + if operation == "rmsnorm": + modulation = (case["shift"], case["scale"], index) + if boundary == "forward": + return extension.h3_rmsnorm_forward(x, case["weight"], 1e-5, *modulation) + return extension.h3_rmsnorm_backward( + case["grad"], x, case["weight"], case["rstd"], *modulation, *tiles + ) + if boundary == "forward": + return extension.h3_gate_residual_forward(x, x, case["gate"], index) + return extension.h3_gate_residual_backward(case["grad"], x, case["gate"], index, *tiles) + + +def _other_cuda(tensor): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + return tensor.to("cuda:1") + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("boundary", ["forward", "backward"]) +@pytest.mark.parametrize("index_case", ["empty", "cpu", "other_cuda"]) +def test_native_rejects_invalid_row_index( + native_extension, native_case, operation, boundary, index_case +): + index = native_case["index"] + if index_case == "empty": + native_case["index"] = index[:0] + elif index_case == "cpu": + native_case["index"] = index.cpu() + else: + native_case["index"] = _other_cuda(index) + with pytest.raises( + RuntimeError, match="index must be a non-empty contiguous int64 tensor on cuda:0" + ): + _native_call(native_extension, operation, boundary, native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +@pytest.mark.parametrize("device", ["cpu", "other_cuda"]) +def test_native_rejects_mixed_device_tiles( + native_extension, native_case, operation, metadata, device +): + tensor = native_case[metadata] + native_case[metadata] = tensor.cpu() if device == "cpu" else _other_cuda(tensor) + with pytest.raises( + RuntimeError, match="tile metadata must be contiguous int64 tensors on cuda:0" + ): + _native_call(native_extension, operation, "backward", native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +def test_native_accepts_same_device_inputs(native_extension, native_case, operation): + x = native_case["x"].detach().clone().requires_grad_(True) + if operation == "rmsnorm": + weight = native_case["weight"].detach().clone().requires_grad_(True) + shift = native_case["shift"].detach().clone().requires_grad_(True) + scale = native_case["scale"].detach().clone().requires_grad_(True) + norm = torch.nn.functional.rms_norm(x, (8,), weight, 1e-5) + ref = norm * (1 + scale) + shift + ref.backward(native_case["grad"]) + expected_grads = (x.grad, weight.grad, shift.grad, scale.grad) + got, _ = _native_call(native_extension, operation, "forward", native_case) + else: + gate = native_case["gate"].detach().clone().requires_grad_(True) + ref = native_case["x"] + gate * x + ref.backward(native_case["grad"]) + expected_grads = (x.grad, gate.grad) + got = _native_call(native_extension, operation, "forward", native_case) + torch.testing.assert_close(got, ref) + actual_grads = _native_call(native_extension, operation, "backward", native_case) + for actual, expected in zip(actual_grads, expected_grads, strict=True): + torch.testing.assert_close(actual, expected) diff --git a/tests/h3/test_h3_report.py b/tests/h3/test_h3_report.py new file mode 100644 index 000000000..b6bb6ec3f --- /dev/null +++ b/tests/h3/test_h3_report.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Optional checkpoint handling preserves standalone operator measurements.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.testing import h3_chain, h3_report +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_adaln_row_gather, provider_norm_modulate +from rl_engine.testing.h3_weights import WEIGHTS_ENV + + +@pytest.fixture +def small_cpu_report(monkeypatch): + """Exercise actual report/autograd logic with small CPU tensors and provider ops.""" + + def dimension(size): + return {5376: 4, 6 * 5376: 24, 257: 160, 777: 160, 4097: 160, 32768: 160}.get(size, size) + + class CpuTorch: + nn = SimpleNamespace( + functional=SimpleNamespace( + rms_norm=lambda x, _shape, weight, eps: F.rms_norm(x, (x.shape[-1],), weight, eps) + ) + ) + + def __getattr__(self, name): + return getattr(torch, name) + + def Generator(self, device): + return torch.Generator(device="cpu") + + def randn(self, *shape, **kwargs): + if len(shape) == 1 and isinstance(shape[0], (tuple, list, torch.Size)): + shape = shape[0] + kwargs["device"] = "cpu" + return torch.randn(*(dimension(size) for size in shape), **kwargs) + + class Gather: + def __call__(self, rows, ti, tags): + return provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + + forward_fp32 = __call__ + + class Norm: + def __call__(self, x, weight): + return F.rms_norm(x, (x.shape[-1],), weight, 1e-5) + + forward_modulated = staticmethod(provider_norm_modulate) + + gather = Gather() + ops = {"adaln_row_gather": gather, "h3_rmsnorm": Norm()} + registry = SimpleNamespace( + get_op=lambda name, device: ops[name], + _priority_map={"cpu": {"adaln_row_gather": ["gather"]}}, + _get_or_create_backend=lambda _backend: gather, + ) + monkeypatch.setattr(h3_report, "torch", CpuTorch()) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: h3_packed_layout(dimension(seq), num, seed=seed, device="cpu"), + ) + + def norm_inputs(seq, seed=0): + generator = torch.Generator().manual_seed(seed) + weight = h3_report.h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0] + shift, scale = [torch.randn(9, 4, generator=generator).bfloat16() for _ in range(2)] + x = torch.randn(1, dimension(seq), 4, generator=generator).bfloat16() + ti, tags = h3_packed_layout(dimension(seq), 3, seed=seed, device="cpu") + return x, weight.bfloat16(), shift, scale, ti * 3 + tags + + monkeypatch.setattr(h3_report, "_norm_inputs", norm_inputs) + return registry + + +def test_gather_without_weights_preserves_op_results(small_cpu_report, monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + + def unexpected_load(*args, **kwargs): + pytest.fail("an unset checkpoint must not be loaded") + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", unexpected_load) + result = h3_report._gather_accuracy(small_cpu_report) + assert all(result["forward_bitwise_vs_index_select"].values()) + assert set(result["op_backward"]) == {"cuda", "provider"} + assert all(row["repeat_bitwise_equal"] for row in result["op_backward"].values()) + assert result["chain_backward"] == [] + assert WEIGHTS_ENV in result["chain_backward_skipped"] + + +@pytest.mark.parametrize("configured", [False, True]) +def test_norm_reports_weight_source(small_cpu_report, monkeypatch, tmp_path, configured): + if configured: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + else: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + calls = [] + + def load_weights(device, names): + assert configured, "an unset checkpoint must not be loaded" + calls.append(names) + return {name: torch.ones(4, dtype=torch.bfloat16) for name in names} + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", load_weights) + result = h3_report._norm_accuracy(small_cpu_report) + assert result["weight_source"] == ("pinned_checkpoint" if configured else "synthetic") + assert len(result["plain_bitwise_vs_nn_rmsnorm"]) == 4 + assert all(result["plain_bitwise_vs_nn_rmsnorm"].values()) + assert result["modulated_bitwise_vs_diffusers"] + assert result["rows_batch_invariant"] + assert set(result["backward"]) == {"cuda", "provider"} + assert bool(calls) is configured + + +def test_gather_with_weights_runs_chain(small_cpu_report, monkeypatch, tmp_path): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + weights = {"weight": torch.ones(1)} + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", lambda device: weights) + + def run_chain(registry, loaded, **case): + assert registry is small_cpu_report + assert loaded is weights + return case + + monkeypatch.setattr(h3_chain, "run_backward_case", run_chain) + result = h3_report._gather_accuracy(small_cpu_report) + assert len(result["chain_backward"]) == 5 + assert result["chain_backward"][-1] == {"num_timesteps": 4, "seq_len": 32768} + assert "chain_backward_skipped" not in result + + +@pytest.mark.parametrize("accuracy", [h3_report._gather_accuracy, h3_report._norm_accuracy]) +def test_configured_missing_weights_raise(small_cpu_report, monkeypatch, tmp_path, accuracy): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + with pytest.raises(FileNotFoundError, match="missing; run scripts/prepare_h3_weights.py"): + accuracy(small_cpu_report) + + +def test_plot_gather_without_chain(small_cpu_report, monkeypatch, tmp_path): + pytest.importorskip("matplotlib") + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + script = Path(__file__).resolve().parents[2] / "scripts" / "plot_h3_evidence.py" + spec = importlib.util.spec_from_file_location("plot_h3_evidence_test", script) + plot = importlib.util.module_from_spec(spec) + spec.loader.exec_module(plot) + accuracy = h3_report._gather_accuracy(small_cpu_report) + report = { + "environment": {"gpu": "CPU fixture", "torch": torch.__version__, "cuda": "none"}, + "rl_kernel_commit": "testcommit", + "model_revision": "pinned_revision", + "accuracy": accuracy, + "perf": [ + { + "case": "S=160", + "candidate_us": 1, + "provider_us": 2, + "candidate_backward_us": 3, + "provider_backward_us": 4, + } + ], + } + figure = plot.plot_gather(report) + output = tmp_path / "figure.png" + figure.savefig(output) + assert output.stat().st_size > 0 + assert figure.axes[1].texts[0].get_text() == accuracy["chain_backward_skipped"] + plot.plt.close(figure) From dd35c14482d7a7a5693b78848230acccd0e555ca Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:16:48 +0000 Subject: [PATCH 04/14] docs(h3): gate_residual vs existing implementations (B200, 4010854) 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-gate-residual.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 2ef31ba54..6b923d740 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -81,6 +81,29 @@ The data is in [`report.json`](../usage/evidence/h3-adaln-gate-residual-b200/rep written by `scripts/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. +## Existing implementations (RFC #420 reuse rule) + +![gate_residual vs existing implementations](../usage/evidence/h3-prior-art-b200/gate_residual.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers residual + gate.index_select(...) * y | **no** (param/table grads not repeatable) | 3.0e-03 / 4.3e-02 / 637 µs | 2.9e-03 / 1.4e-01 / 2502 µs | +| rl-kernel H3GateResidualCudaOp | yes | 3.0e-03 / 2.6e-03 / 1210 µs | 2.9e-03 / 2.6e-03 / 1596 µ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. [`gate_residual.json`](../usage/evidence/h3-prior-art-b200/gate_residual.json) +was written from a clean tree at `4010854` by + +```bash +python scripts/h3_prior_art.py --op gate_residual --out docs/usage/evidence/h3-prior-art-b200/gate_residual.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/gate_residual.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 3bbcca42f7e8a5a8a7b4d8753a46c108cd1bd653 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:30:51 +0000 Subject: [PATCH 05/14] evidence(h3): gate_residual direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 910d614a. There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index e69de13cd..e003fd5b8 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -104,6 +104,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. + +![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) + +There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. +[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/gate_residual.json \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.png +``` + ## Tests ```bash From a8e5901b5d194bf846c5ad670ea9780850087c04 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:18 +0000 Subject: [PATCH 06/14] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 12 ++++++------ docs/operators/h3-adaln-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 6 files changed, 37 insertions(+), 37 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index e003fd5b8..8847d6161 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -75,15 +75,15 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](../../reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-gate-residual-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) -![gate_residual vs existing implementations](../../reports/experiments/h3-prior-art-b200/gate_residual.png) +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -94,7 +94,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../../reports/experiments/h3-prior-art-b200/gate_residual.json) +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) was written from a clean tree at `4010854` by ```bash @@ -111,10 +111,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a76cc0417..a9c8832e2 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -124,9 +124,9 @@ dependence on T. ## Evidence -![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also @@ -135,7 +135,7 @@ or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![adaln_projection vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_projection.png) +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../../reports/experiments/h3-prior-art-b200/adaln_projection.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) was written from a clean tree at `c88691c` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 88fee5f3d..94e32b711 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -122,13 +122,13 @@ alternates each iteration and is recorded in the report. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../../reports/experiments/h3-adaln-row-gather-b200/figure.png) +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward bitwise checks, op-level backward, and the chain-backward cases plotted above. -- [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. `report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only @@ -141,7 +141,7 @@ model does; every CUDA pass/fail flag is unchanged. ## Existing implementations (RFC #420 reuse rule) -![adaln_row_gather vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.png) +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) was written from a clean tree at `6c900ae` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_row_gather bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 740f16232..75dd3d477 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -92,9 +92,9 @@ each iteration and is recorded in the report. ## Evidence -![h3_rmsnorm on B200: latency and backward accuracy](../../reports/experiments/h3-rmsnorm-b200/figure.png) +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report.json), written by +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; @@ -103,7 +103,7 @@ The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report. ## Existing implementations (RFC #420 reuse rule) -![norm_modulate vs existing implementations](../../reports/experiments/h3-prior-art-b200/norm_modulate.png) +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -119,7 +119,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](../../reports/experiments/h3-prior-art-b200/norm_modulate.json) +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) was written from a clean tree at `ee83dec` by ```bash @@ -136,10 +136,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![norm_modulate bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 8049a87df..cc5809cee 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -109,15 +109,15 @@ kernel time. ## Evidence -![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../../reports/experiments/h3-timestep-mlp-b200/figure.png) +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-mlp-b200/report.json), written +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![timestep_mlp vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_mlp.png) +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -134,7 +134,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../../reports/experiments/h3-prior-art-b200/timestep_mlp.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) was written from a clean tree at `ffd958a` by ```bash @@ -151,10 +151,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index c5f69a15b..75edf172b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -88,15 +88,15 @@ about 1 KB per timestep. ## Evidence -![timestep_sinusoid_h3 on B200: latency and error vs FP64](../../reports/experiments/h3-timestep-sinusoid-b200/figure.png) +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-sinusoid-b200/report.json). +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). `tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. ## Existing implementations (RFC #420 reuse rule) -![timestep_sinusoid vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -108,7 +108,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) was written from a clean tree at `7a82917` by ```bash @@ -126,11 +126,11 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). -![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. -[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) was written from a clean tree at `0eb5ddd5` by ```bash From 38018081392ba4c1222a1bb32e0261a6122cf804 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:58 +0000 Subject: [PATCH 07/14] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 4 ++-- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 6 files changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 8847d6161..b35bc41c1 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -78,7 +78,7 @@ tile setup. ![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `e7bc53a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) @@ -114,7 +114,7 @@ gradient with diffusers' path and with this kernel. ![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `a746bff0` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a9c8832e2..3f79732cd 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -127,7 +127,7 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `579775f1` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `50d1e924` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 94e32b711..923ad2f77 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,11 +131,11 @@ There are two data files: - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +`report.json` was regenerated from a clean tree at commit `29b07c5d` with backward-only timings and alternating execution order, and with the op-backward check's rows and gradients drawn from a seeded generator. Before, they came from the process's default CUDA generator, -so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). -`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +so the fraction depended on how many draws the runner made first (99.99965% at `d2a6546`). +`chain_replay.json` was written from a clean tree at commit `61d23ade`. Both were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is unchanged. @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `62f8009e` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 75dd3d477..ac625ae35 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -95,7 +95,7 @@ each iteration and is recorded in the report. ![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by -`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -139,7 +139,7 @@ gradient with diffusers' path and with this kernel. ![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `83797833` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index cc5809cee..b65b19fca 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -112,7 +112,7 @@ kernel time. ![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written -by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `418998c`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) @@ -154,7 +154,7 @@ gradient with diffusers' path and with this kernel. ![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `73939f2d` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 75edf172b..ff6579cd8 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -131,7 +131,7 @@ kernel compiled with `--use_fast_math`, which the build excludes (see Known Limi SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) -was written from a clean tree at `0eb5ddd5` by +was written from a clean tree at `be6a94ae` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ From 669513d223d33fbc8d0af14b1b00733a50c858fd Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:07:41 +0000 Subject: [PATCH 08/14] fix(h3): keep the gate-residual forward out of the timed backward _gate_perf had no candidate_backward_setup/provider_backward_setup, so each timed backward call also created the leaves, ran the forward and built the graph while measure() reported backward_timing_scope = backward_only. The forward moves into the setup like in _gather_perf and _norm_perf. A test checks that every perf case with a timed backward has its setup. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 13 +++++++++---- .../models/minimax_h3/test_h3_report_timing.py | 18 ++++++++++++++++++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 492bfc1b6..aa1b53498 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -606,9 +606,12 @@ def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: gate = _gate_view(table) grad = torch.randn_like(residual) - def backward(fn, residual=residual, y=y, table=table, grad=grad): + def prepare_backward(fn, residual=residual, y=y, table=table): leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] - fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + return fn(leaves[0], leaves[1], _gate_view(leaves[2])) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -622,10 +625,12 @@ def backward(fn, residual=residual, y=y, table=table, grad=grad): "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( r, g, i, y ), - "candidate_backward": lambda b=backward, i=index: b( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: provider_gate_residual(r, g, i, y_) ), } diff --git a/tests/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py index 853680a67..9b77b16cf 100644 --- a/tests/models/minimax_h3/test_h3_report_timing.py +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -52,3 +52,21 @@ 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 + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (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 e5318ad2ab1a75100336762b08d057b65d5a853d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:11:52 +0000 Subject: [PATCH 09/14] fix(h3): do not ask for the checkpoint when the gate-residual evidence never reads it h3_evidence.py treated every op but timestep_sinusoid_h3 as weighted. adaln_gate_residual runs on seeded inputs only, so the CLI refused to run without the 4.8 GB shard and a run with it would have been labelled pinned_checkpoint. The weight-free ops are now an explicit set. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_cli.py | 1 + tools/validation/models/h3_evidence.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_cli.py b/tests/models/minimax_h3/test_h3_cli.py index 8e97fc6a2..3f9b726ed 100644 --- a/tests/models/minimax_h3/test_h3_cli.py +++ b/tests/models/minimax_h3/test_h3_cli.py @@ -145,6 +145,7 @@ def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, we "op, weight_source", [ ("timestep_sinusoid_h3", "not_applicable"), + ("adaln_gate_residual", "not_applicable"), ("timestep_mlp_fp32", "pinned_checkpoint"), ("adaln_projection_3mod", "pinned_checkpoint"), ], diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index ed8267a94..892366786 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -37,6 +37,9 @@ load_h3_manifest, ) +# Operators whose evidence runs on seeded inputs only and never reads the checkpoint. +WEIGHT_FREE_OPS = frozenset({"timestep_sinusoid_h3", "adaln_gate_residual"}) + _NEEDS_PINNED_MLP = ( "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " "RL_KERNEL_H3_WEIGHTS (tools/weights/prepare_h3_weights.py) " @@ -85,7 +88,7 @@ def main() -> None: ) args = parser.parse_args() require_pinned_weights(args.op, args.allow_random_weights) - uses_weights = args.op != "timestep_sinusoid_h3" + uses_weights = args.op not in WEIGHT_FREE_OPS smoke_run = args.allow_random_weights and args.op == "timestep_mlp_fp32" if uses_weights and h3_weights_dir() is None and not smoke_run: raise SystemExit(f"set {WEIGHTS_ENV} to the tools/weights/prepare_h3_weights.py output dir") From 239ef808e279b6f62dec4592c4311c7cc1b0c32f Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:27:36 +0000 Subject: [PATCH 10/14] fix(h3): assert the gate row in the native kernels and check the backward position repeat - gate_residual.cu: Gate carries the table row count and the row kernel asserts 0 <= index < rows before using it as an offset, like the RMSNorm and row-gather kernels (check_range=False callers and direct _C calls get a device-side assert instead of an out-of-range read). - The backward uses _repeat_positions() instead of index.repeat(rows // len(index)). - Test: rows -1 and R in forward and backward expect the device-side assert (subprocess). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 7 ++- .../minimax_h3/gate_residual.py | 9 ++-- .../minimax_h3/test_h3_adaln_gate_residual.py | 52 +++++++++++++++++++ 3 files changed, 63 insertions(+), 5 deletions(-) diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index c2bd7921a..6e0b5a85b 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -50,6 +50,7 @@ struct Gate { int64_t stride; const int64_t* index; // (S,) int64_t seq; + int64_t num_rows; }; // One block per row (grid-strided); 16-byte vectors when the row allows it. @@ -60,7 +61,9 @@ __global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __re Gate gate, bool vectorized) { constexpr int kVecElems = 16 / sizeof(T); for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { - const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; const T* ar = a + r * n; const T* br = kDy ? nullptr : b + r * n; T* orow = out + r * n; @@ -148,7 +151,7 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); - return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; } void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index 2c6c6be55..ad9bb6d8f 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -7,9 +7,12 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _repeat_positions, + _segment_tiles, +) from rl_engine.backends.extension import _C, _EXT_AVAILABLE -from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.gate_residual import validate_h3_gate_residual KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" @@ -40,7 +43,7 @@ def forward(ctx, residual, y, gate, index): def backward(ctx, grad): y2, gate, index = ctx.saved_tensors g2 = grad.contiguous().view_as(y2) - positions = index.repeat(y2.shape[0] // index.shape[0]) + positions = _repeat_positions(index, y2.shape[0]) tiles = _segment_tiles(positions, gate.shape[0]) dy, dgate = _C.h3_gate_residual_backward(g2, y2, gate, index, *tiles) record_backward( diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index e4c303e10..4dd9521f8 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -11,6 +11,12 @@ from __future__ import annotations +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest import torch @@ -113,3 +119,49 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("adaln_gate_residual", device="cuda") assert type(op).__name__ == "H3GateResidualCudaOp" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +@pytest.mark.parametrize("bad_row", [-1, 2], ids=["negative", "past-table"]) +def test_gate_row_outside_the_table_asserts_in_the_kernel(entrypoint, bad_row): + """check_range=False callers get a device-side assert, not a read past the gate table.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + if not hasattr(_C, "h3_gate_residual_forward"): + sys.exit(77) + entrypoint, bad_row = sys.argv[1], int(sys.argv[2]) + x = torch.randn(4, 8, device="cuda") + gate = torch.ones(2, 8, device="cuda") # a table of two rows + index = torch.tensor([0, bad_row], dtype=torch.int64, device="cuda") + if entrypoint == "forward": + _C.h3_gate_residual_forward(x, x, gate, index) + else: + i64 = dict(dtype=torch.int64, device="cuda") + _C.h3_gate_residual_backward( + torch.ones_like(x), x, gate, index, + torch.tensor([0, 2, 1, 3], **i64), torch.tensor([0, 2], **i64), + torch.tensor([2, 4], **i64), torch.tensor([0, 1, 2], **i64), + ) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(bad_row)], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + if completed.returncode == 77: + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + assert completed.returncode != 0, "the kernel accepted a gate row outside the table" + assert ( + "device-side assert triggered" in completed.stderr + ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" From a2d0674bd8b8e975c3493abff6351f6429b7b8e2 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:16:56 +0000 Subject: [PATCH 11/14] docs(h3): gate-residual latencies from the backward-only rerun The table and the report links now come from the B200 report rerun at 239ef808, where the timed backward no longer includes the forward pass: backward 1.02 / 1.08 / 2.21 ms for S = 4097 / 32768 / 131072 against 0.55 / 2.24 / 8.92 ms for diffusers (the speedup at S = 131072 grows from 3.2x to 4.0x). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index b35bc41c1..e869f4363 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -65,9 +65,9 @@ B200, B = 1, H = 5376, BF16: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | -| 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | -| 131072 | 0.92 ms | 1.59 ms | 3.29 ms | 10.6 ms | +| 4097 | 0.06 ms | 0.07 ms | 1.02 ms | 0.55 ms | +| 32768 | 0.24 ms | 0.43 ms | 1.08 ms | 2.24 ms | +| 131072 | 0.88 ms | 1.61 ms | 2.21 ms | 8.92 ms | The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and @@ -75,10 +75,11 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `e7bc53a`. It also records +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `239ef808` on an otherwise idle B200 (backward timings exclude the leaf +creation, the forward pass and the graph construction; the execution order alternates and is recorded). It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) From 1857fa69aed91b8333b9ba7cb16e985d0bfd1eba Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:58:21 +0000 Subject: [PATCH 12/14] fix(h3): check the gate backward's tile metadata and lengths, and scan the reference's rows once gate_residual.cu: the backward runs the shared tile-metadata pre-pass (tile_metadata.cuh, #486) before the segment kernels, so sorted positions, tiles and segment boundaries that do not describe the rows end in a device-side assert instead of an out-of-bounds read, including the case of no tiles with a non-trivial seg_first_tile (the partial buffer then holds one uninitialized row). Host checks: sorted_pos has M entries, tile_begin and tile_end have the same length, seg_first_tile has at least one entry, and the column blocks fit grid.y (65535); the fold launch is skipped for zero segments. Native gate backward, idle B200, min of 400: 86.6 -> 90.7 us at S = 4097, 362.7 -> 367.5 us at S = 32768, 1315.0 -> 1319.7 us at S = 131072 (one extra launch of about 4 us). Tests run corrupted metadata in a child process and the host checks. gate_residual.py: dy is viewed as the shape of y (validate_h3_gate_residual makes it the residual's shape, so nothing changes for valid calls) and the saved shapes are named for the tensor they belong to. Reference: the gate gradient takes the positions of all rows from positions_by_row (one sort and one host read instead of one torch.nonzero scan and sync per row); a test compares the gradient bitwise with the loop it replaces. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 22 ++++- .../minimax_h3/gate_residual.py | 6 +- .../reference/minimax_h3/gate_residual.py | 5 +- .../minimax_h3/test_h3_adaln_gate_residual.py | 87 +++++++++++++++++++ 4 files changed, 112 insertions(+), 8 deletions(-) diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 6e0b5a85b..123b434fa 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -27,6 +27,8 @@ #include #include +#include "tile_metadata.cuh" + namespace { __device__ __forceinline__ float to_f(float v) { return v; } @@ -216,8 +218,11 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: } const int64_t rows = grad.size(0); const int64_t n = grad.size(1); + TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); const int64_t segments = seg_first_tile.numel() - 1; TORCH_CHECK(segments == gate.size(0), "segments must equal the gate rows"); + TORCH_CHECK(sorted_pos.numel() == rows, "sorted_pos must have M entries"); + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); const int64_t tiles = tile_begin.numel(); const c10::cuda::CUDAGuard guard(grad.device()); auto stream = at::cuda::getCurrentCUDAStream(); @@ -225,7 +230,14 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: auto partial = torch::empty({std::max(tiles, 1), n}, grad.options().dtype(at::kFloat)); auto dgate = torch::empty({segments, n}, grad.options()); const int threads = 256; - const unsigned col_blocks = static_cast((n + threads - 1) / threads); + const int64_t col_blocks64 = (n + threads - 1) / threads; + TORCH_CHECK(col_blocks64 <= 65535, "N=", n, " needs ", col_blocks64, + " column blocks, more than the grid.y limit of 65535"); + const unsigned col_blocks = static_cast(col_blocks64); + // The tile list is caller-built metadata: check it before the kernels read through it. + h3_tiles::check_tile_metadata(stream, sorted_pos.data_ptr(), rows, + tile_begin.data_ptr(), tile_end.data_ptr(), tiles, + seg_first_tile.data_ptr(), segments); H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward", [&] { const T* gp = reinterpret_cast(grad.data_ptr()); const T* yp = reinterpret_cast(y.data_ptr()); @@ -237,9 +249,11 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: gp, yp, sorted_pos.data_ptr(), tile_begin.data_ptr(), tile_end.data_ptr(), partial.data_ptr(), n); } - gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, stream>>>( - partial.data_ptr(), seg_first_tile.data_ptr(), - reinterpret_cast(dgate.data_ptr()), n); + if (segments > 0) { // a zero-sized grid is not a valid launch + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(dgate.data_ptr()), n); + } }); C10_CUDA_KERNEL_LAUNCH_CHECK(); return {dy, dgate}; diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index ad9bb6d8f..04f6bdf7d 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -36,7 +36,9 @@ def forward(ctx, residual, y, gate, index): y2 = y.contiguous().view(-1, hidden) out = _C.h3_gate_residual_forward(res2, y2, gate, index) ctx.save_for_backward(y2, gate, index) - ctx.shape = residual.shape + # validate_h3_gate_residual makes y and residual one shape; each gradient is still viewed + # as the shape of its own input. + ctx.residual_shape, ctx.y_shape = residual.shape, y.shape return out.view(residual.shape) @staticmethod @@ -49,7 +51,7 @@ def backward(ctx, grad): record_backward( "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) - return grad, dy.view(ctx.shape), dgate, None + return grad, dy.view(ctx.y_shape), dgate, None class H3GateResidualCudaOp: diff --git a/rl_engine/reference/minimax_h3/gate_residual.py b/rl_engine/reference/minimax_h3/gate_residual.py index f68a0fdf4..0ad2c1b0e 100644 --- a/rl_engine/reference/minimax_h3/gate_residual.py +++ b/rl_engine/reference/minimax_h3/gate_residual.py @@ -18,6 +18,8 @@ import torch +from rl_engine.reference.minimax_h3.adaln_row_gather import positions_by_row + _DTYPES = (torch.float32, torch.bfloat16, torch.float16) @@ -69,8 +71,7 @@ def backward(ctx, grad): contrib = (grad.float() * y.float()).reshape(-1, index.shape[0], y.shape[-1]) per_position = contrib.sum(dim=0) # (S, H): fixed reduction over the batch d_gate = per_position.new_zeros((ctx.gate_rows, y.shape[-1])) - for row in range(ctx.gate_rows): - positions = torch.nonzero(index == row).flatten() + for row, positions in enumerate(positions_by_row(index, ctx.gate_rows)): if positions.numel(): d_gate[row] = per_position.index_select(0, positions).sum(dim=0) return grad, grad * gathered, d_gate.to(gathered.dtype), None diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index 4dd9521f8..b26ab8a54 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -52,6 +52,25 @@ def _gate(table, chunk=2): class TestReference: + def test_gate_backward_is_bitwise_the_per_row_scan_it_replaced(self): + from rl_engine.reference.minimax_h3.gate_residual import _DeterministicGateResidual + + g = torch.Generator(device="cpu").manual_seed(5) + residual, y = torch.randn(2, 300, 16, generator=g), torch.randn(2, 300, 16, generator=g) + gate = torch.randn(7, 16, generator=g) + index = torch.randint(0, 5, (300,), generator=g) # rows 5 and 6 are never referenced + grad = torch.randn(2, 300, 16, generator=g) + leaf = gate.clone().requires_grad_(True) + _DeterministicGateResidual.apply(residual, y, leaf, index).backward(grad) + + per_position = (grad * y).reshape(-1, 300, 16).sum(dim=0) + expected = torch.zeros(7, 16) + for row in range(7): # the loop of the earlier implementation + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + expected[row] = per_position.index_select(0, positions).sum(dim=0) + assert torch.equal(leaf.grad, expected) + def test_rejects_bad_inputs(self): op = NativeH3GateResidualOp() res, y, gate = torch.zeros(4, 8), torch.zeros(4, 8), torch.zeros(3, 8) @@ -165,3 +184,71 @@ def test_gate_row_outside_the_table_asserts_in_the_kernel(entrypoint, bad_row): assert ( "device-side assert triggered" in completed.stderr ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + + +def _gate_backward_in_child(mutation: str) -> subprocess.CompletedProcess: + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + if not hasattr(_C, "h3_gate_residual_backward"): + sys.exit(77) + mutation = sys.argv[1] + x = torch.randn(4, 8, device="cuda") + gate = torch.ones(2, 8, device="cuda") + index = torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda") + i64 = dict(dtype=torch.int64, device="cuda") + sorted_pos = torch.tensor([0, 2, 1, 3], **i64) + begin, end = torch.tensor([0, 2], **i64), torch.tensor([2, 4], **i64) + first = torch.tensor([0, 1, 2], **i64) + if mutation == "tile_end_past_rows": + end = torch.tensor([2, 9], **i64) + elif mutation == "position_past_rows": + sorted_pos = torch.tensor([0, 2, 1, 4], **i64) + elif mutation == "segments_past_tiles": + first = torch.tensor([0, 1, 5], **i64) + elif mutation == "no_tiles": + begin, end = begin[:0], end[:0] + elif mutation == "short_sorted_pos": + sorted_pos = sorted_pos[:3] + elif mutation == "no_boundary": + first = first[:0] + grad = torch.ones_like(x) + _C.h3_gate_residual_backward(grad, x, gate, index, sorted_pos, begin, end, first) + 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, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_the_gate_backward_stops_on_tile_metadata_that_does_not_describe_the_rows(): + control = _gate_backward_in_child("none") + if control.returncode == 77: + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + assert control.returncode == 0, control.stderr + for mutation in ( + "tile_end_past_rows", + "position_past_rows", + "segments_past_tiles", + "no_tiles", + ): + done = _gate_backward_in_child(mutation) + assert done.returncode != 0, mutation + assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) + # Lengths are checked on the host, before any launch. + for mutation, message in ( + ("short_sorted_pos", "sorted_pos must have M entries"), + ("no_boundary", "seg_first_tile must hold R + 1 entries"), + ): + done = _gate_backward_in_child(mutation) + assert done.returncode != 0 and message in done.stderr, (mutation, done.stderr) From ab407b39ccc7f09066ce7300a00c8ff17c9c4401 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:38:27 +0000 Subject: [PATCH 13/14] fix(h3): reject gate tables whose rows overlap, and state the residual gradient shape * `make_gate` refuses a table whose row stride is below N (an expanded table would make every row alias the same memory), as the modulation tables already did; `H3GateResidualCudaOp` copies such a table to dense rows first, so an expanded gate still works and gives the bytes of its dense copy. Views of a wider AdaLN table (unit column stride, larger row stride) go to the kernel unchanged. * The residual's VJP is the incoming cotangent, which already has the residual's shape because forward returns `out.view(residual.shape)`; `ctx.residual_shape` was saved and unused, which read like a missing reshape. The return is now `grad.reshape(ctx.residual_shape)` (a no-op for equal shapes, and it accepts a non-contiguous gradient where `view` would not), and a test feeds a (B, S, H) residual with a transposed cotangent and compares every gradient with the eager VJP bitwise. * `_C.pyi`: contracts for `h3_gate_residual_forward` / `_backward`. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 4 ++ .../minimax_h3/gate_residual.py | 10 ++++- .../minimax_h3/test_h3_adaln_gate_residual.py | 37 +++++++++++++++++++ 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 123b434fa..81051fcc6 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -149,6 +149,10 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc "gate must be a CUDA tensor with the activations' dtype"); TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, "gate must be (R, N) with unit column stride"); + // Rows are read as dense N-element spans at `stride(0)`: a view of a wider table is fine, an + // expanded or overlapping one (stride(0) < N) is not a table of rows. + TORCH_CHECK(gate.size(0) <= 1 || gate.stride(0) >= gate.size(1), "gate row stride ", + gate.stride(0), " is smaller than N=", gate.size(1)); TORCH_CHECK(index.is_cuda() && index.device() == x.device() && index.scalar_type() == at::kLong && index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index 04f6bdf7d..da7319bfb 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -51,7 +51,9 @@ def backward(ctx, grad): record_backward( "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) - return grad, dy.view(ctx.y_shape), dgate, None + # The cotangent has the output's shape, which is the residual's; the residual's VJP is that + # cotangent itself. reshape (not view) so a non-contiguous gradient is accepted. + return grad.reshape(ctx.residual_shape), dy.view(ctx.y_shape), dgate, None class H3GateResidualCudaOp: @@ -80,7 +82,11 @@ def forward(self, residual, y, gate, index, *, check_range: bool = True): validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) if not residual.is_cuda: raise ValueError("H3GateResidualCudaOp needs CUDA tensors") - if gate.stride(1) != 1: + # A view of a wider table (``table.view(R, 6, H)[:, 2]``: unit column stride, row stride 6H) + # goes to the kernel as it is, which indexes rows by the stride. Anything that is not a + # table of dense rows (non-unit columns, or rows that overlap because the stride is below + # H, as for an expanded table) is copied first. + if gate.stride(1) != 1 or (gate.shape[0] > 1 and gate.stride(0) < gate.shape[1]): gate = gate.contiguous() return _H3GateResidualCuda.apply(residual, y, gate, index.contiguous()) diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index b26ab8a54..193abce60 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -125,6 +125,43 @@ def grads(fn, dtype=None): assert rel < 5e-3 # only the BF16 rounding of the table gradient assert torch.count_nonzero(ours[2].view(9, 6, -1)[:, [0, 1, 3, 4, 5]]) == 0 + def test_residual_gradient_keeps_the_residual_shape_for_batched_and_strided_gradients(self): + """A (B, S, H) residual gets a (B, S, H) gradient even when the incoming one is a view.""" + + table, residual, y, index = _case(300, hidden=64, seed=5) + leaves = [t.detach().clone().requires_grad_(True) for t in (residual, y, table)] + out = _cuda_op()(leaves[0], leaves[1], _gate(leaves[2]), index) + # a non-contiguous cotangent: the transpose of a (B, H, S) tensor + grad = torch.randn(2, 64, 300, device="cuda").bfloat16().transpose(1, 2) + assert not grad.is_contiguous() and grad.shape == out.shape + out.backward(grad) + ref = [t.detach().clone().requires_grad_(True) for t in (residual, y, table)] + provider_gate_residual(ref[0], _gate(ref[2]), index, ref[1]).backward(grad) + assert leaves[0].grad.shape == residual.shape and leaves[1].grad.shape == y.shape + assert torch.equal(leaves[0].grad, ref[0].grad) + assert torch.equal(leaves[1].grad, ref[1].grad) + + def test_a_gate_table_whose_rows_overlap_is_copied_not_read_in_place(self): + """An expanded table (row stride 0) is a legal gate: it is copied to dense rows first.""" + + from rl_engine.backends.extension import _C + + hidden = 64 + row = (torch.randn(1, hidden, device="cuda") * 0.5).bfloat16() + expanded = row.expand(9, hidden) # stride (0, 1): all nine rows alias one + assert expanded.stride(0) == 0 + _, residual, y, index = _case(100, hidden=hidden, seed=6) + index = index % 9 + dense = expanded.contiguous() + got = _cuda_op()(residual, y, expanded, index) + assert torch.equal(got, _cuda_op()(residual, y, dense, index)) + assert torch.equal(got, provider_gate_residual(residual, dense, index, y)) + # the native entry point refuses the layout instead of reading overlapping rows + res2, y2 = residual.reshape(-1, hidden), y.reshape(-1, hidden) + with pytest.raises(RuntimeError, match="row stride"): + _C.h3_gate_residual_forward(res2, y2, expanded, index) + _C.h3_gate_residual_forward(res2, y2, dense, index) # the dense copy is accepted + def test_rejects_cpu_and_bad_index(self): table, residual, y, index = _case(10, hidden=16) with pytest.raises(IndexError): From 60e7785c754a2f08b4f2e207977430063a181dd1 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:05:10 +0000 Subject: [PATCH 14/14] fix(h3): return no residual gradient for a residual that needs none The gate residual backward handed back the cotangent as the residual's gradient whether or not the residual required one. It now follows needs_input_grad[0], like the other ports of this group. The tile-metadata mutation test gains a gap, an overlap and a short last tile, and expects the empty tile list to be refused on the host. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/gate_residual.py | 6 ++-- .../minimax_h3/test_h3_adaln_gate_residual.py | 31 ++++++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index da7319bfb..f39cf351b 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -52,8 +52,10 @@ def backward(ctx, grad): "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) # The cotangent has the output's shape, which is the residual's; the residual's VJP is that - # cotangent itself. reshape (not view) so a non-contiguous gradient is accepted. - return grad.reshape(ctx.residual_shape), dy.view(ctx.y_shape), dgate, None + # cotangent itself. reshape (not view) so a non-contiguous gradient is accepted. A residual + # that needs no gradient (a frozen skip path) gets none. + d_residual = grad.reshape(ctx.residual_shape) if ctx.needs_input_grad[0] else None + return d_residual, dy.view(ctx.y_shape), dgate, None class H3GateResidualCudaOp: diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index 193abce60..f93561263 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -141,6 +141,26 @@ def test_residual_gradient_keeps_the_residual_shape_for_batched_and_strided_grad assert torch.equal(leaves[0].grad, ref[0].grad) assert torch.equal(leaves[1].grad, ref[1].grad) + def test_no_residual_gradient_is_returned_for_a_residual_that_needs_none(self): + """The backward node hands back ``None`` for a frozen residual, the cotangent otherwise.""" + + table, residual, y, index = _case(100, hidden=64, seed=7) + grad = torch.randn_like(residual) + for needs_residual in (False, True): + leaves = [ + residual.detach().clone().requires_grad_(needs_residual), + y.detach().clone().requires_grad_(True), + table.detach().clone().requires_grad_(True), + ] + out = _cuda_op()(leaves[0], leaves[1], _gate(leaves[2]), index) + # grad_fn.apply returns what backward returned, before autograd drops the unused ones + d_residual, d_y, *_ = out.grad_fn.apply(grad) + if needs_residual: + assert torch.equal(d_residual, grad) + else: + assert d_residual is None + assert d_y.shape == y.shape + def test_a_gate_table_whose_rows_overlap_is_copied_not_read_in_place(self): """An expanded table (row stride 0) is a legal gate: it is copied to dense rows first.""" @@ -248,6 +268,12 @@ def _gate_backward_in_child(mutation: str) -> subprocess.CompletedProcess: first = torch.tensor([0, 1, 5], **i64) elif mutation == "no_tiles": begin, end = begin[:0], end[:0] + elif mutation == "gap_between_tiles": + begin = torch.tensor([0, 3], **i64) + elif mutation == "overlapping_tiles": + begin = torch.tensor([0, 1], **i64) + elif mutation == "last_tile_stops_short": + end = torch.tensor([2, 3], **i64) elif mutation == "short_sorted_pos": sorted_pos = sorted_pos[:3] elif mutation == "no_boundary": @@ -277,13 +303,16 @@ def test_the_gate_backward_stops_on_tile_metadata_that_does_not_describe_the_row "tile_end_past_rows", "position_past_rows", "segments_past_tiles", - "no_tiles", + "gap_between_tiles", + "overlapping_tiles", + "last_tile_stops_short", ): done = _gate_backward_in_child(mutation) assert done.returncode != 0, mutation assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) # Lengths are checked on the host, before any launch. for mutation, message in ( + ("no_tiles", "the tile list is empty but there are 4 positions to sum"), ("short_sorted_pos", "sorted_pos must have M entries"), ("no_boundary", "seg_first_tile must hold R + 1 entries"), ):