From 503305ad83907c4a5724f292eb2dc599e9d34273 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:36:25 +0000 Subject: [PATCH 01/24] [WS1][CUDA][MiniMax-H3] adaln_projection_3mod: three-modality AdaLN projection RFC #420 step 3, row `adaln_projection_3mod` (stacked on timestep_mlp_fp32). - SiLU at temb's FP32 precision and exactly one cast to the projection dtype (bitwise equal to the provider activation); a BF16 temb is rejected (RFC probe H7). - 2688 -> 96768 projection on a BF16 tensor-core path in det_linear.cu, contract h3-det-linear-bf16-mma-v1: warp = 16 output columns x 8 input rows (zero-padded), K in ascending 16-wide groups, each mma.sync m16n8k16 from a zero accumulator added into an FP32 running sum. Same instruction sequence for every T <= 8, so rows are bitwise batch/order/tile invariant; 99.98% of outputs correctly rounded vs 99.84% for cuBLAS; kernel time flat at ~78 us for T = 1..8. SM80+. - Six (3T, 5376) views in diffusers' view(-1, 6H).chunk(6) layout, row t*3+m; a channel-tagging test pins chunk/modality placement (probe H4). - Backward: FP32 d_temb through the identity-VJP cast, ascending-row dW/db folds, no cuBLAS or atomics. The FP64 golden applies the cast straight-through so its own gradient is not rounded to BF16. - Chain stage (end-to-end test), report/figure hooks, registry, gtest, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 124 ++++++++- docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-adaln-projection.md | 142 ++++++++++ rl_engine/kernels/gtest/operator_inputs.py | 26 ++ rl_engine/kernels/gtest/operator_specs.py | 13 + .../kernels/ops/cuda/h3/adaln_projection.py | 106 ++++++++ .../ops/pytorch/h3/adaln_projection.py | 98 +++++++ rl_engine/kernels/registry.py | 13 + rl_engine/testing/h3_chain.py | 21 +- rl_engine/testing/h3_provider.py | 10 + rl_engine/testing/h3_report.py | 84 +++++- rl_engine/tests/test_dispatch.py | 2 +- scripts/plot_h3_evidence.py | 38 +++ tests/h3/test_h3_adaln_projection.py | 249 ++++++++++++++++++ tests/test_ws1_gtest_gpu.py | 1 + 16 files changed, 915 insertions(+), 14 deletions(-) create mode 100644 docs/operators/h3-adaln-projection.md create mode 100644 rl_engine/kernels/ops/cuda/h3/adaln_projection.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/adaln_projection.py create mode 100644 tests/h3/test_h3_adaln_projection.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index cc82acafa..04f7060fb 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -20,6 +20,8 @@ // whose result is identical on every lane (FP add is commutative); // * the bias is added once after the butterfly; the activation runs in FP32 // (SiLU as v / (1 + expf(-v)), PyTorch's formula); one cast at the store. +// * this FP32 forward is used for FP32 weights; BF16 weights take the +// tensor-core forward of contract h3-det-linear-bf16-mma-v1 below. // d_input dx[t, k] = sum_n g[t, n] * W[n, k] // * N is split into fixed chunks of 64 rows; each chunk is an ascending // fmaf chain into FP32 from 0; the chunk partials are then left-folded in @@ -27,6 +29,20 @@ // d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] // * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one // cross-row reduction; its order is the logical row order. +// +// Reduction contract h3-det-linear-bf16-mma-v1 (BF16 x and W, FP32 accumulate) +// ----------------------------------------------------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * a warp owns 16 output columns (MMA rows) and 8 input rows (MMA columns; +// rows past T are zero), so every launch runs the same instruction +// sequence for every T <= 8, and an MMA column never sees another's data; +// * K is visited in groups of 16 in ascending order. Within a 32-wide step, +// quad q supplies k = 32s + 8q + {0..3} to the first mma.sync m16n8k16 and +// k = 32s + 8q + {4..7} to the second (the same k map for W and x); +// * every mma.sync starts from a zero accumulator, so the tensor core only +// sums 16 products; its FP32 result is added to the running FP32 sum with +// an IEEE add, group by group in ascending k; +// * then bias, FP32 activation and one cast, as above. #include #include @@ -42,12 +58,9 @@ namespace { #ifndef H3_FWD_F32_COLS #define H3_FWD_F32_COLS 2 #define H3_FWD_F32_AHEAD 6 -#define H3_FWD_BF16_COLS 4 -#define H3_FWD_BF16_AHEAD 2 #endif #ifndef H3_FWD_CONFIGS -#define H3_FWD_CONFIGS H3_FWD_CASE(H3_FWD_F32_COLS, H3_FWD_F32_AHEAD) \ - H3_FWD_CASE(H3_FWD_BF16_COLS, H3_FWD_BF16_AHEAD) +#define H3_FWD_CONFIGS H3_FWD_CASE(H3_FWD_F32_COLS, H3_FWD_F32_AHEAD) #endif constexpr int kWarp = 32; @@ -177,6 +190,93 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) } } +__device__ __forceinline__ void mma_m16n8k16_bf16(float (&c)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + // BF16 mma.sync needs SM80+; the host refuses to launch below that. + __trap(); +#else + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +#endif +} + +constexpr int kMmaCols = 16; // output columns per warp (MMA M) +constexpr int kMmaRows = 8; // input rows per pass (MMA N) +constexpr int kMmaAhead = 4; // 32-wide K steps loaded before use + +__device__ __forceinline__ void mma_group_accumulate(float (&acc)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { + float part[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + mma_m16n8k16_bf16(part, a, b0, b1); +#pragma unroll + for (int i = 0; i < 4; ++i) acc[i] += part[i]; +} + +// Contract h3-det-linear-bf16-mma-v1. Lane = 4 * g + q: g picks the MMA row +// (output columns n0 + g and n0 + g + 8) and the MMA column (input row t0 + g). +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_bf16_mma_kernel(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, float* __restrict__ pre_act, + int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation) { + const int lane = threadIdx.x % kWarp; + const int g = lane >> 2; + const int q = lane & 3; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + + for (int64_t n0 = warp_global * kMmaCols; n0 < n_out; n0 += warp_stride * kMmaCols) { + const bool lo_ok = n0 + g < n_out; + const bool hi_ok = n0 + g + 8 < n_out; + const __nv_bfloat16* w_lo = w + (n0 + g) * k_in; + const __nv_bfloat16* w_hi = w + (n0 + g + 8) * k_in; + for (int64_t t0 = 0; t0 < rows; t0 += kMmaRows) { + const bool x_ok = t0 + g < rows; + const __nv_bfloat16* x_row = x + (t0 + g) * k_in; + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int64_t kb = 0; kb < k_in; kb += 32 * kMmaAhead) { + uint4 a_lo[kMmaAhead], a_hi[kMmaAhead], bx[kMmaAhead]; +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + const int64_t k = kb + 32 * u + 8 * q; + const bool k_ok = k < k_in; + a_lo[u] = (k_ok && lo_ok) ? __ldcs(reinterpret_cast(w_lo + k)) : zero; + a_hi[u] = (k_ok && hi_ok) ? __ldcs(reinterpret_cast(w_hi + k)) : zero; + bx[u] = (k_ok && x_ok) ? *reinterpret_cast(x_row + k) : zero; + } +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + if (kb + 32 * u >= k_in) break; + const uint32_t first[4] = {a_lo[u].x, a_hi[u].x, a_lo[u].y, a_hi[u].y}; + const uint32_t second[4] = {a_lo[u].z, a_hi[u].z, a_lo[u].w, a_hi[u].w}; + mma_group_accumulate(acc, first, bx[u].x, bx[u].y); + mma_group_accumulate(acc, second, bx[u].z, bx[u].w); + } + } + // acc = {(n0+g, t0+2q), (n0+g, t0+2q+1), (n0+g+8, t0+2q), (n0+g+8, t0+2q+1)} +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int64_t n = n0 + g + (i >= 2 ? 8 : 0); + const int64_t t = t0 + 2 * q + (i & 1); + if (n < n_out && t < rows) { + float v = acc[i] + (bias != nullptr ? __bfloat162float(bias[n]) : 0.0f); + if (pre_act != nullptr) pre_act[t * n_out + n] = v; + if (activation == kActSilu) v = v / (1.0f + expf(-v)); + out[t * n_out + n] = __float2bfloat16(v); + } + } + } + } +} + // partial[c, t, k] = sum_{n in chunk c} g[t, n] * w[n, k] template __global__ void det_linear_dinput_partial_kernel(const g_t* __restrict__ grad, @@ -287,9 +387,8 @@ struct ForwardConfig { int ahead; }; -ForwardConfig pick_forward_config(int64_t rows, int64_t elem_size) { +ForwardConfig pick_forward_config(int64_t rows) { const int row_tile = rows <= 1 ? 1 : (rows <= 2 ? 2 : 4); - if (elem_size == 2) return {row_tile, H3_FWD_BF16_COLS, H3_FWD_BF16_AHEAD}; return {row_tile, H3_FWD_F32_COLS, H3_FWD_F32_AHEAD}; } @@ -365,18 +464,21 @@ std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor torch::Tensor pre; if (save_pre_activation) pre = torch::empty({rows, n_out}, x.options().dtype(at::kFloat)); auto stream = at::cuda::getCurrentCUDAStream(); - const ForwardConfig cfg = pick_forward_config(rows, x.element_size()); if (x.scalar_type() == at::kFloat) { + const ForwardConfig cfg = pick_forward_config(rows); launch_forward(cfg, cptr(x), cptr(weight), bias.has_value() ? cptr(*bias) : nullptr, mptr(out), save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation, stream); } else { - launch_forward<__nv_bfloat16>( - cfg, cptr(x), cptr(weight), + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 8, + "the BF16 h3 det_linear forward uses mma.sync and needs SM80 or newer"); + const int64_t blocks = grid_for((n_out + kMmaCols - 1) / kMmaCols, kWarpsPerBlock); + det_linear_forward_bf16_mma_kernel<<(blocks), kWarpsPerBlock * kWarp, + 0, stream>>>( + cptr(x), cptr(weight), bias.has_value() ? cptr(*bias) : nullptr, mptr(out), - save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation, - stream); + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation); } C10_CUDA_KERNEL_LAUNCH_CHECK(); if (save_pre_activation) return {out, pre}; diff --git a/docs/.nav.yml b/docs/.nav.yml index 0dbcaec7d..b5a3c718b 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -32,6 +32,7 @@ nav: - MiniMax-H3: - operators/h3-timestep-sinusoid.md - operators/h3-timestep-mlp.md + - operators/h3-adaln-projection.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index ac127f3c4..d713f52d1 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -34,4 +34,5 @@ Every operator page should include: - [Token Embedding](embedding.md) - [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) - [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) +- [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md new file mode 100644 index 000000000..8e91cf9bf --- /dev/null +++ b/docs/operators/h3-adaln-projection.md @@ -0,0 +1,142 @@ +# MiniMax-H3 Three-Modality AdaLN Projection + +## Summary + +`adaln_projection_3mod` is the per-block `MiniMaxH3AdaLayerNormModulation` +(RFC #420, WS1 step 3). It turns the shared FP32 timestep embedding into the six +modulation tensors of one transformer block, for all three modalities: + +```text +act = silu(temb).to(bf16) SiLU in FP32, one declared cast +table = act @ W.T + b (T, 96768) BF16, W: (96768, 2688) +rows = table.view(3T, 32256) row t * 3 + m, m: 0 video, 1 text, 2 audio +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) # (3T, 5376) each +``` + +Projection output channel `o = m * 32256 + c * 5376 + h` is chunk `c`, hidden index `h` +of modality `m`. The table rows are laid out `[t0m0, t0m1, t0m2, t1m0, ...]`. That is the +layout `adaln_row_gather` addresses with `timestep_indices * 3 + token_tags`. + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227`, `transformer_blocks.0.adaln_proj.linear` +(BF16 weight and bias). All 50 blocks have the same shape. + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("adaln_projection_3mod", device="cuda") +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, weight, bias) +table = op.forward_table(temb, weight, bias) # raw (T, 96768) for a fused gather +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+, validated on SM100) | `rl_engine.kernels.ops.cuda.h3.adaln_projection.H3AdaLNProjectionCudaOp` | `rl_engine._C.h3_det_linear_*` | BF16 weights: contract `h3-det-linear-bf16-mma-v1`; FP32 weights: `h3-det-linear-v1` | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.adaln_projection.NativeH3AdaLNProjectionOp` | n/a | `forward`: provider path; `forward_fp32`: declared-cast FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `temb` | `(T, D)`, `T >= 1` | float32 | H3: `D = 2688`. BF16 is rejected (RFC probe H7) | +| `weight` | `(6 * H * 3, D)` | bfloat16 (checkpoint) or float32 | First dim must be a multiple of 18. H3: `H = 5376` | +| `bias` | `(6 * H * 3,)` | same as `weight` | | +| outputs | six `(3T, H)` views | weight dtype | Views of one `(T, 18H)` table, as in diffusers | + +The CUDA kernel needs `D` to be a multiple of 8 for BF16 or 4 for FP32, and BF16 needs SM80+. These inputs fail +closed: a BF16 `temb`, mismatched dtypes or shapes, an empty `temb`, and mixed devices. + +## Numerics + +- **Mixed-precision boundary.** SiLU runs in FP32 at `temb`'s precision, and the result is + rounded to BF16 exactly once. These are the provider's own elementwise ops, so the + activation is bitwise equal to diffusers. Casting before the SiLU (probe H7) is + rejected at the API. It is also numerically visible: on block 0 it changes about 52% of + the BF16 outputs. +- **Projection (BF16 weights).** It follows contract `h3-det-linear-bf16-mma-v1` in + `csrc/cuda/h3/det_linear.cu`: + - Each warp owns 16 output columns and 8 input rows. Rows past T are zero, so every + T <= 8 runs the same instruction sequence. A tensor-core column never sees another + column's data. + - K is visited in groups of 16 in ascending order. Each `mma.sync m16n8k16` (BF16 in, + FP32 out) starts from zero, so the tensor core only sums 16 products. Its result is + added to an FP32 running sum with an IEEE add. + - Then the bias is added and the result is rounded to BF16 once. + + A timestep's 18 modulation rows are bitwise independent of the other timesteps in the + call and of their order. Restarting the accumulator for every 16-wide group is what keeps + accuracy at the level of an FP32 FMA chain. Letting the tensor core accumulate across + the whole of K, as cuBLAS does, loses about 0.15% of correctly rounded outputs. FP32 + weights use the FMA path of `h3-det-linear-v1` instead, because FP32 on tensor cores + would mean TF32, which the contract forbids. +- **Golden.** `forward_fp32` computes the SiLU in FP64, rounds it to BF16 at the declared + boundary (the cast is model semantics), and runs the projection in FP64. It returns + FP32 without the final rounding. Its gradient stays in FP64 end to end, because the + cast is applied straight-through (identity VJP). Writing the cast as + `.to(bf16).double()` would make autograd round the golden's own `d_temb` to BF16. +- **Backward.** No cuBLAS and no atomics. `d_act` is computed with the chunked + deterministic `grad @ W`, kept in FP32 through the cast (whose VJP is the identity), + and passed through the FP32 SiLU VJP; it is row-local and batch-invariant. + `dW` and `db` are ascending-row FP32 folds, rounded to BF16 once. Unlike the provider, + `d_temb` is not rounded to BF16 on its way through the cast. + +Measured on a B200 with the pinned block-0 weights and torch 2.13.0+cu130. Fractions are +medians over 20 random draws of `temb`, with T = 4: + +| Comparison | CUDA | provider (cuBLAS) | +| --- | --- | --- | +| outputs equal to the correctly rounded golden | 99.982% (worst draw 99.976%) | 99.842% (worst draw 99.787%) | +| CUDA equal to provider | 99.84% | | +| early-cast golden (probe H7) equal to CUDA | 43% | | +| `d_temb` vs FP64 golden (gtest, T = 3) | max abs 2.3e-5 | 6.1e-2 (`d_act` rounded to BF16) | +| `dW`, `db` vs FP64 golden rounded to BF16 | > 99.9% bitwise equal | | + +The contract tolerance is `reduction` / `bfloat16`: atol 5e-2, rtol 2e-2, and atol 1e-1 +for gradients. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op adaln_projection_3mod +``` + +B200, pinned BF16 weights (520 MB per block, streamed once per call): + +| T | CUDA op | provider (SiLU + cast + cuBLAS) | CUDA GEMV kernel | cuBLAS kernel | +| --- | --- | --- | --- | --- | +| 1 | 102 µs | 112 µs | 78.3 µs (6.6 TB/s) | 93 µs | +| 2 | 102 µs | 101 µs | 78.4 µs | 79 µs | +| 4 | 102 µs | 101 µs | 79.0 µs | 79 µs | + +The kernel takes the same time for every T <= 8, because it always computes 8 rows. It +matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 24 µs of the op time is +the SiLU and cast kernels plus the Python wrapper. + +An earlier FMA-only BF16 kernel ran at 81, 98 and 166 µs for T = 1, 2 and 4. Its time +grew with T because every extra row adds FMAs per weight element. Packed FP32 FMA +(`__ffma2_rn`) gave identical bits but no speedup. Moving to tensor cores removed the +dependence on T. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_adaln_projection.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end: sinusoid -> MLP -> projection +python scripts/check_operator.py --op adaln_projection_3mod --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op adaln_projection_3mod \ + --out docs/usage/evidence/h3-adaln-projection-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-adaln-projection-b200/report.json +``` + +## Known Limitations + +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- The BF16 path needs SM80 or newer (`mma.sync`). It is rejected at launch on older GPUs. +- T > 8 runs one more weight pass per 8 rows. +- Sharding the 96768-wide projection (`tp_adaln_3mod`) is a separate WS2 row. diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 435d42023..255a48c0c 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -19,6 +19,7 @@ H3_FREQ_DIM = 256 H3_TIME_HIDDEN = 5376 H3_TIME_EMBED = 2688 +H3_HIDDEN = 5376 def make_operator_inputs( @@ -47,6 +48,7 @@ def make_operator_inputs( "kv_cache_attention": _make_kv_cache_attention_inputs, "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, + "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, } try: return builders[op_name](args, dtype, device) @@ -80,6 +82,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "timestep_sinusoid_h3": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}", "timestep_mlp_fp32": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}x{H3_TIME_HIDDEN}" f"x{H3_TIME_EMBED}", + "adaln_projection_3mod": f"{_h3_num_timesteps(args)}x{H3_TIME_EMBED}" + f"x{6 * 3 * _h3_hidden(args)}", } try: return names[op_name] @@ -385,6 +389,28 @@ def _make_timestep_mlp_fp32_inputs( } +def _h3_hidden(args: argparse.Namespace) -> int: + return _arg_int(args, "normalized_dim", H3_HIDDEN) + + +def _make_adaln_projection_3mod_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # temb is FP32 by contract (the SiLU runs before the cast); ``dtype`` is the + # projection's weight dtype, BF16 in the checkpoint. + n_out = 6 * 3 * _h3_hidden(args) + temb = _floating_tensor( + (_h3_num_timesteps(args), H3_TIME_EMBED), args, torch.float32, device, 0 + ) + weight = _floating_tensor((n_out, H3_TIME_EMBED), args, torch.float32, device, 1) + bias = _floating_tensor((n_out,), args, torch.float32, device, 2) + return { + "temb": temb * 2.0, + "weight": (weight * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (bias * 0.1).to(dtype), + } + + 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 cefebd46f..91d877b62 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -272,6 +272,19 @@ def _load_object(path: str) -> Any: }, grad_input_names=("x", "w1", "b1", "w2", "b2"), ), + "adaln_projection_3mod": OperatorSpec( + name="adaln_projection_3mod", + op_class="reduction", + gold_path=("rl_engine.kernels.ops.pytorch.h3.adaln_projection.NativeH3AdaLNProjectionOp"), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.kernels.ops.pytorch.h3.adaln_projection.NativeH3AdaLNProjectionOp" + ), + "cuda": "rl_engine.kernels.ops.cuda.h3.adaln_projection.H3AdaLNProjectionCudaOp", + }, + grad_input_names=("temb", "weight", "bias"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py new file mode 100644 index 000000000..c8a9b42cc --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 three-modality AdaLN projection (RFC #420 ``adaln_projection_3mod``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + ACT_NONE, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.pytorch.h3.adaln_projection import ( + split_adaln_table, + validate_h3_adaln_projection, +) + +KERNEL_ID = "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.kernels.ops.cuda.h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNProjectionCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias): + temb = temb.contiguous() + # Declared mixed-precision boundary: SiLU at temb's FP32 precision, + # then exactly one rounding to the projection dtype. These are the + # provider's own elementwise ops, so the activation is bitwise equal. + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias, activation=ACT_NONE) + ctx.save_for_backward(temb, act, weight) + return table + + @staticmethod + def backward(ctx, grad_table): + temb, act, weight = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad + grad_table = grad_table.float().contiguous() + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(grad_table, act, weight.dtype) + if need_temb: + # The cast's VJP is the identity; the SiLU VJP runs in FP32. + d_act = det_linear_backward_input(grad_table, weight, torch.float32) + d_temb = silu_backward_fp32(d_act, temb) + record_backward( + "adaln_projection_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return d_temb, dw if need_weight else None, db if need_bias else None + + +class H3AdaLNProjectionCudaOp: + """CUDA candidate: FP32 SiLU, one BF16 cast, deterministic GEMV, six views. + + The projection weight (96768 x 2688 BF16, 520 MB per block) is streamed + once per call; every output element follows contract ``h3-det-linear-v1``, + so a timestep's 3 x 6 modulation rows do not depend on the other + timesteps in the call. The six outputs are views of one ``(T, 6H*3)`` + table, laid out exactly like diffusers' ``view(-1, 6H).chunk(6)``. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, temb, weight, bias): + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") + table = _H3AdaLNProjectionCuda.apply(temb, weight, bias) + return split_adaln_table(table, hidden) + + def forward_table(self, temb, weight, bias) -> torch.Tensor: + """The raw ``(T, 6H*3)`` projection, for consumers that gather it directly.""" + + validate_h3_adaln_projection(temb, weight, bias) + return _H3AdaLNProjectionCuda.apply(temb, weight, bias) + + def forward_fp32(self, temb, weight, bias): + return self.forward(temb, weight, bias) diff --git a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py new file mode 100644 index 000000000..75ec2b8bd --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 three-modality AdaLN projection (RFC #420 row ``adaln_projection_3mod``). + +One ``MiniMaxH3AdaLayerNormModulation`` per transformer block:: + + act = silu(temb).to(weight.dtype) SiLU in FP32, one cast to BF16 + table = act @ W.T + b (T, 6 * H * 3), BF16 + rows = table.view(3 * T, 6 * H) row t * 3 + m, m in {video, text, audio} + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) + +Output channel ``o = m * 6H + c * H + h`` of the projection is chunk ``c``, hidden +index ``h`` of modality ``m``. Each of the six outputs has shape ``(3T, H)``, +and row ``t * 3 + m`` is what ``timestep_indices * 3 + token_tags`` addresses. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_WEIGHT_DTYPES = (torch.bfloat16, torch.float32) + + +def adaln_hidden_size(weight: torch.Tensor) -> int: + rows_per_hidden = H3_ADALN_CHUNKS * H3_MODALITY_NUM + if weight.dim() != 2 or weight.shape[0] % rows_per_hidden != 0: + raise ValueError( + f"AdaLN weight must be (6 * H * 3, D); got {tuple(weight.shape)}, whose first " + f"dim is not a multiple of {rows_per_hidden}" + ) + return weight.shape[0] // rows_per_hidden + + +def validate_h3_adaln_projection(temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.device != temb.device: + raise ValueError(f"{name} is on {tensor.device}, temb is on {temb.device}") + if temb.dtype != torch.float32: + # RFC probe H7: the activation must run at the timestep-embedding + # precision; a BF16 temb means the cast happened before the SiLU. + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in _WEIGHT_DTYPES or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share a dtype in {_WEIGHT_DTYPES}, got {weight.dtype}/{bias.dtype}" + ) + hidden = adaln_hidden_size(weight) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + if bias.shape != (weight.shape[0],): + raise ValueError(f"bias must be ({weight.shape[0]},), got {tuple(bias.shape)}") + return hidden + + +def split_adaln_table(table: torch.Tensor, hidden: int) -> tuple[torch.Tensor, ...]: + """(T, 6H*3) projection output -> six (3T, H) views, diffusers order.""" + + return table.view(-1, H3_ADALN_CHUNKS * hidden).chunk(H3_ADALN_CHUNKS, dim=-1) + + +class NativeH3AdaLNProjectionOp: + """PyTorch reference for the H3 AdaLN projection. + + ``forward`` is the provider path (``F.silu`` in FP32, cast, ``F.linear``). + ``forward_fp32`` is the golden: the SiLU in FP64, rounded once to the + weight dtype at the declared boundary (the cast is model semantics, not + an implementation detail), then the projection in FP64. The six outputs are + returned in FP32 without the final BF16 rounding. Its gradient is FP64 + end to end: the cast is applied straight-through (identity VJP). + """ + + op_class = "reduction" + + def __call__(self, temb, weight, bias): + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + hidden = validate_h3_adaln_projection(temb, weight, bias) + table = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + return split_adaln_table(table, hidden) + + def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + hidden = validate_h3_adaln_projection(temb, weight, bias) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + # The declared cast rounds the value; its VJP is the identity. Written + # straight-through so autograd does not round the FP64 gradient to the + # weight dtype on its way back (``.to(bf16)`` would). + act = act + (act.to(weight.dtype).double() - act).detach() + table = F.linear(act, weight.double(), bias.double()) + return tuple(chunk.float() for chunk in split_adaln_table(table, hidden)) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index e3b3765d5..22bbab259 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -209,6 +209,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ) PYTORCH_H3_TIMESTEP_MLP = "rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp" CUDA_H3_TIMESTEP_MLP = "rl_engine.kernels.ops.cuda.h3.timestep_mlp.H3TimestepMLPCudaOp" + PYTORCH_H3_ADALN_PROJECTION = ( + "rl_engine.kernels.ops.pytorch.h3.adaln_projection.NativeH3AdaLNProjectionOp" + ) + CUDA_H3_ADALN_PROJECTION = ( + "rl_engine.kernels.ops.cuda.h3.adaln_projection.H3AdaLNProjectionCudaOp" + ) def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -627,6 +633,10 @@ def __init__(self): OpBackend.CUDA_H3_TIMESTEP_MLP, OpBackend.PYTORCH_H3_TIMESTEP_MLP, ], + "adaln_projection_3mod": [ + OpBackend.CUDA_H3_ADALN_PROJECTION, + OpBackend.PYTORCH_H3_ADALN_PROJECTION, + ], }, "rocm": { "logp": [ @@ -669,6 +679,7 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -715,6 +726,7 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -742,6 +754,7 @@ def __init__(self): "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], }, # 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 a2c6a80c5..8671ef141 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -30,7 +30,11 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps -from rl_engine.testing.h3_provider import provider_time_embedder, provider_time_proj +from rl_engine.testing.h3_provider import ( + provider_adaln_modulation, + provider_time_embedder, + provider_time_proj, +) REPO_ROOT = Path(__file__).resolve().parents[2] @@ -54,6 +58,11 @@ def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: return [weights[f"time_embedder.linear_{i}.{p}"] for i in (1, 2) for p in ("weight", "bias")] +def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + prefix = "transformer_blocks.0.adaln_proj.linear" + return ctx["weights"][f"{prefix}.weight"], ctx["weights"][f"{prefix}.bias"] + + STAGES: list[Stage] = [ Stage( name="timestep_sinusoid_h3", @@ -75,6 +84,16 @@ def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: golden_atol=1e-4, # reduction / float32 golden_rtol=1e-4, ), + Stage( + name="adaln_projection_3mod", + op_type="adaln_projection_3mod", + candidate=lambda op, ctx, up: op(up, *adaln_params(ctx)), + provider=lambda ctx, up: provider_adaln_modulation(up, *adaln_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *adaln_params(ctx)), + provider_bitwise_isolated=False, # reduction: tensor-core tree differs from cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ), ] diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index 6ae2c5790..f64f7174f 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -46,3 +46,13 @@ def provider_time_embedder( sample = F.linear(features.to(w1.dtype), w1, b1) sample = F.silu(sample) return F.linear(sample, w2, b2) + + +def provider_adaln_modulation( + temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, hidden_size: int = 5376 +) -> tuple[torch.Tensor, ...]: + """``MiniMaxH3AdaLayerNormModulation.forward``: six (3T, H) views.""" + + temb = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + temb = temb.view(-1, 6 * hidden_size) + return temb.chunk(6, dim=-1) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 175547a21..3f84eb520 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -17,7 +17,11 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing.h3_cases import h3_timesteps -from rl_engine.testing.h3_provider import provider_time_embedder, provider_time_proj +from rl_engine.testing.h3_provider import ( + provider_adaln_modulation, + provider_time_embedder, + provider_time_proj, +) from rl_engine.testing.h3_weights import h3_weights_dir, load_h3_conditioning_weights # Keys of a perf case that hold a timed callable, in display order. @@ -178,11 +182,89 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: } +# --------------------------------------------------------------------------- # +# adaln_projection_3mod +# --------------------------------------------------------------------------- # + +ADALN_NAMES = [ + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +] + + +def _adaln_params() -> tuple[torch.Tensor, torch.Tensor]: + weight, bias = h3_params(ADALN_NAMES, [(96768, 2688), (96768,)]) + return weight.bfloat16(), bias.bfloat16() + + +def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_projection_3mod", device="cuda") + weight, bias = _adaln_params() + cases = [] + for num in (1, 2, 3, 4): + temb = torch.randn(num, 2688, device="cuda") + cases.append( + { + "op": "adaln_projection_3mod", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight.numel() * weight.element_size(), + "candidate": lambda temb=temb: op(temb, weight, bias), + "provider": lambda temb=temb: provider_adaln_modulation(temb, weight, bias), + } + ) + return cases + + +def _flat_cat(outputs) -> torch.Tensor: + return torch.cat([out.reshape(-1) for out in outputs]) + + +def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, Any]: + """Fraction of BF16 outputs equal to the correctly rounded FP64 golden, per draw.""" + + op = registry.get_op("adaln_projection_3mod", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["adaln_projection_3mod"][-1] + ) + weight, bias = _adaln_params() + cuda_frac, provider_frac, early_frac = [], [], [] + for seed in range(draws): + g = torch.Generator(device="cuda").manual_seed(seed) + temb = torch.randn(4, 2688, device="cuda", generator=g) * 2 + gold = _flat_cat(golden.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat_cat(golden.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + ours = _flat_cat(op(temb, weight, bias)) + theirs = _flat_cat(provider_adaln_modulation(temb, weight, bias)) + cuda_frac.append(float((ours == gold).float().mean())) + provider_frac.append(float((theirs == gold).float().mean())) + early_frac.append(float((ours == early).float().mean())) + temb = torch.randn(9, 2688, device="cuda") + full = op(temb, weight, bias) + 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 i in range(9) + ) + return { + "draws": draws, + "num_timesteps": 4, + "cuda_correctly_rounded": cuda_frac, + "provider_correctly_rounded": provider_frac, + "early_cast_golden_match": early_frac, + "rows_batch_invariant": invariant, + } + + 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, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, "timestep_mlp_fp32": _mlp_accuracy, + "adaln_projection_3mod": _projection_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index b436e3f2d..c20e157e4 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -226,7 +226,7 @@ def fake_load_backend(backend): ] -H3_OPS = ("timestep_sinusoid_h3", "timestep_mlp_fp32") +H3_OPS = ("timestep_sinusoid_h3", "timestep_mlp_fp32", "adaln_projection_3mod") @pytest.mark.parametrize("op_name", H3_OPS) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 2baa21bb6..6385b4aa4 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -224,9 +224,47 @@ def plot_mlp(report: dict[str, Any]): return fig +def plot_projection(report: dict[str, Any]): + fig, (left, right) = _figure(report, "adaln_projection_3mod (pinned block-0 AdaLN weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _style( + right, + f"Outputs equal to the correctly rounded FP64 golden, {acc['draws']} draws", + "% of BF16 outputs", + ) + series = [ + ("RL-Kernel CUDA", RL_KERNEL, [100 * v for v in acc["cuda_correctly_rounded"]]), + ("diffusers (cuBLAS)", PROVIDER, [100 * v for v in acc["provider_correctly_rounded"]]), + ] + for i, (name, color, values) in enumerate(series): + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax_vals = sorted(values) + median = ax_vals[len(ax_vals) // 2] + right.scatter(jitter, values, s=36, color=color, alpha=0.8, linewidths=0, zorder=3) + right.hlines(median, i - 0.25, i + 0.25, color=color, lw=2.5, zorder=4) + right.annotate( + f"median {median:.3f}%", (i + 0.28, median), fontsize=8, color=INK, va="center" + ) + right.set_xticks(range(len(series)), [name for name, _, _ in series]) + right.set_xlim(-0.6, len(series) - 0.1) + early = sorted(acc["early_cast_golden_match"])[len(acc["early_cast_golden_match"]) // 2] + right.text( + 0.98, + 0.62, + f"an early BF16 cast (probe H7)\nwould match only {100 * early:.0f}%", + transform=right.transAxes, + fontsize=8, + color=MUTED, + ha="right", + ) + 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, } diff --git a/tests/h3/test_h3_adaln_projection.py b/tests/h3/test_h3_adaln_projection.py new file mode 100644 index 000000000..c9aac92d1 --- /dev/null +++ b/tests/h3/test_h3_adaln_projection.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_projection_3mod``: SiLU + 2688 -> 96768 projection, 6 x 3 table. + +* layout: the six outputs and the three modality rows per timestep land + exactly where diffusers puts them (RFC probe H4); +* mixed precision: SiLU in FP32, one cast to BF16; an early cast (probe H7) + is rejected at the API and is measurably different numerically; +* accuracy: forward and gradients against the declared-cast golden on the + pinned block-0 weights; +* invariance and determinism: a timestep's 18 modulation rows do not depend + on the other timesteps, and repeated runs are bitwise equal. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.kernels.ops.pytorch.h3.adaln_projection import NativeH3AdaLNProjectionOp +from rl_engine.testing.h3_provider import provider_adaln_modulation + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3AdaLNProjectionCudaOp() + + +def _temb(num, dim=2688, seed=0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(num, dim, generator=g) * 2).to(device) + + +def _synthetic(hidden, dim, dtype=torch.bfloat16, seed=0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).to(dtype) + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype) + return weight.to(device), bias.to(device) + + +def _flat(outputs): + return torch.cat([out.reshape(-1) for out in outputs]) + + +class TestReference: + def test_provider_layout_matches_diffusers(self): + temb = _temb(2, dim=16, device="cpu") + weight, bias = _synthetic(8, 16, device="cpu") + ours = NativeH3AdaLNProjectionOp().forward(temb, weight, bias) + theirs = provider_adaln_modulation(temb, weight, bias, hidden_size=8) + assert len(ours) == 6 + for a, b in zip(ours, theirs): + assert a.shape == (6, 8) and torch.equal(a, b) + + def test_rejects_bf16_temb(self): + # RFC probe H7: casting temb before the SiLU. + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(TypeError, match="float32"): + NativeH3AdaLNProjectionOp().forward(_temb(2, 16, device="cpu").bfloat16(), weight, bias) + + def test_rejects_bad_shapes(self): + weight, bias = _synthetic(8, 16, device="cpu") + op = NativeH3AdaLNProjectionOp() + with pytest.raises(ValueError): # not a multiple of 6 * 3 rows + op.forward(_temb(2, 16, device="cpu"), weight[:-1], bias[:-1]) + with pytest.raises(ValueError): + op.forward(_temb(2, 15, device="cpu"), weight, bias) + with pytest.raises(ValueError): + op.forward(_temb(0, 16, device="cpu"), weight, bias) + with pytest.raises(TypeError): + op.forward(_temb(2, 16, device="cpu"), weight, bias.float()) + + +@requires_cuda +class TestLayout: + def test_channel_identity_lands_in_its_chunk_and_modality_row(self): + """Probe H4: bias o = m*6H + c*H + h tagged (m, c); zero weight passes it through.""" + + hidden, dim, num = 16, 32, 3 + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + o = torch.arange(n_out) + modality, chunk = o // (H3_ADALN_CHUNKS * hidden), (o // hidden) % H3_ADALN_CHUNKS + bias = (modality * 10 + chunk).to(torch.bfloat16).cuda() + weight = torch.zeros(n_out, dim, dtype=torch.bfloat16, device="cuda") + outputs = _cuda_op()(_temb(num, dim), weight, bias) + assert len(outputs) == 6 + for c, out in enumerate(outputs): + assert out.shape == (num * 3, hidden) + for t in range(num): + for m in range(3): + expected = torch.full((hidden,), float(m * 10 + c), device="cuda") + assert torch.equal(out[t * 3 + m].float(), expected), (c, t, m) + + def test_outputs_are_views_of_one_table_like_diffusers(self): + weight, bias = _synthetic(16, 32) + outputs = _cuda_op()(_temb(2, 32), weight, bias) + base = outputs[0].untyped_storage().data_ptr() + assert all(out.untyped_storage().data_ptr() == base for out in outputs) + assert outputs[1].data_ptr() - outputs[0].data_ptr() == 16 * outputs[0].element_size() + + +@pytest.fixture(scope="module") +def block0(h3_weights_cpu): + return h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4]) + def test_forward_is_correctly_rounded_declared_golden(self, block0, num): + weight, bias = block0 + temb = _temb(num, seed=num) + outputs = _cuda_op()(temb, weight, bias) + golden = NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias) + for out, gold in zip(outputs, golden): + assert out.dtype == torch.bfloat16 and out.shape == (3 * num, 5376) + # tolerance_contract.json forward_accuracy / reduction / bfloat16. + torch.testing.assert_close(out.float(), gold, atol=5e-2, rtol=2e-2) + # Beyond the contract: almost every element is the correctly rounded + # golden value; the rest are 1-ULP ties of the FP32 accumulation. + match = (_flat(outputs) == _flat(golden).bfloat16()).float().mean().item() + assert match > 0.999 + + def test_early_cast_is_numerically_detectable(self, block0): + """Probe H7: a BF16 cast before the SiLU moves about half the outputs.""" + + weight, bias = block0 + temb = _temb(3, seed=7) + op = NativeH3AdaLNProjectionOp() + ours = _flat(_cuda_op()(temb, weight, bias)) + declared = _flat(op.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat(op.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + assert (ours == declared).float().mean() > 0.999 + assert (ours == early).float().mean() < 0.8 + + def test_rows_are_batch_and_position_invariant(self, block0): + weight, bias = block0 + op = _cuda_op() + temb = _temb(5, seed=3) + full = op(temb, weight, bias) + for i in range(5): + single = op(temb[i : i + 1], weight, bias) + for s, f in zip(single, full): + assert torch.equal(s, f[3 * i : 3 * i + 3]) + swapped = op(temb.flip(0), weight, bias) + for s, f in zip(swapped, full): + assert torch.equal(s.view(5, 3, -1).flip(0), f.view(5, 3, -1)) + + def test_backward_against_golden(self, block0): + weight, bias = block0 + temb = _temb(2, seed=5) + g = torch.Generator(device="cuda").manual_seed(0) + grads = [torch.randn(6, 5376, device="cuda", generator=g).bfloat16() for _ in range(6)] + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(_cuda_op()(*leaves)), grads) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + act = act + (act.to(torch.bfloat16).double() - act).detach() # declared cast, identity VJP + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376).chunk(6, dim=-1) + torch.autograd.backward(list(table), [gr.double() for gr in grads]) + # tolerance_contract.json gradient_accuracy / reduction: FP32 temb grad + # and BF16 weight/bias grads (rounded once from an FP32 accumulation). + torch.testing.assert_close(leaves[0].grad.double(), ref[0].grad, atol=1e-4, rtol=1e-4) + for leaf, r in zip(leaves[1:], ref[1:]): + torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) + assert (leaf.grad == r.grad.to(torch.bfloat16)).float().mean() > 0.999 + + +@requires_cuda +class TestCudaSynthetic: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) + @pytest.mark.parametrize("hidden, dim", [(24, 16), (40, 72), (8, 2688)]) + def test_other_sizes(self, dtype, hidden, dim): + weight, bias = _synthetic(hidden, dim, dtype=dtype, seed=hidden) + temb = _temb(3, dim, seed=hidden) + tol = dict(atol=5e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-4, rtol=1e-4) + for out, gold in zip( + _cuda_op()(temb, weight, bias), + NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias), + ): + torch.testing.assert_close(out.float(), gold, **tol) + + def test_repeat_forward_backward_bitwise(self): + weight, bias = _synthetic(64, 2688, seed=9) + temb = _temb(3, seed=9) + grads = [torch.randn(9, 64, device="cuda").bfloat16() for _ in range(6)] + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + outs = _cuda_op()(*leaves) + torch.autograd.backward(list(outs), grads) + runs.append([*[o.detach() for o in outs], *[leaf.grad for leaf in leaves]]) + for later in runs[1:]: + for a, b in zip(runs[0], later): + assert torch.equal(a, b) + + def test_temb_grad_rows_are_batch_invariant(self): + weight, bias = _synthetic(64, 2688, seed=11) + temb = _temb(4, seed=11) + grads = [torch.randn(12, 64, device="cuda").bfloat16() for _ in range(6)] + + def d_temb(rows): + leaf = temb[rows].detach().clone().requires_grad_(True) + outs = _cuda_op()(leaf, weight, bias) + torch.autograd.backward( + list(outs), [g.view(4, 3, -1)[rows].reshape(-1, 64) for g in grads] + ) + return leaf.grad + + full = d_temb(slice(0, 4)) + for i in range(4): + assert torch.equal(d_temb(slice(i, i + 1))[0], full[i]) + + @pytest.mark.parametrize("num", [8, 9, 17]) + def test_rows_invariant_across_tensor_core_row_tiles(self, num): + """The BF16 path computes rows in tiles of 8; a row's bits must not depend on its tile.""" + + weight, bias = _synthetic(64, 2688, seed=13) + temb = _temb(num, seed=num) + full = _cuda_op()(temb, weight, bias) + for i in (0, 7, num - 1): + single = _cuda_op()(temb[i : i + 1], weight, bias) + for s_out, f_out in zip(single, full): + assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) + + def test_rejects_cpu(self): + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(ValueError): + _cuda_op()(_temb(2, 16, device="cpu"), weight, bias) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_projection_3mod", device="cuda") + assert type(op).__name__ == "H3AdaLNProjectionCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 6dde75c53..fdef637f7 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -51,6 +51,7 @@ def test_all_ws1_single_ops_are_registered(): "prefix_shared_attention", "timestep_sinusoid_h3", "timestep_mlp_fp32", + "adaln_projection_3mod", } <= names From 404e766bb1d5131776b03a44edc12bd93c4878a6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:36:46 +0000 Subject: [PATCH 02/24] [WS1][MiniMax-H3] adaln_projection_3mod: B200 evidence report and figure Produced by scripts/h3_evidence.py at d06e120 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-projection.md | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 8e91cf9bf..0e27a5761 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -108,19 +108,29 @@ B200, pinned BF16 weights (520 MB per block, streamed once per call): | T | CUDA op | provider (SiLU + cast + cuBLAS) | CUDA GEMV kernel | cuBLAS kernel | | --- | --- | --- | --- | --- | -| 1 | 102 µs | 112 µs | 78.3 µs (6.6 TB/s) | 93 µs | -| 2 | 102 µs | 101 µs | 78.4 µs | 79 µs | -| 4 | 102 µs | 101 µs | 79.0 µs | 79 µs | +| 1 | 105 µs | 115 µs | 78.3 µs (6.6 TB/s) | 93 µs | +| 2 | 105 µs | 103 µs | 78.4 µs | 79 µs | +| 4 | 106 µs | 104 µs | 79.0 µs | 79 µs | The kernel takes the same time for every T <= 8, because it always computes 8 rows. It matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 24 µs of the op time is -the SiLU and cast kernels plus the Python wrapper. +the SiLU and cast kernels plus the Python wrapper. With those included, the op is within +about 2% of the provider for T >= 2. An earlier FMA-only BF16 kernel ran at 81, 98 and 166 µs for T = 1, 2 and 4. Its time grew with T because every extra row adds FMAs per weight element. Packed FP32 FMA (`__ffma2_rn`) gave identical bits but no speedup. Moving to tensor cores removed the dependence on T. +## Evidence + +![adaln_projection_3mod on B200: latency and correctly rounded outputs](../usage/evidence/h3-adaln-projection-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-adaln-projection-b200/report.json), +written by `scripts/h3_evidence.py` from a clean tree at commit `d06e120`. The report also +records that a timestep's 18 modulation rows are bitwise identical whether it runs alone +or in a batch of 9. + ## Tests ```bash From 41ecf74b6300e83179bb020bf8f72331f4cd08ed Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:53:53 +0000 Subject: [PATCH 03/24] fix(h3): address AdaLN projection review findings Validate pinned tensor contents at preparation and load time. Execute chain prerequisites for selected stages and require checkpoint weights for weighted evidence. Reject invalid CUDA tile and backward inputs, align the CUDA wrapper guards, skip unsupported devices, and rotate comparative timing with raw samples. Validation: 143 H3 tests passed on B200 after a fresh CUDA build; illegal FP32 tiles fail compilation; isolated AdaLN replay and pinned-weight evidence generation passed; Black, isort, Ruff, and Flake8 passed. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 3 + .../kernels/ops/cuda/h3/adaln_projection.py | 2 + rl_engine/testing/h3_manifest.json | 30 +++- rl_engine/testing/h3_report.py | 37 +++- rl_engine/testing/h3_weights.py | 22 ++- scripts/h3_chain_replay.py | 12 +- scripts/h3_evidence.py | 10 +- scripts/prepare_h3_weights.py | 10 +- tests/h3/test_h3_cli.py | 160 ++++++++++++++++++ tests/h3/test_h3_conditioning_e2e.py | 5 +- tests/h3/test_h3_det_linear.py | 34 ++++ tests/h3/test_h3_report.py | 68 ++++++++ tests/h3/test_h3_weights.py | 70 ++++++++ tests/h3/test_prepare_h3_weights.py | 5 +- 14 files changed, 443 insertions(+), 25 deletions(-) create mode 100644 tests/h3/test_h3_cli.py create mode 100644 tests/h3/test_h3_det_linear.py create mode 100644 tests/h3/test_h3_report.py create mode 100644 tests/h3/test_h3_weights.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 95e16b234..246d27d60 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -124,6 +124,7 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) using WV = Vec16; using XV = Vec16; constexpr int kVec = WV::kN; + static_assert(kCols * kRows <= kWarp, "one warp cannot store more than kWarp outputs per tile"); static_assert(XV::kN == kVec, "x and weight share a dtype"); constexpr int64_t kStride = static_cast(kWarp) * kVec; @@ -499,6 +500,7 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, "out_dtype must be float32 or bfloat16"); const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); const int64_t n_out = weight.size(0); const int64_t k_in = weight.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); @@ -544,6 +546,7 @@ std::vector h3_det_linear_backward_weight(torch::Tensor grad, tor TORCH_CHECK(w_dtype == at::kFloat || w_dtype == at::kBFloat16, "w_dtype must be float32 or bfloat16"); const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); const int64_t n_out = grad.size(1); const int64_t k_in = x.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py index c8a9b42cc..aee0ea503 100644 --- a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py +++ b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py @@ -100,6 +100,8 @@ def forward_table(self, temb, weight, bias) -> torch.Tensor: """The raw ``(T, 6H*3)`` projection, for consumers that gather it directly.""" validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") return _H3AdaLNProjectionCuda.apply(temb, weight, bias) def forward_fp32(self, temb, weight, bias): diff --git a/rl_engine/testing/h3_manifest.json b/rl_engine/testing/h3_manifest.json index 52b1989b1..5d08b87f3 100644 --- a/rl_engine/testing/h3_manifest.json +++ b/rl_engine/testing/h3_manifest.json @@ -30,12 +30,30 @@ } }, "tensors": { - "time_embedder.linear_1.weight": {"dtype": "float32", "shape": [5376, 256]}, - "time_embedder.linear_1.bias": {"dtype": "float32", "shape": [5376]}, - "time_embedder.linear_2.weight": {"dtype": "float32", "shape": [2688, 5376]}, - "time_embedder.linear_2.bias": {"dtype": "float32", "shape": [2688]}, - "transformer_blocks.0.adaln_proj.linear.weight": {"dtype": "bfloat16", "shape": [96768, 2688]}, - "transformer_blocks.0.adaln_proj.linear.bias": {"dtype": "bfloat16", "shape": [96768]} + "time_embedder.linear_1.weight": { + "dtype": "float32", "shape": [5376, 256], + "sha256": "4c50eabe3c434dc05422e6501cdf7baa5cecbb1a7706ae9599c7a5b0e34e2913" + }, + "time_embedder.linear_1.bias": { + "dtype": "float32", "shape": [5376], + "sha256": "64763a5e39cb078eb0d8c5be5d56ad36cac3dec37be93a87e106d82ea6c071b7" + }, + "time_embedder.linear_2.weight": { + "dtype": "float32", "shape": [2688, 5376], + "sha256": "697819dd9413e0fec3a50987a964f1fd5db568ca545267414238f452294a8154" + }, + "time_embedder.linear_2.bias": { + "dtype": "float32", "shape": [2688], + "sha256": "096a3c7539f0248698cc65376c8da260e7365bf4ccef847c09bb9a36782a743c" + }, + "transformer_blocks.0.adaln_proj.linear.weight": { + "dtype": "bfloat16", "shape": [96768, 2688], + "sha256": "a0000193b30ea36a1996f53fc0d6c4c37ce5fd4585284f0f9bae570d25c6b8cf" + }, + "transformer_blocks.0.adaln_proj.linear.bias": { + "dtype": "bfloat16", "shape": [96768], + "sha256": "d59f9de9e07e6b4e3b57009627e276a26971ce379555fc99dbb06d3a237c8fc0" + } }, "reference_implementation": { "repo": "huggingface/diffusers", diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 3f84eb520..d4b45e09b 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -63,12 +63,35 @@ def peak_mib(fn: Callable[[], Any]) -> float: def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} - for key in TIMED_KEYS: - if key in case: - us = time_us(case[key], warmup, iters) - row[f"{key}_us"] = us - row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 - row[f"{key}_peak_mib"] = peak_mib(case[key]) + keys = [key for key in TIMED_KEYS if key in case] + if not keys: + return row + for _ in range(warmup): + for key in keys: + case[key]() + torch.cuda.synchronize() + samples = {key: [] for key in keys} + orders = [] + for index in range(iters): + # Rotate the first implementation so drift does not favor one backend. + offset = index % len(keys) + order = keys[offset:] + keys[:offset] + orders.append(order) + for key in order: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + case[key]() + end.record() + end.synchronize() + samples[key].append(start.elapsed_time(end) * 1e3) + row["timing_samples_us"] = samples + row["timing_order"] = orders + for key in keys: + us = statistics.median(samples[key]) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key]) return row @@ -244,7 +267,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) ) diff --git a/rl_engine/testing/h3_weights.py b/rl_engine/testing/h3_weights.py index 8090c6e83..7d7abdf03 100644 --- a/rl_engine/testing/h3_weights.py +++ b/rl_engine/testing/h3_weights.py @@ -41,6 +41,13 @@ def sha256_file(path: Path, chunk: int = 1 << 24) -> str: return digest.hexdigest() +def sha256_tensor(tensor: torch.Tensor) -> str: + """Hash a CPU tensor's contiguous bytes without safetensors metadata.""" + + tensor_bytes = tensor.contiguous().view(torch.uint8).numpy() + return hashlib.sha256(memoryview(tensor_bytes)).hexdigest() + + def load_h3_conditioning_weights( device: torch.device | str = "cpu", names: list[str] | None = None ) -> dict[str, torch.Tensor]: @@ -55,15 +62,20 @@ def load_h3_conditioning_weights( if not path.is_file(): raise FileNotFoundError(f"{path} missing; run scripts/prepare_h3_weights.py") manifest = load_h3_manifest() - tensors = load_file(str(path), device=str(device)) + tensors = load_file(str(path), device="cpu") wanted = names if names is not None else list(manifest["tensors"]) - out: dict[str, torch.Tensor] = {} - for name in wanted: - spec = manifest["tensors"][name] + for name, spec in manifest["tensors"].items(): tensor = tensors[name] if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: raise ValueError(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") if list(tensor.shape) != spec["shape"]: raise ValueError(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") - out[name] = tensor + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise ValueError(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + out: dict[str, torch.Tensor] = {} + for name in wanted: + if name not in manifest["tensors"]: + raise KeyError(name) + out[name] = tensors[name].to(device=device) return out diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py index e262ca053..512a15f06 100644 --- a/scripts/h3_chain_replay.py +++ b/scripts/h3_chain_replay.py @@ -35,7 +35,11 @@ def _ints(text: str) -> list[int]: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--stages", default=",".join(s.name for s in h3_chain.STAGES)) + parser.add_argument( + "--stages", + default=",".join(s.name for s in h3_chain.STAGES), + help="comma list of stages to report; preceding stages run automatically", + ) parser.add_argument("--timesteps", type=_ints, default=[1, 2, 4], help="comma list of T") parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") parser.add_argument("--seed", type=int, default=0) @@ -49,6 +53,8 @@ def main() -> None: if unknown: raise SystemExit(f"unknown stages: {unknown}") stages = [s for s in h3_chain.STAGES if s.name in wanted] + last_stage = max(i for i, stage in enumerate(h3_chain.STAGES) if stage.name in wanted) + execution_stages = h3_chain.STAGES[: last_stage + 1] torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False @@ -65,8 +71,9 @@ def main() -> None: num_timesteps=num_timesteps, seq_len=seq_len, seed=args.seed, - stages=stages, + stages=execution_stages, ) + case["stages"] = [entry for entry in case["stages"] if entry["stage"] in wanted] cases.append(case) summary = ", ".join( f"{e['stage']}=" @@ -85,6 +92,7 @@ def main() -> None: **h3_chain.git_state(), "environment": h3_chain.environment(), "stages": [s.name for s in stages], + "executed_stages": [s.name for s in execution_stages], "cases": cases, } if args.out is not None: diff --git a/scripts/h3_evidence.py b/scripts/h3_evidence.py index c20778234..8351da245 100644 --- a/scripts/h3_evidence.py +++ b/scripts/h3_evidence.py @@ -26,7 +26,11 @@ from rl_engine.kernels.registry import KernelRegistry # noqa: E402 from rl_engine.testing.h3_chain import environment, git_state # noqa: E402 from rl_engine.testing.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 -from rl_engine.testing.h3_weights import load_h3_manifest # noqa: E402 +from rl_engine.testing.h3_weights import ( # noqa: E402 + WEIGHTS_ENV, + h3_weights_dir, + load_h3_manifest, +) def main() -> None: @@ -34,6 +38,9 @@ def main() -> None: parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) parser.add_argument("--out", type=Path, required=True) args = parser.parse_args() + uses_weights = args.op != "timestep_sinusoid_h3" + if uses_weights and h3_weights_dir() is None: + raise SystemExit(f"set {WEIGHTS_ENV} to the scripts/prepare_h3_weights.py output dir") if not torch.cuda.is_available(): raise SystemExit("needs a CUDA device") torch.backends.cuda.matmul.allow_tf32 = False @@ -46,6 +53,7 @@ def main() -> None: "op": args.op, "rfc": manifest["rfc"], "model_revision": manifest["model_identity"]["revision"], + "weight_source": "pinned_checkpoint" if uses_weights else "not_applicable", "reference_commit": manifest["reference_implementation"]["commit"], **git_state(), "environment": environment(), diff --git a/scripts/prepare_h3_weights.py b/scripts/prepare_h3_weights.py index 910f34584..262408736 100644 --- a/scripts/prepare_h3_weights.py +++ b/scripts/prepare_h3_weights.py @@ -23,7 +23,12 @@ REPO_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(REPO_ROOT)) -from rl_engine.testing.h3_weights import EXTRACTED_FILE, load_h3_manifest, sha256_file # noqa: E402 +from rl_engine.testing.h3_weights import ( # noqa: E402 + EXTRACTED_FILE, + load_h3_manifest, + sha256_file, + sha256_tensor, +) def _check(path: Path, expected: str) -> None: @@ -86,6 +91,9 @@ def fetch(filename: str) -> Path: raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") if list(tensor.shape) != spec["shape"]: raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise SystemExit(f"{name}: sha256 {actual} != manifest {spec['sha256']}") target = out_dir / EXTRACTED_FILE save_file(tensors, str(target), metadata={"revision": identity["revision"]}) diff --git a/tests/h3/test_h3_cli.py b/tests/h3/test_h3_cli.py new file mode 100644 index 000000000..7712fc41a --- /dev/null +++ b/tests/h3/test_h3_cli.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The H3 CLIs preserve chain dependencies and require pinned report weights.""" + +from __future__ import annotations + +import json +import sys +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.testing import h3_chain +from rl_engine.testing.h3_weights import WEIGHTS_ENV +from scripts import h3_chain_replay, h3_evidence + + +@pytest.fixture +def cpu_chain_cli(monkeypatch): + executed = [] + state = {"provider_shift": 0} + + class CpuOp: + def __call__(self, value): + return value + 1 + + forward_fp32 = __call__ + + stages = [] + for index, stage in enumerate(h3_chain.STAGES): + + def stage_input(ctx, upstream, index=index): + if index == 0: + return ctx["timestep"] + assert isinstance(upstream, torch.Tensor), "a preceding stage must run first" + return upstream + + stages.append( + replace( + stage, + candidate=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + golden=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + provider=lambda ctx, up, stage_input=stage_input, index=index: ( + stage_input(ctx, up) + 1 + (state["provider_shift"] if index == 0 else 0) + ), + ) + ) + + def get_op(op_type, *, device): + executed.append(op_type) + return CpuOp() + + monkeypatch.setattr(h3_chain, "STAGES", stages) + monkeypatch.setattr(h3_chain, "make_context", lambda *a, **kw: {"timestep": torch.zeros(1)}) + monkeypatch.setattr(h3_chain, "golden_op", lambda *a: CpuOp()) + monkeypatch.setattr(h3_chain_replay, "KernelRegistry", lambda: SimpleNamespace(get_op=get_op)) + monkeypatch.setattr(h3_chain_replay, "load_h3_conditioning_weights", lambda *a: {}) + monkeypatch.setattr(h3_chain, "environment", lambda: {}) + monkeypatch.setattr(h3_chain, "git_state", lambda: {}) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + return executed, state + + +@pytest.mark.parametrize("provider_shift", [0, 1]) +@pytest.mark.parametrize( + "requested, execution_count", + [ + ("timestep_mlp_fp32", 2), + ("adaln_projection_3mod", 3), + ("timestep_sinusoid_h3,adaln_projection_3mod", 3), + ("timestep_mlp_fp32,timestep_sinusoid_h3", 2), + ], +) +def test_chain_selection_runs_dependencies( + cpu_chain_cli, monkeypatch, tmp_path, capsys, requested, execution_count, provider_shift +): + executed, state = cpu_chain_cli + state["provider_shift"] = provider_shift + out = tmp_path / "chain.json" + monkeypatch.setattr( + sys, + "argv", + [ + "h3_chain_replay.py", + "--stages", + requested, + "--timesteps", + "1", + "--seq-lens", + "3", + "--out", + str(out), + ], + ) + h3_chain_replay.main() + + report = json.loads(out.read_text()) + stage_names = [stage.name for stage in h3_chain.STAGES] + expected_report = [name for name in stage_names if name in requested.split(",")] + assert executed == stage_names[:execution_count] + assert report["executed_stages"] == executed + assert report["stages"] == expected_report + case = report["cases"][0] + assert [entry["stage"] for entry in case["stages"]] == expected_report + assert all(entry["chained_vs_golden"]["bitwise_equal"] for entry in case["stages"]) + assert case["first_drift"] == (stage_names[0] if provider_shift else None) + assert case["first_isolated_drift"] == (stage_names[0] if provider_shift else None) + summary = capsys.readouterr().out.split("; first_drift=", 1)[0] + for name in stage_names: + assert (f"{name}=" in summary) == (name in expected_report) + + +@pytest.mark.parametrize("op", ["timestep_mlp_fp32", "adaln_projection_3mod"]) +@pytest.mark.parametrize("weights_env", [None, " "]) +def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, weights_env): + if weights_env is None: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, weights_env) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + with pytest.raises(SystemExit, match=WEIGHTS_ENV): + h3_evidence.main() + assert not out.exists() + + +@pytest.mark.parametrize( + "op, weight_source", + [ + ("timestep_sinusoid_h3", "not_applicable"), + ("timestep_mlp_fp32", "pinned_checkpoint"), + ("adaln_projection_3mod", "pinned_checkpoint"), + ], +) +def test_evidence_records_weight_source(monkeypatch, tmp_path, op, weight_source): + if weight_source == "not_applicable": + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + monkeypatch.setattr(h3_evidence, "KernelRegistry", lambda: object()) + monkeypatch.setattr( + h3_evidence, + "git_state", + lambda: {"rl_kernel_commit": "test_commit", "tracked_tree_dirty": False}, + ) + monkeypatch.setattr(h3_evidence, "environment", lambda: {}) + monkeypatch.setitem(h3_evidence.ACCURACY, op, lambda registry: {}) + monkeypatch.setitem(h3_evidence.PERF_CASES, op, lambda registry: []) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + h3_evidence.main() + assert json.loads(out.read_text())["weight_source"] == weight_source diff --git a/tests/h3/test_h3_conditioning_e2e.py b/tests/h3/test_h3_conditioning_e2e.py index edd75a3b2..2bd1c47bd 100644 --- a/tests/h3/test_h3_conditioning_e2e.py +++ b/tests/h3/test_h3_conditioning_e2e.py @@ -22,7 +22,10 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing import h3_chain -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) # (T distinct timesteps, S packed rows): single row, short packed, realistic-size packed. CASES = [(1, 3), (3, 257), (4, 4097)] diff --git a/tests/h3/test_h3_det_linear.py b/tests/h3/test_h3_det_linear.py new file mode 100644 index 000000000..f8f90e83c --- /dev/null +++ b/tests/h3/test_h3_det_linear.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Native input validation for the H3 deterministic linear kernels.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.base import _C +from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), + pytest.mark.skipif(not det_linear_available(), reason="rl_engine._C lacks h3_det_linear_*"), +] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_backward_input_rejects_zero_rows(dtype): + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + weight = torch.empty(16, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_input(grad, weight, dtype) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_rejects_zero_rows(dtype, with_bias): + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + x = torch.empty(0, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_weight(grad, x, dtype, with_bias) diff --git a/tests/h3/test_h3_report.py b/tests/h3/test_h3_report.py new file mode 100644 index 000000000..eab1c8807 --- /dev/null +++ b/tests/h3/test_h3_report.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Comparative timing balances backend order and preserves its raw evidence.""" + +from __future__ import annotations + +import pytest + +from rl_engine.testing import h3_report + + +def test_measure_interleaves_backends_and_records_samples(monkeypatch): + elapsed_us = 0.0 + calls = [] + + class Event: + def __init__(self, *, enable_timing): + assert enable_timing + + def record(self): + self.timestamp = elapsed_us + + def synchronize(self): + pass + + def elapsed_time(self, end): + return (end.timestamp - self.timestamp) / 1e3 + + def candidate(): + nonlocal elapsed_us + calls.append("candidate") + elapsed_us += 2.0 + + def provider(): + nonlocal elapsed_us + calls.append("provider") + elapsed_us += 4.0 + + monkeypatch.setattr(h3_report.torch.cuda, "Event", Event) + monkeypatch.setattr(h3_report.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn: 0.0) + report = h3_report.measure( + {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 + ) + assert calls == [ + "candidate", + "provider", # warmup + "candidate", + "provider", + "provider", + "candidate", + "candidate", + "provider", + "provider", + "candidate", + ] + assert report["timing_order"] == [ + ["candidate", "provider"], + ["provider", "candidate"], + ["candidate", "provider"], + ["provider", "candidate"], + ] + assert report["timing_samples_us"] == {"candidate": [2.0] * 4, "provider": [4.0] * 4} + assert report["candidate_us"] == pytest.approx(2.0) + assert report["provider_us"] == pytest.approx(4.0) + assert report["candidate_gbps"] == pytest.approx(0.5) + assert report["provider_gbps"] == pytest.approx(0.25) diff --git a/tests/h3/test_h3_weights.py b/tests/h3/test_h3_weights.py new file mode 100644 index 000000000..015f78954 --- /dev/null +++ b/tests/h3/test_h3_weights.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Extracted H3 weights must retain the pinned tensor contents.""" + +from __future__ import annotations + +import hashlib + +import pytest +import torch + +from rl_engine.testing import h3_weights + + +@pytest.mark.parametrize( + "dtype,raw_bytes", + [ + (torch.float32, bytes.fromhex("000000000000803f")), + (torch.bfloat16, bytes.fromhex("0000803f")), + ], +) +def test_tensor_checksum_hashes_raw_bytes(dtype, raw_bytes): + tensor = torch.tensor([0.0, 1.0], dtype=dtype) + assert h3_weights.sha256_tensor(tensor) == hashlib.sha256(raw_bytes).hexdigest() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) +@pytest.mark.parametrize("names", [None, ["bias"]]) +def test_loading_rejects_modified_weight_contents(tmp_path, monkeypatch, dtype, mismatch, names): + from safetensors.torch import save_file + + tensors = { + "weight": torch.tensor([[0.0, 1.0], [2.0, 3.0]], dtype=dtype), + "bias": torch.tensor([0.0, 1.0], dtype=dtype), + } + manifest = { + "tensors": { + name: { + "dtype": str(tensor.dtype).removeprefix("torch."), + "shape": list(tensor.shape), + "sha256": h3_weights.sha256_tensor(tensor), + } + for name, tensor in tensors.items() + } + } + if mismatch == "dtype": + tensors["weight"] = tensors["weight"].to(torch.float64) + elif mismatch == "shape": + tensors["weight"] = tensors["weight"].reshape(4) + elif mismatch == "sha256": + tensors["weight"][0, 0] = 1.0 + # Tensor checksums do not depend on serializer metadata or tensor key order. + save_file( + dict(reversed(list(tensors.items()))), + str(tmp_path / h3_weights.EXTRACTED_FILE), + metadata={"revision": "a different serializer metadata value"}, + ) + monkeypatch.setenv(h3_weights.WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(h3_weights, "load_h3_manifest", lambda: manifest) + + if mismatch: + with pytest.raises(ValueError, match=f"weight: {mismatch} .* != manifest"): + h3_weights.load_h3_conditioning_weights("cpu", names=names) + else: + actual = h3_weights.load_h3_conditioning_weights("cpu", names=names) + assert list(actual) == (names if names is not None else list(tensors)) + for name, tensor in actual.items(): + assert torch.equal(tensor, tensors[name]) diff --git a/tests/h3/test_prepare_h3_weights.py b/tests/h3/test_prepare_h3_weights.py index 826fbd95e..8a1c94298 100644 --- a/tests/h3/test_prepare_h3_weights.py +++ b/tests/h3/test_prepare_h3_weights.py @@ -12,11 +12,11 @@ import pytest import torch -from rl_engine.testing.h3_weights import EXTRACTED_FILE, sha256_file +from rl_engine.testing.h3_weights import EXTRACTED_FILE, sha256_file, sha256_tensor @pytest.mark.parametrize("optimize", [0, 1]) -@pytest.mark.parametrize("mismatch", [None, "dtype", "shape"]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, mismatch): import huggingface_hub from safetensors.torch import load_file, save_file @@ -42,6 +42,7 @@ def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, misma "weight": { "dtype": "bfloat16" if mismatch == "dtype" else "float32", "shape": [3] if mismatch == "shape" else [2], + "sha256": "0" * 64 if mismatch == "sha256" else sha256_tensor(tensor), } }, } From 50c97dbd5e6b95b9e0222ee7b7d5cdbfa66fac42 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:56:13 +0000 Subject: [PATCH 04/24] ci: enable CodeRabbit auto review for test-h3 CodeRabbit skips automatic reviews targeting non-default branches by default. Allow the test-h3 base branch so pushes to the H3 stack receive incremental reviews. Validation: configuration validated against the official CodeRabbit JSON schema. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .coderabbit.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .coderabbit.yaml diff --git a/.coderabbit.yaml b/.coderabbit.yaml new file mode 100644 index 000000000..cc811057e --- /dev/null +++ b/.coderabbit.yaml @@ -0,0 +1,5 @@ +# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json +reviews: + auto_review: + base_branches: + - "^test-h3$" From d1b71b6b8ac6ba0139f74fe97055127258e9edf1 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:30:47 +0000 Subject: [PATCH 05/24] docs(h3): document conditioning operators and review regression tests Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- benchmarks/benchmark_h3_conditioning.py | 2 + rl_engine/_C.pyi | 42 +++++++++++++++++-- rl_engine/kernels/gtest/operator_inputs.py | 8 ++++ .../kernels/ops/cuda/h3/adaln_projection.py | 24 +++++++++++ rl_engine/kernels/ops/cuda/h3/det_linear.py | 25 +++++++++++ rl_engine/kernels/ops/cuda/h3/timestep_mlp.py | 17 ++++++++ .../kernels/ops/cuda/h3/timestep_sinusoid.py | 19 +++++++++ .../ops/pytorch/h3/adaln_projection.py | 22 ++++++++++ .../kernels/ops/pytorch/h3/timestep_mlp.py | 16 +++++++ .../ops/pytorch/h3/timestep_sinusoid.py | 15 +++++++ rl_engine/kernels/registry.py | 1 + rl_engine/testing/h3_cases.py | 13 ++++-- rl_engine/testing/h3_chain.py | 24 ++++++++++- rl_engine/testing/h3_provider.py | 17 ++++++-- rl_engine/testing/h3_report.py | 26 +++++++++++- rl_engine/testing/h3_weights.py | 6 +++ rl_engine/tests/test_dispatch.py | 2 + scripts/h3_chain_replay.py | 4 ++ scripts/h3_evidence.py | 2 + scripts/plot_h3_evidence.py | 18 +++++++- scripts/prepare_h3_weights.py | 6 +++ setup.py | 1 + tests/h3/conftest.py | 2 + tests/h3/test_h3_adaln_projection.py | 36 ++++++++++++++++ tests/h3/test_h3_cli.py | 14 +++++++ tests/h3/test_h3_conditioning_e2e.py | 6 +++ tests/h3/test_h3_det_linear.py | 4 ++ tests/h3/test_h3_report.py | 14 +++++++ tests/h3/test_h3_timestep_mlp.py | 42 +++++++++++++++++++ tests/h3/test_h3_timestep_sinusoid.py | 42 +++++++++++++++++++ tests/h3/test_h3_weights.py | 4 ++ tests/h3/test_prepare_h3_weights.py | 2 + tests/test_ws1_gtest_gpu.py | 2 + 33 files changed, 463 insertions(+), 15 deletions(-) diff --git a/benchmarks/benchmark_h3_conditioning.py b/benchmarks/benchmark_h3_conditioning.py index de1aaceb7..dc5aed949 100644 --- a/benchmarks/benchmark_h3_conditioning.py +++ b/benchmarks/benchmark_h3_conditioning.py @@ -27,6 +27,8 @@ def main() -> None: + """Benchmark selected CUDA operators, print summaries, and optionally save JSON.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--op", default="all", choices=["all", *PERF_CASES]) parser.add_argument("--warmup", type=int, default=20) diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fbe7558a8..29a889f27 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -330,17 +330,51 @@ def h3_timestep_sinusoid_forward( num_channels: int = 256, max_period: float = 10000.0, check_range: bool = True, -) -> torch.Tensor: ... +) -> torch.Tensor: + """Return FP32 CUDA cosine-then-sine features of shape ``(T, num_channels)``. + + Require a nonempty 1-D FP32 CUDA timestep tensor, positive even channels + and positive ``max_period``. Check finite values in ``[0, 1]`` when enabled; + the Python operator supplies the analytic backward for this native forward. + """ + ... + def h3_det_linear_forward( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None, activation: int = 0, save_pre_activation: bool = False, -) -> list[torch.Tensor]: ... +) -> list[torch.Tensor]: + """Return a deterministic CUDA projection and optional FP32 pre-activation. + + Require same-device/dtype FP32 or BF16 contiguous, 16-byte-aligned matrices + ``x (T, K)`` and ``weight (N, K)``, ``T > 0`` and ``K`` divisible by the + 16-byte vector width. Optional contiguous bias is ``(N,)`` in that dtype. + Activation 0 is identity and 1 is SiLU; BF16 forward requires SM80 or newer. + Return ``[out]`` in the input dtype or ``[out, pre]`` when saving activation. + """ + ... + def h3_det_linear_backward_input( grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype -) -> torch.Tensor: ... +) -> torch.Tensor: + """Return a deterministic ``(T, K)`` CUDA input gradient in ``out_dtype``. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``weight (N, K)``. Output dtype + must be FP32 or BF16; output-column chunks accumulate in fixed order. + """ + ... + def h3_det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True -) -> list[torch.Tensor]: ... +) -> list[torch.Tensor]: + """Return CUDA weight and optional bias gradients using an ascending row fold. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``x (T, K)``. Return ``[dweight]`` + with shape ``(N, K)``, or ``[dweight, dbias]`` with ``dbias (N,)`` when + requested, in FP32 or BF16 as specified by ``w_dtype``. + """ + ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 255a48c0c..ced084383 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -28,6 +28,7 @@ def make_operator_inputs( dtype: torch.dtype, device: torch.device, ) -> dict[str, Any]: + """Build operator keyword inputs on the requested device from CLI shape and seed options.""" builders = { "rms_norm": _make_rms_norm_inputs, "qk_norm": _make_qk_norm_inputs, @@ -57,6 +58,7 @@ def make_operator_inputs( def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: + """Return the operator's dimension label for benchmark and evidence reports.""" batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) names = { @@ -344,6 +346,7 @@ def _make_kv_cache_attention_inputs( def _h3_num_timesteps(args: argparse.Namespace) -> int: + """Read the packed H3 timestep count, falling back to the CLI batch size.""" # H3 packs a handful of distinct timesteps; reuse --batch as their count. return _arg_int(args, "num_timesteps", _arg_int(args, "batch", 2)) @@ -351,6 +354,7 @@ def _h3_num_timesteps(args: argparse.Namespace) -> int: def _h3_timesteps( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> torch.Tensor: + """Build seeded timesteps in [0, 1] with endpoint cases in the requested storage dtype.""" num = _h3_num_timesteps(args) mode = _arg_str(args, "input_mode", "random") if mode == "constant": @@ -366,12 +370,14 @@ def _h3_timesteps( def _make_timestep_sinusoid_h3_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: + """Supply a one-dimensional packed timestep tensor to the H3 sinusoid operator.""" return {"timestep": _h3_timesteps(args, dtype, device)} def _make_timestep_mlp_fp32_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: + """Build [T, 256] sinusoid features and seeded FP32 256→5376→2688 MLP parameters.""" # The H3 time_embedder is declared FP32; ``dtype`` does not apply to it. del dtype from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import NativeH3TimestepSinusoidOp @@ -390,12 +396,14 @@ def _make_timestep_mlp_fp32_inputs( def _h3_hidden(args: argparse.Namespace) -> int: + """Read the AdaLN channel width, defaulting to the checkpoint's 5376 channels.""" return _arg_int(args, "normalized_dim", H3_HIDDEN) def _make_adaln_projection_3mod_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: + """Build FP32 [T, 2688] embeddings and 18H projection parameters in the chosen dtype.""" # temb is FP32 by contract (the SiLU runs before the cast); ``dtype`` is the # projection's weight dtype, BF16 in the checkpoint. n_out = 6 * 3 * _h3_hidden(args) diff --git a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py index aee0ea503..fde0480e6 100644 --- a/rl_engine/kernels/ops/cuda/h3/adaln_projection.py +++ b/rl_engine/kernels/ops/cuda/h3/adaln_projection.py @@ -35,6 +35,11 @@ class _H3AdaLNProjectionCuda(torch.autograd.Function): @staticmethod def forward(ctx, temb, weight, bias): + """Project FP32 ``(T, D)`` embeddings to a weight-dtype ``(T, 18H)`` table. + + Save the embedding, rounded SiLU activation and weight for the VJP. + """ + temb = temb.contiguous() # Declared mixed-precision boundary: SiLU at temb's FP32 precision, # then exactly one rounding to the projection dtype. These are the @@ -46,6 +51,12 @@ def forward(ctx, temb, weight, bias): @staticmethod def backward(ctx, grad_table): + """Return requested embedding, weight and bias gradients in their input dtypes. + + Accumulate projection gradients deterministically and evaluate the + straight-through cast and SiLU VJP in FP32 for the embedding gradient. + """ + temb, act, weight = ctx.saved_tensors need_temb, need_weight, need_bias = ctx.needs_input_grad grad_table = grad_table.float().contiguous() @@ -81,15 +92,26 @@ class H3AdaLNProjectionCudaOp: contract = CONTRACT def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + if not det_linear_available(): raise RuntimeError( "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" ) def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` modulation views for CUDA inputs.""" + return self.forward(temb, weight, bias) def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six ``(3T, H)`` views after FP32 SiLU and a deterministic projection. + + Require CUDA FP32 ``temb`` of shape ``(T, D)`` with ``T > 0`` and + same-device FP32 or BF16 weight/bias of shapes ``(18H, D)``/``(18H,)``. + Outputs share one table in the weight dtype; autograd covers all inputs. + """ + hidden = validate_h3_adaln_projection(temb, weight, bias) if not temb.is_cuda: raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") @@ -105,4 +127,6 @@ def forward_table(self, temb, weight, bias) -> torch.Tensor: return _H3AdaLNProjectionCuda.apply(temb, weight, bias) def forward_fp32(self, temb, weight, bias): + """Run the CUDA projection with its declared weight-dtype output boundary.""" + return self.forward(temb, weight, bias) diff --git a/rl_engine/kernels/ops/cuda/h3/det_linear.py b/rl_engine/kernels/ops/cuda/h3/det_linear.py index f5dab854e..f12acabba 100644 --- a/rl_engine/kernels/ops/cuda/h3/det_linear.py +++ b/rl_engine/kernels/ops/cuda/h3/det_linear.py @@ -24,10 +24,14 @@ def det_linear_available() -> bool: + """Return whether the extension exposes all three H3 deterministic linear APIs.""" + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) def _require() -> None: + """Raise when a deterministic linear entry point is unavailable.""" + if not det_linear_available(): raise RuntimeError( "rl_engine._C lacks the h3_det_linear_* symbols; rebuild the CUDA extension " @@ -43,6 +47,14 @@ def det_linear_forward( activation: int = ACT_NONE, save_pre_activation: bool = False, ) -> list[torch.Tensor]: + """Project CUDA ``x`` of shape ``(T, K)`` with ``(N, K)`` weights and optional bias. + + Inputs share an FP32 or BF16 dtype/device, ``T > 0``, and ``K`` supports + 16-byte vector loads; BF16 forward needs SM80 or newer. Return ``[out]`` + in the input dtype, plus an FP32 pre-activation when requested. Activation + is identity or SiLU; contiguous copies preserve the reduction contract. + """ + _require() return _C.h3_det_linear_forward( x.contiguous(), @@ -56,6 +68,12 @@ def det_linear_forward( def det_linear_backward_input( grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype ) -> torch.Tensor: + """Compute a deterministic ``(T, K)`` input VJP on the inputs' CUDA device. + + Convert nonempty ``(T, N)`` gradients to FP32 and combine them with FP32 + or BF16 ``(N, K)`` weights; return FP32 or BF16 as specified by ``out_dtype``. + """ + _require() return _C.h3_det_linear_backward_input( grad.float().contiguous(), weight.contiguous(), out_dtype @@ -65,6 +83,13 @@ def det_linear_backward_input( def det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True ) -> list[torch.Tensor]: + """Fold nonempty CUDA ``(T, N)`` gradients and ``(T, K)`` inputs in row order. + + Convert gradients to FP32; ``x`` is FP32 or BF16 on the same device. + Return ``[dweight]`` of shape ``(N, K)``, optionally followed by ``dbias`` + of shape ``(N,)``, both in the requested FP32 or BF16 ``w_dtype``. + """ + _require() return _C.h3_det_linear_backward_weight( grad.float().contiguous(), x.contiguous(), w_dtype, bool(with_bias) diff --git a/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py b/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py index da85e87d8..041116c70 100644 --- a/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py +++ b/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py @@ -32,6 +32,8 @@ class _H3TimestepMLPCuda(torch.autograd.Function): @staticmethod def forward(ctx, x, w1, b1, w2, b2): + """Evaluate the FP32 CUDA linear-SiLU-linear graph and save VJP inputs.""" + x = x.contiguous() hidden, pre = det_linear_forward(x, w1, b1, activation=ACT_SILU, save_pre_activation=True) (out,) = det_linear_forward(hidden, w2, b2, activation=ACT_NONE) @@ -40,6 +42,8 @@ def forward(ctx, x, w1, b1, w2, b2): @staticmethod def backward(ctx, grad_out): + """Return requested FP32 input and parameter VJPs with fixed reduction order.""" + x, w1, w2, pre, hidden = ctx.saved_tensors need_x, need_w1, need_b1, need_w2, need_b2 = ctx.needs_input_grad grad_out = grad_out.float().contiguous() @@ -81,19 +85,32 @@ class H3TimestepMLPCudaOp: contract = CONTRACT def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + if not det_linear_available(): raise RuntimeError( "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" ) def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` timestep embeddings for same-device CUDA inputs.""" + return self.forward(x, w1, b1, w2, b2) def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Apply two deterministic FP32 CUDA projections with an intervening SiLU. + + Require same-device FP32 tensors: ``x`` is nonempty ``(T, K)``, + weights are ``(H, K)`` and ``(D, H)``, and biases are ``(H,)`` and + ``(D,)``. Return ``(T, D)`` with autograd support for all five inputs. + """ + validate_h3_timestep_mlp(x, w1, b1, w2, b2) if not x.is_cuda: raise ValueError("H3TimestepMLPCudaOp needs CUDA tensors") return _H3TimestepMLPCuda.apply(x, w1, b1, w2, b2) def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Use the standard CUDA MLP, whose inputs and output are already FP32.""" + return self.forward(x, w1, b1, w2, b2) diff --git a/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py index 72669d7c0..8601ef8f8 100644 --- a/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py +++ b/rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py @@ -21,12 +21,16 @@ def h3_sinusoid_cuda_available() -> bool: + """Return whether the extension provides the CUDA H3 sinusoid entry point.""" + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) class _H3TimestepSinusoidCuda(torch.autograd.Function): @staticmethod def forward(ctx, timestep: torch.Tensor, num_channels: int, check_range: bool): + """Compute FP32 CUDA features and save them for the analytic timestep VJP.""" + t32 = timestep.detach().float().contiguous() out = _C.h3_timestep_sinusoid_forward( t32, int(num_channels), float(H3_MAX_PERIOD), check_range @@ -38,6 +42,8 @@ def forward(ctx, timestep: torch.Tensor, num_channels: int, check_range: bool): @staticmethod def backward(ctx, grad_out: torch.Tensor): + """Reduce channel derivatives with a fixed FP32 tree into the timestep dtype.""" + (out,) = ctx.saved_tensors half = ctx.num_channels // 2 freq = NativeH3TimestepSinusoidOp.frequencies_fp32(ctx.num_channels, out.device) @@ -70,6 +76,8 @@ class H3TimestepSinusoidCudaOp: backward_impl = BACKWARD_IMPL def __init__(self) -> None: + """Require the native H3 CUDA sinusoid symbol before creating this operator.""" + if not h3_sinusoid_cuda_available(): raise RuntimeError( "rl_engine._C.h3_timestep_sinusoid_forward is unavailable; rebuild the " @@ -77,6 +85,8 @@ def __init__(self) -> None: ) def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 CUDA features for timesteps in ``[0, 1]`` with gradients.""" + return self.forward(timestep, num_channels=num_channels) def forward( @@ -86,6 +96,13 @@ def forward( num_channels: int = H3_FREQ_DIM, check_range: bool = True, ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on CUDA. + + Require nonempty 1-D FP32, FP16 or BF16 timesteps and a positive even + channel count. Native validation rejects nonfinite/out-of-range values + when ``check_range`` is true; backward returns the original input dtype. + """ + # Native code owns value validation, avoiding a second host sync here. validate_h3_timesteps(timestep, num_channels, check_range=False) if not timestep.is_cuda: @@ -93,5 +110,7 @@ def forward( return _H3TimestepSinusoidCuda.apply(timestep, num_channels, check_range) def forward_fp32(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Use the standard CUDA sinusoid path, which already returns FP32 features.""" + # The op is FP32 end to end; the FP32 path is the op itself. return self.forward(timestep, num_channels=num_channels) diff --git a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py index 75ec2b8bd..b6d013961 100644 --- a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py +++ b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py @@ -26,6 +26,8 @@ def adaln_hidden_size(weight: torch.Tensor) -> int: + """Infer ``H`` from a 2-D projection weight whose row count is divisible by 18.""" + rows_per_hidden = H3_ADALN_CHUNKS * H3_MODALITY_NUM if weight.dim() != 2 or weight.shape[0] % rows_per_hidden != 0: raise ValueError( @@ -36,6 +38,12 @@ def adaln_hidden_size(weight: torch.Tensor) -> int: def validate_h3_adaln_projection(temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): + """Check the AdaLN input shapes, shared device and precision boundary; return ``H``. + + Require nonempty FP32 ``temb`` of shape ``(T, D)`` and FP32 or BF16 weight + of shape ``(18H, D)`` with matching-dtype bias of shape ``(18H,)``. + """ + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): if not isinstance(tensor, torch.Tensor): raise TypeError(f"{name} must be a torch.Tensor") @@ -79,14 +87,28 @@ class NativeH3AdaLNProjectionOp: op_class = "reduction" def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` views using the PyTorch provider path.""" + return self.forward(temb, weight, bias) def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Apply FP32 SiLU, cast to the weight dtype and project on the input device. + + Validate ``(T, D)`` embeddings and ``(18H, D)``/``(18H,)`` parameters; + return six weight-dtype ``(3T, H)`` views with PyTorch autograd support. + """ + hidden = validate_h3_adaln_projection(temb, weight, bias) table = F.linear(F.silu(temb).to(weight.dtype), weight, bias) return split_adaln_table(table, hidden) def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six FP32 golden views from FP64 SiLU and projection arithmetic. + + Preserve the declared activation rounding to the weight dtype while + using an identity VJP at that cast boundary for the FP64 golden graph. + """ + hidden = validate_h3_adaln_projection(temb, weight, bias) t64 = temb.double() act = t64 * torch.sigmoid(t64) diff --git a/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py b/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py index ac54d295f..51ef48db7 100644 --- a/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py +++ b/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py @@ -27,6 +27,12 @@ def validate_h3_timestep_mlp( w2: torch.Tensor, b2: torch.Tensor, ) -> None: + """Check a same-device FP32 ``(T, K) -> (T, H) -> (T, D)`` MLP contract. + + Require ``T > 0``, weights of shapes ``(H, K)``/``(D, H)`` and biases + of shapes ``(H,)``/``(D,)``; reject lower-precision parameters or inputs. + """ + named = {"x": x, "w1": w1, "b1": b1, "w2": w2, "b2": b2} for name, tensor in named.items(): if not isinstance(tensor, torch.Tensor): @@ -64,13 +70,23 @@ class NativeH3TimestepMLPOp: op_class = "reduction" def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` embeddings using the PyTorch provider graph.""" + return self.forward(x, w1, b1, w2, b2) def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Validate and evaluate linear-SiLU-linear in FP32 on the input device. + + Inputs follow the ``(T, K) -> (T, H) -> (T, D)`` contract and retain + PyTorch autograd support; CUDA callers must disable TF32 for this path. + """ + validate_h3_timestep_mlp(x, w1, b1, w2, b2) return F.linear(F.silu(F.linear(x, w1, b1)), w2, b2) def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Evaluate the validated MLP in FP64 and round its ``(T, D)`` output to FP32.""" + validate_h3_timestep_mlp(x, w1, b1, w2, b2) z = F.linear(x.double(), w1.double(), b1.double()) h = z * torch.sigmoid(z) diff --git a/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py b/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py index 8ca0c0ccb..54da13527 100644 --- a/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py +++ b/rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py @@ -68,6 +68,8 @@ class NativeH3TimestepSinusoidOp: op_class = "elementwise" def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 sinusoidal features for nonempty timesteps in ``[0, 1]``.""" + return self.forward(timestep, num_channels=num_channels) def forward( @@ -77,6 +79,13 @@ def forward( num_channels: int = H3_FREQ_DIM, check_range: bool = True, ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on the input device. + + Require nonempty 1-D FP32, BF16 or FP16 timesteps and positive even + channels. Validate finite values in ``[0, 1]`` when ``check_range`` is + true and preserve PyTorch autograd through the provider formula. + """ + validate_h3_timesteps(timestep, num_channels, check_range=check_range) half = num_channels // 2 # Same op sequence and dtypes as diffusers get_timestep_embedding with @@ -96,6 +105,12 @@ def forward_fp32( num_channels: int = H3_FREQ_DIM, check_range: bool = True, ) -> torch.Tensor: + """Evaluate the validated sinusoid formula in FP64 and round once to FP32. + + Return ``(T, num_channels)`` on the timestep device with cosine channels + followed by sine channels and gradients through the FP64 golden graph. + """ + validate_h3_timesteps(timestep, num_channels, check_range=check_range) half = num_channels // 2 k = torch.arange(half, dtype=torch.float64, device=timestep.device) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 22bbab259..880363b6c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -456,6 +456,7 @@ class KernelRegistry: """ def __init__(self): + """Initialize backend caches, capability contracts, and device-specific priorities.""" self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() self.semantic = SemanticOperatorCatalog(_default_semantic_descriptors()) diff --git a/rl_engine/testing/h3_cases.py b/rl_engine/testing/h3_cases.py index 527a820ec..8eb509d33 100644 --- a/rl_engine/testing/h3_cases.py +++ b/rl_engine/testing/h3_cases.py @@ -9,7 +9,11 @@ def h3_timesteps(num: int, *, seed: int = 0, device: str = "cuda") -> torch.Tensor: - """Distinct-looking timesteps in [0, 1] with both endpoints present.""" + """Return seeded FP32 timesteps of shape ``(num,)`` on ``device``. + + Requires positive ``num``; the first value is 0 and the last is 1 when + ``num > 1``, with all remaining values sampled uniformly in [0, 1]. + """ generator = torch.Generator(device="cpu").manual_seed(seed) t = torch.rand(num, generator=generator) @@ -22,10 +26,11 @@ def h3_timesteps(num: int, *, seed: int = 0, device: str = "cuda") -> torch.Tens def h3_packed_layout( seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" ) -> tuple[torch.Tensor, torch.Tensor]: - """(timestep_indices, token_tags) for a packed sequence with every modality. + """Return seeded int64 timestep indices and modality tags, each shape ``(seq_len,)``. - Rows are video, text and audio blocks in random order, each block mapped - to a random timestep, like H3's conditioning/target packing. + Each row independently samples a timestep in ``[0, num_timesteps)`` and + a video/text/audio tag in ``[0, 3)``. The first three rows cover all tags + when available; both tensors are returned on ``device``. """ generator = torch.Generator(device="cpu").manual_seed(seed) diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 8671ef141..11ca1c0fb 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -54,11 +54,15 @@ class Stage: def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: + """Return the context's MLP weights and biases in linear_1, linear_2 call order.""" + weights = ctx["weights"] return [weights[f"time_embedder.linear_{i}.{p}"] for i in (1, 2) for p in ("weight", "bias")] def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + """Return the context's block-0 AdaLN projection weight and bias.""" + prefix = "transformer_blocks.0.adaln_proj.linear" return ctx["weights"][f"{prefix}.weight"], ctx["weights"][f"{prefix}.bias"] @@ -98,12 +102,16 @@ def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: def _flat(value: Any) -> list[torch.Tensor]: + """Collect tensors from nested sequences in their original order.""" + if isinstance(value, torch.Tensor): return [value] return [tensor for item in value for tensor in _flat(item)] def compare(lhs: Any, rhs: Any, atol: float = 0.0, rtol: float = 0.0) -> dict[str, Any]: + """Report tensor equality, maximum FP32 absolute error, and tolerance checks.""" + lhs_t, rhs_t = _flat(lhs), _flat(rhs) bitwise = all(torch.equal(a, b) for a, b in zip(lhs_t, rhs_t, strict=True)) max_abs = 0.0 @@ -117,6 +125,8 @@ def compare(lhs: Any, rhs: Any, atol: float = 0.0, rtol: float = 0.0) -> dict[st def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dict[str, Any]: + """Build seeded CUDA timesteps and packed row labels with the supplied weights.""" + timestep_indices, token_tags = h3_packed_layout(seq_len, num_timesteps, seed=seed) return { "timestep": h3_timesteps(num_timesteps, seed=seed), @@ -127,7 +137,7 @@ def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dic def golden_op(registry: KernelRegistry, op_type: str): - """The PyTorch reference: the last entry of the CPU priority list.""" + """Return the PyTorch reference at the end of the operator's CPU priority list.""" return registry._get_or_create_backend(registry._priority_map["cpu"][op_type][-1]) @@ -141,6 +151,12 @@ def run_case( seed: int = 0, stages: list[Stage] | None = None, ) -> dict[str, Any]: + """Replay CUDA candidate, provider, and golden paths and report each stage's drift. + + Supplied stages must include their prerequisites in execution order. + Omitting ``stages`` runs the full registered chain without gradients. + """ + stages = STAGES if stages is None else stages ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) report: dict[str, Any] = { @@ -182,7 +198,11 @@ def run_case( def git_state() -> dict[str, Any]: + """Return HEAD and tracked-tree dirtiness, defaulting to unknown/dirty if Git fails.""" + def run(*args: str) -> str: + """Run Git in the repository root and return stripped stdout.""" + return subprocess.check_output(["git", *args], cwd=REPO_ROOT, text=True).strip() try: @@ -194,6 +214,8 @@ def run(*args: str) -> str: def environment() -> dict[str, Any]: + """Record the active CUDA device, software versions, and matmul TF32 setting.""" + return { "gpu": torch.cuda.get_device_name(), "capability": list(torch.cuda.get_device_capability()), diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index f64f7174f..eaf31ba92 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -19,7 +19,11 @@ def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch.Tensor: - """``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0)``.""" + """Replay provider FP32 cosine/sine features for ``(T,)`` timesteps. + + Matches ``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0)`` + with the default channel count, returning shape ``(T, 256)``. + """ half_dim = num_channels // 2 exponent = -math.log(10000) * torch.arange( @@ -41,7 +45,10 @@ def provider_time_embedder( w2: torch.Tensor, b2: torch.Tensor, ) -> torch.Tensor: - """``TimestepEmbedding(256, 5376, out_dim=2688)``: linear_1 -> SiLU -> linear_2, FP32.""" + """Replay linear_1, SiLU, and linear_2 using the checkpoint weight dtype. + + H3's FP32 weights map features of shape ``(T, 256)`` to ``(T, 2688)``. + """ sample = F.linear(features.to(w1.dtype), w1, b1) sample = F.silu(sample) @@ -51,7 +58,11 @@ def provider_time_embedder( def provider_adaln_modulation( temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, hidden_size: int = 5376 ) -> tuple[torch.Tensor, ...]: - """``MiniMaxH3AdaLayerNormModulation.forward``: six (3T, H) views.""" + """Replay provider SiLU, the weight-dtype cast, and the AdaLN projection. + + Return six ``(3T, hidden_size)`` views in timestep-major modality order + for checkpoint weights shaped ``(18 * hidden_size, 2688)``. + """ temb = F.linear(F.silu(temb).to(weight.dtype), weight, bias) temb = temb.view(-1, 6 * hidden_size) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index d4b45e09b..ad9eb87fc 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -53,6 +53,8 @@ def time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 200) -> float: def peak_mib(fn: Callable[[], Any]) -> float: + """Synchronize one CUDA call and report peak allocated memory above baseline in MiB.""" + torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() base = torch.cuda.memory_allocated() @@ -62,6 +64,12 @@ def peak_mib(fn: Callable[[], Any]) -> float: def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + """Time case callables in rotating order after warmup and return their evidence. + + Record raw CUDA-event samples, executed orders, median microseconds, + decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. + """ + row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} keys = [key for key in TIMED_KEYS if key in case] if not keys: @@ -101,6 +109,8 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st def _sinusoid_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA sinusoid timing cases with checked, unchecked, and provider calls.""" + op = registry.get_op("timestep_sinusoid_h3", device="cuda") cases = [] for num in (1, 2, 4, 64): @@ -120,6 +130,8 @@ def _sinusoid_perf(registry: KernelRegistry) -> list[dict[str, Any]]: def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: + """Compare sinusoid outputs to provider and FP64 golden across timestep counts.""" + op = registry.get_op("timestep_sinusoid_h3", device="cuda") golden = registry._get_or_create_backend( registry._priority_map["cpu"]["timestep_sinusoid_h3"][-1] @@ -140,7 +152,7 @@ def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Tensor]: - """Pinned checkpoint tensors when available, otherwise same-shape random ones.""" + """Return CUDA checkpoint parameters, or seeded random tensors when weights are unset.""" if h3_weights_dir() is not None: weights = load_h3_conditioning_weights("cuda", names) @@ -161,6 +173,8 @@ def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Ten def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA MLP candidate/provider timing cases for one to four timesteps.""" + op = registry.get_op("timestep_mlp_fp32", device="cuda") params = h3_params(MLP_NAMES, MLP_SHAPES) weight_bytes = sum(p.numel() * p.element_size() for p in params) @@ -181,6 +195,8 @@ def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: + """Measure per-draw MLP error against FP64 and test single-row/batched equality.""" + op = registry.get_op("timestep_mlp_fp32", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["timestep_mlp_fp32"][-1]) sinusoid = registry.get_op("timestep_sinusoid_h3", device="cuda") @@ -216,11 +232,15 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: def _adaln_params() -> tuple[torch.Tensor, torch.Tensor]: + """Return block-0 AdaLN projection weights and bias as BF16 CUDA tensors.""" + weight, bias = h3_params(ADALN_NAMES, [(96768, 2688), (96768,)]) return weight.bfloat16(), bias.bfloat16() def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA AdaLN candidate/provider timing cases for one to four timesteps.""" + op = registry.get_op("adaln_projection_3mod", device="cuda") weight, bias = _adaln_params() cases = [] @@ -240,11 +260,13 @@ def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: def _flat_cat(outputs) -> torch.Tensor: + """Concatenate flattened modulation tensors in their output order.""" + return torch.cat([out.reshape(-1) for out in outputs]) def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, Any]: - """Fraction of BF16 outputs equal to the correctly rounded FP64 golden, per draw.""" + """Report per-draw BF16 equality to FP64 goldens and single-row/batch invariance.""" op = registry.get_op("adaln_projection_3mod", device="cuda") golden = registry._get_or_create_backend( diff --git a/rl_engine/testing/h3_weights.py b/rl_engine/testing/h3_weights.py index 7d7abdf03..6b2dc995c 100644 --- a/rl_engine/testing/h3_weights.py +++ b/rl_engine/testing/h3_weights.py @@ -25,15 +25,21 @@ def load_h3_manifest() -> dict[str, Any]: + """Read the bundled checkpoint identity and pinned tensor contracts.""" + return json.loads(MANIFEST_PATH.read_text()) def h3_weights_dir() -> Path | None: + """Expand the configured extraction directory, or return None when unset or blank.""" + value = os.environ.get(WEIGHTS_ENV, "").strip() return Path(value).expanduser() if value else None def sha256_file(path: Path, chunk: int = 1 << 24) -> str: + """Return a file SHA256 while reading bounded chunks.""" + digest = hashlib.sha256() with path.open("rb") as handle: while block := handle.read(chunk): diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index c20e157e4..7f810d815 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -231,6 +231,8 @@ def fake_load_backend(backend): @pytest.mark.parametrize("op_name", H3_OPS) def test_h3_ops_prefer_cuda_and_fall_back_to_pytorch_reference(op_name): + """Prefer H3 CUDA backends on CUDA and use their PyTorch fallbacks elsewhere.""" + registry = KernelRegistry() cuda_chain = registry._priority_map["cuda"][op_name] assert cuda_chain[0].name.startswith("CUDA_H3_") diff --git a/scripts/h3_chain_replay.py b/scripts/h3_chain_replay.py index 512a15f06..395d3e17c 100644 --- a/scripts/h3_chain_replay.py +++ b/scripts/h3_chain_replay.py @@ -30,10 +30,14 @@ def _ints(text: str) -> list[int]: + """Parse comma-separated integers for the replay CLI's case dimensions.""" + return [int(v) for v in text.split(",")] def main() -> None: + """Replay requested CUDA stages with prerequisites and optionally save JSON evidence.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--stages", diff --git a/scripts/h3_evidence.py b/scripts/h3_evidence.py index 8351da245..426b0848d 100644 --- a/scripts/h3_evidence.py +++ b/scripts/h3_evidence.py @@ -34,6 +34,8 @@ def main() -> None: + """Collect CUDA accuracy and performance evidence, requiring weights for weighted ops.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) parser.add_argument("--out", type=Path, required=True) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 64edd56fe..7ebd228d0 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -35,6 +35,8 @@ def _style(ax, title: str, ylabel: str) -> None: + """Apply the shared report axis styling, title, and vertical-axis label.""" + ax.set_facecolor(SURFACE) ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) ax.set_ylabel(ylabel, color=MUTED, fontsize=9) @@ -47,7 +49,7 @@ def _style(ax, title: str, ylabel: str) -> None: def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt: str) -> None: - """Grouped bars with a 2px-style gap and value labels on top.""" + """Draw grouped series on the axis with formatted value labels and a legend.""" n = len(series) width = 0.8 / n @@ -75,10 +77,14 @@ def _bars(ax, labels: list[str], series: list[tuple[str, str, list[float]]], fmt def _log_values(values: list[float]) -> list[float]: + """Replace nonpositive values with the plotting floor for logarithmic axes.""" + return [v if v > 0 else LOG_FLOOR for v in values] def _figure(report: dict[str, Any], suptitle: str): + """Create two report axes with a title and GPU, software, and revision provenance.""" + fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), facecolor=SURFACE) env = report["environment"] fig.suptitle(suptitle, x=0.01, ha="left", fontsize=12, color=INK, fontweight="bold") @@ -94,6 +100,8 @@ def _figure(report: dict[str, Any], suptitle: str): def plot_sinusoid(report: dict[str, Any]): + """Return a sinusoid report figure showing latency and errors against FP64.""" + fig, (left, right) = _figure(report, "timestep_sinusoid_h3") perf = report["perf"] _style(left, "Latency per call (lower is better)", "µs") @@ -155,6 +163,8 @@ def plot_sinusoid(report: dict[str, Any]): def _latency_panel(ax, perf: list[dict[str, Any]], provider_name: str) -> None: + """Plot candidate and provider latency summaries in microseconds for each case.""" + _style(ax, "Latency per call (lower is better)", "µs") _bars( ax, @@ -209,6 +219,8 @@ def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], con def plot_mlp(report: dict[str, Any]): + """Return an MLP report figure showing latency and per-draw FP64 error distributions.""" + fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") _latency_panel(left, report["perf"], "diffusers (cuBLAS)") acc = report["accuracy"] @@ -225,6 +237,8 @@ def plot_mlp(report: dict[str, Any]): def plot_projection(report: dict[str, Any]): + """Return an AdaLN figure with latency, BF16 rounding accuracy, and the early-cast probe.""" + fig, (left, right) = _figure(report, "adaln_projection_3mod (pinned block-0 AdaLN weights)") _latency_panel(left, report["perf"], "diffusers (cuBLAS)") acc = report["accuracy"] @@ -269,6 +283,8 @@ def plot_projection(report: dict[str, Any]): def main() -> None: + """Read an operator report and save its figure as a PNG beside it or at ``--out``.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("report", type=Path) parser.add_argument("--out", type=Path, default=None) diff --git a/scripts/prepare_h3_weights.py b/scripts/prepare_h3_weights.py index 262408736..a9af15c37 100644 --- a/scripts/prepare_h3_weights.py +++ b/scripts/prepare_h3_weights.py @@ -32,6 +32,8 @@ def _check(path: Path, expected: str) -> None: + """Verify a downloaded file's SHA-256, exiting on mismatch or printing success.""" + actual = sha256_file(path) if actual != expected: raise SystemExit(f"sha256 mismatch for {path.name}: {actual} != {expected}") @@ -39,6 +41,8 @@ def _check(path: Path, expected: str) -> None: def main() -> None: + """Download pinned model artifacts, verify tensor identities, and write extracted weights.""" + parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--out", type=Path, required=True, help="output directory") parser.add_argument( @@ -60,6 +64,8 @@ def main() -> None: out_dir.mkdir(parents=True, exist_ok=True) def fetch(filename: str) -> Path: + """Download one file from the manifest's model subfolder and pinned revision.""" + return Path( hf_hub_download( identity["hf_repo"], diff --git a/setup.py b/setup.py index b388bdd29..2211604c5 100644 --- a/setup.py +++ b/setup.py @@ -96,6 +96,7 @@ def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: def get_extensions(): + """Configure the optional PyTorch native extension for the active CUDA or ROCm toolchain.""" torch, _, CUDAExtension = _load_torch_extension_tools() if torch is None: message = ( diff --git a/tests/h3/conftest.py b/tests/h3/conftest.py index 44c8fe117..1c3ba8caa 100644 --- a/tests/h3/conftest.py +++ b/tests/h3/conftest.py @@ -17,6 +17,8 @@ @pytest.fixture(scope="session") def h3_weights_cpu() -> dict[str, torch.Tensor]: + """Load pinned CPU weights once per session, skipping unavailable artifacts.""" + if h3_weights_dir() is None: pytest.skip(f"{WEIGHTS_ENV} not set; run scripts/prepare_h3_weights.py") try: diff --git a/tests/h3/test_h3_adaln_projection.py b/tests/h3/test_h3_adaln_projection.py index c9aac92d1..0b031fa4e 100644 --- a/tests/h3/test_h3_adaln_projection.py +++ b/tests/h3/test_h3_adaln_projection.py @@ -29,6 +29,8 @@ def _cuda_op(): + """Construct the CUDA projection operator, skipping missing native linear kernels.""" + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available @@ -38,11 +40,15 @@ def _cuda_op(): def _temb(num, dim=2688, seed=0, device="cuda"): + """Generate seeded FP32 timestep embeddings on CPU and move them to the test device.""" + g = torch.Generator(device="cpu").manual_seed(seed) return (torch.randn(num, dim, generator=g) * 2).to(device) def _synthetic(hidden, dim, dtype=torch.bfloat16, seed=0, device="cuda"): + """Create seeded projection weights for six chunks and three modalities per channel.""" + g = torch.Generator(device="cpu").manual_seed(seed) n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).to(dtype) @@ -51,11 +57,15 @@ def _synthetic(hidden, dim, dtype=torch.bfloat16, seed=0, device="cuda"): def _flat(outputs): + """Flatten and concatenate all modulation chunks for elementwise comparisons.""" + return torch.cat([out.reshape(-1) for out in outputs]) class TestReference: def test_provider_layout_matches_diffusers(self): + """Match the provider values and timestep-major modality layout on CPU.""" + temb = _temb(2, dim=16, device="cpu") weight, bias = _synthetic(8, 16, device="cpu") ours = NativeH3AdaLNProjectionOp().forward(temb, weight, bias) @@ -66,11 +76,15 @@ def test_provider_layout_matches_diffusers(self): def test_rejects_bf16_temb(self): # RFC probe H7: casting temb before the SiLU. + """Reject an early BF16 cast so SiLU receives the declared FP32 embeddings.""" + weight, bias = _synthetic(8, 16, device="cpu") with pytest.raises(TypeError, match="float32"): NativeH3AdaLNProjectionOp().forward(_temb(2, 16, device="cpu").bfloat16(), weight, bias) def test_rejects_bad_shapes(self): + """Reject invalid projection dimensions, empty batches, and mismatched bias dtype.""" + weight, bias = _synthetic(8, 16, device="cpu") op = NativeH3AdaLNProjectionOp() with pytest.raises(ValueError): # not a multiple of 6 * 3 rows @@ -104,6 +118,8 @@ def test_channel_identity_lands_in_its_chunk_and_modality_row(self): assert torch.equal(out[t * 3 + m].float(), expected), (c, t, m) def test_outputs_are_views_of_one_table_like_diffusers(self): + """Require all six modulation chunks to share storage with the expected offsets.""" + weight, bias = _synthetic(16, 32) outputs = _cuda_op()(_temb(2, 32), weight, bias) base = outputs[0].untyped_storage().data_ptr() @@ -113,6 +129,8 @@ def test_outputs_are_views_of_one_table_like_diffusers(self): @pytest.fixture(scope="module") def block0(h3_weights_cpu): + """Move the pinned block-zero projection weight and bias to CUDA for this module.""" + return h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() @@ -120,6 +138,8 @@ def block0(h3_weights_cpu): class TestCudaRealWeights: @pytest.mark.parametrize("num", [1, 2, 3, 4]) def test_forward_is_correctly_rounded_declared_golden(self, block0, num): + """Check checkpoint accuracy and near-exact BF16 rounding of the declared golden.""" + weight, bias = block0 temb = _temb(num, seed=num) outputs = _cuda_op()(temb, weight, bias) @@ -146,6 +166,8 @@ def test_early_cast_is_numerically_detectable(self, block0): assert (ours == early).float().mean() < 0.8 def test_rows_are_batch_and_position_invariant(self, block0): + """Keep each timestep modulation row identical alone and under batch reordering.""" + weight, bias = block0 op = _cuda_op() temb = _temb(5, seed=3) @@ -159,6 +181,8 @@ def test_rows_are_batch_and_position_invariant(self, block0): assert torch.equal(s.view(5, 3, -1).flip(0), f.view(5, 3, -1)) def test_backward_against_golden(self, block0): + """Check embedding and parameter gradients against a high-precision declared-cast VJP.""" + weight, bias = block0 temb = _temb(2, seed=5) g = torch.Generator(device="cuda").manual_seed(0) @@ -184,6 +208,8 @@ class TestCudaSynthetic: @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("hidden, dim", [(24, 16), (40, 72), (8, 2688)]) def test_other_sizes(self, dtype, hidden, dim): + """Match the projection golden at non-checkpoint dimensions in FP32 and BF16.""" + weight, bias = _synthetic(hidden, dim, dtype=dtype, seed=hidden) temb = _temb(3, dim, seed=hidden) tol = dict(atol=5e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-4, rtol=1e-4) @@ -194,6 +220,8 @@ def test_other_sizes(self, dtype, hidden, dim): torch.testing.assert_close(out.float(), gold, **tol) def test_repeat_forward_backward_bitwise(self): + """Require repeated projection outputs and all input gradients to match bitwise.""" + weight, bias = _synthetic(64, 2688, seed=9) temb = _temb(3, seed=9) grads = [torch.randn(9, 64, device="cuda").bfloat16() for _ in range(6)] @@ -208,11 +236,15 @@ def test_repeat_forward_backward_bitwise(self): assert torch.equal(a, b) def test_temb_grad_rows_are_batch_invariant(self): + """Keep embedding gradients bitwise equal in full batches and single-row calls.""" + weight, bias = _synthetic(64, 2688, seed=11) temb = _temb(4, seed=11) grads = [torch.randn(12, 64, device="cuda").bfloat16() for _ in range(6)] def d_temb(rows): + """Compute selected embedding-row gradients with their matching modality gradients.""" + leaf = temb[rows].detach().clone().requires_grad_(True) outs = _cuda_op()(leaf, weight, bias) torch.autograd.backward( @@ -237,11 +269,15 @@ def test_rows_invariant_across_tensor_core_row_tiles(self, num): assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) def test_rejects_cpu(self): + """Reject CPU tensors at the CUDA projection boundary.""" + weight, bias = _synthetic(8, 16, device="cpu") with pytest.raises(ValueError): _cuda_op()(_temb(2, 16, device="cpu"), weight, bias) def test_registry_dispatches_cuda(self): + """Resolve the projection registry entry to its dedicated CUDA implementation.""" + from rl_engine.kernels.registry import KernelRegistry _cuda_op() diff --git a/tests/h3/test_h3_cli.py b/tests/h3/test_h3_cli.py index 7712fc41a..7288d795a 100644 --- a/tests/h3/test_h3_cli.py +++ b/tests/h3/test_h3_cli.py @@ -20,11 +20,15 @@ @pytest.fixture def cpu_chain_cli(monkeypatch): + """Replace CUDA chain stages with CPU stubs that track dependencies and provider drift.""" + executed = [] state = {"provider_shift": 0} class CpuOp: def __call__(self, value): + """Advance the synthetic chain value by one for candidate and golden replay.""" + return value + 1 forward_fp32 = __call__ @@ -33,6 +37,8 @@ def __call__(self, value): for index, stage in enumerate(h3_chain.STAGES): def stage_input(ctx, upstream, index=index): + """Use raw timesteps for stage zero and require predecessor outputs thereafter.""" + if index == 0: return ctx["timestep"] assert isinstance(upstream, torch.Tensor), "a preceding stage must run first" @@ -50,6 +56,8 @@ def stage_input(ctx, upstream, index=index): ) def get_op(op_type, *, device): + """Record each dispatched stage and return its deterministic CPU stub.""" + executed.append(op_type) return CpuOp() @@ -79,6 +87,8 @@ def get_op(op_type, *, device): def test_chain_selection_runs_dependencies( cpu_chain_cli, monkeypatch, tmp_path, capsys, requested, execution_count, provider_shift ): + """Execute prerequisite stages but report requested stages and preserve first drift.""" + executed, state = cpu_chain_cli state["provider_shift"] = provider_shift out = tmp_path / "chain.json" @@ -118,6 +128,8 @@ def test_chain_selection_runs_dependencies( @pytest.mark.parametrize("op", ["timestep_mlp_fp32", "adaln_projection_3mod"]) @pytest.mark.parametrize("weights_env", [None, " "]) def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, weights_env): + """Reject weighted evidence without a configured checkpoint before writing a report.""" + if weights_env is None: monkeypatch.delenv(WEIGHTS_ENV, raising=False) else: @@ -138,6 +150,8 @@ def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, we ], ) def test_evidence_records_weight_source(monkeypatch, tmp_path, op, weight_source): + """Label weighted evidence as pinned and sinusoid evidence as weight-independent.""" + if weight_source == "not_applicable": monkeypatch.delenv(WEIGHTS_ENV, raising=False) else: diff --git a/tests/h3/test_h3_conditioning_e2e.py b/tests/h3/test_h3_conditioning_e2e.py index 2bd1c47bd..a24ceebb9 100644 --- a/tests/h3/test_h3_conditioning_e2e.py +++ b/tests/h3/test_h3_conditioning_e2e.py @@ -33,16 +33,22 @@ @pytest.fixture(scope="module") def chain_weights(h3_weights_cpu): + """Move all pinned conditioning tensors to CUDA once for the end-to-end module.""" + return {name: tensor.cuda() for name, tensor in h3_weights_cpu.items()} @pytest.fixture(scope="module") def registry(): + """Provide the production registry used to dispatch every conditioning stage.""" + return KernelRegistry() @pytest.mark.parametrize("num_timesteps, seq_len", CASES) def test_chain_forward(registry, chain_weights, num_timesteps, seq_len): + """Check CUDA dispatch, golden tolerance, repeatability, and declared provider parity.""" + report = h3_chain.run_case( registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len ) diff --git a/tests/h3/test_h3_det_linear.py b/tests/h3/test_h3_det_linear.py index f8f90e83c..63ae03a8e 100644 --- a/tests/h3/test_h3_det_linear.py +++ b/tests/h3/test_h3_det_linear.py @@ -19,6 +19,8 @@ @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_backward_input_rejects_zero_rows(dtype): + """Reject empty gradient batches in native input-gradient kernels for both dtypes.""" + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) weight = torch.empty(16, 8, device="cuda", dtype=dtype) with pytest.raises(RuntimeError, match="grad must have at least one row"): @@ -28,6 +30,8 @@ def test_backward_input_rejects_zero_rows(dtype): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @pytest.mark.parametrize("with_bias", [False, True]) def test_backward_weight_rejects_zero_rows(dtype, with_bias): + """Reject empty gradient batches before parameter-gradient or bias-gradient reduction.""" + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) x = torch.empty(0, 8, device="cuda", dtype=dtype) with pytest.raises(RuntimeError, match="grad must have at least one row"): diff --git a/tests/h3/test_h3_report.py b/tests/h3/test_h3_report.py index eab1c8807..be62daa38 100644 --- a/tests/h3/test_h3_report.py +++ b/tests/h3/test_h3_report.py @@ -11,28 +11,42 @@ def test_measure_interleaves_backends_and_records_samples(monkeypatch): + """Balance candidate/provider timing order and preserve samples and derived bandwidth.""" + elapsed_us = 0.0 calls = [] class Event: def __init__(self, *, enable_timing): + """Require timing-enabled events in the deterministic timing stub.""" + assert enable_timing def record(self): + """Capture the synthetic elapsed clock at this event.""" + self.timestamp = elapsed_us def synchronize(self): + """Leave synchronization inert because the timing stub has no pending GPU work.""" + pass def elapsed_time(self, end): + """Return elapsed synthetic event time in CUDA milliseconds.""" + return (end.timestamp - self.timestamp) / 1e3 def candidate(): + """Record a candidate call and advance the synthetic clock by two microseconds.""" + nonlocal elapsed_us calls.append("candidate") elapsed_us += 2.0 def provider(): + """Record a provider call and advance the synthetic clock by four microseconds.""" + nonlocal elapsed_us calls.append("provider") elapsed_us += 4.0 diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py index 71f1e694f..f378a3b35 100644 --- a/tests/h3/test_h3_timestep_mlp.py +++ b/tests/h3/test_h3_timestep_mlp.py @@ -26,6 +26,8 @@ def _cuda_op(): + """Construct the CUDA MLP operator, skipping missing native linear kernels.""" + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available from rl_engine.kernels.ops.cuda.h3.timestep_mlp import H3TimestepMLPCudaOp @@ -35,10 +37,14 @@ def _cuda_op(): def _features(num, device="cuda", seed=0): + """Generate seeded H3 sinusoidal features on the requested test device.""" + return NativeH3TimestepSinusoidOp().forward(h3_timesteps(num, seed=seed, device=device)) def _synthetic_params(k_in, hidden, out, device="cuda", seed=0): + """Create seeded FP32 parameters for a two-layer SiLU MLP of the requested sizes.""" + g = torch.Generator(device="cpu").manual_seed(seed) w1 = torch.randn(hidden, k_in, generator=g) / k_in**0.5 b1 = torch.randn(hidden, generator=g) * 0.1 @@ -48,12 +54,16 @@ def _synthetic_params(k_in, hidden, out, device="cuda", seed=0): def _real_params(weights): + """Return pinned timestep MLP weights and biases on CUDA in layer order.""" + return [ weights[f"time_embedder.linear_{i}.{p}"].cuda() for i in (1, 2) for p in ("weight", "bias") ] def _fp64_grads(x, params, grad_out): + """Compute gradients of both linear layers and SiLU using FP64 autograd.""" + leaves = [t.detach().double().requires_grad_(True) for t in (x, *params)] z = F.linear(leaves[0], leaves[1], leaves[2]) F.linear(z * torch.sigmoid(z), leaves[3], leaves[4]).backward(grad_out.double()) @@ -62,6 +72,8 @@ def _fp64_grads(x, params, grad_out): class TestReference: def test_provider_and_golden_agree(self): + """Keep the FP32 provider close to the high-precision MLP replay on CPU.""" + x = _features(5, device="cpu") params = _synthetic_params(256, 64, 32, device="cpu") op = NativeH3TimestepMLPOp() @@ -72,12 +84,16 @@ def test_provider_and_golden_agree(self): @pytest.mark.parametrize("which", range(5)) def test_rejects_bf16_anywhere(self, which): # RFC probe H7: casting a declared FP32 path to BF16 early. + """Reject BF16 in every argument of the declared FP32 MLP path.""" + args = [_features(2, device="cpu"), *_synthetic_params(256, 64, 32, device="cpu")] args[which] = args[which].bfloat16() with pytest.raises(TypeError, match="float32"): NativeH3TimestepMLPOp().forward(*args) def test_rejects_shape_mismatches(self): + """Reject incompatible layer dimensions, malformed biases, and empty feature batches.""" + x = _features(2, device="cpu") w1, b1, w2, b2 = _synthetic_params(256, 64, 32, device="cpu") op = NativeH3TimestepMLPOp() @@ -95,6 +111,8 @@ def test_rejects_shape_mismatches(self): class TestCudaRealWeights: @pytest.mark.parametrize("num", [1, 2, 3, 4, 7]) def test_forward_against_fp64_golden(self, h3_weights_cpu, num): + """Check checkpoint forward accuracy, FP32 dtype, and the timestep output shape.""" + params = _real_params(h3_weights_cpu) x = _features(num, seed=num) out = _cuda_op()(x, *params) @@ -105,6 +123,8 @@ def test_forward_against_fp64_golden(self, h3_weights_cpu, num): assert (out - gold).abs().max().item() < 1e-5 def test_rows_are_batch_and_position_invariant(self, h3_weights_cpu): + """Keep each checkpoint embedding row bitwise unchanged across batch sizes and positions.""" + op = _cuda_op() params = _real_params(h3_weights_cpu) x_all = _features(9, seed=1) @@ -118,6 +138,8 @@ def test_rows_are_batch_and_position_invariant(self, h3_weights_cpu): assert torch.equal(op(x, *params)[pos], full[4]), (num, pos) def test_backward_against_fp64_golden(self, h3_weights_cpu): + """Match all five checkpoint input gradients to the FP64 autograd golden.""" + params = _real_params(h3_weights_cpu) x = _features(3, seed=2) leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] @@ -134,6 +156,8 @@ def test_backward_against_fp64_golden(self, h3_weights_cpu): class TestCudaSynthetic: @pytest.mark.parametrize("k_in, hidden, out", [(8, 24, 16), (12, 40, 20), (256, 5376, 2688)]) def test_other_sizes(self, k_in, hidden, out): + """Match the high-precision forward golden for small and checkpoint-sized MLPs.""" + params = _synthetic_params(k_in, hidden, out, seed=k_in) x = torch.randn(5, k_in, device="cuda") torch.testing.assert_close( @@ -144,11 +168,15 @@ def test_other_sizes(self, k_in, hidden, out): ) def test_rejects_unaligned_k(self): + """Reject inner dimensions that violate the native FP32 four-element alignment.""" + params = _synthetic_params(6, 8, 8) with pytest.raises(RuntimeError, match="multiple of 4"): _cuda_op()(torch.randn(2, 6, device="cuda"), *params) def test_rejects_cpu_and_mixed_devices(self): + """Reject CPU-only and mixed-device inputs at the CUDA MLP boundary.""" + params = _synthetic_params(8, 24, 16) with pytest.raises(ValueError): _cuda_op()(torch.randn(2, 8), *[p.cpu() for p in params]) @@ -156,6 +184,8 @@ def test_rejects_cpu_and_mixed_devices(self): _cuda_op()(torch.randn(2, 8), *params) def test_repeat_and_backward_are_deterministic(self): + """Require repeated MLP outputs and all input gradients to be bitwise equal.""" + op = _cuda_op() params = _synthetic_params(256, 5376, 2688, seed=3) x = _features(4, seed=3) @@ -171,12 +201,16 @@ def test_repeat_and_backward_are_deterministic(self): assert torch.equal(a, b) def test_dx_rows_are_batch_invariant(self): + """Keep feature gradients bitwise unchanged when rows are evaluated independently.""" + op = _cuda_op() params = _synthetic_params(256, 5376, 2688, seed=4) x = _features(6, seed=4) grad_out = torch.randn(6, 2688, device="cuda") def dx(rows): + """Compute selected feature-row gradients using matching upstream rows.""" + leaf = x[rows].detach().clone().requires_grad_(True) op(leaf, *params).backward(grad_out[rows]) return leaf.grad @@ -186,6 +220,8 @@ def dx(rows): assert torch.equal(dx(slice(i, i + 1))[0], full[i]) def test_registry_dispatches_cuda(self): + """Resolve the timestep MLP registry entry to its dedicated CUDA implementation.""" + from rl_engine.kernels.registry import KernelRegistry _cuda_op() @@ -198,6 +234,8 @@ def test_registry_dispatches_cuda(self): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) class TestCudaDetLinearDeviceContract: def test_forward_rejects_bias_on_another_device(self, dtype): + """Reject a projection bias on a different CUDA device from its weight.""" + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_forward _cuda_op() @@ -208,6 +246,8 @@ def test_forward_rejects_bias_on_another_device(self, dtype): det_linear_forward(x, weight, bias) def test_backward_input_rejects_weight_on_another_device(self, dtype): + """Reject native input-gradient operands on different CUDA devices.""" + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_backward_input _cuda_op() @@ -217,6 +257,8 @@ def test_backward_input_rejects_weight_on_another_device(self, dtype): det_linear_backward_input(grad, weight, dtype) def test_backward_weight_rejects_input_on_another_device(self, dtype): + """Reject native parameter-gradient operands on different CUDA devices.""" + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_backward_weight _cuda_op() diff --git a/tests/h3/test_h3_timestep_sinusoid.py b/tests/h3/test_h3_timestep_sinusoid.py index bcd567dc3..2f9747a00 100644 --- a/tests/h3/test_h3_timestep_sinusoid.py +++ b/tests/h3/test_h3_timestep_sinusoid.py @@ -28,6 +28,8 @@ def _cuda_op(): + """Construct the CUDA sinusoid operator, skipping a missing native forward entrypoint.""" + from rl_engine.kernels.ops.cuda.h3.timestep_sinusoid import ( H3TimestepSinusoidCudaOp, h3_sinusoid_cuda_available, @@ -40,6 +42,8 @@ def _cuda_op(): class TestReference: def test_layout_is_cos_then_sin(self): + """Check FP32 cosine-then-sine layout and the unit-frequency boundary values.""" + t = torch.tensor([0.0, 1.0]) out = NativeH3TimestepSinusoidOp().forward(t) assert out.shape == (2, 256) and out.dtype == torch.float32 @@ -50,11 +54,15 @@ def test_layout_is_cos_then_sin(self): assert out[1, 128].item() == pytest.approx(math.sin(1.0), abs=1e-7) def test_provider_path_matches_diffusers_cpu(self): + """Require CPU sinusoid values to match the diffusers provider bitwise.""" + t = h3_timesteps(9, device="cpu") ours = NativeH3TimestepSinusoidOp().forward(t) assert torch.equal(ours, provider_time_proj(t, 256)) def test_golden_close_to_provider_path(self): + """Keep the CPU provider within FP32 error of the high-precision sinusoid golden.""" + t = h3_timesteps(33, device="cpu") op = NativeH3TimestepSinusoidOp() torch.testing.assert_close(op.forward(t), op.forward_fp32(t), atol=1e-6, rtol=0) @@ -71,10 +79,14 @@ def test_golden_close_to_provider_path(self): ], ) def test_rejects_invalid_timesteps(self, bad, error): + """Reject empty, malformed, non-floating, non-finite, or out-of-range timesteps.""" + with pytest.raises(error): NativeH3TimestepSinusoidOp().forward(bad) def test_rejects_odd_channels(self): + """Reject odd feature counts that cannot split into equal cosine and sine halves.""" + with pytest.raises(ValueError): NativeH3TimestepSinusoidOp().forward(torch.tensor([0.5]), num_channels=255) @@ -83,6 +95,8 @@ def test_rejects_odd_channels(self): class TestCuda: @pytest.mark.parametrize("num", [1, 2, 3, 7, 64, 1000, 4097]) def test_bitwise_equal_to_diffusers_cuda_path(self, num): + """Require CUDA sinusoid values to match both provider replays for varied batch sizes.""" + t = h3_timesteps(num, seed=num) out = _cuda_op()(t) assert out.dtype == torch.float32 and out.shape == (num, 256) @@ -91,12 +105,16 @@ def test_bitwise_equal_to_diffusers_cuda_path(self, num): @pytest.mark.parametrize("num_channels", [2, 8, 96, 256, 320]) def test_other_channel_counts_match_reference(self, num_channels): + """Preserve bitwise reference parity across supported even channel counts.""" + t = h3_timesteps(5) out = _cuda_op()(t, num_channels=num_channels) ref = NativeH3TimestepSinusoidOp().forward(t, num_channels=num_channels) assert torch.equal(out, ref) def test_fp32_contract_against_fp64_golden(self): + """Bound CUDA sinusoid error against the high-precision golden by FP32 precision.""" + t = h3_timesteps(257) out = _cuda_op()(t) gold = NativeH3TimestepSinusoidOp().forward_fp32(t) @@ -106,10 +124,14 @@ def test_fp32_contract_against_fp64_golden(self): @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) def test_low_precision_timesteps_are_upcast_first(self, dtype): + """Match provider FP32 upcasting for BF16 and FP16 timesteps.""" + t = h3_timesteps(6).to(dtype) assert torch.equal(_cuda_op()(t), provider_time_proj(t, 256)) def test_batch_position_and_repeat_invariance(self): + """Keep row bits unchanged across batch positions, repeated calls, and permutations.""" + op = _cuda_op() base = torch.tensor(0.3141592, device="cuda") single = op(base.view(1)) @@ -126,21 +148,29 @@ def test_batch_position_and_repeat_invariance(self): assert torch.equal(op(t[perm]), first[perm]) def test_non_contiguous_input(self): + """Treat strided timesteps identically to their contiguous copies.""" + t = h3_timesteps(20) strided = t[::2] assert not strided.is_contiguous() assert torch.equal(_cuda_op()(strided), _cuda_op()(strided.contiguous())) def test_rejects_cpu_input(self): + """Reject CPU timesteps at the CUDA sinusoid boundary.""" + with pytest.raises(ValueError): _cuda_op()(torch.tensor([0.5])) def test_rejects_out_of_range(self): + """Reject scaled timesteps outside the declared unit interval.""" + with pytest.raises(ValueError): _cuda_op()(torch.tensor([0.0, 999.0], device="cuda")) @pytest.mark.parametrize("bad", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]) def test_native_entrypoint_rejects_invalid_values(self, bad): + """Reject invalid values in native and wrapper calls even when neighbors are valid.""" + from rl_engine.kernels.ops.base import _C _cuda_op() @@ -152,6 +182,8 @@ def test_native_entrypoint_rejects_invalid_values(self, bad): _cuda_op()(t) def test_native_entrypoint_accepts_boundaries(self): + """Accept both interval endpoints and match provider values in the native call.""" + from rl_engine.kernels.ops.base import _C _cuda_op() @@ -159,6 +191,8 @@ def test_native_entrypoint_accepts_boundaries(self): assert torch.equal(_C.h3_timestep_sinusoid_forward(t), provider_time_proj(t, 256)) def test_trusted_path_supports_cuda_graph(self): + """Capture and replay the prevalidated sinusoid path without changing its output.""" + op = _cuda_op() t = h3_timesteps(7) expected = op(t) @@ -174,6 +208,8 @@ def test_trusted_path_supports_cuda_graph(self): assert torch.equal(out, expected) def test_backward_matches_autograd_reference(self): + """Match timestep gradients to the high-precision autograd sinusoid replay.""" + t = h3_timesteps(17).requires_grad_(True) t_ref = t.detach().clone().requires_grad_(True) grad = torch.randn(17, 256, device="cuda", generator=torch.Generator("cuda").manual_seed(1)) @@ -182,10 +218,14 @@ def test_backward_matches_autograd_reference(self): torch.testing.assert_close(t.grad, t_ref.grad, atol=1e-4, rtol=1e-4) def test_backward_is_row_invariant(self): + """Keep a timestep gradient bitwise unchanged when unrelated rows join the batch.""" + op = _cuda_op() grad_row = torch.randn(256, device="cuda", generator=torch.Generator("cuda").manual_seed(2)) def grad_of(t, pos): + """Differentiate one selected timestep using a fixed upstream feature-gradient row.""" + t = t.clone().requires_grad_(True) grad = torch.zeros(t.shape[0], 256, device="cuda") grad[pos] = grad_row @@ -198,6 +238,8 @@ def grad_of(t, pos): assert torch.equal(grad_of(t, 13), single) def test_registry_dispatches_cuda(self): + """Resolve the sinusoid registry entry to its dedicated CUDA implementation.""" + from rl_engine.kernels.registry import KernelRegistry _cuda_op() diff --git a/tests/h3/test_h3_weights.py b/tests/h3/test_h3_weights.py index 015f78954..c081234e2 100644 --- a/tests/h3/test_h3_weights.py +++ b/tests/h3/test_h3_weights.py @@ -21,6 +21,8 @@ ], ) def test_tensor_checksum_hashes_raw_bytes(dtype, raw_bytes): + """Hash the exact FP32 and BF16 bit patterns without numerical conversion.""" + tensor = torch.tensor([0.0, 1.0], dtype=dtype) assert h3_weights.sha256_tensor(tensor) == hashlib.sha256(raw_bytes).hexdigest() @@ -29,6 +31,8 @@ def test_tensor_checksum_hashes_raw_bytes(dtype, raw_bytes): @pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) @pytest.mark.parametrize("names", [None, ["bias"]]) def test_loading_rejects_modified_weight_contents(tmp_path, monkeypatch, dtype, mismatch, names): + """Reject altered shape, dtype, or bytes even outside the requested tensor subset.""" + from safetensors.torch import save_file tensors = { diff --git a/tests/h3/test_prepare_h3_weights.py b/tests/h3/test_prepare_h3_weights.py index 8a1c94298..1405e044e 100644 --- a/tests/h3/test_prepare_h3_weights.py +++ b/tests/h3/test_prepare_h3_weights.py @@ -18,6 +18,8 @@ @pytest.mark.parametrize("optimize", [0, 1]) @pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, mismatch): + """Enforce tensor contracts before extraction writes, including under Python optimization.""" + import huggingface_hub from safetensors.torch import load_file, save_file diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index fdef637f7..8e77797ce 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -33,6 +33,8 @@ def _run(script: str, *args: str, timeout: int = 300) -> None: def test_all_ws1_single_ops_are_registered(): + """Require gtest specs for all WS1 single operators and the H3 conditioning stages.""" + names = set(operator_names()) assert { "rms_norm", From 59476c9788984e2a257e639ed69d3d75d95690ee Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:32:23 +0000 Subject: [PATCH 06/24] test(h3): enforce complete tensor pairings in regression checks Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- scripts/plot_h3_evidence.py | 4 ++-- tests/h3/test_h3_adaln_projection.py | 15 ++++++++------- tests/h3/test_h3_timestep_mlp.py | 4 ++-- 3 files changed, 12 insertions(+), 11 deletions(-) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 7ebd228d0..988f68c7d 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -181,7 +181,7 @@ def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], con """Per-draw max |error| as boxes with every draw overlaid, log scale.""" _style(ax, title, "max abs error per draw") - for i, (name, color, values) in enumerate(series): + for i, (_name, color, values) in enumerate(series): ax.boxplot( [values], positions=[i], @@ -251,7 +251,7 @@ def plot_projection(report: dict[str, Any]): ("RL-Kernel CUDA", RL_KERNEL, [100 * v for v in acc["cuda_correctly_rounded"]]), ("diffusers (cuBLAS)", PROVIDER, [100 * v for v in acc["provider_correctly_rounded"]]), ] - for i, (name, color, values) in enumerate(series): + for i, (_name, color, values) in enumerate(series): jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] ax_vals = sorted(values) median = ax_vals[len(ax_vals) // 2] diff --git a/tests/h3/test_h3_adaln_projection.py b/tests/h3/test_h3_adaln_projection.py index 0b031fa4e..d052651a3 100644 --- a/tests/h3/test_h3_adaln_projection.py +++ b/tests/h3/test_h3_adaln_projection.py @@ -71,7 +71,7 @@ def test_provider_layout_matches_diffusers(self): ours = NativeH3AdaLNProjectionOp().forward(temb, weight, bias) theirs = provider_adaln_modulation(temb, weight, bias, hidden_size=8) assert len(ours) == 6 - for a, b in zip(ours, theirs): + for a, b in zip(ours, theirs, strict=True): assert a.shape == (6, 8) and torch.equal(a, b) def test_rejects_bf16_temb(self): @@ -144,7 +144,7 @@ def test_forward_is_correctly_rounded_declared_golden(self, block0, num): temb = _temb(num, seed=num) outputs = _cuda_op()(temb, weight, bias) golden = NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias) - for out, gold in zip(outputs, golden): + for out, gold in zip(outputs, golden, strict=True): assert out.dtype == torch.bfloat16 and out.shape == (3 * num, 5376) # tolerance_contract.json forward_accuracy / reduction / bfloat16. torch.testing.assert_close(out.float(), gold, atol=5e-2, rtol=2e-2) @@ -174,10 +174,10 @@ def test_rows_are_batch_and_position_invariant(self, block0): full = op(temb, weight, bias) for i in range(5): single = op(temb[i : i + 1], weight, bias) - for s, f in zip(single, full): + for s, f in zip(single, full, strict=True): assert torch.equal(s, f[3 * i : 3 * i + 3]) swapped = op(temb.flip(0), weight, bias) - for s, f in zip(swapped, full): + for s, f in zip(swapped, full, strict=True): assert torch.equal(s.view(5, 3, -1).flip(0), f.view(5, 3, -1)) def test_backward_against_golden(self, block0): @@ -198,7 +198,7 @@ def test_backward_against_golden(self, block0): # tolerance_contract.json gradient_accuracy / reduction: FP32 temb grad # and BF16 weight/bias grads (rounded once from an FP32 accumulation). torch.testing.assert_close(leaves[0].grad.double(), ref[0].grad, atol=1e-4, rtol=1e-4) - for leaf, r in zip(leaves[1:], ref[1:]): + for leaf, r in zip(leaves[1:], ref[1:], strict=True): torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) assert (leaf.grad == r.grad.to(torch.bfloat16)).float().mean() > 0.999 @@ -216,6 +216,7 @@ def test_other_sizes(self, dtype, hidden, dim): for out, gold in zip( _cuda_op()(temb, weight, bias), NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias), + strict=True, ): torch.testing.assert_close(out.float(), gold, **tol) @@ -232,7 +233,7 @@ def test_repeat_forward_backward_bitwise(self): torch.autograd.backward(list(outs), grads) runs.append([*[o.detach() for o in outs], *[leaf.grad for leaf in leaves]]) for later in runs[1:]: - for a, b in zip(runs[0], later): + for a, b in zip(runs[0], later, strict=True): assert torch.equal(a, b) def test_temb_grad_rows_are_batch_invariant(self): @@ -265,7 +266,7 @@ def test_rows_invariant_across_tensor_core_row_tiles(self, num): full = _cuda_op()(temb, weight, bias) for i in (0, 7, num - 1): single = _cuda_op()(temb[i : i + 1], weight, bias) - for s_out, f_out in zip(single, full): + for s_out, f_out in zip(single, full, strict=True): assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) def test_rejects_cpu(self): diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py index f378a3b35..c61d4faac 100644 --- a/tests/h3/test_h3_timestep_mlp.py +++ b/tests/h3/test_h3_timestep_mlp.py @@ -147,7 +147,7 @@ def test_backward_against_fp64_golden(self, h3_weights_cpu): 3, 2688, device="cuda", generator=torch.Generator("cuda").manual_seed(0) ) _cuda_op()(*leaves).backward(grad_out) - for name, leaf, ref in zip(NAMES, leaves, _fp64_grads(x, params, grad_out)): + for name, leaf, ref in zip(NAMES, leaves, _fp64_grads(x, params, grad_out), strict=True): # tolerance_contract.json gradient_accuracy / reduction / float32. torch.testing.assert_close(leaf.grad.double(), ref, atol=1e-4, rtol=1e-4, msg=name) @@ -197,7 +197,7 @@ def test_repeat_and_backward_are_deterministic(self): out.backward(grad_out) runs.append([out.detach(), *[leaf.grad for leaf in leaves]]) for later in runs[1:]: - for a, b in zip(runs[0], later): + for a, b in zip(runs[0], later, strict=True): assert torch.equal(a, b) def test_dx_rows_are_batch_invariant(self): From 5f5cf945b42556653474e1070a8389b5b2450b03 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:44:51 +0000 Subject: [PATCH 07/24] fix(h3): preserve gradients for empty linear dimensions Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 8 ++++- .../ops/pytorch/h3/adaln_projection.py | 2 +- tests/h3/test_h3_det_linear.py | 30 +++++++++++++++++++ tests/h3/test_h3_timestep_mlp.py | 24 +++++++++++++++ 4 files changed, 62 insertions(+), 2 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 246d27d60..636a00414 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -504,9 +504,11 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei const int64_t n_out = weight.size(0); const int64_t k_in = weight.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); + auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); + if (k_in == 0) return out; + if (n_out == 0) return out.zero_(); const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; auto partial = torch::empty({chunks, rows, k_in}, grad.options()); - auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); auto stream = at::cuda::getCurrentCUDAStream(); const int threads = 256; dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); @@ -551,6 +553,10 @@ std::vector h3_det_linear_backward_weight(torch::Tensor grad, tor const int64_t k_in = x.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); auto dw = torch::empty({n_out, k_in}, grad.options().dtype(w_dtype)); + if (n_out == 0) { + if (!with_bias) return {dw}; + return {dw, torch::empty({n_out}, grad.options().dtype(w_dtype))}; + } auto stream = at::cuda::getCurrentCUDAStream(); const int threads = 256; const int64_t total = n_out * k_in; diff --git a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py index b6d013961..e9cac5ff0 100644 --- a/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py +++ b/rl_engine/kernels/ops/pytorch/h3/adaln_projection.py @@ -103,7 +103,7 @@ def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: return split_adaln_table(table, hidden) def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: - """Return six FP32 golden views from FP64 SiLU and projection arithmetic. + """Return six FP32 golden outputs from FP64 SiLU and projection arithmetic. Preserve the declared activation rounding to the weight dtype while using an identity VJP at that cast boundary for the FP64 golden graph. diff --git a/tests/h3/test_h3_det_linear.py b/tests/h3/test_h3_det_linear.py index 63ae03a8e..a84184811 100644 --- a/tests/h3/test_h3_det_linear.py +++ b/tests/h3/test_h3_det_linear.py @@ -36,3 +36,33 @@ def test_backward_weight_rejects_zero_rows(dtype, with_bias): x = torch.empty(0, 8, device="cuda", dtype=dtype) with pytest.raises(RuntimeError, match="grad must have at least one row"): _C.h3_det_linear_backward_weight(grad, x, dtype, with_bias) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +def test_backward_input_handles_empty_columns(dtype, out_dtype, k_in, n_out): + """Match the input VJP of empty linear dimensions in the requested output dtype.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + weight = torch.empty(n_out, k_in, device="cuda", dtype=dtype) + expected = (grad @ weight.float()).to(out_dtype) + actual = _C.h3_det_linear_backward_input(grad, weight, out_dtype) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_handles_empty_columns(dtype, w_dtype, k_in, n_out, with_bias): + """Return empty weight gradients while preserving any nonempty bias reduction.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + x = torch.arange(3 * k_in, device="cuda", dtype=dtype).reshape(3, k_in) + expected_weight = (grad.T @ x.float()).to(w_dtype) + actual = _C.h3_det_linear_backward_weight(grad, x, w_dtype, with_bias) + assert len(actual) == (2 if with_bias else 1) + torch.testing.assert_close(actual[0], expected_weight, atol=0, rtol=0) + if with_bias: + torch.testing.assert_close(actual[1], grad.sum(dim=0).to(w_dtype), atol=0, rtol=0) diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py index c61d4faac..fa74344ab 100644 --- a/tests/h3/test_h3_timestep_mlp.py +++ b/tests/h3/test_h3_timestep_mlp.py @@ -167,6 +167,30 @@ def test_other_sizes(self, k_in, hidden, out): rtol=1e-4, ) + def test_zero_hidden_size_matches_cpu_forward_backward(self): + """Match CPU forward values and all input gradients with a zero-width hidden layer.""" + + generator = torch.Generator(device="cpu").manual_seed(0) + cpu_inputs = [ + torch.randn(3, 8, generator=generator), + torch.empty(0, 8), + torch.empty(0), + torch.empty(16, 0), + torch.linspace(-1.0, 1.0, 16), + ] + cpu_inputs = [tensor.requires_grad_(True) for tensor in cpu_inputs] + cuda_inputs = [tensor.detach().cuda().requires_grad_(True) for tensor in cpu_inputs] + cpu_out = NativeH3TimestepMLPOp()(*cpu_inputs) + cuda_out = _cuda_op()(*cuda_inputs) + torch.testing.assert_close(cuda_out.cpu(), cpu_out, atol=0, rtol=0) + grad = torch.randn(3, 16, generator=generator) + cpu_out.backward(grad) + cuda_out.backward(grad.cuda()) + for name, cpu_input, cuda_input in zip(NAMES, cpu_inputs, cuda_inputs, strict=True): + torch.testing.assert_close( + cuda_input.grad.cpu(), cpu_input.grad, atol=0, rtol=0, msg=name + ) + def test_rejects_unaligned_k(self): """Reject inner dimensions that violate the native FP32 four-element alignment.""" From 28cba8702a7e5219ee6360aef899e8b578a2ab6a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:30:08 +0000 Subject: [PATCH 08/24] docs(h3): adaln_projection vs existing implementations (B200, c88691c) 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-projection.md | 29 +++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 0e27a5761..910165a3b 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -131,6 +131,35 @@ written by `scripts/h3_evidence.py` from a clean tree at commit `d06e120`. The r records that a timestep's 18 modulation rows are bitwise identical whether it runs alone or in a batch of 9. +## Existing implementations (RFC #420 reuse rule) + +![adaln_projection vs existing implementations](../usage/evidence/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 | +|---|---|---|---|---| +| diffusers AdaLN projection, BF16 F.linear [plain] | **no** (14214 rows, 3080 sub-batches, 2232 sweep cases) | 3.6e-03 / 3.9e-03 / 716 µs | 3.4e-03 / 3.2e-03 / 728 µs | 3.2e-03 / 3.5e-03 / 2346 µs | +| rl-kernel H3AdaLNProjectionCudaOp | yes | 3.6e-03 / 3.9e-03 / 2276 µs | 3.4e-03 / 3.2e-03 / 77446 µs | 3.2e-03 / 3.5e-03 / 626992 µs | +| diffusers AdaLN projection, BF16 F.linear [vllm] | yes | 3.6e-03 / 3.9e-03 / 681 µs | 3.4e-03 / 3.2e-03 / 722 µs | 3.2e-03 / 3.5e-03 / 2330 µs | +| diffusers AdaLN projection, BF16 F.linear [sglang] | yes | 3.6e-03 / 3.9e-03 / 3096 µs | 3.4e-03 / 3.2e-03 / 3690 µs | 3.2e-03 / 3.5e-03 / 14296 µs | +| diffusers AdaLN projection, BF16 F.linear [sglang_ieee] | yes | 3.6e-03 / 3.9e-03 / 3138 µs | 3.4e-03 / 3.2e-03 / 3740 µs | 3.2e-03 / 3.5e-03 / 14325 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_te_native] | **no** (7107 rows, 2056 sub-batches, 1104 sweep cases) | 3.6e-03 / 3.9e-03 / 702 µs | 3.4e-03 / 3.2e-03 / 706 µs | 3.2e-03 / 3.5e-03 / 2330 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_triton] | yes | 3.6e-03 / 3.9e-03 / 2532 µs | 3.4e-03 / 3.2e-03 / 3204 µs | 3.2e-03 / 3.5e-03 / 13771 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_triton_ieee] | yes | 3.6e-03 / 3.9e-03 / 2533 µs | 3.4e-03 / 3.2e-03 / 3206 µs | 3.2e-03 / 3.5e-03 / 13773 µ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 +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`](../usage/evidence/h3-prior-art-b200/adaln_projection.json) +was written from a clean tree at `c88691c` by + +```bash +python scripts/h3_prior_art.py --op adaln_projection --out docs/usage/evidence/h3-prior-art-b200/adaln_projection.json --megatron-src +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/adaln_projection.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From c4f6145ffa65a114b6a9e2b2977954626eeb1ef9 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:46:47 +0000 Subject: [PATCH 09/24] evidence(h3): adaln_projection direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 0e476483. 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%. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index e93cdee6b..92dca1376 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -160,6 +160,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![adaln_projection bitwise comparison](../../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 + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png +``` + ## Tests ```bash From 49fa04a0d191092c8dd1dda2fea3671f64423d3d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:46:51 +0000 Subject: [PATCH 10/24] test(h3): round the projection golden's activation from the FP32 SiLU test_backward_against_golden rounded an FP64 SiLU to BF16 for its declared cast, while the kernel and diffusers round the FP32 SiLU. Near a BF16 tie the two pick different neighbours and shift a whole dW column. The test was seeded (seed 5 scores 1.0), so it never failed, but over 300 temb draws 8 fell below 0.9999 (worst 0.999646). With the FP32-SiLU golden every dW/db element is correctly rounded for all 300 draws; d_temb is unchanged. Same fix as test_h3_adaln_modulation.py. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_adaln_projection.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 6f9323ed5..91732f6d2 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -194,7 +194,11 @@ def test_backward_against_golden(self, block0): ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] act = ref[0] * torch.sigmoid(ref[0]) - act = act + (act.to(torch.bfloat16).double() - act).detach() # declared cast, identity VJP + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column. + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # declared cast, identity VJP table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376).chunk(6, dim=-1) torch.autograd.backward(list(table), [gr.double() for gr in grads]) # tolerance_contract.json gradient_accuracy / reduction: FP32 temb grad From 55370d6c2aaa79ae477634932c1856f452f914cb Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:20:09 +0000 Subject: [PATCH 11/24] test(h3): require exact dW/db equality in the projection golden check With the FP32-SiLU golden (8bb2e092), the BF16 weight and bias gradients equal the correctly rounded golden at every element for all 500 temb draws tried (T = 2, two products per element). Replace the '> 0.999 of elements' threshold with torch.equal; the FP64 closeness check stays. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_adaln_projection.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 91732f6d2..5f0e7713e 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -206,7 +206,10 @@ def test_backward_against_golden(self, block0): torch.testing.assert_close(leaves[0].grad.double(), ref[0].grad, atol=1e-4, rtol=1e-4) for leaf, r in zip(leaves[1:], ref[1:], strict=True): torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) - assert (leaf.grad == r.grad.to(torch.bfloat16)).float().mean() > 0.999 + # With T = 2 each dW/db element sums two products, and the one BF16 cast + # of the FP32 sum is the correctly rounded value of every element + # (checked on 500 of 500 temb draws). + assert torch.equal(leaf.grad, r.grad.to(torch.bfloat16)) @requires_cuda From 0d433f6f46072ab73818ed40bcfb63b235015f3d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:36:32 +0000 Subject: [PATCH 12/24] test(h3): round the projection forward golden from the FP32 SiLU, require 0.9995 The two forward checks compared the kernel with forward_fp32, which rounds an FP64 SiLU to BF16 while the kernel and diffusers round the FP32 SiLU. Near a BF16 tie that shifts a whole output column: over 300 temb draws the old golden matched as little as 95.7% of the outputs, so the '> 0.999' checks only passed for their fixed seeds. A test-local _declared_golden now rounds F.silu(temb.float()). Exact equality is not reachable: the FP32 tensor-core accumulation over K = 2688 leaves 32-104 outputs off on every draw (ties, and a few ULPs where near-zero sums cancel). Over 300 draws per T = 1..4 the worst match is 0.999618-0.999731, so both checks now require >= 0.9995. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/test_h3_adaln_projection.py | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 5f0e7713e..29269b061 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -20,7 +20,10 @@ import torch.nn.functional as F from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM -from rl_engine.reference.minimax_h3.adaln_projection import NativeH3AdaLNProjectionOp +from rl_engine.reference.minimax_h3.adaln_projection import ( + NativeH3AdaLNProjectionOp, + split_adaln_table, +) from rl_engine.validation.models.h3_provider import provider_adaln_modulation requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") @@ -64,6 +67,19 @@ def _flat(outputs): return torch.cat([out.reshape(-1) for out in outputs]) +def _declared_golden(temb, weight, bias): + """FP64 projection of the declared BF16 activation, rounded from the FP32 SiLU. + + The kernel and diffusers round the FP32 SiLU; ``forward_fp32`` rounds an FP64 + SiLU, which can pick the other BF16 neighbour near a tie and shift a whole + output column (one input in 300 then matched only 95.7% of the outputs). + """ + + act = F.silu(temb.float()).to(weight.dtype).double() + table = F.linear(act, weight.double(), bias.double()) + return tuple(chunk.float() for chunk in split_adaln_table(table, weight.shape[0] // 18)) + + class TestReference: def test_provider_layout_matches_diffusers(self): """Match the provider values and timestep-major modality layout on CPU.""" @@ -151,9 +167,12 @@ def test_forward_is_correctly_rounded_declared_golden(self, block0, num): # tolerance_contract.json forward_accuracy / reduction / bfloat16. torch.testing.assert_close(out.float(), gold, atol=5e-2, rtol=2e-2) # Beyond the contract: almost every element is the correctly rounded - # golden value; the rest are 1-ULP ties of the FP32 accumulation. - match = (_flat(outputs) == _flat(golden).bfloat16()).float().mean().item() - assert match > 0.999 + # golden value. The rest come from the FP32 tensor-core accumulation over + # K = 2688 (ties, and a few ULPs where near-zero sums cancel), so exact + # equality is not reachable: over 300 temb draws per T the worst match was + # 0.999618 (T = 1), 0.999690, 0.999728 and 0.999731 (T = 4). + declared = _flat(_declared_golden(temb, weight, bias)).bfloat16() + assert (_flat(outputs) == declared).float().mean().item() >= 0.9995 def test_early_cast_is_numerically_detectable(self, block0): """Probe H7: a BF16 cast before the SiLU moves about half the outputs.""" @@ -162,9 +181,9 @@ def test_early_cast_is_numerically_detectable(self, block0): temb = _temb(3, seed=7) op = NativeH3AdaLNProjectionOp() ours = _flat(_cuda_op()(temb, weight, bias)) - declared = _flat(op.forward_fp32(temb, weight, bias)).bfloat16() + declared = _flat(_declared_golden(temb, weight, bias)).bfloat16() early = _flat(op.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() - assert (ours == declared).float().mean() > 0.999 + assert (ours == declared).float().mean() >= 0.9995 assert (ours == early).float().mean() < 0.8 def test_rows_are_batch_and_position_invariant(self, block0): From 579775f1306a3afb4103de10925db79a67cdc503 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:50:58 +0000 Subject: [PATCH 13/24] fix(h3): round the adaln_projection golden's activation from the FP32 SiLU NativeH3AdaLNProjectionOp.forward_fp32 rounded an FP64 SiLU to the weight dtype for the declared cast, while the model and the kernel round the FP32 SiLU. Near a BF16 tie the two pick different neighbours and shift a whole output column: over 300 temb draws the old golden matched the kernel's BF16 outputs as little as 95.7% of the time. The golden now rounds F.silu(temb.float()); the VJP stays the FP64 SiLU derivative (identity through the cast). The test-local golden added in 42943618 is no longer needed and the forward checks use forward_fp32 again. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../reference/minimax_h3/adaln_projection.py | 11 ++++++---- .../minimax_h3/test_h3_adaln_projection.py | 22 +++---------------- 2 files changed, 10 insertions(+), 23 deletions(-) diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py index 03c1ad683..be83239c5 100644 --- a/rl_engine/reference/minimax_h3/adaln_projection.py +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -112,9 +112,12 @@ def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: hidden = validate_h3_adaln_projection(temb, weight, bias) t64 = temb.double() act = t64 * torch.sigmoid(t64) - # The declared cast rounds the value; its VJP is the identity. Written - # straight-through so autograd does not round the FP64 gradient to the - # weight dtype on its way back (``.to(bf16)`` would). - act = act + (act.to(weight.dtype).double() - act).detach() + # The declared cast rounds the FP32 SiLU, as the model (and the kernel) + # computes it: an FP64 SiLU can round to the other weight-dtype neighbour + # near a tie and shift a whole output column. Its VJP is the identity, + # written straight-through so autograd does not round the FP64 gradient + # to the weight dtype on its way back (``.to(bf16)`` would). + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() table = F.linear(act, weight.double(), bias.double()) return tuple(chunk.float() for chunk in split_adaln_table(table, hidden)) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 29269b061..5fb10a73f 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -20,10 +20,7 @@ import torch.nn.functional as F from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM -from rl_engine.reference.minimax_h3.adaln_projection import ( - NativeH3AdaLNProjectionOp, - split_adaln_table, -) +from rl_engine.reference.minimax_h3.adaln_projection import NativeH3AdaLNProjectionOp from rl_engine.validation.models.h3_provider import provider_adaln_modulation requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") @@ -67,19 +64,6 @@ def _flat(outputs): return torch.cat([out.reshape(-1) for out in outputs]) -def _declared_golden(temb, weight, bias): - """FP64 projection of the declared BF16 activation, rounded from the FP32 SiLU. - - The kernel and diffusers round the FP32 SiLU; ``forward_fp32`` rounds an FP64 - SiLU, which can pick the other BF16 neighbour near a tie and shift a whole - output column (one input in 300 then matched only 95.7% of the outputs). - """ - - act = F.silu(temb.float()).to(weight.dtype).double() - table = F.linear(act, weight.double(), bias.double()) - return tuple(chunk.float() for chunk in split_adaln_table(table, weight.shape[0] // 18)) - - class TestReference: def test_provider_layout_matches_diffusers(self): """Match the provider values and timestep-major modality layout on CPU.""" @@ -171,7 +155,7 @@ def test_forward_is_correctly_rounded_declared_golden(self, block0, num): # K = 2688 (ties, and a few ULPs where near-zero sums cancel), so exact # equality is not reachable: over 300 temb draws per T the worst match was # 0.999618 (T = 1), 0.999690, 0.999728 and 0.999731 (T = 4). - declared = _flat(_declared_golden(temb, weight, bias)).bfloat16() + declared = _flat(golden).bfloat16() assert (_flat(outputs) == declared).float().mean().item() >= 0.9995 def test_early_cast_is_numerically_detectable(self, block0): @@ -181,7 +165,7 @@ def test_early_cast_is_numerically_detectable(self, block0): temb = _temb(3, seed=7) op = NativeH3AdaLNProjectionOp() ours = _flat(_cuda_op()(temb, weight, bias)) - declared = _flat(_declared_golden(temb, weight, bias)).bfloat16() + declared = _flat(op.forward_fp32(temb, weight, bias)).bfloat16() early = _flat(op.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() assert (ours == declared).float().mean() >= 0.9995 assert (ours == early).float().mean() < 0.8 From 068fcaa378ff05dc11b20f33b947912317079304 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:19:59 +0000 Subject: [PATCH 14/24] evidence(h3): regenerate adaln_projection_3mod report after the golden fix Written by tools/validation/models/h3_evidence.py from a clean tree at a0834552 on an idle B200. The accuracy fields are identical to the previous report (none of its 20 draws has an activation on a BF16 tie); the op latencies were re-measured: CUDA 100/100/101 us and provider 112/100/101 us for T = 1/2/4. The doc's table and evidence pointer follow. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 92dca1376..a76cc0417 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -108,14 +108,14 @@ B200, pinned BF16 weights (520 MB per block, streamed once per call): | T | CUDA op | provider (SiLU + cast + cuBLAS) | CUDA GEMV kernel | cuBLAS kernel | | --- | --- | --- | --- | --- | -| 1 | 105 µs | 115 µs | 78.3 µs (6.6 TB/s) | 93 µs | -| 2 | 105 µs | 103 µs | 78.4 µs | 79 µs | -| 4 | 106 µs | 104 µs | 79.0 µs | 79 µs | +| 1 | 100 µs | 112 µs | 78.3 µs (6.6 TB/s) | 93 µs | +| 2 | 100 µs | 100 µs | 78.4 µs | 79 µs | +| 4 | 101 µs | 101 µs | 79.0 µs | 79 µs | The kernel takes the same time for every T <= 8, because it always computes 8 rows. It -matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 24 µs of the op time is +matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 22 µs of the op time is the SiLU and cast kernels plus the Python wrapper. With those included, the op is within -about 2% of the provider for T >= 2. +about 1% of the provider for T >= 2. An earlier FMA-only BF16 kernel ran at 81, 98 and 166 µs for T = 1, 2 and 4. Its time grew with T because every extra row adds FMAs per weight element. Packed FP32 FMA @@ -127,7 +127,9 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `d06e120`. The report also +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 records that a timestep's 18 modulation rows are bitwise identical whether it runs alone or in a batch of 9. From 7f2fa858bcfe42e51130f20a95352f002a66c931 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:56:42 +0000 Subject: [PATCH 15/24] fix(h3): seed the adaln_projection timing and invariance inputs in the report _projection_perf and the batch-invariance part of _projection_accuracy drew temb from the process's default CUDA generator. Seed both. The per-draw accuracy inputs were already seeded; the full accuracy section from this commit is identical to the committed report's, so it needs no regeneration. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 92e8bd632..2dff72a8a 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -246,7 +246,8 @@ def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: weight, bias = _adaln_params() cases = [] for num in (1, 2, 3, 4): - temb = torch.randn(num, 2688, device="cuda") + gen = torch.Generator(device="cuda").manual_seed(num) + temb = torch.randn(num, 2688, device="cuda", generator=gen) cases.append( { "op": "adaln_projection_3mod", @@ -285,7 +286,8 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, cuda_frac.append(float((ours == gold).float().mean())) provider_frac.append(float((theirs == gold).float().mean())) early_frac.append(float((ours == early).float().mean())) - temb = torch.randn(9, 2688, device="cuda") + gen = torch.Generator(device="cuda").manual_seed(draws) + temb = torch.randn(9, 2688, device="cuda", generator=gen) full = op(temb, weight, bias) invariant = all( all( From 06bf572c136b2529b26f2ca871e2aaa0bf069814 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:12 +0000 Subject: [PATCH 16/24] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 3 files changed, 18 insertions(+), 18 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a76cc0417..a9c8832e2 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -124,9 +124,9 @@ dependence on T. ## Evidence -![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also @@ -135,7 +135,7 @@ or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![adaln_projection vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_projection.png) +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../../reports/experiments/h3-prior-art-b200/adaln_projection.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) was written from a clean tree at `c88691c` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-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 441e00004bc856365e8734ca607e2839a8811c4c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:54 +0000 Subject: [PATCH 17/24] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a9c8832e2..3f79732cd 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -127,7 +127,7 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `579775f1` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `50d1e924` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-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 a5d155784564a902fe6e3afe164da5d432afcb88 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:43:48 +0000 Subject: [PATCH 18/24] fix(h3): use one timing harness across the stacked PRs The stacked PRs carried three versions of measure(). #487-#494 have the most recent one (alternating order, setup outside the timed region, backward-only scope); take that block and its test here so the whole stack agrees and merging the PRs in order does not conflict. Also move the noqa: E402 of prepare_h3_weights.py to the first line of the import. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 91 ++++++++++++++--------- tests/models/minimax_h3/test_h3_report.py | 2 +- 2 files changed, 56 insertions(+), 37 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 2dff72a8a..6069de83a 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -34,37 +34,59 @@ ) -def time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 200) -> float: - """Median CUDA-event time of ``fn`` in microseconds (includes the Python wrapper).""" +def _prepared_call( + fn: Callable[..., Any], setup: Callable[[], Any] | None = None +) -> Callable[[], Any]: + if setup is None: + return fn + prepared = setup() + return lambda: fn(prepared) + + +def _sample_us(fn: Callable[..., Any], setup: Callable[[], Any] | None = None) -> float: + run = _prepared_call(fn, setup) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1e3 + + +def time_us( + fn: Callable[..., Any], + warmup: int = 20, + iters: int = 200, + *, + setup: Callable[[], Any] | None = None, +) -> float: + """Median CUDA-event microseconds; optional setup runs outside the timed region. + + When supplied, ``setup`` runs before every call and its return value is passed + to ``fn``. Backward measurements use it to build a fresh forward graph. + """ for _ in range(warmup): - fn() + _prepared_call(fn, setup)() torch.cuda.synchronize() - samples = [] - for _ in range(iters): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - fn() - end.record() - end.synchronize() - samples.append(start.elapsed_time(end) * 1e3) - return statistics.median(samples) + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) -def peak_mib(fn: Callable[[], Any]) -> float: - """Synchronize one CUDA call and report peak allocated memory above baseline in MiB.""" +def peak_mib(fn: Callable[..., Any], *, setup: Callable[[], Any] | None = None) -> float: + """Incremental peak CUDA memory of the call, excluding optional setup.""" + run = _prepared_call(fn, setup) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() base = torch.cuda.memory_allocated() - fn() + run() torch.cuda.synchronize() return (torch.cuda.max_memory_allocated() - base) / 2**20 def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: - """Time case callables in rotating order after warmup and return their evidence. + """Interleave timed keys, reversing their execution order on alternate iterations. Record raw CUDA-event samples, executed orders, median microseconds, decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. @@ -74,32 +96,29 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st keys = [key for key in TIMED_KEYS if key in case] if not keys: return row - for _ in range(warmup): - for key in keys: - case[key]() + orders = (keys, list(reversed(keys))) + row["execution_order"] = { + "policy": "alternating", + "iteration_0": orders[0], + "iteration_1": orders[1], + } + if any(key.endswith("_backward") for key in keys): + row["backward_timing_scope"] = "backward_only" + for iteration in range(warmup): + for key in orders[iteration % 2]: + _prepared_call(case[key], case.get(f"{key}_setup"))() torch.cuda.synchronize() samples = {key: [] for key in keys} - orders = [] - for index in range(iters): - # Rotate the first implementation so drift does not favor one backend. - offset = index % len(keys) - order = keys[offset:] + keys[:offset] - orders.append(order) - for key in order: - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - case[key]() - end.record() - end.synchronize() - samples[key].append(start.elapsed_time(end) * 1e3) + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) row["timing_samples_us"] = samples - row["timing_order"] = orders + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] for key in keys: us = statistics.median(samples[key]) row[f"{key}_us"] = us row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 - row[f"{key}_peak_mib"] = peak_mib(case[key]) + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) return row diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py index 98e6ab714..b79973eb1 100644 --- a/tests/models/minimax_h3/test_h3_report.py +++ b/tests/models/minimax_h3/test_h3_report.py @@ -53,7 +53,7 @@ def provider(): monkeypatch.setattr(h3_report.torch.cuda, "Event", Event) monkeypatch.setattr(h3_report.torch.cuda, "synchronize", lambda: None) - monkeypatch.setattr(h3_report, "peak_mib", lambda fn: 0.0) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) report = h3_report.measure( {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 ) From 07d319996358ac2c1a89f6c6bdd82c314f62eecf Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:28:37 +0000 Subject: [PATCH 19/24] fix(h3): count every output in the bitwise heatmap and keep results when a helper fails - plot_h3_bitwise_reuse: a cell now adds up differ/numel over all forward outputs (adaln_row_gather has six) and shows a shape or dtype mismatch as a full difference instead of reading output 0 only. - h3_prior_art: a mistyped --megatron-src no longer raises after the modes have finished and their part files were deleted; the Megatron commit is left empty. - h3_bitwise_reuse: a failing fast-math sinusoid build is recorded as unavailable instead of ending the child before the other candidates are written. - tools/migration/layout.json: map scripts/h3_bitwise_reuse.py and scripts/plot_h3_bitwise_reuse.py. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/test_h3_bitwise_reuse_plot.py | 36 +++++++++++++++++++ tools/migration/layout.json | 2 ++ tools/validation/models/h3_bitwise_reuse.py | 7 +++- tools/validation/models/h3_prior_art.py | 21 ++++++----- .../models/plot_h3_bitwise_reuse.py | 9 +++-- 5 files changed, 62 insertions(+), 13 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py diff --git a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py new file mode 100644 index 000000000..8fa02de7e --- /dev/null +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""A heatmap cell counts every forward output of the operator, not only the first.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("matplotlib") +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "tools/validation/models")) + +from plot_h3_bitwise_reuse import _cell # noqa: E402 + + +def _out(differ: int, numel: int = 100) -> dict: + return {"equal": differ == 0, "differ": differ, "numel": numel} + + +def test_cell_sums_over_all_outputs(): + assert _cell({"out": [_out(0), _out(0), _out(0)]}) == 0.0 + # a difference that is only in the fifth of six outputs is not hidden by output 0 + assert _cell({"out": [_out(0)] * 4 + [_out(30)] + [_out(0)]}) == pytest.approx(30 / 600) + + +def test_cell_marks_shape_or_dtype_mismatch_as_full_difference(): + assert _cell({"out": [_out(0), {"note": "shape mismatch"}]}) == 1.0 + + +def test_cell_without_data_is_not_applicable(): + assert _cell(None) is None + assert _cell({"out": []}) is None + assert _cell({"out": [None]}) is None diff --git a/tools/migration/layout.json b/tools/migration/layout.json index 206cbda6e..0b04ce765 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -819,9 +819,11 @@ "scripts/check_rocm_env.py": "tools/env/check_rocm_env.py", "scripts/check_stateful_kv.py": "tools/validation/operators/check_stateful_kv.py", "scripts/ci_smoke.py": "tools/checks/ci_smoke.py", + "scripts/h3_bitwise_reuse.py": "tools/validation/models/h3_bitwise_reuse.py", "scripts/h3_chain_replay.py": "tools/validation/models/h3_chain_replay.py", "scripts/h3_evidence.py": "tools/validation/models/h3_evidence.py", "scripts/h3_prior_art.py": "tools/validation/models/h3_prior_art.py", + "scripts/plot_h3_bitwise_reuse.py": "tools/validation/models/plot_h3_bitwise_reuse.py", "scripts/plot_h3_evidence.py": "tools/validation/models/plot_h3_evidence.py", "scripts/plot_h3_prior_art.py": "tools/validation/models/plot_h3_prior_art.py", "scripts/prepare_h3_weights.py": "tools/weights/prepare_h3_weights.py", diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py index d43e5c70e..345aa18cc 100644 --- a/tools/validation/models/h3_bitwise_reuse.py +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -118,7 +118,12 @@ def dump(op: str, mode: str, parts: Path, megatron_src: str | None, fast_math: b torch.backends.cudnn.allow_tf32 = False cands = PA.candidates(op, mode) if fast_math and op == "timestep_sinusoid" and mode == "plain": - cands.append({"name": "rl_kernel_fast_math", **_fast_math_sinusoid(parts / "build")}) + try: + cands.append({"name": "rl_kernel_fast_math", **_fast_math_sinusoid(parts / "build")}) + except Exception as exc: # noqa: BLE001 + (parts / "rl_kernel_fast_math.unavailable").write_text( + f"{type(exc).__name__}: {exc}"[:300] + ) for cand in cands: name = cand.pop("name") if "unavailable" in cand: diff --git a/tools/validation/models/h3_prior_art.py b/tools/validation/models/h3_prior_art.py index 9012542e2..07555ad26 100644 --- a/tools/validation/models/h3_prior_art.py +++ b/tools/validation/models/h3_prior_art.py @@ -911,15 +911,18 @@ def main() -> None: part.unlink(missing_ok=True) megatron_commit = None if args.megatron_src: - megatron_commit = ( - subprocess.run( - ["git", "rev-parse", "--short", "HEAD"], - cwd=args.megatron_src, - capture_output=True, - text=True, - ).stdout.strip() - or None - ) + try: + megatron_commit = ( + subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + cwd=args.megatron_src, + capture_output=True, + text=True, + ).stdout.strip() + or None + ) + except OSError: # a bad path must not discard the finished modes + megatron_commit = None report = { "kind": "h3_prior_art", "rfc": "RL-Align/RL-Kernel#420", diff --git a/tools/validation/models/plot_h3_bitwise_reuse.py b/tools/validation/models/plot_h3_bitwise_reuse.py index 85983fc0a..6da21441f 100644 --- a/tools/validation/models/plot_h3_bitwise_reuse.py +++ b/tools/validation/models/plot_h3_bitwise_reuse.py @@ -29,10 +29,13 @@ def _cell(entry: dict | None) -> float | None: if not entry: return None - out = entry["out"][0] - if out is None or "differ" not in out: + outs = [o for o in entry["out"] if o is not None] + if not outs: return None - return 0.0 if out["equal"] else out["differ"] / out["numel"] + if any("differ" not in o for o in outs): # shape or dtype mismatch + return 1.0 + numel = sum(o["numel"] for o in outs) + return sum(o["differ"] for o in outs) / numel if numel else None def _text(frac: float | None) -> str: From 90990842c2046bc65d6c6e276d60a9776389f1d4 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:35:00 +0000 Subject: [PATCH 20/24] fix(h3): refuse the projection op below SM80 and state the activation cast of the weight dtype - H3AdaLNProjectionCudaOp checks the device in __init__ (det_linear_device_supported), so on a pre-SM80 GPU the registry marks the backend failed and falls back to the PyTorch reference instead of failing at every BF16 call. - The formula in the operator doc casts the activation to the weight dtype like the wrapper does (BF16 for the checkpoint weights, FP32 is allowed by the contract). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-projection.md | 2 +- .../cuda/model_specific/minimax_h3/adaln_projection.py | 5 ++++- .../cuda/model_specific/minimax_h3/det_linear.py | 6 ++++++ tests/models/minimax_h3/test_h3_adaln_projection.py | 9 +++++++++ 4 files changed, 20 insertions(+), 2 deletions(-) diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 3f79732cd..5de612b3f 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -7,7 +7,7 @@ modulation tensors of one transformer block, for all three modalities: ```text -act = silu(temb).to(bf16) SiLU in FP32, one declared cast +act = silu(temb).to(W.dtype) SiLU in FP32, one declared cast (BF16 for the checkpoint weights) table = act @ W.T + b (T, 96768) BF16, W: (96768, 2688) rows = table.view(3T, 32256) row t * 3 + m, m: 0 video, 1 text, 2 audio shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) # (3T, 5376) each diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py index 6c9f5bf49..1363d8a44 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py @@ -8,16 +8,17 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( ACT_NONE, CONTRACT, det_linear_available, det_linear_backward_input, det_linear_backward_weight, + det_linear_device_supported, det_linear_forward, silu_backward_fp32, ) +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.adaln_projection import ( split_adaln_table, validate_h3_adaln_projection, @@ -98,6 +99,8 @@ def __init__(self) -> None: raise RuntimeError( "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" ) + if not det_linear_device_supported(): + raise RuntimeError("the BF16 h3 det_linear forward needs SM80 or newer") def __call__(self, temb, weight, bias): """Return six weight-dtype ``(3T, H)`` modulation views for CUDA inputs.""" diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py index 34d523464..d288d2fa9 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py @@ -29,6 +29,12 @@ def det_linear_available() -> bool: return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) +def det_linear_device_supported() -> bool: + """Return whether the current CUDA device can run the BF16 forward (``mma.sync``, SM80+).""" + + return not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] >= 8 + + def _require() -> None: """Raise when a deterministic linear entry point is unavailable.""" diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 5fb10a73f..f35b13a79 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -41,6 +41,15 @@ def _cuda_op(): return H3AdaLNProjectionCudaOp() +def test_pre_sm80_device_is_refused_at_construction_so_the_registry_can_fall_back(monkeypatch): + from rl_engine.backends.cuda.model_specific.minimax_h3 import adaln_projection as module + + monkeypatch.setattr(module, "det_linear_available", lambda: True) + monkeypatch.setattr(module, "det_linear_device_supported", lambda: False) + with pytest.raises(RuntimeError, match="SM80"): + module.H3AdaLNProjectionCudaOp() + + def _temb(num, dim=2688, seed=0, device="cuda"): """Generate seeded FP32 timestep embeddings on CPU and move them to the test device.""" From e6ee9327d26129f8a1fd4dfe32af7041fbc83ccc Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:12:13 +0000 Subject: [PATCH 21/24] test(h3): give the stubbed MLP accuracy a weight source check_weight_source reads accuracy["weights_source"] for timestep_mlp_fp32 since the early check of #484 was merged in; the stub returned an empty dict and the evidence CLI test failed with a KeyError. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_cli.py b/tests/models/minimax_h3/test_h3_cli.py index 8e97fc6a2..0732760b6 100644 --- a/tests/models/minimax_h3/test_h3_cli.py +++ b/tests/models/minimax_h3/test_h3_cli.py @@ -166,7 +166,7 @@ def test_evidence_records_weight_source(monkeypatch, tmp_path, op, weight_source lambda: {"rl_kernel_commit": "test_commit", "tracked_tree_dirty": False}, ) monkeypatch.setattr(h3_evidence, "environment", lambda: {}) - monkeypatch.setitem(h3_evidence.ACCURACY, op, lambda registry: {}) + monkeypatch.setitem(h3_evidence.ACCURACY, op, lambda registry: {"weights_source": "pinned"}) monkeypatch.setitem(h3_evidence.PERF_CASES, op, lambda registry: []) out = tmp_path / "report.json" monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) From ca14fadec7502be0b2d626c0279007102676b381 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:16:35 +0000 Subject: [PATCH 22/24] fix(h3): check the SM80 requirement of the BF16 forward for the device of the tensors det_linear_device_supported() took the current CUDA device and returned True without CUDA. It now takes the device in question (default: the current one) and is False when CUDA is unavailable or the device is not a CUDA device, so a constructor that calls it before any tensor exists fails and the registry falls back to the PyTorch reference. det_linear_forward() repeats the check for the device of the BF16 inputs and raises a RuntimeError naming SM80 instead of letting the launch fail inside the native call; FP32 is not affected. Tests pin that the capability is read for the requested device, that no CUDA means no support, and that a BF16 forward on a pre-SM80 device is refused while FP32 runs. The odd-width tests skip their BF16 forward cases on such a device, and their docstring now describes the 16-column MMA tile of the BF16 forward. Also drop the duplicated static_assert(kCols * kRows <= kWarp) of the FP32 forward kernel. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 1 - .../model_specific/minimax_h3/det_linear.py | 19 +++++++-- .../minimax_h3/test_h3_adaln_projection.py | 2 +- tests/models/minimax_h3/test_h3_det_linear.py | 40 ++++++++++++++++++- .../test_h3_det_linear_odd_width.py | 19 +++++++-- 5 files changed, 71 insertions(+), 10 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index d0ed82851..fb514abcc 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -134,7 +134,6 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) static_assert(kCols * kRows <= kWarp, "the epilogue stores through lane == c * kRows + r, so kCols * kRows must fit a warp"); constexpr int kVec = WV::kN; - static_assert(kCols * kRows <= kWarp, "one warp cannot store more than kWarp outputs per tile"); static_assert(XV::kN == kVec, "x and weight share a dtype"); constexpr int64_t kStride = static_cast(kWarp) * kVec; diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py index d288d2fa9..646685dc2 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py @@ -29,10 +29,19 @@ def det_linear_available() -> bool: return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) -def det_linear_device_supported() -> bool: - """Return whether the current CUDA device can run the BF16 forward (``mma.sync``, SM80+).""" +def det_linear_device_supported(device: torch.device | int | str | None = None) -> bool: + """Return whether ``device`` can run the BF16 forward (``mma.sync``, SM80+). - return not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] >= 8 + ``device`` defaults to the current CUDA device. Without CUDA, or for a non-CUDA device, the + answer is False: constructors use this as an early check so the registry falls back to the + PyTorch reference, and ``det_linear_forward`` repeats it for the device of the actual tensors. + """ + + if not torch.cuda.is_available(): + return False + if device is not None and torch.device(device).type != "cuda": + return False + return torch.cuda.get_device_capability(device)[0] >= 8 def _require() -> None: @@ -62,6 +71,10 @@ def det_linear_forward( """ _require() + if x.dtype == torch.bfloat16 and not det_linear_device_supported(x.device): + raise RuntimeError( + f"the BF16 h3 det_linear forward needs SM80 or newer (tensors are on {x.device})" + ) return _C.h3_det_linear_forward( x.contiguous(), weight.contiguous(), diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index f35b13a79..b22c5dea1 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -45,7 +45,7 @@ def test_pre_sm80_device_is_refused_at_construction_so_the_registry_can_fall_bac from rl_engine.backends.cuda.model_specific.minimax_h3 import adaln_projection as module monkeypatch.setattr(module, "det_linear_available", lambda: True) - monkeypatch.setattr(module, "det_linear_device_supported", lambda: False) + monkeypatch.setattr(module, "det_linear_device_supported", lambda *args: False) with pytest.raises(RuntimeError, match="SM80"): module.H3AdaLNProjectionCudaOp() diff --git a/tests/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py index f9bc79829..f1a2e3175 100644 --- a/tests/models/minimax_h3/test_h3_det_linear.py +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -8,8 +8,12 @@ import pytest import torch +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_device_supported, + det_linear_forward, +) from rl_engine.backends.extension import _C -from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), @@ -66,3 +70,37 @@ def test_backward_weight_handles_empty_columns(dtype, w_dtype, k_in, n_out, with torch.testing.assert_close(actual[0], expected_weight, atol=0, rtol=0) if with_bias: torch.testing.assert_close(actual[1], grad.sum(dim=0).to(w_dtype), atol=0, rtol=0) + + +def test_device_support_follows_the_device_that_is_asked_about(monkeypatch): + """The capability is read for the given device, not for the current one.""" + + seen = [] + + def capability(device=None): + seen.append(device) + return (7, 5) if device == torch.device("cuda:0") else (9, 0) + + monkeypatch.setattr(torch.cuda, "get_device_capability", capability) + assert not det_linear_device_supported(torch.device("cuda:0")) + assert det_linear_device_supported(torch.device("cuda:1")) + assert not det_linear_device_supported(torch.device("cpu")) + assert seen == [torch.device("cuda:0"), torch.device("cuda:1")] + + +def test_no_cuda_means_no_support(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + assert not det_linear_device_supported() + + +def test_bf16_forward_is_refused_on_a_pre_sm80_device_with_a_clear_error(monkeypatch): + """The wrapper checks the tensors' device, so no opaque launch error reaches the caller.""" + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (7, 5)) + x = torch.zeros(1, 64, device="cuda", dtype=torch.bfloat16) + w = torch.zeros(8, 64, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="SM80"): + det_linear_forward(x, w, None) + # FP32 does not use mma.sync and is not refused. + out = det_linear_forward(x.float(), w.float(), None)[0] + assert out.shape == (1, 8) diff --git a/tests/models/minimax_h3/test_h3_det_linear_odd_width.py b/tests/models/minimax_h3/test_h3_det_linear_odd_width.py index 2b0ee1aac..08036db8f 100644 --- a/tests/models/minimax_h3/test_h3_det_linear_odd_width.py +++ b/tests/models/minimax_h3/test_h3_det_linear_odd_width.py @@ -3,10 +3,12 @@ """Output widths that are not a multiple of the column tile of the deterministic linear. -Each warp of the forward kernel owns kCols adjacent output columns (2 for FP32, 4 for BF16), so the -last tile of an odd or short N has columns that are not loaded and must not be consumed. The -results for such widths must match an FP64 golden, and must be bitwise the leading columns of the -same product with zero columns padded on (the tail does not influence the real columns). +Each warp of the FP32 forward kernel owns kCols = 2 adjacent output columns, and each warp of the +BF16 forward (an ``mma.sync`` kernel) owns a 16-column tile (kMmaCols), so the last tile of an odd +or short N has columns that are not loaded and must not be consumed. The results for such widths +must match an FP64 golden, and must be bitwise the leading columns of the same product with zero +columns padded on (the tail does not influence the real columns). The BF16 forward needs SM80, so +its cases are skipped on older devices. """ from __future__ import annotations @@ -18,6 +20,7 @@ det_linear_available, det_linear_backward_input, det_linear_backward_weight, + det_linear_device_supported, det_linear_forward, ) @@ -32,6 +35,11 @@ ROWS = [1, 3, 5] +def _skip_unsupported(dtype): + if dtype == torch.bfloat16 and not det_linear_device_supported(): + pytest.skip("the BF16 forward needs SM80 or newer") + + def _inputs(dtype, rows, n_out, seed=0): g = torch.Generator(device="cpu").manual_seed(seed + n_out) x = torch.randn(rows, K_IN, generator=g).to(dtype).cuda() @@ -49,6 +57,7 @@ def _pad(weight, bias, extra): @pytest.mark.parametrize("n_out", WIDTHS) @pytest.mark.parametrize("rows", ROWS) def test_forward_matches_golden_and_ignores_the_tail(dtype, n_out, rows): + _skip_unsupported(dtype) x, weight, bias = _inputs(dtype, rows, n_out) out = det_linear_forward(x, weight, bias)[0] assert out.shape == (rows, n_out) and torch.isfinite(out.float()).all() @@ -95,6 +104,8 @@ def test_backward_matches_golden_and_ignores_the_tail(dtype, n_out): def test_k_that_is_not_a_multiple_of_the_lane_stride_matches_the_golden(dtype, k_in, n_out): """Row lengths that are a vector multiple but not a multiple of the 32-lane stride.""" + _skip_unsupported(dtype) + g = torch.Generator(device="cpu").manual_seed(k_in + n_out) x = torch.randn(3, k_in, generator=g).to(dtype).cuda() weight = (torch.randn(n_out, k_in, generator=g) / k_in**0.5).to(dtype).cuda() From 9c9789fd97f6a99d286e2615c89198230b3955ca Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:24:18 +0000 Subject: [PATCH 23/24] fix(h3): keep the bitwise-reuse report when a reference lacks a case, mark missing cells, reject zero-width AdaLN input h3_bitwise_reuse.compare indexed runs[ref][case] for every case of every candidate. A reference child that died after saving only some of its cases (an out-of-memory at the largest size) raised a KeyError before the report was written, losing the comparisons that had completed and leaving the .parts directory behind. A case the reference never saved is now recorded as unavailable. The plot treats an entry without outputs as not applicable, writes "n/a" in the cell and draws it grey instead of the white of "bitwise equal". validate_h3_adaln_projection rejects a weight with zero input width: F.linear over an empty reduction returns the broadcast bias instead of failing. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../reference/minimax_h3/adaln_projection.py | 3 ++ .../minimax_h3/test_h3_adaln_projection.py | 5 +++ .../minimax_h3/test_h3_bitwise_reuse_plot.py | 39 ++++++++++++++++++- tools/validation/models/h3_bitwise_reuse.py | 26 +++++++++---- .../models/plot_h3_bitwise_reuse.py | 8 ++-- 5 files changed, 69 insertions(+), 12 deletions(-) diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py index be83239c5..a6b9fd714 100644 --- a/rl_engine/reference/minimax_h3/adaln_projection.py +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -58,6 +58,9 @@ def validate_h3_adaln_projection(temb: torch.Tensor, weight: torch.Tensor, bias: f"weight/bias must share a dtype in {_WEIGHT_DTYPES}, got {weight.dtype}/{bias.dtype}" ) hidden = adaln_hidden_size(weight) + if weight.shape[1] == 0: + # F.linear over an empty reduction returns the broadcast bias instead of failing. + raise ValueError(f"AdaLN weight must have a nonzero input width, got {tuple(weight.shape)}") if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: raise ValueError( f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index b22c5dea1..23aff80f6 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -107,6 +107,11 @@ def test_rejects_bad_shapes(self): with pytest.raises(TypeError): op.forward(_temb(2, 16, device="cpu"), weight, bias.float()) + def test_rejects_a_zero_width_input_instead_of_returning_the_bias(self): + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(ValueError, match="nonzero input width"): + NativeH3AdaLNProjectionOp().forward(_temb(2, 0, device="cpu"), weight[:, :0], bias) + @requires_cuda class TestLayout: diff --git a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py index 8fa02de7e..ffb87bd65 100644 --- a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -13,7 +13,8 @@ pytest.importorskip("matplotlib") sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "tools/validation/models")) -from plot_h3_bitwise_reuse import _cell # noqa: E402 +from h3_bitwise_reuse import _compare_case # noqa: E402 +from plot_h3_bitwise_reuse import _cell, _panel # noqa: E402 def _out(differ: int, numel: int = 100) -> dict: @@ -34,3 +35,39 @@ def test_cell_without_data_is_not_applicable(): assert _cell(None) is None assert _cell({"out": []}) is None assert _cell({"out": [None]}) is None + + +def test_a_case_the_reference_never_saved_is_unavailable_and_does_not_abort_the_comparison(): + import torch + + run = {"out": [torch.ones(3)], "grads": {"x": torch.zeros(2)}} + ref = {"out": [torch.ones(3)], "grads": {"x": torch.zeros(2)}} + done = _compare_case(run, ref) + assert done["out"][0]["equal"] and done["grads"]["x"]["equal"] + missing = _compare_case(run, None) + assert "unavailable" in missing and _cell(missing) is None + + +def test_missing_comparisons_are_drawn_grey_not_as_bitwise_equal(): + import matplotlib.pyplot as plt + + def result(frac_differ): + return {"out": [_out(frac_differ)], "grads": {}} + + report = { + "op": "probe", + "rl_kernel_commit": "0123456789", + "cases": ["a", "b"], + "results": { + "rl_kernel": {"vs diffusers": {"a": result(0), "b": {"unavailable": "no run"}}}, + "diffusers": {}, + }, + } + fig, ax = plt.subplots() + _panel(ax, report) + grid = ax.images[0].get_array() + assert grid.mask[0, 1] and not grid.mask[0, 0] # the made comparison is data, the other masked + assert ax.images[0].cmap(ax.images[0].norm(0.0)) != ax.images[0].cmap._rgba_bad + texts = [t.get_text() for t in ax.texts] + assert "n/a" in texts and "0" in texts + plt.close(fig) diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py index 345aa18cc..81a80b75c 100644 --- a/tools/validation/models/h3_bitwise_reuse.py +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -161,6 +161,23 @@ def _same(a: torch.Tensor | None, b: torch.Tensor | None) -> dict[str, Any] | No } +def _compare_case(run: dict[str, Any], ref_run: dict[str, Any] | None) -> dict[str, Any]: + """Compare one case with the reference's run of it. + + A case the reference never saved is unavailable. + + A reference child can die after saving only some of its cases (an out-of-memory at the largest + size, say); the comparisons that did complete must still reach the report. + """ + + if ref_run is None: + return {"unavailable": "the reference run has no saved result for this case"} + return { + "out": [_same(a, b) for a, b in zip(run["out"], ref_run["out"])], + "grads": {g: _same(run["grads"].get(g), ref_run["grads"].get(g)) for g in run["grads"]}, + } + + def compare(parts: Path) -> dict[str, Any]: runs: dict[str, dict[str, Any]] = {} for f in sorted(parts.glob("*.pt")): @@ -174,14 +191,7 @@ def compare(parts: Path) -> dict[str, Any]: if ref == name: continue entry[f"vs {ref}"] = { - case: { - "out": [_same(a, b) for a, b in zip(run["out"], runs[ref][case]["out"])], - "grads": { - g: _same(run["grads"].get(g), runs[ref][case]["grads"].get(g)) - for g in run["grads"] - }, - } - for case, run in by_case.items() + case: _compare_case(run, runs[ref].get(case)) for case, run in by_case.items() } out[name] = entry for f in parts.glob("*.unavailable"): diff --git a/tools/validation/models/plot_h3_bitwise_reuse.py b/tools/validation/models/plot_h3_bitwise_reuse.py index 6da21441f..777b1edbf 100644 --- a/tools/validation/models/plot_h3_bitwise_reuse.py +++ b/tools/validation/models/plot_h3_bitwise_reuse.py @@ -27,7 +27,7 @@ def _cell(entry: dict | None) -> float | None: - if not entry: + if not entry or "out" not in entry: # no entry, or a case the reference never saved return None outs = [o for o in entry["out"] if o is not None] if not outs: @@ -62,14 +62,16 @@ def _panel(ax, report: dict) -> None: ax.text( j, i, - "—" if row == ref else _text(frac), + "—" if row == ref else ("n/a" if frac is None else _text(frac)), ha="center", va="center", fontsize=7.5, color="white" if frac is not None and frac > 0.3 else "black", fontweight="bold" if frac == 0 else None, ) - ax.imshow(np.sqrt(np.nan_to_num(grid)), cmap="Reds", vmin=0, vmax=1, aspect="auto") + # A comparison that was not made is grey, so it cannot be read as the white of "bitwise equal". + cmap = plt.get_cmap("Reds").with_extremes(bad="0.82") + ax.imshow(np.ma.masked_invalid(np.sqrt(grid)), cmap=cmap, vmin=0, vmax=1, aspect="auto") ax.set_yticks(range(len(rows)), rows, fontsize=8) ax.set_xticks(range(len(cols)), [str(c) for _, c in cols], fontsize=7, rotation=45) transform = ax.get_xaxis_transform() From 984f287df703c02abbd616a7a2cd0de59639b34c Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:15:02 +0000 Subject: [PATCH 24/24] fix(h3): fail closed on a device the probe cannot query, skip the projection bias fold, reject a weight without rows * `det_linear_device_supported` returns False where torch raises (an index outside this host's devices, a device string it cannot parse, a CUDA runtime that fails to start), so a constructor that uses it falls back to the PyTorch reference instead of failing; integer indices keep working (`torch.device(0)` is `cuda:0`). Tests cover each case. * `adaln_projection` backward passes `with_bias=` to `det_linear_backward_weight`, like the fused modulation op, so the bias fold does not run when only the weight needs a gradient. A test freezes each input in turn, compares the remaining gradients bitwise with the all-trainable run and records the argument of every call. * `adaln_hidden_size` rejects a (0, D) weight: it divided by 18 and gave H = 0 and six empty outputs. * `det_linear.cu`: the two FP32 launch-shape knobs (`H3_FWD_F32_COLS`, `H3_FWD_F32_AHEAD`) must be overridden together; a build that sets one fails in the preprocessor with that message instead of with an undeclared identifier. The test runs the preprocessor on stand-ins for the torch headers. * `h3_bitwise_reuse._compare_case` records a different number of outputs as a difference (a bare zip stopped at the shorter list and could call the case bitwise equal) and zips strictly otherwise. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 6 ++ .../minimax_h3/adaln_projection.py | 6 +- .../model_specific/minimax_h3/det_linear.py | 16 +++-- .../reference/minimax_h3/adaln_projection.py | 2 + .../minimax_h3/test_h3_adaln_projection.py | 47 +++++++++++++++ .../minimax_h3/test_h3_bitwise_reuse_plot.py | 16 +++++ tests/models/minimax_h3/test_h3_det_linear.py | 60 +++++++++++++++++++ tools/validation/models/h3_bitwise_reuse.py | 10 +++- 8 files changed, 156 insertions(+), 7 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index fb514abcc..c8625d0bf 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -63,6 +63,12 @@ namespace { +// The FP32 forward launch shape is one pair of knobs, and the single compiled case below is built +// from the same pair, so the shape picked at run time is always one that was compiled. A build that +// overrides one knob has to override the other (and H3_FWD_CONFIGS for a multi-case sweep). +#if defined(H3_FWD_F32_COLS) != defined(H3_FWD_F32_AHEAD) +#error "define H3_FWD_F32_COLS and H3_FWD_F32_AHEAD together" +#endif #ifndef H3_FWD_F32_COLS #define H3_FWD_F32_COLS 2 #define H3_FWD_F32_AHEAD 6 diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py index 1363d8a44..1946cf6a9 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py @@ -63,7 +63,11 @@ def backward(ctx, grad_table): grad_table = grad_table.float().contiguous() d_temb = dw = db = None if need_weight or need_bias: - dw, db = det_linear_backward_weight(grad_table, act, weight.dtype) + # The bias fold is skipped when only the weight needs a gradient. + dw, *rest = det_linear_backward_weight( + grad_table, act, weight.dtype, with_bias=need_bias + ) + db = rest[0] if need_bias else None if need_temb: # The cast's VJP is the identity; the SiLU VJP runs in FP32. d_act = det_linear_backward_input(grad_table, weight, torch.float32) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py index 646685dc2..ff35ee019 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py @@ -32,16 +32,22 @@ def det_linear_available() -> bool: def det_linear_device_supported(device: torch.device | int | str | None = None) -> bool: """Return whether ``device`` can run the BF16 forward (``mma.sync``, SM80+). - ``device`` defaults to the current CUDA device. Without CUDA, or for a non-CUDA device, the - answer is False: constructors use this as an early check so the registry falls back to the - PyTorch reference, and ``det_linear_forward`` repeats it for the device of the actual tensors. + ``device`` defaults to the current CUDA device. Without CUDA, for a non-CUDA device, for an + index this host does not have, or when the CUDA runtime cannot be queried, the answer is False: + constructors use this as an early check so the registry falls back to the PyTorch reference, + and ``det_linear_forward`` repeats it for the device of the actual tensors. """ if not torch.cuda.is_available(): return False - if device is not None and torch.device(device).type != "cuda": + try: + if device is not None and torch.device(device).type != "cuda": + return False + return torch.cuda.get_device_capability(device)[0] >= 8 + except (AssertionError, RuntimeError, TypeError, ValueError): + # torch raises AssertionError for an index outside the host's devices and RuntimeError for + # a device string it cannot parse or a CUDA runtime that fails to start. return False - return torch.cuda.get_device_capability(device)[0] >= 8 def _require() -> None: diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py index a6b9fd714..62ed98c90 100644 --- a/rl_engine/reference/minimax_h3/adaln_projection.py +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -34,6 +34,8 @@ def adaln_hidden_size(weight: torch.Tensor) -> int: f"AdaLN weight must be (6 * H * 3, D); got {tuple(weight.shape)}, whose first " f"dim is not a multiple of {rows_per_hidden}" ) + if weight.shape[0] == 0: + raise ValueError(f"AdaLN weight must have a nonzero row count, got {tuple(weight.shape)}") return weight.shape[0] // rows_per_hidden diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py index 23aff80f6..c464fd28b 100644 --- a/tests/models/minimax_h3/test_h3_adaln_projection.py +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -107,6 +107,14 @@ def test_rejects_bad_shapes(self): with pytest.raises(TypeError): op.forward(_temb(2, 16, device="cpu"), weight, bias.float()) + def test_rejects_a_weight_without_rows_instead_of_hidden_size_zero(self): + """A (0, D) weight divides by 18 and would give H = 0 and six empty outputs.""" + + temb = _temb(2, 16, device="cpu") + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(ValueError, match="nonzero row count"): + NativeH3AdaLNProjectionOp().forward(temb, weight[:0], bias[:0]) + def test_rejects_a_zero_width_input_instead_of_returning_the_bias(self): weight, bias = _synthetic(8, 16, device="cpu") with pytest.raises(ValueError, match="nonzero input width"): @@ -295,6 +303,45 @@ def test_rows_invariant_across_tensor_core_row_tiles(self, num): for s_out, f_out in zip(single, full, strict=True): assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) + @pytest.mark.parametrize( + "frozen", [("bias",), ("weight",), ("temb",), ("weight", "bias"), ("temb", "bias")] + ) + def test_frozen_inputs_get_no_gradient_and_do_not_change_the_others(self, frozen, monkeypatch): + """Only a bias that needs a gradient is folded; the other gradients stay bitwise equal.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3 import adaln_projection + + op = _cuda_op() + weight, bias = _synthetic(64, 2688, seed=11) + temb = _temb(3, seed=11) + grads = [torch.randn(9, 64, device="cuda").bfloat16() for _ in range(6)] + names = ("temb", "weight", "bias") + + def run(frozen_names): + leaves = [t.detach().clone() for t in (temb, weight, bias)] + for name, leaf in zip(names, leaves, strict=True): + leaf.requires_grad_(name not in frozen_names) + torch.autograd.backward(list(op(*leaves)), grads) + return leaves + + full = run(()) + calls = [] + real = adaln_projection.det_linear_backward_weight + + def recording(*args, **kwargs): + calls.append(kwargs.get("with_bias", True)) + return real(*args, **kwargs) + + monkeypatch.setattr(adaln_projection, "det_linear_backward_weight", recording) + part = run(frozen) + for name, a, b in zip(names, full, part, strict=True): + if name in frozen: + assert b.grad is None, name + else: + assert torch.equal(a.grad, b.grad), name + wants_params = "weight" not in frozen or "bias" not in frozen + assert calls == ([("bias" not in frozen)] if wants_params else []), (frozen, calls) + def test_rejects_cpu(self): """Reject CPU tensors at the CUDA projection boundary.""" diff --git a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py index ffb87bd65..ae039698b 100644 --- a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -48,6 +48,22 @@ def test_a_case_the_reference_never_saved_is_unavailable_and_does_not_abort_the_ assert "unavailable" in missing and _cell(missing) is None +def test_a_different_number_of_outputs_is_a_difference_not_a_shorter_comparison(): + import torch + + ones = torch.ones(3) + run = {"out": [ones, ones], "grads": {"x": torch.zeros(2)}} + ref = {"out": [ones], "grads": {"x": torch.zeros(2)}} + done = _compare_case(run, ref) + # the extra output is not dropped (a zip would call this case bitwise equal) ... + assert done["out"] == [{"equal": False, "note": "2 vs 1 outputs"}] + assert _cell(done) == 1.0 + # ... and the gradients are still compared. + assert done["grads"]["x"]["equal"] + same = _compare_case(run, {"out": [ones, ones], "grads": {"x": torch.zeros(2)}}) + assert [o["equal"] for o in same["out"]] == [True, True] + + def test_missing_comparisons_are_drawn_grey_not_as_bitwise_equal(): import matplotlib.pyplot as plt diff --git a/tests/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py index f1a2e3175..1bd78a42b 100644 --- a/tests/models/minimax_h3/test_h3_det_linear.py +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -5,6 +5,10 @@ from __future__ import annotations +import shutil +import subprocess +from pathlib import Path + import pytest import torch @@ -88,6 +92,26 @@ def capability(device=None): assert seen == [torch.device("cuda:0"), torch.device("cuda:1")] +@pytest.mark.parametrize("device", [torch.cuda.device_count(), 99, "cuda:99", "bogus", 1.5]) +def test_a_device_this_host_does_not_have_is_unsupported_not_an_error(device): + """The probe answers False where torch raises, so a constructor falls back, not fails.""" + + assert det_linear_device_supported(device) is False + + +def test_an_integer_index_names_a_cuda_device(): + assert det_linear_device_supported(0) == det_linear_device_supported(torch.device("cuda", 0)) + assert det_linear_device_supported("cuda:0") == det_linear_device_supported(0) + + +def test_a_cuda_runtime_that_cannot_be_queried_is_unsupported(monkeypatch): + def broken(device=None): + raise RuntimeError("CUDA driver initialization failed") + + monkeypatch.setattr(torch.cuda, "get_device_capability", broken) + assert det_linear_device_supported() is False + + def test_no_cuda_means_no_support(monkeypatch): monkeypatch.setattr(torch.cuda, "is_available", lambda: False) assert not det_linear_device_supported() @@ -104,3 +128,39 @@ def test_bf16_forward_is_refused_on_a_pre_sm80_device_with_a_clear_error(monkeyp # FP32 does not use mma.sync and is not refused. out = det_linear_forward(x.float(), w.float(), None)[0] assert out.shape == (1, 8) + + +@pytest.mark.skipif(shutil.which("g++") is None, reason="needs a C++ preprocessor") +@pytest.mark.parametrize( + ("defines", "builds"), + [ + ([], True), + (["-DH3_FWD_F32_COLS=2", "-DH3_FWD_F32_AHEAD=6"], True), + (["-DH3_FWD_F32_COLS=4"], False), + (["-DH3_FWD_F32_AHEAD=4"], False), + ], +) +def test_the_fp32_launch_shape_knobs_are_overridden_together(tmp_path, defines, builds): + """One knob alone would leave the compiled case and the run-time shape apart: say so. + + Only the preprocessor runs, on empty stand-ins for the torch and CUDA headers. + """ + + for header in ( + "torch/extension.h", + "ATen/cuda/CUDAContext.h", + "c10/cuda/CUDAGuard.h", + "c10/cuda/CUDAException.h", + "cuda_bf16.h", + ): + (tmp_path / header).parent.mkdir(parents=True, exist_ok=True) + (tmp_path / header).write_text("") + source = Path(__file__).resolve().parents[3] / "csrc/cuda/h3/det_linear.cu" + done = subprocess.run( + ["g++", "-E", "-x", "c++", "-std=c++17", f"-I{tmp_path}", *defines, str(source)], + capture_output=True, + text=True, + ) + assert (done.returncode == 0) is builds, done.stderr[-400:] + if not builds: + assert "define H3_FWD_F32_COLS and H3_FWD_F32_AHEAD together" in done.stderr diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py index 81a80b75c..134c82a07 100644 --- a/tools/validation/models/h3_bitwise_reuse.py +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -172,8 +172,16 @@ def _compare_case(run: dict[str, Any], ref_run: dict[str, Any] | None) -> dict[s if ref_run is None: return {"unavailable": "the reference run has no saved result for this case"} + if len(run["out"]) != len(ref_run["out"]): + # Outputs that have no counterpart are a difference, not something to drop: a zip would + # stop at the shorter list and could report the case as bitwise equal. + outs: list[dict[str, Any] | None] = [ + {"equal": False, "note": f"{len(run['out'])} vs {len(ref_run['out'])} outputs"} + ] + else: + outs = [_same(a, b) for a, b in zip(run["out"], ref_run["out"], strict=True)] return { - "out": [_same(a, b) for a, b in zip(run["out"], ref_run["out"])], + "out": outs, "grads": {g: _same(run["grads"].get(g), ref_run["grads"].get(g)) for g in run["grads"]}, }