diff --git a/benchmarks/models/benchmark_h3_conditioning.py b/benchmarks/models/benchmark_h3_conditioning.py index b739f34ab..df977b46e 100644 --- a/benchmarks/models/benchmark_h3_conditioning.py +++ b/benchmarks/models/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/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index d9a7d5653..c8625d0bf 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 @@ -30,10 +32,24 @@ // chain with the multiplier 1, which is a plain FP32 add (fmaf(g, 1, a) == a + g). // // Tails: the host requires K % VEC == 0 (VEC = 16 / sizeof(w)), and every 16-byte -// load in the forward kernels starts at a multiple of VEC, so `k < K` implies -// `k + VEC <= K`: no load straddles the end of a row, and every row start stays -// 16-byte aligned. N and the row count have no such requirement; their tails are -// masked per column / per row. +// load in the forward kernels starts at a multiple of VEC (the MMA kernel's at +// k = kb + 32u + 8q), so `k < K` implies `k + VEC <= K`: no load straddles the end +// of a row, and every row start stays 16-byte aligned. N and the row count have no +// such requirement; their tails are masked per column / per row. +// +// 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 @@ -47,15 +63,18 @@ 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 -#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; @@ -193,6 +212,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, @@ -303,9 +409,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}; } @@ -383,18 +488,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}; @@ -413,9 +521,13 @@ 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()); + 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; // chunks is blockIdx.y and rows * k_in / threads is blockIdx.x of the fold: fail with a clear message // instead of "invalid configuration argument" when a launch dimension would overflow. @@ -424,7 +536,6 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei TORCH_CHECK((rows * k_in + 255) / 256 <= std::numeric_limits::max(), "rows * K = ", rows * k_in, " is too large for one input-gradient fold launch"); 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)); @@ -464,10 +575,15 @@ 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()); 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/docs/.nav.yml b/docs/.nav.yml index 395041926..bd35a2f85 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -33,6 +33,7 @@ nav: - MiniMax-H3: - operators/h3-timestep-sinusoid.md - operators/h3-timestep-mlp.md + - operators/h3-adaln-projection.md - Developer Guide: - contributing/README.md - Contributor Guide: contributing/contributor-guide.md diff --git a/docs/operators/README.md b/docs/operators/README.md index 57b97f32c..2a1a0967a 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -34,5 +34,6 @@ 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) - [H3 Flash Attention (optional fast backend)](h3-flash-attention.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..5de612b3f --- /dev/null +++ b/docs/operators/h3-adaln-projection.md @@ -0,0 +1,202 @@ +# 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(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 +``` + +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.runtime.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.backends.cuda.model_specific.minimax_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.reference.minimax_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/models/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 | 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 22 µs of the op time is +the SiLU and cast kernels plus the Python wrapper. With those included, the op is within +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 +(`__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](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 `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 +or in a batch of 9. + +## Existing implementations (RFC #420 reuse rule) + +![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 | +|---|---|---|---|---| +| 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`](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 +python tools/validation/models/h3_prior_art.py --op adaln_projection --out reports/experiments/h3-prior-art-b200/adaln_projection.json --megatron-src +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/adaln_projection.json +``` + +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](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 `50d1e924` 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 +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_adaln_projection.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end: sinusoid -> MLP -> projection +python tools/validation/operators/check_operator.py --op adaln_projection_3mod --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op adaln_projection_3mod \ + --out reports/experiments/h3-adaln-projection-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/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/_C.pyi b/rl_engine/_C.pyi index d4275fcc0..efb12a121 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -337,17 +337,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/backends/cuda/model_specific/minimax_h3/adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py new file mode 100644 index 000000000..1946cf6a9 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py @@ -0,0 +1,139 @@ +# 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.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, +) + +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.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +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 + # 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): + """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() + d_temb = dw = db = None + if need_weight or need_bias: + # 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) + 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: + """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" + ) + 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.""" + + 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") + 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) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") + 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/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py index 9391114c2..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 @@ -24,10 +24,35 @@ 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 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, 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 + 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 + + 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,7 +68,19 @@ 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() + 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(), @@ -56,6 +93,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 +108,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/backends/cuda/model_specific/minimax_h3/timestep_mlp.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py index 87763d308..565d42111 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py +++ b/rl_engine/backends/cuda/model_specific/minimax_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() @@ -86,19 +90,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/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py index 86b838403..f44d03e44 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py +++ b/rl_engine/backends/cuda/model_specific/minimax_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 # The VJP slices the native output as [cos | sin] over frequencies_fp32 order. @@ -78,6 +84,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 " @@ -86,6 +94,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( @@ -95,6 +105,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. + """ + # Dtype, rank, emptiness and even num_channels are checked here on every call (they are # what backward's [cos | sin] split relies on); only the value check is left to the native # entry point, which enforces check_range itself and avoids a second host sync here. @@ -110,6 +127,8 @@ def forward_fp32( num_channels: int = H3_FREQ_DIM, check_range: bool = True, ) -> torch.Tensor: + """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. The signature matches the # reference's, so validation can call either one with the same arguments. return self.forward(timestep, num_channels=num_channels, check_range=check_range) diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py new file mode 100644 index 000000000..62ed98c90 --- /dev/null +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -0,0 +1,128 @@ +# 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.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_WEIGHT_DTYPES = (torch.bfloat16, torch.float32) + + +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( + 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 + + +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") + 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 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)}" + ) + 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 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 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. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + # 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/rl_engine/reference/minimax_h3/timestep_mlp.py b/rl_engine/reference/minimax_h3/timestep_mlp.py index f8348d99c..a5429ce69 100644 --- a/rl_engine/reference/minimax_h3/timestep_mlp.py +++ b/rl_engine/reference/minimax_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): @@ -66,13 +72,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/reference/minimax_h3/timestep_sinusoid.py b/rl_engine/reference/minimax_h3/timestep_sinusoid.py index 2f374c81a..ef782a4c6 100644 --- a/rl_engine/reference/minimax_h3/timestep_sinusoid.py +++ b/rl_engine/reference/minimax_h3/timestep_sinusoid.py @@ -75,6 +75,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( @@ -84,6 +86,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 @@ -103,6 +112,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/runtime/registry.py b/rl_engine/runtime/registry.py index 556101696..fb9466ae1 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -216,6 +216,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_TIMESTEP_MLP = ( "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp.H3TimestepMLPCudaOp" ) + PYTORCH_H3_ADALN_PROJECTION = ( + "rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp" + ) + CUDA_H3_ADALN_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection.H3AdaLNProjectionCudaOp" + ) def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -515,6 +521,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()) @@ -694,6 +701,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": [ @@ -737,6 +748,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": [ @@ -788,6 +800,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], @@ -816,6 +829,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/validation/models/h3_cases.py b/rl_engine/validation/models/h3_cases.py index 527a820ec..8eb509d33 100644 --- a/rl_engine/validation/models/h3_cases.py +++ b/rl_engine/validation/models/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/validation/models/h3_chain.py b/rl_engine/validation/models/h3_chain.py index 90489566d..70978a3e0 100644 --- a/rl_engine/validation/models/h3_chain.py +++ b/rl_engine/validation/models/h3_chain.py @@ -30,7 +30,11 @@ from rl_engine.runtime.registry import KernelRegistry from rl_engine.validation.models.h3_cases import h3_packed_layout, h3_timesteps -from rl_engine.validation.models.h3_provider import provider_time_embedder, provider_time_proj +from rl_engine.validation.models.h3_provider import ( + provider_adaln_modulation, + provider_time_embedder, + provider_time_proj, +) REPO_ROOT = Path(__file__).resolve().parents[3] @@ -50,10 +54,19 @@ 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"] + + STAGES: list[Stage] = [ Stage( name="timestep_sinusoid_h3", @@ -75,16 +88,30 @@ 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, + ), ] 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 @@ -98,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), @@ -108,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]) @@ -122,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] = { @@ -163,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: @@ -175,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/validation/models/h3_manifest.json b/rl_engine/validation/models/h3_manifest.json index 52b1989b1..5d08b87f3 100644 --- a/rl_engine/validation/models/h3_manifest.json +++ b/rl_engine/validation/models/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/validation/models/h3_provider.py b/rl_engine/validation/models/h3_provider.py index 2e906d235..87a63ae0e 100644 --- a/rl_engine/validation/models/h3_provider.py +++ b/rl_engine/validation/models/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,8 +45,25 @@ 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) 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, ...]: + """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) + return temb.chunk(6, dim=-1) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 5d9437a01..aa4bdba56 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -17,7 +17,11 @@ from rl_engine.runtime.registry import KernelRegistry from rl_engine.validation.models.h3_cases import h3_timesteps -from rl_engine.validation.models.h3_provider import provider_time_embedder, provider_time_proj +from rl_engine.validation.models.h3_provider import ( + provider_adaln_modulation, + provider_time_embedder, + provider_time_proj, +) from rl_engine.validation.models.h3_weights import h3_weights_dir, load_h3_conditioning_weights # Keys of a perf case that hold a timed callable, in display order. @@ -124,6 +128,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): @@ -143,6 +149,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] @@ -163,7 +171,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) @@ -190,6 +198,8 @@ def h3_weights_source() -> str: 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) @@ -212,6 +222,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") @@ -237,11 +249,97 @@ 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]: + """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 = [] + for num in (1, 2, 3, 4): + gen = torch.Generator(device="cuda").manual_seed(num) + temb = torch.randn(num, 2688, device="cuda", generator=gen) + 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: + """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]: + """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( + 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())) + 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( + torch.equal(s, f[3 * i : 3 * i + 3]) + for s, f in zip(op(temb[i : i + 1], weight, bias), full, strict=True) + ) + 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/validation/models/h3_weights.py b/rl_engine/validation/models/h3_weights.py index 261b4385e..c6e5942c2 100644 --- a/rl_engine/validation/models/h3_weights.py +++ b/rl_engine/validation/models/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): @@ -41,6 +47,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 +68,20 @@ def load_h3_conditioning_weights( if not path.is_file(): raise FileNotFoundError(f"{path} missing; run tools/weights/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/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index 6ca195306..b19c65b35 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/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( @@ -27,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, @@ -49,6 +51,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) @@ -57,6 +60,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.""" if op_name == "video_patch_pack_unpack": shape = "x".join(str(dim) for dim in _video_patch_shape(args)) return f"{_arg_str(args, 'patch_direction', 'pack')}-{shape}" @@ -86,6 +90,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] @@ -379,6 +385,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)) @@ -386,6 +393,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": @@ -401,12 +409,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.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp @@ -424,6 +434,30 @@ 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) + 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/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index d3f940390..f4e7bf58a 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -305,6 +305,22 @@ 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.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp"), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp" + ), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "adaln_projection.H3AdaLNProjectionCudaOp" + ), + }, + grad_input_names=("temb", "weight", "bias"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/tests/models/minimax_h3/conftest.py b/tests/models/minimax_h3/conftest.py index c1c9296a1..8bc18d0df 100644 --- a/tests/models/minimax_h3/conftest.py +++ b/tests/models/minimax_h3/conftest.py @@ -22,6 +22,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 tools/weights/prepare_h3_weights.py") try: diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py new file mode 100644 index 000000000..c464fd28b --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -0,0 +1,359 @@ +# 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.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +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") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _cuda_op(): + """Construct the CUDA projection operator, skipping missing native linear kernels.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + 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 *args: 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.""" + + 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) + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype) + return weight.to(device), bias.to(device) + + +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) + theirs = provider_adaln_modulation(temb, weight, bias, hidden_size=8) + assert len(ours) == 6 + 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): + # 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 + 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()) + + 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"): + NativeH3AdaLNProjectionOp().forward(_temb(2, 0, device="cpu"), weight[:, :0], bias) + + +@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): + """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() + 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): + """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() + + +@requires_cuda +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) + golden = NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias) + 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) + # Beyond the contract: almost every element is the correctly rounded + # 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(golden).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.""" + + 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.9995 + 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) + 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, 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, strict=True): + 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) + 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]) + # 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 + # 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:], strict=True): + torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) + # 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 +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) + 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) + + 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)] + 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, strict=True): + 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( + 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, 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.""" + + 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.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_projection_3mod", device="cuda") + assert type(op).__name__ == "H3AdaLNProjectionCudaOp" 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..ae039698b --- /dev/null +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -0,0 +1,89 @@ +# 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 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: + 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 + + +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_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 + + 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/tests/models/minimax_h3/test_h3_cli.py b/tests/models/minimax_h3/test_h3_cli.py new file mode 100644 index 000000000..0732760b6 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_cli.py @@ -0,0 +1,174 @@ +# 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.validation.models import h3_chain +from rl_engine.validation.models.h3_weights import WEIGHTS_ENV +from tools.validation.models import h3_chain_replay, h3_evidence + + +@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__ + + stages = [] + 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" + 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): + """Record each dispatched stage and return its deterministic CPU stub.""" + + 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 +): + """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" + 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): + """Reject weighted evidence without a configured checkpoint before writing a report.""" + + 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): + """Label weighted evidence as pinned and sinusoid evidence as weight-independent.""" + + 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: {"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)]) + h3_evidence.main() + assert json.loads(out.read_text())["weight_source"] == weight_source diff --git a/tests/models/minimax_h3/test_h3_conditioning_e2e.py b/tests/models/minimax_h3/test_h3_conditioning_e2e.py index 104090d10..4ea90a5f7 100644 --- a/tests/models/minimax_h3/test_h3_conditioning_e2e.py +++ b/tests/models/minimax_h3/test_h3_conditioning_e2e.py @@ -22,7 +22,10 @@ from rl_engine.runtime.registry import KernelRegistry from rl_engine.validation.models 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)] @@ -30,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(require_h3_kernels, 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/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py new file mode 100644 index 000000000..1bd78a42b --- /dev/null +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -0,0 +1,166 @@ +# 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 shutil +import subprocess +from pathlib import Path + +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 + +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): + """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"): + _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): + """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"): + _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) + + +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")] + + +@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() + + +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) + + +@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/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() diff --git a/tests/models/minimax_h3/test_h3_timestep_mlp.py b/tests/models/minimax_h3/test_h3_timestep_mlp.py index 28359f0c2..f5a055727 100644 --- a/tests/models/minimax_h3/test_h3_timestep_mlp.py +++ b/tests/models/minimax_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.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available from rl_engine.backends.cuda.model_specific.minimax_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() @@ -102,6 +118,8 @@ def test_a_w1_that_is_not_a_matrix_gets_a_dimension_error(self, bad_w1): 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) @@ -112,6 +130,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) @@ -125,6 +145,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)] @@ -132,7 +154,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) @@ -141,6 +163,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( @@ -150,6 +174,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 + ) + @pytest.mark.parametrize("with_bias", [True, False]) def test_weight_gradient_with_no_input_columns_keeps_the_bias_gradient(self, with_bias): # K = 0: dW is an (N, 0) tensor, but the bias gradient is still the column sums of grad, so @@ -170,11 +218,15 @@ def test_weight_gradient_with_no_input_columns_keeps_the_bias_gradient(self, wit ) 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]) @@ -182,6 +234,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) @@ -193,16 +247,20 @@ 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): + """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 @@ -257,6 +315,8 @@ def recording(*args, **kwargs): assert calls == expected, (frozen, calls) def test_registry_dispatches_cuda(self): + """Resolve the timestep MLP registry entry to its dedicated CUDA implementation.""" + from rl_engine.runtime.registry import KernelRegistry _cuda_op() @@ -269,6 +329,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.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_forward _cuda_op() @@ -279,6 +341,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.backends.cuda.model_specific.minimax_h3.det_linear import ( det_linear_backward_input, ) @@ -290,6 +354,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.backends.cuda.model_specific.minimax_h3.det_linear import ( det_linear_backward_weight, ) diff --git a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py index 6adbdd208..6d7ad06e5 100644 --- a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py +++ b/tests/models/minimax_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.backends.cuda.model_specific.minimax_h3.timestep_sinusoid import ( H3TimestepSinusoidCudaOp, h3_sinusoid_cuda_available, @@ -46,6 +48,8 @@ def _dense_unit_sweep() -> torch.Tensor: 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 @@ -56,11 +60,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) @@ -77,10 +85,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) @@ -97,6 +109,8 @@ def test_frequency_table_rejects_what_the_validator_rejects(self, num_channels): 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) @@ -120,12 +134,16 @@ def test_sglang_timestep_embedding_is_bitwise_equal(self): @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) @@ -135,10 +153,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)) @@ -155,21 +177,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.backends.extension import _C _cuda_op() @@ -181,6 +211,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.backends.extension import _C _cuda_op() @@ -188,6 +220,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) @@ -203,6 +237,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)) @@ -237,10 +273,14 @@ def test_odd_num_channels_is_rejected_before_the_native_call(self): _cuda_op().forward(t, num_channels=255, check_range=False) 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 @@ -253,6 +293,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.runtime.registry import KernelRegistry _cuda_op() diff --git a/tests/models/minimax_h3/test_h3_weights.py b/tests/models/minimax_h3/test_h3_weights.py new file mode 100644 index 000000000..13b9b2489 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_weights.py @@ -0,0 +1,74 @@ +# 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.validation.models 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): + """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() + + +@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): + """Reject altered shape, dtype, or bytes even outside the requested tensor subset.""" + + 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/models/minimax_h3/test_prepare_h3_weights.py b/tests/models/minimax_h3/test_prepare_h3_weights.py index 0c470e8d2..5c9a966b9 100644 --- a/tests/models/minimax_h3/test_prepare_h3_weights.py +++ b/tests/models/minimax_h3/test_prepare_h3_weights.py @@ -12,12 +12,14 @@ import pytest import torch -from rl_engine.validation.models.h3_weights import EXTRACTED_FILE, sha256_file +from rl_engine.validation.models.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): + """Enforce tensor contracts before extraction writes, including under Python optimization.""" + import huggingface_hub from safetensors.torch import load_file, save_file @@ -42,6 +44,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), } }, } diff --git a/tests/runtime/test_dispatch.py b/tests/runtime/test_dispatch.py index 02f642b87..9566ef6c2 100644 --- a/tests/runtime/test_dispatch.py +++ b/tests/runtime/test_dispatch.py @@ -232,11 +232,13 @@ 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) 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/tests/validation/operators/test_ws1_gtest_gpu.py b/tests/validation/operators/test_ws1_gtest_gpu.py index dae43dd92..9fd9f2595 100644 --- a/tests/validation/operators/test_ws1_gtest_gpu.py +++ b/tests/validation/operators/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", @@ -51,6 +53,7 @@ def test_all_ws1_single_ops_are_registered(): "prefix_shared_attention", "timestep_sinusoid_h3", "timestep_mlp_fp32", + "adaln_projection_3mod", } <= names diff --git a/tools/migration/layout.json b/tools/migration/layout.json index 7d37d4ef5..0b04ce765 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -467,6 +467,10 @@ "docs/design/ws2-attention-single-gpu-harness.md": "docs/archive/design/ws2-attention-single-gpu-harness.md", "docs/design/ws2-cp-attention-contract.md": "docs/contracts/ws2-cp-attention-contract.md", "docs/design/ws2-module-debug-matrix.md": "docs/archive/design/ws2-module-debug-matrix.md", + "docs/usage/evidence/h3-adaln-projection-b200/figure.png": "reports/experiments/h3-adaln-projection-b200/figure.png", + "docs/usage/evidence/h3-adaln-projection-b200/report.json": "reports/experiments/h3-adaln-projection-b200/report.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_projection.json": "reports/experiments/h3-prior-art-b200/adaln_projection.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_projection.png": "reports/experiments/h3-prior-art-b200/adaln_projection.png", "docs/usage/evidence/h3-prior-art-b200/timestep_mlp.json": "reports/experiments/h3-prior-art-b200/timestep_mlp.json", "docs/usage/evidence/h3-prior-art-b200/timestep_mlp.png": "reports/experiments/h3-prior-art-b200/timestep_mlp.png", "docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.json": "reports/experiments/h3-prior-art-b200/timestep_sinusoid.json", @@ -685,6 +689,7 @@ "rl_engine/kernels/ops/cuda/attention/strict_runtime.py": "rl_engine/backends/cuda/attention/strict_runtime.py", "rl_engine/kernels/ops/cuda/ffn.py": "rl_engine/backends/cuda/ffn/ffn.py", "rl_engine/kernels/ops/cuda/h3/__init__.py": "rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py", + "rl_engine/kernels/ops/cuda/h3/adaln_projection.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py", "rl_engine/kernels/ops/cuda/h3/det_linear.py": "rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py", "rl_engine/kernels/ops/cuda/h3/timestep_mlp.py": "rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py", "rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py": "rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py", @@ -719,6 +724,7 @@ "rl_engine/kernels/ops/pytorch/ffn/__init__.py": "rl_engine/ops/ffn/__init__.py", "rl_engine/kernels/ops/pytorch/ffn/ffn.py": "rl_engine/ops/ffn/qwen3.py", "rl_engine/kernels/ops/pytorch/h3/__init__.py": "rl_engine/reference/minimax_h3/__init__.py", + "rl_engine/kernels/ops/pytorch/h3/adaln_projection.py": "rl_engine/reference/minimax_h3/adaln_projection.py", "rl_engine/kernels/ops/pytorch/h3/fixed_order.py": "rl_engine/reference/minimax_h3/fixed_order.py", "rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py": "rl_engine/reference/minimax_h3/timestep_mlp.py", "rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py": "rl_engine/reference/minimax_h3/timestep_sinusoid.py", @@ -813,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", @@ -850,9 +858,14 @@ "tests/distributed/test_rocm_strict_attention_cp.py": "tests/distributed/cp/test_rocm_strict_attention_cp.py", "tests/distributed/test_transport_deterministic_collective.py": "tests/distributed/collectives/test_transport_deterministic_collective.py", "tests/h3/conftest.py": "tests/models/minimax_h3/conftest.py", + "tests/h3/test_h3_adaln_projection.py": "tests/models/minimax_h3/test_h3_adaln_projection.py", + "tests/h3/test_h3_cli.py": "tests/models/minimax_h3/test_h3_cli.py", "tests/h3/test_h3_conditioning_e2e.py": "tests/models/minimax_h3/test_h3_conditioning_e2e.py", + "tests/h3/test_h3_det_linear.py": "tests/models/minimax_h3/test_h3_det_linear.py", + "tests/h3/test_h3_report.py": "tests/models/minimax_h3/test_h3_report.py", "tests/h3/test_h3_timestep_mlp.py": "tests/models/minimax_h3/test_h3_timestep_mlp.py", "tests/h3/test_h3_timestep_sinusoid.py": "tests/models/minimax_h3/test_h3_timestep_sinusoid.py", + "tests/h3/test_h3_weights.py": "tests/models/minimax_h3/test_h3_weights.py", "tests/h3/test_prepare_h3_weights.py": "tests/models/minimax_h3/test_prepare_h3_weights.py", "tests/linear_logp_tp.py": "tests/distributed/tp/linear_logp_tp.py", "tests/test_alignment_model_wrappers.py": "tests/models/qwen3/test_alignment_model_wrappers.py", diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py index d43e5c70e..134c82a07 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: @@ -156,6 +161,31 @@ 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"} + 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": outs, + "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")): @@ -169,14 +199,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/h3_chain_replay.py b/tools/validation/models/h3_chain_replay.py index 217db85e1..5c3d8f110 100644 --- a/tools/validation/models/h3_chain_replay.py +++ b/tools/validation/models/h3_chain_replay.py @@ -30,12 +30,20 @@ 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", 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 +57,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 +75,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 +96,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/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index 5aa5e4f73..ed8267a94 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -31,7 +31,11 @@ h3_weights_source, measure, ) -from rl_engine.validation.models.h3_weights import load_h3_manifest # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + WEIGHTS_ENV, + h3_weights_dir, + load_h3_manifest, +) _NEEDS_PINNED_MLP = ( "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " @@ -68,6 +72,8 @@ def require_pinned_weights(op: str, allow_random: bool) -> None: 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) @@ -79,6 +85,10 @@ def main() -> None: ) args = parser.parse_args() require_pinned_weights(args.op, args.allow_random_weights) + uses_weights = args.op != "timestep_sinusoid_h3" + smoke_run = args.allow_random_weights and args.op == "timestep_mlp_fp32" + if uses_weights and h3_weights_dir() is None and not smoke_run: + raise SystemExit(f"set {WEIGHTS_ENV} to the tools/weights/prepare_h3_weights.py output dir") if not torch.cuda.is_available(): raise SystemExit("needs a CUDA device") torch.backends.cuda.matmul.allow_tf32 = False @@ -91,6 +101,11 @@ def main() -> None: "op": args.op, "rfc": manifest["rfc"], "model_revision": manifest["model_identity"]["revision"], + "weight_source": ( + ("pinned_checkpoint" if h3_weights_source() == "pinned" else "random") + if uses_weights + else "not_applicable" + ), "reference_commit": manifest["reference_implementation"]["commit"], **git_state(), "environment": environment(), 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..777b1edbf 100644 --- a/tools/validation/models/plot_h3_bitwise_reuse.py +++ b/tools/validation/models/plot_h3_bitwise_reuse.py @@ -27,12 +27,15 @@ 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 - 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: @@ -59,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() diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py index ecea9ab1e..ab9af46d9 100644 --- a/tools/validation/models/plot_h3_evidence.py +++ b/tools/validation/models/plot_h3_evidence.py @@ -36,6 +36,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) @@ -48,7 +50,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 @@ -76,10 +78,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") @@ -95,6 +101,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") @@ -156,6 +164,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, @@ -172,7 +182,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], @@ -221,6 +231,8 @@ def mlp_title(report: dict[str, Any]) -> str: 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)") # name the weights the report was measured with instead of always claiming the pinned ones fig.suptitle(mlp_title(report), x=0.01, ha="left", fontsize=12, color=INK, fontweight="bold") @@ -238,13 +250,55 @@ def plot_mlp(report: dict[str, Any]): return fig +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"] + _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, } 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/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py index c21f56c2d..5c6f163f6 100644 --- a/tools/weights/prepare_h3_weights.py +++ b/tools/weights/prepare_h3_weights.py @@ -27,10 +27,13 @@ EXTRACTED_FILE, load_h3_manifest, sha256_file, + sha256_tensor, ) 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}") @@ -38,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( @@ -59,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"], @@ -90,6 +97,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"]})