From 418998c1a8c76731fd9e70efe7625e5912708608 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:32:38 +0000 Subject: [PATCH 01/19] [WS1][CUDA][MiniMax-H3] timestep_mlp_fp32: deterministic FP32 timestep MLP RFC #420 step 3, row `timestep_mlp_fp32` (stacked on timestep_sinusoid_h3). - New CUDA module csrc/cuda/h3/det_linear.cu, contract h3-det-linear-v1: warp per output-column group, lane-strided 16-byte K chunks, ascending fmaf chains, xor butterfly, bias then fused SiLU. The launch shape never changes a column's summation order, so rows are batch/position invariant and runs are repeat-bitwise. - Backward without cuBLAS or atomics: dinput via fixed 64-row N chunks folded in order (row-local), dweight/dbias as ascending-row FP32 folds. - PyTorch reference (provider path + FP64 golden); BF16 anywhere is rejected (RFC probe H7: the time_embedder is a declared FP32 module). - On the pinned weights both CUDA and cuBLAS are ~1e-6 from the FP64 golden (100x inside the contract); the op is 2x faster for T >= 2. - Chain stage (covered by the end-to-end test), report/figure hooks, registry, gtest, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 476 ++++++++++++++++++ csrc/ops.cpp | 18 + docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-timestep-mlp.md | 128 +++++ rl_engine/_C.pyi | 13 + rl_engine/kernels/gtest/operator_inputs.py | 25 + rl_engine/kernels/gtest/operator_specs.py | 11 + rl_engine/kernels/ops/cuda/h3/det_linear.py | 79 +++ rl_engine/kernels/ops/cuda/h3/timestep_mlp.py | 99 ++++ .../kernels/ops/pytorch/h3/timestep_mlp.py | 77 +++ rl_engine/kernels/registry.py | 9 + rl_engine/testing/h3_chain.py | 17 +- rl_engine/testing/h3_provider.py | 15 + rl_engine/testing/h3_report.py | 71 ++- rl_engine/tests/test_dispatch.py | 2 +- scripts/plot_h3_evidence.py | 71 +++ setup.py | 1 + tests/h3/test_h3_timestep_mlp.py | 193 +++++++ tests/test_ws1_gtest_gpu.py | 1 + 20 files changed, 1305 insertions(+), 3 deletions(-) create mode 100644 csrc/cuda/h3/det_linear.cu create mode 100644 docs/operators/h3-timestep-mlp.md create mode 100644 rl_engine/kernels/ops/cuda/h3/det_linear.py create mode 100644 rl_engine/kernels/ops/cuda/h3/timestep_mlp.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py create mode 100644 tests/h3/test_h3_timestep_mlp.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu new file mode 100644 index 000000000..cc82acafa --- /dev/null +++ b/csrc/cuda/h3/det_linear.cu @@ -0,0 +1,476 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Deterministic, batch-invariant row-wise linear layers for the MiniMax-H3 +// conditioning path (RFC #420: `timestep_mlp_fp32`, `adaln_projection_3mod`). +// +// These layers run on a handful of rows (one per distinct timestep) against +// large weights, so they are GEMVs bound by weight bandwidth. The reduction +// order is fixed and depends only on K (and N for the input gradient), never +// on the number of rows or their position. +// +// Reduction contract h3-det-linear-v1 +// ----------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * one warp per output column n; lane l owns the 16-byte K chunks +// c = l, l + 32, l + 64, ... (VEC = 16 / sizeof(w) elements each); +// * each lane accumulates its chunks in ascending order and the elements of +// a chunk in ascending order with fmaf into an FP32 accumulator from 0; +// * the 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1), +// 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. +// 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 +// ascending chunk order and cast once. +// d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] +// * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one +// cross-row reduction; its order is the logical row order. + +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +#ifndef H3_FWD_F32_COLS +#define H3_FWD_F32_COLS 2 +#define H3_FWD_F32_AHEAD 6 +#define H3_FWD_BF16_COLS 4 +#define H3_FWD_BF16_AHEAD 2 +#endif +#ifndef H3_FWD_CONFIGS +#define H3_FWD_CONFIGS H3_FWD_CASE(H3_FWD_F32_COLS, H3_FWD_F32_AHEAD) \ + H3_FWD_CASE(H3_FWD_BF16_COLS, H3_FWD_BF16_AHEAD) +#endif + +constexpr int kWarp = 32; +constexpr int kWarpsPerBlock = 8; +constexpr int kRowTile = 4; +constexpr int kDInputChunk = 64; + +enum Activation : int64_t { kActNone = 0, kActSilu = 1 }; + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ float from_float(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); // round to nearest even, like at::BFloat16 +} + +// 16-byte vector load converted to FP32. +template +struct Vec16 { + static constexpr int kN = 16 / sizeof(T); + __device__ __forceinline__ static void load(const T* ptr, float (&out)[kN]) { + unpack(load_raw(ptr), out); + } + __device__ __forceinline__ static uint4 load_raw(const T* ptr) { + return *reinterpret_cast(ptr); + } + // Streaming weights are read once: bypass L1 allocation. + __device__ __forceinline__ static uint4 load_stream(const T* ptr) { + return __ldcs(reinterpret_cast(ptr)); + } + __device__ __forceinline__ static void unpack(const uint4& raw, float (&out)[kN]) { + const T* vals = reinterpret_cast(&raw); +#pragma unroll + for (int i = 0; i < kN; ++i) out[i] = to_float(vals[i]); + } +}; + +__device__ __forceinline__ float warp_butterfly_sum(float v) { +#pragma unroll + for (int offset = kWarp / 2; offset > 0; offset >>= 1) { + v += __shfl_xor_sync(0xffffffffu, v, offset); + } + return v; +} + +// Each warp owns kCols adjacent output columns so that kLoadAhead x kCols +// 16-byte weight loads are in flight per lane. Grouping columns and chunks +// changes only when loads are issued: every column still accumulates its own +// chunks in ascending order, exactly as a one-column-per-warp kernel would. +template +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_kernel(const x_t* __restrict__ x, const w_t* __restrict__ w, + const w_t* __restrict__ bias, out_t* __restrict__ out, + float* __restrict__ pre_act, int64_t rows, int64_t n_out, + int64_t k_in, int64_t activation) { + using WV = Vec16; + using XV = Vec16; + constexpr int kVec = WV::kN; + static_assert(XV::kN == kVec, "x and weight share a dtype"); + constexpr int64_t kStride = static_cast(kWarp) * kVec; + + const int lane = threadIdx.x % kWarp; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + + for (int64_t n0 = warp_global * kCols; n0 < n_out; n0 += warp_stride * kCols) { + for (int64_t t0 = 0; t0 < rows; t0 += kRows) { + float acc[kCols][kRows]; +#pragma unroll + for (int c = 0; c < kCols; ++c) { +#pragma unroll + for (int r = 0; r < kRows; ++r) acc[c][r] = 0.0f; + } + for (int64_t g0 = static_cast(lane) * kVec; g0 < k_in; + g0 += kStride * kLoadAhead) { + uint4 raw[kLoadAhead][kCols]; +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; +#pragma unroll + for (int c = 0; c < kCols; ++c) { + if (k0 < k_in && n0 + c < n_out) raw[u][c] = WV::load_stream(w + (n0 + c) * k_in + k0); + } + } +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; + if (k0 >= k_in) break; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + if (t0 + r < rows) { + float xv[kVec]; + XV::load(x + (t0 + r) * k_in + k0, xv); +#pragma unroll + for (int c = 0; c < kCols; ++c) { + float wv[kVec]; + WV::unpack(raw[u][c], wv); +#pragma unroll + for (int j = 0; j < kVec; ++j) acc[c][r] = fmaf(xv[j], wv[j], acc[c][r]); + } + } + } + } + } +#pragma unroll + for (int c = 0; c < kCols; ++c) { + const int64_t n = n0 + c; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + const int64_t t = t0 + r; + const float sum = warp_butterfly_sum(acc[c][r]); + if (lane == c * kRows + r && t < rows && n < n_out) { + float v = sum + (bias != nullptr ? to_float(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] = from_float(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, + const w_t* __restrict__ w, + float* __restrict__ partial, int64_t rows, + int64_t n_out, int64_t k_in) { + const int64_t k = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t chunk = blockIdx.y; + if (k >= k_in) return; + const int64_t n_begin = chunk * kDInputChunk; + const int64_t n_end = min(n_begin + kDInputChunk, n_out); + for (int64_t t0 = 0; t0 < rows; t0 += kRowTile) { + float acc[kRowTile]; +#pragma unroll + for (int r = 0; r < kRowTile; ++r) acc[r] = 0.0f; + for (int64_t n = n_begin; n < n_end; ++n) { + const float wv = to_float(w[n * k_in + k]); +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) acc[r] = fmaf(to_float(grad[t * n_out + n]), wv, acc[r]); + } + } +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) partial[(chunk * rows + t) * k_in + k] = acc[r]; + } + } +} + +template +__global__ void det_fold_chunks_kernel(const float* __restrict__ partial, out_t* __restrict__ out, + int64_t chunks, int64_t elems) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= elems) return; + float acc = partial[i]; + for (int64_t c = 1; c < chunks; ++c) acc += partial[c * elems + i]; + out[i] = from_float(acc); +} + +// dW[n, k] = sum_t g[t, n] * x[t, k] ; ascending t, FP32 from 0, one cast. +template +__global__ void det_linear_dweight_kernel(const g_t* __restrict__ grad, const x_t* __restrict__ x, + w_t* __restrict__ dw, int64_t rows, int64_t n_out, + int64_t k_in) { + const int64_t total = n_out * k_in; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < total; + i += static_cast(gridDim.x) * blockDim.x) { + const int64_t n = i / k_in; + const int64_t k = i - n * k_in; + float acc = 0.0f; + for (int64_t t = 0; t < rows; ++t) { + acc = fmaf(to_float(grad[t * n_out + n]), to_float(x[t * k_in + k]), acc); + } + dw[i] = from_float(acc); + } +} + +template +__global__ void det_linear_dbias_kernel(const g_t* __restrict__ grad, w_t* __restrict__ db, + int64_t rows, int64_t n_out) { + const int64_t n = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (n >= n_out) return; + float acc = 0.0f; + for (int64_t t = 0; t < rows; ++t) acc += to_float(grad[t * n_out + n]); + db[n] = from_float(acc); +} + +template +struct CudaType { + using type = T; +}; +template <> +struct CudaType { + using type = __nv_bfloat16; +}; + +template +const typename CudaType::type* cptr(const torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} +template +typename CudaType::type* mptr(torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} + +void check_matrix(const torch::Tensor& t, const char* name) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.dim() == 2, name, " must be 2-D"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(t.scalar_type() == at::kFloat || t.scalar_type() == at::kBFloat16, name, + " must be float32 or bfloat16, got ", t.scalar_type()); + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, name, + " must be 16-byte aligned"); +} + +int64_t grid_for(int64_t work, int64_t per_block) { + const int64_t blocks = (work + per_block - 1) / per_block; + return std::max(1, std::min(blocks, 1 << 20)); +} + +// Launch shape only: the row tile, the columns per warp and the load depth +// change which loads are in flight, never a column's accumulation order. +struct ForwardConfig { + int rows; + int cols; + int ahead; +}; + +ForwardConfig pick_forward_config(int64_t rows, int64_t elem_size) { + 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}; +} + +template +void launch_forward_impl(const T* x, const T* w, const T* bias, T* out, float* pre, + int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + const int64_t blocks = grid_for((n_out + kCols - 1) / kCols, kWarpsPerBlock); + det_linear_forward_kernel + <<(blocks), kWarpsPerBlock * kWarp, 0, stream>>>( + x, w, bias, out, pre, rows, n_out, k_in, activation); +} + +template +void launch_forward_rows(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation, cudaStream_t stream) { +#define H3_FWD_CASE(C, A) \ + if (cfg.cols == C && cfg.ahead == A) { \ + launch_forward_impl(x, w, bias, out, pre, rows, n_out, k_in, activation, \ + stream); \ + return; \ + } + H3_FWD_CONFIGS +#undef H3_FWD_CASE + TORCH_CHECK(false, "unsupported det_linear forward config cols=", cfg.cols, + " ahead=", cfg.ahead); +} + +template +void launch_forward(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + if (cfg.rows == 1) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else if (cfg.rows == 2) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } +} + +} // namespace + +std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor weight, + c10::optional bias, + int64_t activation, bool save_pre_activation) { + check_matrix(x, "x"); + check_matrix(weight, "weight"); + TORCH_CHECK(x.device() == weight.device(), "x and weight must be on the same device"); + TORCH_CHECK(x.scalar_type() == weight.scalar_type(), + "x and weight must share a dtype (cast at the declared boundary first), got ", + x.scalar_type(), " and ", weight.scalar_type()); + TORCH_CHECK(activation == kActNone || activation == kActSilu, "unknown activation ", + activation); + const int64_t rows = x.size(0); + const int64_t k_in = x.size(1); + const int64_t n_out = weight.size(0); + TORCH_CHECK(rows > 0, "x must have at least one row"); + TORCH_CHECK(weight.size(1) == k_in, "weight is [", n_out, ", ", weight.size(1), + "] but x has K=", k_in); + const int64_t vec = 16 / x.element_size(); + TORCH_CHECK(k_in % vec == 0, "K=", k_in, " must be a multiple of ", vec, + " for 16-byte vector loads"); + if (bias.has_value()) { + TORCH_CHECK(bias->is_cuda() && bias->dim() == 1 && bias->size(0) == n_out && + bias->is_contiguous() && bias->scalar_type() == weight.scalar_type(), + "bias must be a contiguous [N] CUDA tensor with the weight dtype"); + } + + const c10::cuda::CUDAGuard device_guard(x.device()); + auto out = torch::empty({rows, n_out}, x.options()); + 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) { + 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), + bias.has_value() ? cptr(*bias) : nullptr, mptr(out), + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation, + stream); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (save_pre_activation) return {out, pre}; + return {out}; +} + +// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype. +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(weight, "weight"); + TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N"); + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + const int64_t rows = grad.size(0); + const int64_t n_out = weight.size(0); + const int64_t k_in = weight.size(1); + const c10::cuda::CUDAGuard device_guard(grad.device()); + const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; + auto partial = torch::empty({chunks, rows, k_in}, grad.options()); + auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); + if (weight.scalar_type() == at::kFloat) { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, n_out, + k_in); + } else { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, + n_out, k_in); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const int64_t elems = rows * k_in; + const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); + if (out_dtype == at::kFloat) { + det_fold_chunks_kernel<<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } else { + det_fold_chunks_kernel<__nv_bfloat16><<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// grad [T, N] float32, x [T, K] (float32 or bfloat16) -> (dW [N, K], dbias [N]) in w_dtype. +std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, + c10::ScalarType w_dtype, + bool with_bias) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(x, "x"); + TORCH_CHECK(grad.size(0) == x.size(0), "grad rows != x rows"); + TORCH_CHECK(w_dtype == at::kFloat || w_dtype == at::kBFloat16, + "w_dtype must be float32 or bfloat16"); + const int64_t rows = grad.size(0); + 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)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const int64_t total = n_out * k_in; + const unsigned blocks = static_cast(grid_for(total, threads)); + const float* g = grad.data_ptr(); +#define H3_DW_LAUNCH(XT, XC, WT, WC) \ + det_linear_dweight_kernel<<>>( \ + g, cptr(x), mptr(dw), rows, n_out, k_in) + if (x.scalar_type() == at::kFloat && w_dtype == at::kFloat) { + H3_DW_LAUNCH(float, float, float, float); + } else if (x.scalar_type() == at::kFloat) { + H3_DW_LAUNCH(float, float, at::BFloat16, __nv_bfloat16); + } else if (w_dtype == at::kFloat) { + H3_DW_LAUNCH(at::BFloat16, __nv_bfloat16, float, float); + } else { + H3_DW_LAUNCH(at::BFloat16, __nv_bfloat16, at::BFloat16, __nv_bfloat16); + } +#undef H3_DW_LAUNCH + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!with_bias) return {dw}; + auto db = torch::empty({n_out}, grad.options().dtype(w_dtype)); + const unsigned bias_blocks = static_cast((n_out + threads - 1) / threads); + if (w_dtype == at::kFloat) { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } else { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dw, db}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 71b5b5502..f245f8e15 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -502,6 +502,14 @@ at::Tensor prefix_shared_attention( (defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)) torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, double max_period); +std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor weight, + c10::optional bias, + int64_t activation, bool save_pre_activation); +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype); +std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, + c10::ScalarType w_dtype, + bool with_bias); #endif // PyBind11 Module Registration @@ -765,5 +773,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("h3_timestep_sinusoid_forward", &h3_timestep_sinusoid_forward, "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0); + m.def("h3_det_linear_forward", &h3_det_linear_forward, + "Batch-invariant warp-per-column linear (contract h3-det-linear-v1), optional SiLU", + py::arg("x"), py::arg("weight"), py::arg("bias") = py::none(), + py::arg("activation") = 0, py::arg("save_pre_activation") = false); + m.def("h3_det_linear_backward_input", &h3_det_linear_backward_input, + "Deterministic grad @ weight with fixed 64-row N chunks folded in order", + py::arg("grad"), py::arg("weight"), py::arg("out_dtype")); + m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight, + "Deterministic dW/dbias as ascending-row FP32 folds", + py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index d4c01f6c5..0dbcaec7d 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -31,6 +31,7 @@ nav: - operators/embedding.md - MiniMax-H3: - operators/h3-timestep-sinusoid.md + - operators/h3-timestep-mlp.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index dee8796e4..ac127f3c4 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -33,4 +33,5 @@ Every operator page should include: - [Sampling](sampling.md) - [Token Embedding](embedding.md) - [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) +- [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md new file mode 100644 index 000000000..d5077f3f9 --- /dev/null +++ b/docs/operators/h3-timestep-mlp.md @@ -0,0 +1,128 @@ +# MiniMax-H3 FP32 Timestep MLP + +## Summary + +`timestep_mlp_fp32` is H3's `TimestepEmbedding(in_channels=256, time_embed_dim=5376, +out_dim=2688)`. It runs on the sinusoidal features of the distinct timesteps +(RFC #420, WS1 step 3): + +```text +temb = linear_2(silu(linear_1(features))) (T, 256) -> (T, 5376) -> (T, 2688) +``` + +`time_embedder` is a `_keep_in_fp32_modules` module. Its weights, biases, +activations and output are all FP32. `temb` stays FP32 because every AdaLN +projection applies its own SiLU before casting to BF16. + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227`. The tensors are +`time_embedder.linear_{1,2}.{weight,bias}` (see `rl_engine/testing/h3_manifest.json`). + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("timestep_mlp_fp32", device="cuda") +temb = op(features, w1, b1, w2, b2) # all float32 +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.kernels.ops.cuda.h3.timestep_mlp.H3TimestepMLPCudaOp` | `rl_engine._C.h3_det_linear_{forward,backward_input,backward_weight}` | Contract `h3-det-linear-v1` | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp` | n/a | `forward`: provider path (`F.linear`/`F.silu`); `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `x` | `(T, K)`, `T >= 1` | float32 | H3: `K = 256` | +| `w1`, `b1` | `(H, K)`, `(H,)` | float32 | H3: `H = 5376` | +| `w2`, `b2` | `(D, H)`, `(D,)` | float32 | H3: `D = 2688` | +| output | `(T, D)` | float32 | | + +The CUDA kernel needs `K` and `H` to be multiples of 4, for 16-byte vector loads. +These inputs fail closed: BF16 in any argument (RFC probe H7), mismatched shapes, +an empty `x`, mixed devices, and unaligned `K`. + +## Numerics: contract `h3-det-linear-v1` + +Forward, for each output `y[t, n]`: + +- One warp owns a group of output columns. Lane `l` reads the 16-byte K chunks + `l, l + 32, l + 64, ...`. +- Each lane accumulates its chunks in ascending order, and the elements inside a chunk + in ascending order, using `fmaf` into an FP32 accumulator that starts at 0. +- The 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1). +- The bias is added once, then SiLU runs in FP32 as `v / (1 + expf(-v))` + (PyTorch's formula). The result is stored as FP32. + +How many columns a warp owns, the row tile (1, 2 or 4), and how many chunks are +loaded ahead change only when loads are issued, never the order in which a column +is summed. A timestep's `temb` is therefore bitwise identical however many other +timesteps share the call, and wherever it sits among them. + +Backward is deterministic. There is no cuBLAS and there are no atomics: + +- `d_hidden = g @ W2` and `dx = d_pre @ W1`: N is split into fixed 64-row chunks. + Each chunk is an ascending `fmaf` chain, and the chunks are then left-folded in + ascending order. These gradients are row-local and batch-invariant. +- `d_pre = d_hidden * s * (1 + z * (1 - s))`, elementwise in FP32, using the saved + pre-activation `z`. +- `dW[n, k] = sum_t g[t, n] * x[t, k]` and `db[n] = sum_t g[t, n]`: an ascending-`t` + FP32 fold that starts from 0, in logical row order. + +Accuracy, measured on a B200 with the pinned checkpoint weights, TF32 disabled and 200 +draws of 4 timesteps (each draw includes t = 0 and t = 1): + +| Comparison | CUDA | provider (cuBLAS) | +| --- | --- | --- | +| max abs error vs FP64 golden, median over draws | 1.1e-6 | 5.4e-7 | +| gradients vs FP64 autograd (`x`, `w1`, `b1`, `w2`, `b2`) | <= 1.2e-5 | | + +Both outputs are about 100× inside the contract (`reduction` / `float32`, atol and rtol +1e-4). Neither is more accurate in general: cuBLAS's tree happens to do slightly better +on these weights. What the CUDA kernel adds is a fixed summation order, which makes each +row batch- and position-invariant and every run repeat-bitwise, and it is 2× faster +for T >= 2. The two outputs are not bitwise equal, so `scripts/h3_chain_replay.py` +reports this stage as the chain's `first_drift`, which is expected for a reduction. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op timestep_mlp_fp32 +``` + +B200, pinned weights (63.3 MB FP32). End-to-end op times include the Python wrapper: + +| T | CUDA op | provider | kernels only (CUDA) | +| --- | --- | --- | --- | +| 1 | 32.8 µs | 36.5 µs | 3.8 + 11.6 µs | +| 2 | 38.8 µs | 79.9 µs | 3.9 + 13.7 µs | +| 4 | 42.7 µs | 83.1 µs | 5.0 + 15.8 µs | + +The 5376→2688 layer reads its weights at about 5.0 TB/s at T = 1. For T >= 2 the +provider switches from cuBLAS GEMV to an SGEMM path, which takes about 52 µs of +kernel time. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_timestep_mlp.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end: sinusoid -> MLP +python scripts/check_operator.py --op timestep_mlp_fp32 --candidate cuda --device cuda \ + --dtype fp32 --batch 3 --check-grad +python scripts/h3_evidence.py --op timestep_mlp_fp32 \ + --out docs/usage/evidence/h3-timestep-mlp-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-timestep-mlp-b200/report.json +``` + +## Known Limitations + +- FP32 only, by contract. The BF16 and FP16 gtest rows do not apply to this op. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- Weight gradients are reductions across timesteps, so they follow logical row + order. They are deterministic, but not batch-invariant (no cross-row reduction can be). diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 970df8515..b8cf1d58a 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -328,3 +328,16 @@ def deterministic_collective_rocm_ipc_all_gather_input( def h3_timestep_sinusoid_forward( timestep: torch.Tensor, num_channels: int = 256, max_period: float = 10000.0 ) -> torch.Tensor: ... +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]: ... +def h3_det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: ... +def h3_det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 4562f2aad..435d42023 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -17,6 +17,8 @@ DEFAULT_ROPE_THETA = 1.0e6 DEFAULT_RMS_EPS = 1.0e-6 H3_FREQ_DIM = 256 +H3_TIME_HIDDEN = 5376 +H3_TIME_EMBED = 2688 def make_operator_inputs( @@ -44,6 +46,7 @@ def make_operator_inputs( "lm_head": _make_lm_head_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, } try: return builders[op_name](args, dtype, device) @@ -75,6 +78,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", "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}", } try: return names[op_name] @@ -360,6 +365,26 @@ def _make_timestep_sinusoid_h3_inputs( 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]: + # The H3 time_embedder is declared FP32; ``dtype`` does not apply to it. + del dtype + from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import NativeH3TimestepSinusoidOp + + features = NativeH3TimestepSinusoidOp().forward(_h3_timesteps(args, torch.float32, device)) + scale1, scale2 = H3_FREQ_DIM**-0.5, H3_TIME_HIDDEN**-0.5 + return { + "x": features, + "w1": _floating_tensor((H3_TIME_HIDDEN, H3_FREQ_DIM), args, torch.float32, device, 1) + * scale1, + "b1": _floating_tensor((H3_TIME_HIDDEN,), args, torch.float32, device, 2) * 0.1, + "w2": _floating_tensor((H3_TIME_EMBED, H3_TIME_HIDDEN), args, torch.float32, device, 3) + * scale2, + "b2": _floating_tensor((H3_TIME_EMBED,), args, torch.float32, device, 4) * 0.1, + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index cc740d0ff..cefebd46f 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -261,6 +261,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("timestep",), ), + "timestep_mlp_fp32": OperatorSpec( + name="timestep_mlp_fp32", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp", + "cuda": "rl_engine.kernels.ops.cuda.h3.timestep_mlp.H3TimestepMLPCudaOp", + }, + grad_input_names=("x", "w1", "b1", "w2", "b2"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/det_linear.py b/rl_engine/kernels/ops/cuda/h3/det_linear.py new file mode 100644 index 000000000..f5dab854e --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/det_linear.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Python surface of the H3 deterministic row-wise linear (contract h3-det-linear-v1). + +See ``csrc/cuda/h3/det_linear.cu`` for the reduction order. Every helper +raises instead of falling back to cuBLAS when the extension is missing. +""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + +ACT_NONE = 0 +ACT_SILU = 1 +CONTRACT = "h3-det-linear-v1" +_SYMBOLS = ( + "h3_det_linear_forward", + "h3_det_linear_backward_input", + "h3_det_linear_backward_weight", +) + + +def det_linear_available() -> bool: + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +def _require() -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks the h3_det_linear_* symbols; rebuild the CUDA extension " + "with csrc/cuda/h3/det_linear.cu (no cuBLAS fallback is allowed)" + ) + + +def det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + *, + activation: int = ACT_NONE, + save_pre_activation: bool = False, +) -> list[torch.Tensor]: + _require() + return _C.h3_det_linear_forward( + x.contiguous(), + weight.contiguous(), + None if bias is None else bias.contiguous(), + int(activation), + bool(save_pre_activation), + ) + + +def det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: + _require() + return _C.h3_det_linear_backward_input( + grad.float().contiguous(), weight.contiguous(), out_dtype + ) + + +def det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True +) -> list[torch.Tensor]: + _require() + return _C.h3_det_linear_backward_weight( + grad.float().contiguous(), x.contiguous(), w_dtype, bool(with_bias) + ) + + +def silu_backward_fp32(grad: torch.Tensor, pre_activation: torch.Tensor) -> torch.Tensor: + """Elementwise SiLU VJP in FP32: g * s * (1 + z * (1 - s)), s = sigmoid(z).""" + + z = pre_activation.float() + s = torch.sigmoid(z) + return grad.float() * (s * (1.0 + z * (1.0 - s))) diff --git a/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py b/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py new file mode 100644 index 000000000..da85e87d8 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/timestep_mlp.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 FP32 timestep MLP (RFC #420 ``timestep_mlp_fp32``).""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + ACT_NONE, + ACT_SILU, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.pytorch.h3.timestep_mlp import validate_h3_timestep_mlp + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[silu]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.kernels.ops.cuda.h3.det_linear.silu_backward_fp32" +) + + +class _H3TimestepMLPCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w1, b1, w2, b2): + 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) + ctx.save_for_backward(x, w1, w2, pre, hidden) + return out + + @staticmethod + def backward(ctx, grad_out): + 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() + dw2 = db2 = dw1 = db1 = dx = None + if need_w2 or need_b2: + dw2, db2 = det_linear_backward_weight(grad_out, hidden, torch.float32) + d_pre = None + if need_x or need_w1 or need_b1: + d_hidden = det_linear_backward_input(grad_out, w2, torch.float32) + d_pre = silu_backward_fp32(d_hidden, pre).contiguous() + if need_w1 or need_b1: + dw1, db1 = det_linear_backward_weight(d_pre, x, torch.float32) + if need_x: + dx = det_linear_backward_input(d_pre, w1, torch.float32) + record_backward( + "timestep_mlp_fp32", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return ( + dx, + dw1 if need_w1 else None, + db1 if need_b1 else None, + dw2 if need_w2 else None, + db2 if need_b2 else None, + ) + + +class H3TimestepMLPCudaOp: + """CUDA candidate: two warp-per-column GEMVs, SiLU fused into the first. + + The weights dominate the traffic (5.5 MB + 57.8 MB FP32), the rows are the + handful of distinct timesteps, and every output element follows contract + ``h3-det-linear-v1``, so a timestep's ``temb`` bytes do not depend on which + other timesteps share the call. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, x, w1, b1, w2, b2): + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + 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: + return self.forward(x, w1, b1, w2, b2) diff --git a/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py b/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py new file mode 100644 index 000000000..ac54d295f --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 FP32 timestep MLP (RFC #420 row ``timestep_mlp_fp32``). + +``TimestepEmbedding(in_channels=256, time_embed_dim=5376, out_dim=2688)``:: + + temb = linear_2(silu(linear_1(features))) 256 -> 5376 -> 2688 + +``time_embedder`` is in ``_keep_in_fp32_modules``: weights, biases, +activations and the output are FP32, and ``temb`` stays FP32 because every +AdaLN projection applies its own SiLU to it before casting (RFC #420 §4). +A BF16 argument is a contract violation (RFC probe H7), not a lower-precision +variant, so it is rejected. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def validate_h3_timestep_mlp( + x: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> None: + named = {"x": x, "w1": w1, "b1": b1, "w2": w2, "b2": b2} + for name, tensor in named.items(): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype != torch.float32: + raise TypeError( + f"{name} must be float32: the H3 time_embedder is an FP32 module " + f"(got {tensor.dtype})" + ) + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if x.dim() != 2 or x.shape[0] == 0: + raise ValueError(f"x must be a non-empty (T, K) matrix, got {tuple(x.shape)}") + hidden, k_in = w1.shape if w1.dim() == 2 else (None, None) + if k_in != x.shape[1]: + raise ValueError(f"w1 must be (H, {x.shape[1]}), got {tuple(w1.shape)}") + if b1.shape != (hidden,): + raise ValueError(f"b1 must be ({hidden},), got {tuple(b1.shape)}") + if w2.dim() != 2 or w2.shape[1] != hidden: + raise ValueError(f"w2 must be (D, {hidden}), got {tuple(w2.shape)}") + if b2.shape != (w2.shape[0],): + raise ValueError(f"b2 must be ({w2.shape[0]},), got {tuple(b2.shape)}") + + +class NativeH3TimestepMLPOp: + """PyTorch reference for the H3 timestep MLP. + + ``forward`` is the provider path (``nn.Linear`` / ``F.silu`` in FP32; run + it with TF32 disabled, as the RL-Kernel contract requires). + ``forward_fp32`` is the independent golden: the same graph in FP64, + rounded once to FP32, so its own rounding error (~1e-16 relative) is far + below the FP32 contract and its reduction order does not matter. + """ + + op_class = "reduction" + + def __call__(self, x, w1, b1, w2, b2): + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + 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: + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + z = F.linear(x.double(), w1.double(), b1.double()) + h = z * torch.sigmoid(z) + return F.linear(h, w2.double(), b2.double()).float() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 533536f57..e3b3765d5 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -207,6 +207,8 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_TIMESTEP_SINUSOID = ( "rl_engine.kernels.ops.cuda.h3.timestep_sinusoid.H3TimestepSinusoidCudaOp" ) + PYTORCH_H3_TIMESTEP_MLP = "rl_engine.kernels.ops.pytorch.h3.timestep_mlp.NativeH3TimestepMLPOp" + CUDA_H3_TIMESTEP_MLP = "rl_engine.kernels.ops.cuda.h3.timestep_mlp.H3TimestepMLPCudaOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -621,6 +623,10 @@ def __init__(self): OpBackend.CUDA_H3_TIMESTEP_SINUSOID, OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID, ], + "timestep_mlp_fp32": [ + OpBackend.CUDA_H3_TIMESTEP_MLP, + OpBackend.PYTORCH_H3_TIMESTEP_MLP, + ], }, "rocm": { "logp": [ @@ -662,6 +668,7 @@ def __init__(self): "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -707,6 +714,7 @@ def __init__(self): "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -733,6 +741,7 @@ def __init__(self): "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index fb1f10943..a2c6a80c5 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -30,7 +30,7 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing.h3_cases import h3_packed_layout, h3_timesteps -from rl_engine.testing.h3_provider import provider_time_proj +from rl_engine.testing.h3_provider import provider_time_embedder, provider_time_proj REPO_ROOT = Path(__file__).resolve().parents[2] @@ -49,6 +49,11 @@ class Stage: golden_rtol: float +def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: + weights = ctx["weights"] + return [weights[f"time_embedder.linear_{i}.{p}"] for i in (1, 2) for p in ("weight", "bias")] + + STAGES: list[Stage] = [ Stage( name="timestep_sinusoid_h3", @@ -60,6 +65,16 @@ class Stage: golden_atol=1e-5, # elementwise / float32 golden_rtol=1e-5, ), + Stage( + name="timestep_mlp_fp32", + op_type="timestep_mlp_fp32", + candidate=lambda op, ctx, up: op(up, *mlp_params(ctx)), + provider=lambda ctx, up: provider_time_embedder(up, *mlp_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *mlp_params(ctx)), + provider_bitwise_isolated=False, # reduction: different tree from cuBLAS + golden_atol=1e-4, # reduction / float32 + golden_rtol=1e-4, + ), ] diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index ff4718c00..6ae2c5790 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -15,6 +15,7 @@ import math import torch +import torch.nn.functional as F def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch.Tensor: @@ -31,3 +32,17 @@ def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) return emb + + +def provider_time_embedder( + features: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> torch.Tensor: + """``TimestepEmbedding(256, 5376, out_dim=2688)``: linear_1 -> SiLU -> linear_2, FP32.""" + + sample = F.linear(features.to(w1.dtype), w1, b1) + sample = F.silu(sample) + return F.linear(sample, w2, b2) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index bd5f9c9a6..175547a21 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -17,7 +17,8 @@ from rl_engine.kernels.registry import KernelRegistry from rl_engine.testing.h3_cases import h3_timesteps -from rl_engine.testing.h3_provider import provider_time_proj +from rl_engine.testing.h3_provider import provider_time_embedder, provider_time_proj +from rl_engine.testing.h3_weights import h3_weights_dir, load_h3_conditioning_weights # Keys of a perf case that hold a timed callable, in display order. TIMED_KEYS = ( @@ -111,9 +112,77 @@ def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: return {"contract_atol": 1e-5, "cases": rows} +def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Tensor]: + """Pinned checkpoint tensors when available, otherwise same-shape random ones.""" + + if h3_weights_dir() is not None: + weights = load_h3_conditioning_weights("cuda", names) + return [weights[name] for name in names] + generator = torch.Generator(device="cuda").manual_seed(0) + return [ + torch.randn(shape, device="cuda", generator=generator) * shape[-1] ** -0.5 + for shape in shapes + ] + + +# --------------------------------------------------------------------------- # +# timestep_mlp_fp32 +# --------------------------------------------------------------------------- # + +MLP_NAMES = [f"time_embedder.linear_{i}.{p}" for i in (1, 2) for p in ("weight", "bias")] +MLP_SHAPES = [(5376, 256), (5376,), (2688, 5376), (2688,)] + + +def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + 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) + cases = [] + for num in (1, 2, 3, 4): + x = torch.rand(num, 256, device="cuda") * 2 - 1 + cases.append( + { + "op": "timestep_mlp_fp32", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight_bytes, + "candidate": lambda x=x: op(x, *params), + "provider": lambda x=x: provider_time_embedder(x, *params), + } + ) + return cases + + +def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: + 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") + params = h3_params(MLP_NAMES, MLP_SHAPES) + cuda_err, provider_err = [], [] + for seed in range(draws): + x = sinusoid(h3_timesteps(4, seed=seed)) + gold = golden.forward_fp32(x, *params) + cuda_err.append(float((op(x, *params) - gold).abs().max())) + provider_err.append(float((provider_time_embedder(x, *params) - gold).abs().max())) + # Row invariance: each timestep alone equals its row in the batch of 9. + x = sinusoid(h3_timesteps(9, seed=1)) + full = op(x, *params) + invariant = all(torch.equal(op(x[i : i + 1], *params)[0], full[i]) for i in range(9)) + return { + "contract_atol": 1e-4, + "draws": draws, + "num_timesteps": 4, + "cuda_max_abs_vs_fp64": cuda_err, + "provider_max_abs_vs_fp64": provider_err, + "rows_batch_invariant": invariant, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, + "timestep_mlp_fp32": _mlp_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, + "timestep_mlp_fp32": _mlp_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index f416295d0..b436e3f2d 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -226,7 +226,7 @@ def fake_load_backend(backend): ] -H3_OPS = ("timestep_sinusoid_h3",) +H3_OPS = ("timestep_sinusoid_h3", "timestep_mlp_fp32") @pytest.mark.parametrize("op_name", H3_OPS) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index b0bcc8eb1..2baa21bb6 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -154,8 +154,79 @@ def plot_sinusoid(report: dict[str, Any]): return fig +def _latency_panel(ax, perf: list[dict[str, Any]], provider_name: str) -> None: + _style(ax, "Latency per call (lower is better)", "µs") + _bars( + ax, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + (provider_name, PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + + +def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], contract: float): + """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): + ax.boxplot( + [values], + positions=[i], + widths=0.45, + showfliers=False, + patch_artist=True, + boxprops={"facecolor": color, "alpha": 0.25, "edgecolor": color}, + medianprops={"color": color, "linewidth": 2}, + whiskerprops={"color": color}, + capprops={"color": color}, + ) + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax.scatter(jitter, values, s=8, color=color, alpha=0.6, linewidths=0, zorder=3) + median = sorted(values)[len(values) // 2] + ax.annotate( + f"median {median:.1e}", + (i + 0.28, median), + fontsize=8, + color=INK, + va="center", + ) + ax.axhline(contract, color=MUTED, lw=1, ls=":") + ax.annotate( + f"contract atol {contract:g}", + (len(series) - 0.5, contract), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + ax.set_yscale("log") + ax.set_xticks(range(len(series)), [name for name, _, _ in series]) + ax.set_xlim(-0.6, len(series) - 0.1) + + +def plot_mlp(report: dict[str, Any]): + fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _error_boxes( + right, + f"Error vs FP64 golden, {acc['draws']} draws of T={acc['num_timesteps']}", + [ + ("RL-Kernel CUDA", RL_KERNEL, acc["cuda_max_abs_vs_fp64"]), + ("diffusers (cuBLAS)", PROVIDER, acc["provider_max_abs_vs_fp64"]), + ], + acc["contract_atol"], + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, + "timestep_mlp_fp32": plot_mlp, } diff --git a/setup.py b/setup.py index 29f870cc9..b388bdd29 100644 --- a/setup.py +++ b/setup.py @@ -158,6 +158,7 @@ def get_extensions(): cuda_sources.append("csrc/cuda/distributed/deterministic_collective.cu") # MiniMax-H3 (RFC #420) conditioning path; plain-arch, runs on SM90 and SM100. cuda_sources.append("csrc/cuda/h3/timestep_sinusoid.cu") + cuda_sources.append("csrc/cuda/h3/det_linear.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py new file mode 100644 index 000000000..7e09da4ba --- /dev/null +++ b/tests/h3/test_h3_timestep_mlp.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_mlp_fp32``: FP32 256 -> 5376 -> 2688 timestep MLP. + +* accuracy: forward and every gradient against an FP64 golden, on the pinned + checkpoint weights and on synthetic non-H3 sizes; +* invariance: a timestep's ``temb`` bytes (and its row-local ``dx``) do not + depend on how many timesteps share the call or where it sits; +* determinism: repeated forward/backward runs are bitwise equal; +* contract: BF16 inputs (RFC probe H7) and malformed shapes fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3.timestep_mlp import NativeH3TimestepMLPOp +from rl_engine.kernels.ops.pytorch.h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.testing.h3_cases import h3_timesteps + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NAMES = ("x", "w1", "b1", "w2", "b2") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available + from rl_engine.kernels.ops.cuda.h3.timestep_mlp import H3TimestepMLPCudaOp + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3TimestepMLPCudaOp() + + +def _features(num, device="cuda", seed=0): + return NativeH3TimestepSinusoidOp().forward(h3_timesteps(num, seed=seed, device=device)) + + +def _synthetic_params(k_in, hidden, out, device="cuda", seed=0): + 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 + w2 = torch.randn(out, hidden, generator=g) / hidden**0.5 + b2 = torch.randn(out, generator=g) * 0.1 + return [t.to(device) for t in (w1, b1, w2, b2)] + + +def _real_params(weights): + 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): + 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()) + return [leaf.grad for leaf in leaves] + + +class TestReference: + def test_provider_and_golden_agree(self): + x = _features(5, device="cpu") + params = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + torch.testing.assert_close( + op.forward(x, *params), op.forward_fp32(x, *params), atol=1e-5, rtol=1e-5 + ) + + @pytest.mark.parametrize("which", range(5)) + def test_rejects_bf16_anywhere(self, which): + # RFC probe H7: casting a declared FP32 path to BF16 early. + 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): + x = _features(2, device="cpu") + w1, b1, w2, b2 = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + with pytest.raises(ValueError): + op.forward(x[:, :128], w1, b1, w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1[:10], w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1, w2[:, :10], b2) + with pytest.raises(ValueError): + op.forward(x[:0], w1, b1, w2, b2) + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4, 7]) + def test_forward_against_fp64_golden(self, h3_weights_cpu, num): + params = _real_params(h3_weights_cpu) + x = _features(num, seed=num) + out = _cuda_op()(x, *params) + gold = NativeH3TimestepMLPOp().forward_fp32(x, *params) + assert out.dtype == torch.float32 and out.shape == (num, 2688) + # tolerance_contract.json forward_accuracy / reduction / float32. + torch.testing.assert_close(out, gold, atol=1e-4, rtol=1e-4) + assert (out - gold).abs().max().item() < 1e-5 + + def test_rows_are_batch_and_position_invariant(self, h3_weights_cpu): + op = _cuda_op() + params = _real_params(h3_weights_cpu) + x_all = _features(9, seed=1) + full = op(x_all, *params) + for i in range(9): + assert torch.equal(op(x_all[i : i + 1], *params)[0], full[i]) + for num in (2, 3, 4, 5, 8): + for pos in (0, num - 1): + x = _features(num, seed=10 + num) + x[pos] = x_all[4] + assert torch.equal(op(x, *params)[pos], full[4]), (num, pos) + + def test_backward_against_fp64_golden(self, h3_weights_cpu): + params = _real_params(h3_weights_cpu) + x = _features(3, seed=2) + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + grad_out = torch.randn( + 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)): + # tolerance_contract.json gradient_accuracy / reduction / float32. + torch.testing.assert_close(leaf.grad.double(), ref, atol=1e-4, rtol=1e-4, msg=name) + + +@requires_cuda +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): + params = _synthetic_params(k_in, hidden, out, seed=k_in) + x = torch.randn(5, k_in, device="cuda") + torch.testing.assert_close( + _cuda_op()(x, *params), + NativeH3TimestepMLPOp().forward_fp32(x, *params), + atol=1e-4, + rtol=1e-4, + ) + + def test_rejects_unaligned_k(self): + 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): + params = _synthetic_params(8, 24, 16) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *[p.cpu() for p in params]) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *params) + + def test_repeat_and_backward_are_deterministic(self): + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=3) + x = _features(4, seed=3) + grad_out = torch.randn(4, 2688, device="cuda") + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + out = op(*leaves) + 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): + assert torch.equal(a, b) + + def test_dx_rows_are_batch_invariant(self): + 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): + leaf = x[rows].detach().clone().requires_grad_(True) + op(leaf, *params).backward(grad_out[rows]) + return leaf.grad + + full = dx(slice(0, 6)) + for i in range(6): + assert torch.equal(dx(slice(i, i + 1))[0], full[i]) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_mlp_fp32", device="cuda") + assert type(op).__name__ == "H3TimestepMLPCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index c688cd3d2..6dde75c53 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -50,6 +50,7 @@ def test_all_ws1_single_ops_are_registered(): "linear_logp", "prefix_shared_attention", "timestep_sinusoid_h3", + "timestep_mlp_fp32", } <= names From e32bb4ab8c946ff27ec40c7f37b5751d7e313b61 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:33:06 +0000 Subject: [PATCH 02/19] [WS1][MiniMax-H3] timestep_mlp_fp32: B200 evidence report and figure Produced by scripts/h3_evidence.py at 65ef7f6 on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index d5077f3f9..1da2434ac 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -107,6 +107,14 @@ The 5376→2688 layer reads its weights at about 5.0 TB/s at T = 1. For T >= 2 t provider switches from cuBLAS GEMV to an SGEMM path, which takes about 52 µs of kernel time. +## Evidence + +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../usage/evidence/h3-timestep-mlp-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-timestep-mlp-b200/report.json), written +by `scripts/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +timestep's row is bitwise identical whether it runs alone or in a batch of 9. + ## Tests ```bash From 74a7772d7dbddfd269263fb3fa8a5b86caded83a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:12:05 +0000 Subject: [PATCH 03/19] fix(h3): reject mixed CUDA devices in deterministic linear helpers Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 4 ++++ docs/operators/h3-timestep-mlp.md | 6 +++--- tests/h3/test_h3_timestep_mlp.py | 33 +++++++++++++++++++++++++++++++ 3 files changed, 40 insertions(+), 3 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index cc82acafa..ddaf199ee 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -358,6 +358,8 @@ std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor TORCH_CHECK(bias->is_cuda() && bias->dim() == 1 && bias->size(0) == n_out && bias->is_contiguous() && bias->scalar_type() == weight.scalar_type(), "bias must be a contiguous [N] CUDA tensor with the weight dtype"); + TORCH_CHECK(bias->device() == weight.device(), + "bias and weight must be on the same device"); } const c10::cuda::CUDAGuard device_guard(x.device()); @@ -390,6 +392,7 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei grad.scalar_type() == at::kFloat, "grad must be a contiguous 2-D float32 CUDA tensor"); check_matrix(weight, "weight"); + TORCH_CHECK(grad.device() == weight.device(), "grad and weight must be on the same device"); TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N"); TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, "out_dtype must be float32 or bfloat16"); @@ -434,6 +437,7 @@ std::vector h3_det_linear_backward_weight(torch::Tensor grad, tor grad.scalar_type() == at::kFloat, "grad must be a contiguous 2-D float32 CUDA tensor"); check_matrix(x, "x"); + TORCH_CHECK(grad.device() == x.device(), "grad and x must be on the same device"); TORCH_CHECK(grad.size(0) == x.size(0), "grad rows != x rows"); TORCH_CHECK(w_dtype == at::kFloat || w_dtype == at::kBFloat16, "w_dtype must be float32 or bfloat16"); diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 1da2434ac..a57f98033 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -99,9 +99,9 @@ B200, pinned weights (63.3 MB FP32). End-to-end op times include the Python wrap | T | CUDA op | provider | kernels only (CUDA) | | --- | --- | --- | --- | -| 1 | 32.8 µs | 36.5 µs | 3.8 + 11.6 µs | -| 2 | 38.8 µs | 79.9 µs | 3.9 + 13.7 µs | -| 4 | 42.7 µs | 83.1 µs | 5.0 + 15.8 µs | +| 1 | 35.5 µs | 40.8 µs | 3.8 + 11.6 µs | +| 2 | 40.4 µs | 84.9 µs | 3.9 + 13.7 µs | +| 4 | 44.9 µs | 88.6 µs | 5.0 + 15.8 µs | The 5376→2688 layer reads its weights at about 5.0 TB/s at T = 1. For T >= 2 the provider switches from cuBLAS GEMV to an SGEMM path, which takes about 52 µs of diff --git a/tests/h3/test_h3_timestep_mlp.py b/tests/h3/test_h3_timestep_mlp.py index 7e09da4ba..71f1e694f 100644 --- a/tests/h3/test_h3_timestep_mlp.py +++ b/tests/h3/test_h3_timestep_mlp.py @@ -191,3 +191,36 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("timestep_mlp_fp32", device="cuda") assert type(op).__name__ == "H3TimestepMLPCudaOp" + + +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +class TestCudaDetLinearDeviceContract: + def test_forward_rejects_bias_on_another_device(self, dtype): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_forward + + _cuda_op() + x = torch.randn(2, 8, device="cuda:0", dtype=dtype) + weight = torch.randn(16, 8, device="cuda:0", dtype=dtype) + bias = torch.randn(16, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="bias and weight must be on the same device"): + det_linear_forward(x, weight, bias) + + def test_backward_input_rejects_weight_on_another_device(self, dtype): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_backward_input + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + weight = torch.randn(16, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and weight must be on the same device"): + det_linear_backward_input(grad, weight, dtype) + + def test_backward_weight_rejects_input_on_another_device(self, dtype): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_backward_weight + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + x = torch.randn(2, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and x must be on the same device"): + det_linear_backward_weight(grad, x, dtype) From 4c040db9a30b8a7052145a7a5850110ab2e09740 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:30:05 +0000 Subject: [PATCH 04/19] docs(h3): timestep_mlp vs existing implementations (B200, ffd958a) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index a57f98033..237c7ae54 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -115,6 +115,35 @@ The data is in [`report.json`](../usage/evidence/h3-timestep-mlp-b200/report.jso by `scripts/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. +## Existing implementations (RFC #420 reuse rule) + +![timestep_mlp vs existing implementations](../usage/evidence/h3-prior-art-b200/timestep_mlp.png) + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers TimestepEmbedding, FP32 F.linear [plain] | **no** (14214 rows, 4108 sub-batches, 2232 sweep cases) | 5.7e-07 / 4.1e-07 / 777 µs | 5.7e-07 / 1.0e-06 / 832 µs | 3.7e-06 / 2.0e-06 / 3609 µs | +| rl-kernel H3TimestepMLPCudaOp | yes | 2.6e-07 / 4.2e-07 / 748 µs | 2.3e-07 / 7.6e-07 / 4595 µs | 2.6e-07 / 2.3e-06 / 42993 µs | +| diffusers TimestepEmbedding, FP32 F.linear [vllm] | **no** (14214 rows, 4096 sub-batches, 2232 sweep cases) | 1.8e-07 / 4.5e-07 / 752 µs | 3.7e-06 / 2.8e-06 / 921 µs | 3.7e-06 / 3.2e-06 / 3826 µs | +| diffusers TimestepEmbedding, FP32 F.linear [sglang] | yes | 1.6e-03 / 1.6e-03 / 1036 µs | 1.5e-03 / 1.5e-03 / 1140 µs | 1.5e-03 / 1.7e-03 / 2636 µs | +| diffusers TimestepEmbedding, FP32 F.linear [sglang_ieee] | yes | 3.3e-06 / 2.3e-06 / 2234 µs | 3.7e-06 / 2.8e-06 / 2428 µs | 3.7e-06 / 3.2e-06 / 7329 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_te_native] | **no** (14214 rows, 4108 sub-batches, 2232 sweep cases) | 1.8e-07 / 4.1e-07 / 729 µs | 3.7e-06 / 1.0e-06 / 814 µs | 3.7e-06 / 2.0e-06 / 3636 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_triton] | yes | 1.6e-03 / 1.6e-03 / 1027 µs | 1.5e-03 / 1.5e-03 / 1137 µs | 1.5e-03 / 1.7e-03 / 2634 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_triton_ieee] | yes | 3.3e-06 / 2.3e-06 / 2309 µs | 3.7e-06 / 2.8e-06 / 2435 µs | 3.7e-06 / 3.2e-06 / 7349 µ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. [`timestep_mlp.json`](../usage/evidence/h3-prior-art-b200/timestep_mlp.json) +was written from a clean tree at `ffd958a` by + +```bash +python scripts/h3_prior_art.py --op timestep_mlp --out docs/usage/evidence/h3-prior-art-b200/timestep_mlp.json --megatron-src +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/timestep_mlp.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 1a21cd6ae3189ac7122cd89054beab00bf77fc83 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:38:43 +0000 Subject: [PATCH 05/19] evidence(h3): timestep_mlp direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 0df83a5a. No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 7f2d477e3..8049a87df 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -144,6 +144,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) + +No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. +[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png +``` + ## Tests ```bash From fb1bd8aa5035051d3a4f98058b19f88d9449cac7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:55:22 +0000 Subject: [PATCH 06/19] fix(h3): seed the timestep_mlp timing inputs in the evidence report _mlp_perf drew its timestep features from the process's default CUDA generator, so the inputs depended on how many draws ran first. Seed them per T. Only timing inputs change; no reported accuracy field depends on them, so the committed report needs no regeneration. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index e42382ca8..f7bd01520 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -139,7 +139,8 @@ def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: weight_bytes = sum(p.numel() * p.element_size() for p in params) cases = [] for num in (1, 2, 3, 4): - x = torch.rand(num, 256, device="cuda") * 2 - 1 + gen = torch.Generator(device="cuda").manual_seed(num) + x = torch.rand(num, 256, device="cuda", generator=gen) * 2 - 1 cases.append( { "op": "timestep_mlp_fp32", From 8132d061c7e21ccfceedfb386e97475612c347f6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:10 +0000 Subject: [PATCH 07/19] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 8049a87df..cc5809cee 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -109,15 +109,15 @@ kernel time. ## Evidence -![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../../reports/experiments/h3-timestep-mlp-b200/figure.png) +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-mlp-b200/report.json), written +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![timestep_mlp vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_mlp.png) +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -134,7 +134,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../../reports/experiments/h3-prior-art-b200/timestep_mlp.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) was written from a clean tree at `ffd958a` by ```bash @@ -151,10 +151,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index c5f69a15b..75edf172b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -88,15 +88,15 @@ about 1 KB per timestep. ## Evidence -![timestep_sinusoid_h3 on B200: latency and error vs FP64](../../reports/experiments/h3-timestep-sinusoid-b200/figure.png) +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-sinusoid-b200/report.json). +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). `tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. ## Existing implementations (RFC #420 reuse rule) -![timestep_sinusoid vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -108,7 +108,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) was written from a clean tree at `7a82917` by ```bash @@ -126,11 +126,11 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). -![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. -[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) was written from a clean tree at `0eb5ddd5` by ```bash From 3b74ce356e99170e561126f337bdb895162e24b2 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:53 +0000 Subject: [PATCH 08/19] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index cc5809cee..b65b19fca 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -112,7 +112,7 @@ kernel time. ![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written -by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `418998c`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) @@ -154,7 +154,7 @@ gradient with diffusers' path and with this kernel. ![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `73939f2d` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 75edf172b..ff6579cd8 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -131,7 +131,7 @@ kernel compiled with `--use_fast_math`, which the build excludes (see Known Limi SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) -was written from a clean tree at `0eb5ddd5` by +was written from a clean tree at `be6a94ae` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ From cb9828f8b2477728e707473d570d33fb9e56445a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Sun, 25 Oct 2026 23:25:43 +0000 Subject: [PATCH 09/19] fix(h3): time competing callables in a shuffled order and fix the sinusoid latency column Review findings on the timing and documentation of the H3 evidence: - measure() timed every sample of the candidate before the provider's, so clock boost, thermal state and L2 contents could bias the gap. Each iteration now runs all callables once in a random order (seeded) and the executed order is stored next to the medians. - The sinusoid table called its latency column fwd+bwd although the op has no backward and every cell is forward-only; the header and the sentence below it now say forward latency. - Move the noqa: E402 of prepare_h3_weights.py to the first line of the multi-line import, where the linters report it. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-sinusoid.md | 4 +- rl_engine/validation/models/h3_report.py | 57 +++++++++++++++--- .../minimax_h3/test_h3_report_timing.py | 59 +++++++++++++++++++ tools/weights/prepare_h3_weights.py | 4 +- 4 files changed, 113 insertions(+), 11 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_report_timing.py diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index ff6579cd8..5696f5df7 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -98,14 +98,14 @@ records that commit and the environment. ![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) -| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd latency | size 256: fwd err / worst grad err / fwd latency | size 2048: fwd err / worst grad err / fwd latency | |---|---|---|---|---| | diffusers get_timestep_embedding (op-for-op replay) | yes | 6.7e-08 / — / 64 µs (fwd only) | 8.6e-08 / — / 64 µs (fwd only) | 9.9e-08 / — / 65 µs (fwd only) | | SGLang 0.5.21 timestep_embedding | yes | 6.7e-08 / — / 13 µs (fwd only) | 8.6e-08 / — / 12 µs (fwd only) | 9.9e-08 / — / 13 µs (fwd only) | | rl-kernel H3TimestepSinusoidCudaOp (check_range=False) | yes | 6.7e-08 / — / 16 µs (fwd only) | 8.6e-08 / — / 16 µs (fwd only) | 9.9e-08 / — / 16 µs (fwd only) | 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 +forward time (the op has no backward) on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A "no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index f7bd01520..5bd5c507b 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -10,6 +10,7 @@ from __future__ import annotations +import random import statistics from typing import Any, Callable @@ -57,14 +58,56 @@ def peak_mib(fn: Callable[[], Any]) -> float: return (torch.cuda.max_memory_allocated() - base) / 2**20 +def time_interleaved( + fns: dict[str, Callable[[], Any]], warmup: int = 20, iters: int = 200, seed: int = 0 +) -> tuple[dict[str, float], list[str]]: + """Median CUDA-event time of every callable, each sample in a freshly shuffled order. + + Clock boost, thermal state and L2 contents drift while a benchmark runs, so timing all + samples of one callable before the other's would bias the comparison. Every iteration runs + each callable once, in a random order from ``seed``. Returns the medians in microseconds and + the executed order of every iteration (the digits are indices into ``list(fns)``). + """ + + names = list(fns) + rng = random.Random(seed) + for fn in fns.values(): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples: dict[str, list[float]] = {name: [] for name in names} + executed: list[str] = [] + for _ in range(iters): + order = names[:] + rng.shuffle(order) + executed.append("".join(str(names.index(name)) for name in order)) + for name in order: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fns[name]() + end.record() + end.synchronize() + samples[name].append(start.elapsed_time(end) * 1e3) + return {name: statistics.median(values) for name, values in samples.items()}, executed + + def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: - row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} - for key in TIMED_KEYS: - if key in case: - us = time_us(case[key], warmup, iters) - row[f"{key}_us"] = us - row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 - row[f"{key}_peak_mib"] = peak_mib(case[key]) + row: dict[str, Any] = { + key: case[key] for key in ("op", "case", "backend", "bytes") if key in case + } + keys = [key for key in TIMED_KEYS if key in case] + medians, executed = time_interleaved({key: case[key] for key in keys}, warmup, iters) + for key in keys: + row[f"{key}_us"] = medians[key] + row[f"{key}_gbps"] = case["bytes"] / (medians[key] * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key]) + row["timing"] = { + "order": "shuffled per sample, all callables timed in every iteration", + "seed": 0, + "callables": keys, + "executed_order": executed, + } return row diff --git a/tests/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py new file mode 100644 index 000000000..6efba02ef --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Timing harness of the H3 evidence: competing callables are timed in a shuffled order. + +Run on a real CUDA device with real kernels (torch matmuls of two sizes); the check is about +the executed order and the medians, not about absolute times. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.validation.models.h3_report import measure, time_interleaved + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _callables(): + a = torch.randn(256, 256, device="cuda") + b = torch.randn(1024, 1024, device="cuda") + return {"small": lambda: a @ a, "large": lambda: b @ b} + + +def test_every_callable_runs_once_per_iteration_in_both_orders(): + fns = _callables() + medians, executed = time_interleaved(fns, warmup=2, iters=40, seed=0) + assert set(medians) == {"small", "large"} + assert all(value > 0 for value in medians.values()) + assert len(executed) == 40 + # digits index list(fns): each iteration runs both callables exactly once + assert {"".join(sorted(order)) for order in executed} == {"01"} + # the order is shuffled, so neither callable always goes first + assert {"01", "10"} <= set(executed) + + +def test_the_executed_order_is_reproducible_from_the_seed(): + fns = _callables() + _, first = time_interleaved(fns, warmup=1, iters=25, seed=3) + _, second = time_interleaved(fns, warmup=1, iters=25, seed=3) + _, other = time_interleaved(fns, warmup=1, iters=25, seed=4) + assert first == second + assert first != other + + +def test_measure_records_the_order_next_to_the_timings(): + case = { + "op": "matmul", + "case": "T=1", + "backend": "torch", + "bytes": 1 << 20, + **{key: fn for key, fn in zip(("candidate", "provider"), _callables().values())}, + } + row = measure(case, warmup=2, iters=20) + assert row["candidate_us"] > 0 and row["provider_us"] > 0 + assert row["timing"]["callables"] == ["candidate", "provider"] + assert len(row["timing"]["executed_order"]) == 20 + assert {"01", "10"} <= set(row["timing"]["executed_order"]) diff --git a/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py index e5470ef18..c21f56c2d 100644 --- a/tools/weights/prepare_h3_weights.py +++ b/tools/weights/prepare_h3_weights.py @@ -23,11 +23,11 @@ REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT)) -from rl_engine.validation.models.h3_weights import ( +from rl_engine.validation.models.h3_weights import ( # noqa: E402 EXTRACTED_FILE, load_h3_manifest, sha256_file, -) # noqa: E402 +) def _check(path: Path, expected: str) -> None: From a74b1a5af85ecbebeedcedb8c6b79e1f4c138f61 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Sun, 25 Oct 2026 23:27:38 +0000 Subject: [PATCH 10/19] fix(h3): record which weights the timestep MLP evidence used h3_params fell back to seeded random tensors when RL_KERNEL_H3_WEIGHTS was unset and nothing recorded it, while the figure title always claimed pinned weights. h3_params now returns the source (pinned or random); the accuracy report and every perf row store it, the title is derived from it, and h3_evidence.py refuses to write a timestep_mlp_fp32 report from random weights unless --allow-random-weights is given (smoke runs). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 20 +++++--- .../minimax_h3/test_h3_weight_source.py | 51 +++++++++++++++++++ tools/validation/models/h3_evidence.py | 22 ++++++++ tools/validation/models/plot_h3_evidence.py | 13 ++++- 4 files changed, 98 insertions(+), 8 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_weight_source.py diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 5bd5c507b..a157688aa 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -94,7 +94,7 @@ def time_interleaved( def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: row: dict[str, Any] = { - key: case[key] for key in ("op", "case", "backend", "bytes") if key in case + key: case[key] for key in ("op", "case", "backend", "bytes", "weights") if key in case } keys = [key for key in TIMED_KEYS if key in case] medians, executed = time_interleaved({key: case[key] for key in keys}, warmup, iters) @@ -155,17 +155,21 @@ def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: return {"contract_atol": 1e-5, "cases": rows} -def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Tensor]: - """Pinned checkpoint tensors when available, otherwise same-shape random ones.""" +def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> tuple[list[torch.Tensor], str]: + """Pinned checkpoint tensors when available, otherwise same-shape random ones. + + Returns the tensors and their source, ``"pinned"`` or ``"random"``; the evidence stores it so a + report made without the pinned checkpoint cannot pass for one that used it. + """ if h3_weights_dir() is not None: weights = load_h3_conditioning_weights("cuda", names) - return [weights[name] for name in names] + return [weights[name] for name in names], "pinned" generator = torch.Generator(device="cuda").manual_seed(0) return [ torch.randn(shape, device="cuda", generator=generator) * shape[-1] ** -0.5 for shape in shapes - ] + ], "random" # --------------------------------------------------------------------------- # @@ -178,7 +182,7 @@ def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Ten def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: op = registry.get_op("timestep_mlp_fp32", device="cuda") - params = h3_params(MLP_NAMES, MLP_SHAPES) + params, weights_source = h3_params(MLP_NAMES, MLP_SHAPES) weight_bytes = sum(p.numel() * p.element_size() for p in params) cases = [] for num in (1, 2, 3, 4): @@ -190,6 +194,7 @@ def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: "case": f"T={num}", "backend": type(op).__name__, "bytes": weight_bytes, + "weights": weights_source, "candidate": lambda x=x: op(x, *params), "provider": lambda x=x: provider_time_embedder(x, *params), } @@ -201,7 +206,7 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: 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") - params = h3_params(MLP_NAMES, MLP_SHAPES) + params, weights_source = h3_params(MLP_NAMES, MLP_SHAPES) cuda_err, provider_err = [], [] for seed in range(draws): x = sinusoid(h3_timesteps(4, seed=seed)) @@ -214,6 +219,7 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: invariant = all(torch.equal(op(x[i : i + 1], *params)[0], full[i]) for i in range(9)) return { "contract_atol": 1e-4, + "weights_source": weights_source, "draws": draws, "num_timesteps": 4, "cuda_max_abs_vs_fp64": cuda_err, diff --git a/tests/models/minimax_h3/test_h3_weight_source.py b/tests/models/minimax_h3/test_h3_weight_source.py new file mode 100644 index 000000000..251a728a2 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_weight_source.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The H3 MLP evidence says which weights it was measured with and refuses unpinned runs.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.validation.models.h3_report import MLP_NAMES, MLP_SHAPES, h3_params +from rl_engine.validation.models.h3_weights import WEIGHTS_ENV, h3_weights_dir +from tools.validation.models.h3_evidence import check_weight_source + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_without_the_pinned_checkpoint_the_params_are_marked_random(monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + params, source = h3_params(MLP_NAMES, MLP_SHAPES) + assert source == "random" + assert [tuple(p.shape) for p in params] == MLP_SHAPES + assert all(p.is_cuda for p in params) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_with_the_pinned_checkpoint_the_params_are_marked_pinned(): + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py") + params, source = h3_params(MLP_NAMES, MLP_SHAPES) + assert source == "pinned" + assert [tuple(p.shape) for p in params] == MLP_SHAPES + + +def test_mlp_evidence_is_refused_without_pinned_weights(): + random_run = {"accuracy": {"weights_source": "random"}} + pinned_run = {"accuracy": {"weights_source": "pinned"}} + with pytest.raises(SystemExit, match="pinned time_embedder weights"): + check_weight_source("timestep_mlp_fp32", random_run, allow_random=False) + check_weight_source("timestep_mlp_fp32", random_run, allow_random=True) + check_weight_source("timestep_mlp_fp32", pinned_run, allow_random=False) + # the sinusoid op has no weights, so it records none + check_weight_source("timestep_sinusoid_h3", {"accuracy": {}}, allow_random=False) + + +def test_the_figure_title_names_the_recorded_weight_source(): + pytest.importorskip("matplotlib") + from tools.validation.models.plot_h3_evidence import mlp_title + + assert "pinned time_embedder weights" in mlp_title({"accuracy": {"weights_source": "pinned"}}) + assert "random" in mlp_title({"accuracy": {"weights_source": "random"}}) + assert "not recorded" in mlp_title({"accuracy": {}}) diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index b4852d61f..41157682a 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -29,10 +29,31 @@ from rl_engine.validation.models.h3_weights import load_h3_manifest # noqa: E402 +def check_weight_source(op: str, report: dict, allow_random: bool) -> None: + """Refuse a timestep_mlp_fp32 report that was measured without the pinned checkpoint.""" + + if ( + op == "timestep_mlp_fp32" + and report["accuracy"]["weights_source"] != "pinned" + and not allow_random + ): + raise SystemExit( + "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " + "RL_KERNEL_H3_WEIGHTS (tools/weights/prepare_h3_weights.py) " + "or pass --allow-random-weights" + ) + + def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) parser.add_argument("--out", type=Path, required=True) + parser.add_argument( + "--allow-random-weights", + action="store_true", + help="write a timestep_mlp_fp32 report even without the pinned checkpoint " + "(smoke runs only)", + ) args = parser.parse_args() if not torch.cuda.is_available(): raise SystemExit("needs a CUDA device") @@ -52,6 +73,7 @@ def main() -> None: "accuracy": ACCURACY[args.op](registry), "perf": [measure(case) for case in PERF_CASES[args.op](registry)], } + check_weight_source(args.op, report, args.allow_random_weights) args.out.parent.mkdir(parents=True, exist_ok=True) args.out.write_text(json.dumps(report, indent=2) + "\n") print( diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py index ddf1ee947..bb8aee771 100644 --- a/tools/validation/models/plot_h3_evidence.py +++ b/tools/validation/models/plot_h3_evidence.py @@ -209,8 +209,19 @@ def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], con ax.set_xlim(-0.6, len(series) - 0.1) +def mlp_title(report: dict[str, Any]) -> str: + """Figure title that names the weights the report was measured with.""" + + source = report["accuracy"].get("weights_source") + label = { + "pinned": "pinned time_embedder weights", + "random": "random same-shape weights, not the checkpoint", + }.get(source, "weight source not recorded") + return f"timestep_mlp_fp32 ({label})" + + def plot_mlp(report: dict[str, Any]): - fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") + fig, (left, right) = _figure(report, mlp_title(report)) _latency_panel(left, report["perf"], "diffusers (cuBLAS)") acc = report["accuracy"] _error_boxes( From 9a748eb25977069d596452132df86b06efc8a586 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:22:38 +0000 Subject: [PATCH 11/19] fix(h3): time candidate and provider with the rotating harness of the later PRs The first fix for the timing-order review finding interleaved the callables in a seeded random order. #485 and #486 already carry a different harness that rotates the first callable each iteration and stores the raw CUDA-event samples and the executed orders, which keeps more of the evidence. Use that same measure() here so the stacked PRs agree and later merges do not conflict: drop time_interleaved, take measure() verbatim from #485, add its test (test_h3_report.py) and retarget the real-GPU test to the rotating harness. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 66 ++++++--------- tests/models/minimax_h3/test_h3_report.py | 82 +++++++++++++++++++ .../minimax_h3/test_h3_report_timing.py | 73 ++++++++--------- 3 files changed, 142 insertions(+), 79 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_report.py diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index a157688aa..ddc58a958 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -10,7 +10,6 @@ from __future__ import annotations -import random import statistics from typing import Any, Callable @@ -58,56 +57,43 @@ def peak_mib(fn: Callable[[], Any]) -> float: return (torch.cuda.max_memory_allocated() - base) / 2**20 -def time_interleaved( - fns: dict[str, Callable[[], Any]], warmup: int = 20, iters: int = 200, seed: int = 0 -) -> tuple[dict[str, float], list[str]]: - """Median CUDA-event time of every callable, each sample in a freshly shuffled order. +def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + """Time case callables in rotating order after warmup and return their evidence. - Clock boost, thermal state and L2 contents drift while a benchmark runs, so timing all - samples of one callable before the other's would bias the comparison. Every iteration runs - each callable once, in a random order from ``seed``. Returns the medians in microseconds and - the executed order of every iteration (the digits are indices into ``list(fns)``). + Record raw CUDA-event samples, executed orders, median microseconds, + decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. """ - names = list(fns) - rng = random.Random(seed) - for fn in fns.values(): - for _ in range(warmup): - fn() + row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} + keys = [key for key in TIMED_KEYS if key in case] + if not keys: + return row + for _ in range(warmup): + for key in keys: + case[key]() torch.cuda.synchronize() - samples: dict[str, list[float]] = {name: [] for name in names} - executed: list[str] = [] - for _ in range(iters): - order = names[:] - rng.shuffle(order) - executed.append("".join(str(names.index(name)) for name in order)) - for name in order: + samples = {key: [] for key in keys} + orders = [] + for index in range(iters): + # Rotate the first implementation so drift does not favor one backend. + offset = index % len(keys) + order = keys[offset:] + keys[:offset] + orders.append(order) + for key in order: start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() - fns[name]() + case[key]() end.record() end.synchronize() - samples[name].append(start.elapsed_time(end) * 1e3) - return {name: statistics.median(values) for name, values in samples.items()}, executed - - -def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: - row: dict[str, Any] = { - key: case[key] for key in ("op", "case", "backend", "bytes", "weights") if key in case - } - keys = [key for key in TIMED_KEYS if key in case] - medians, executed = time_interleaved({key: case[key] for key in keys}, warmup, iters) + samples[key].append(start.elapsed_time(end) * 1e3) + row["timing_samples_us"] = samples + row["timing_order"] = orders for key in keys: - row[f"{key}_us"] = medians[key] - row[f"{key}_gbps"] = case["bytes"] / (medians[key] * 1e-6) / 1e9 + us = statistics.median(samples[key]) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 row[f"{key}_peak_mib"] = peak_mib(case[key]) - row["timing"] = { - "order": "shuffled per sample, all callables timed in every iteration", - "seed": 0, - "callables": keys, - "executed_order": executed, - } return row diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py new file mode 100644 index 000000000..98e6ab714 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Comparative timing balances backend order and preserves its raw evidence.""" + +from __future__ import annotations + +import pytest + +from rl_engine.validation.models import h3_report + + +def test_measure_interleaves_backends_and_records_samples(monkeypatch): + """Balance candidate/provider timing order and preserve samples and derived bandwidth.""" + + elapsed_us = 0.0 + calls = [] + + class Event: + def __init__(self, *, enable_timing): + """Require timing-enabled events in the deterministic timing stub.""" + + assert enable_timing + + def record(self): + """Capture the synthetic elapsed clock at this event.""" + + self.timestamp = elapsed_us + + def synchronize(self): + """Leave synchronization inert because the timing stub has no pending GPU work.""" + + pass + + def elapsed_time(self, end): + """Return elapsed synthetic event time in CUDA milliseconds.""" + + return (end.timestamp - self.timestamp) / 1e3 + + def candidate(): + """Record a candidate call and advance the synthetic clock by two microseconds.""" + + nonlocal elapsed_us + calls.append("candidate") + elapsed_us += 2.0 + + def provider(): + """Record a provider call and advance the synthetic clock by four microseconds.""" + + nonlocal elapsed_us + calls.append("provider") + elapsed_us += 4.0 + + monkeypatch.setattr(h3_report.torch.cuda, "Event", Event) + monkeypatch.setattr(h3_report.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn: 0.0) + report = h3_report.measure( + {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 + ) + assert calls == [ + "candidate", + "provider", # warmup + "candidate", + "provider", + "provider", + "candidate", + "candidate", + "provider", + "provider", + "candidate", + ] + assert report["timing_order"] == [ + ["candidate", "provider"], + ["provider", "candidate"], + ["candidate", "provider"], + ["provider", "candidate"], + ] + assert report["timing_samples_us"] == {"candidate": [2.0] * 4, "provider": [4.0] * 4} + assert report["candidate_us"] == pytest.approx(2.0) + assert report["provider_us"] == pytest.approx(4.0) + assert report["candidate_gbps"] == pytest.approx(0.5) + assert report["provider_gbps"] == pytest.approx(0.25) diff --git a/tests/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py index 6efba02ef..853680a67 100644 --- a/tests/models/minimax_h3/test_h3_report_timing.py +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -1,59 +1,54 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Timing harness of the H3 evidence: competing callables are timed in a shuffled order. +"""``measure`` on a real CUDA device with real kernels. -Run on a real CUDA device with real kernels (torch matmuls of two sizes); the check is about -the executed order and the medians, not about absolute times. +``test_h3_report.py`` checks the bookkeeping with a stubbed clock; this one runs two real matmuls +of different size and checks that every callable is timed in every iteration, that the first +callable rotates, and that the stored samples reproduce the reported medians. """ from __future__ import annotations +import statistics + import pytest import torch -from rl_engine.validation.models.h3_report import measure, time_interleaved +from rl_engine.validation.models.h3_report import measure pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def _callables(): - a = torch.randn(256, 256, device="cuda") - b = torch.randn(1024, 1024, device="cuda") - return {"small": lambda: a @ a, "large": lambda: b @ b} - - -def test_every_callable_runs_once_per_iteration_in_both_orders(): - fns = _callables() - medians, executed = time_interleaved(fns, warmup=2, iters=40, seed=0) - assert set(medians) == {"small", "large"} - assert all(value > 0 for value in medians.values()) - assert len(executed) == 40 - # digits index list(fns): each iteration runs both callables exactly once - assert {"".join(sorted(order)) for order in executed} == {"01"} - # the order is shuffled, so neither callable always goes first - assert {"01", "10"} <= set(executed) - - -def test_the_executed_order_is_reproducible_from_the_seed(): - fns = _callables() - _, first = time_interleaved(fns, warmup=1, iters=25, seed=3) - _, second = time_interleaved(fns, warmup=1, iters=25, seed=3) - _, other = time_interleaved(fns, warmup=1, iters=25, seed=4) - assert first == second - assert first != other - - -def test_measure_records_the_order_next_to_the_timings(): - case = { +def _case(): + small = torch.randn(256, 256, device="cuda") + large = torch.randn(1024, 1024, device="cuda") + return { "op": "matmul", "case": "T=1", "backend": "torch", "bytes": 1 << 20, - **{key: fn for key, fn in zip(("candidate", "provider"), _callables().values())}, + "candidate": lambda: small @ small, + "provider": lambda: large @ large, } - row = measure(case, warmup=2, iters=20) - assert row["candidate_us"] > 0 and row["provider_us"] > 0 - assert row["timing"]["callables"] == ["candidate", "provider"] - assert len(row["timing"]["executed_order"]) == 20 - assert {"01", "10"} <= set(row["timing"]["executed_order"]) + + +def test_every_callable_is_timed_in_every_iteration_with_a_rotating_first_callable(): + row = measure(_case(), warmup=2, iters=20) + assert len(row["timing_samples_us"]["candidate"]) == 20 + assert len(row["timing_samples_us"]["provider"]) == 20 + assert row["timing_order"][0] == ["candidate", "provider"] + assert row["timing_order"][1] == ["provider", "candidate"] + assert all(sorted(order) == ["candidate", "provider"] for order in row["timing_order"]) + assert {order[0] for order in row["timing_order"]} == {"candidate", "provider"} + + +def test_the_reported_medians_come_from_the_stored_samples(): + row = measure(_case(), warmup=2, iters=20) + for key in ("candidate", "provider"): + assert row[f"{key}_us"] == statistics.median(row["timing_samples_us"][key]) + assert row[f"{key}_us"] > 0 + assert row[f"{key}_gbps"] == pytest.approx(row["bytes"] / (row[f"{key}_us"] * 1e-6) / 1e9) + assert row[f"{key}_peak_mib"] >= 0 + # the larger matmul is slower than the smaller one in the same run + assert row["provider_us"] > row["candidate_us"] From 9beb1407eb5ef104ab809984f55cb984f0f2fc93 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:23:01 +0000 Subject: [PATCH 12/19] fix(h3): keep the weight source on every timestep MLP perf row measure() is the shared rotating harness and does not copy the case's weights field, so the evidence writer stores it with the row it gets back. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/test_h3_weight_source.py | 19 ++++++++++++++++++- tools/validation/models/h3_evidence.py | 11 ++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/tests/models/minimax_h3/test_h3_weight_source.py b/tests/models/minimax_h3/test_h3_weight_source.py index 251a728a2..0a44d84ef 100644 --- a/tests/models/minimax_h3/test_h3_weight_source.py +++ b/tests/models/minimax_h3/test_h3_weight_source.py @@ -10,7 +10,7 @@ from rl_engine.validation.models.h3_report import MLP_NAMES, MLP_SHAPES, h3_params from rl_engine.validation.models.h3_weights import WEIGHTS_ENV, h3_weights_dir -from tools.validation.models.h3_evidence import check_weight_source +from tools.validation.models.h3_evidence import check_weight_source, measure_case @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") @@ -49,3 +49,20 @@ def test_the_figure_title_names_the_recorded_weight_source(): assert "pinned time_embedder weights" in mlp_title({"accuracy": {"weights_source": "pinned"}}) assert "random" in mlp_title({"accuracy": {"weights_source": "random"}}) assert "not recorded" in mlp_title({"accuracy": {}}) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_a_perf_row_keeps_the_weight_source_of_its_case(): + a = torch.randn(128, 128, device="cuda") + case = { + "op": "matmul", + "case": "T=1", + "backend": "torch", + "bytes": 1 << 16, + "weights": "random", + "candidate": lambda: a @ a, + } + row = measure_case(case) + assert row["weights"] == "random" + assert row["candidate_us"] > 0 + assert "weights" not in measure_case({k: v for k, v in case.items() if k != "weights"}) diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index 41157682a..e67333b8b 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -29,6 +29,15 @@ from rl_engine.validation.models.h3_weights import load_h3_manifest # noqa: E402 +def measure_case(case: dict) -> dict: + """Time one perf case and keep the weight source the case was built with next to the timings.""" + + row = measure(case) + if "weights" in case: + row["weights"] = case["weights"] + return row + + def check_weight_source(op: str, report: dict, allow_random: bool) -> None: """Refuse a timestep_mlp_fp32 report that was measured without the pinned checkpoint.""" @@ -71,7 +80,7 @@ def main() -> None: **git_state(), "environment": environment(), "accuracy": ACCURACY[args.op](registry), - "perf": [measure(case) for case in PERF_CASES[args.op](registry)], + "perf": [measure_case(case) for case in PERF_CASES[args.op](registry)], } check_weight_source(args.op, report, args.allow_random_weights) args.out.parent.mkdir(parents=True, exist_ok=True) From 31f342d76ed8eb6892420178757c9705c93b1620 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 11:49:09 +0000 Subject: [PATCH 13/19] fix(h3): use one timing harness and record the MLP weight source additively The stacked PRs carried three versions of measure(). #487-#494 have the most recent one (alternating order, setup outside the timed region, backward-only scope); take that block and its test here so the whole stack agrees and merging the PRs in order does not conflict. The later PRs reword the docstrings around h3_params and plot_mlp, so the weight-source change of the previous commits conflicts there. Keep h3_params as it was, name the source with a new h3_weights_source(), keep it on the accuracy report and the perf rows, and set the figure title from it after _figure() instead of editing the line that calls it. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 121 +++++++++++------- tests/models/minimax_h3/test_h3_report.py | 2 +- .../minimax_h3/test_h3_weight_source.py | 17 +-- tools/validation/models/plot_h3_evidence.py | 4 +- 4 files changed, 81 insertions(+), 63 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index ddc58a958..5d9437a01 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -30,35 +30,59 @@ ) -def time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 200) -> float: - """Median CUDA-event time of ``fn`` in microseconds (includes the Python wrapper).""" +def _prepared_call( + fn: Callable[..., Any], setup: Callable[[], Any] | None = None +) -> Callable[[], Any]: + if setup is None: + return fn + prepared = setup() + return lambda: fn(prepared) + + +def _sample_us(fn: Callable[..., Any], setup: Callable[[], Any] | None = None) -> float: + run = _prepared_call(fn, setup) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1e3 + + +def time_us( + fn: Callable[..., Any], + warmup: int = 20, + iters: int = 200, + *, + setup: Callable[[], Any] | None = None, +) -> float: + """Median CUDA-event microseconds; optional setup runs outside the timed region. + + When supplied, ``setup`` runs before every call and its return value is passed + to ``fn``. Backward measurements use it to build a fresh forward graph. + """ for _ in range(warmup): - fn() + _prepared_call(fn, setup)() torch.cuda.synchronize() - samples = [] - for _ in range(iters): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - fn() - end.record() - end.synchronize() - samples.append(start.elapsed_time(end) * 1e3) - return statistics.median(samples) - - -def peak_mib(fn: Callable[[], Any]) -> float: + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) + + +def peak_mib(fn: Callable[..., Any], *, setup: Callable[[], Any] | None = None) -> float: + """Incremental peak CUDA memory of the call, excluding optional setup.""" + + run = _prepared_call(fn, setup) torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() base = torch.cuda.memory_allocated() - fn() + run() torch.cuda.synchronize() return (torch.cuda.max_memory_allocated() - base) / 2**20 def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: - """Time case callables in rotating order after warmup and return their evidence. + """Interleave timed keys, reversing their execution order on alternate iterations. Record raw CUDA-event samples, executed orders, median microseconds, decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. @@ -68,32 +92,29 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st keys = [key for key in TIMED_KEYS if key in case] if not keys: return row - for _ in range(warmup): - for key in keys: - case[key]() + orders = (keys, list(reversed(keys))) + row["execution_order"] = { + "policy": "alternating", + "iteration_0": orders[0], + "iteration_1": orders[1], + } + if any(key.endswith("_backward") for key in keys): + row["backward_timing_scope"] = "backward_only" + for iteration in range(warmup): + for key in orders[iteration % 2]: + _prepared_call(case[key], case.get(f"{key}_setup"))() torch.cuda.synchronize() samples = {key: [] for key in keys} - orders = [] - for index in range(iters): - # Rotate the first implementation so drift does not favor one backend. - offset = index % len(keys) - order = keys[offset:] + keys[:offset] - orders.append(order) - for key in order: - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - case[key]() - end.record() - end.synchronize() - samples[key].append(start.elapsed_time(end) * 1e3) + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) row["timing_samples_us"] = samples - row["timing_order"] = orders + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] for key in keys: us = statistics.median(samples[key]) row[f"{key}_us"] = us row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 - row[f"{key}_peak_mib"] = peak_mib(case[key]) + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) return row @@ -141,21 +162,23 @@ def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: return {"contract_atol": 1e-5, "cases": rows} -def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> tuple[list[torch.Tensor], str]: - """Pinned checkpoint tensors when available, otherwise same-shape random ones. - - Returns the tensors and their source, ``"pinned"`` or ``"random"``; the evidence stores it so a - report made without the pinned checkpoint cannot pass for one that used it. - """ +def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Tensor]: + """Pinned checkpoint tensors when available, otherwise same-shape random ones.""" if h3_weights_dir() is not None: weights = load_h3_conditioning_weights("cuda", names) - return [weights[name] for name in names], "pinned" + return [weights[name] for name in names] generator = torch.Generator(device="cuda").manual_seed(0) return [ torch.randn(shape, device="cuda", generator=generator) * shape[-1] ** -0.5 for shape in shapes - ], "random" + ] + + +def h3_weights_source() -> str: + """Source ``h3_params`` draws from: ``pinned`` checkpoint tensors or ``random`` ones.""" + + return "pinned" if h3_weights_dir() is not None else "random" # --------------------------------------------------------------------------- # @@ -168,7 +191,7 @@ def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> tuple[list[tor def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: op = registry.get_op("timestep_mlp_fp32", device="cuda") - params, weights_source = h3_params(MLP_NAMES, MLP_SHAPES) + params = h3_params(MLP_NAMES, MLP_SHAPES) weight_bytes = sum(p.numel() * p.element_size() for p in params) cases = [] for num in (1, 2, 3, 4): @@ -180,7 +203,7 @@ def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: "case": f"T={num}", "backend": type(op).__name__, "bytes": weight_bytes, - "weights": weights_source, + "weights": h3_weights_source(), "candidate": lambda x=x: op(x, *params), "provider": lambda x=x: provider_time_embedder(x, *params), } @@ -192,7 +215,7 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: 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") - params, weights_source = h3_params(MLP_NAMES, MLP_SHAPES) + params = h3_params(MLP_NAMES, MLP_SHAPES) cuda_err, provider_err = [], [] for seed in range(draws): x = sinusoid(h3_timesteps(4, seed=seed)) @@ -205,7 +228,7 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: invariant = all(torch.equal(op(x[i : i + 1], *params)[0], full[i]) for i in range(9)) return { "contract_atol": 1e-4, - "weights_source": weights_source, + "weights_source": h3_weights_source(), "draws": draws, "num_timesteps": 4, "cuda_max_abs_vs_fp64": cuda_err, diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py index 98e6ab714..b79973eb1 100644 --- a/tests/models/minimax_h3/test_h3_report.py +++ b/tests/models/minimax_h3/test_h3_report.py @@ -53,7 +53,7 @@ def provider(): monkeypatch.setattr(h3_report.torch.cuda, "Event", Event) monkeypatch.setattr(h3_report.torch.cuda, "synchronize", lambda: None) - monkeypatch.setattr(h3_report, "peak_mib", lambda fn: 0.0) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) report = h3_report.measure( {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 ) diff --git a/tests/models/minimax_h3/test_h3_weight_source.py b/tests/models/minimax_h3/test_h3_weight_source.py index 0a44d84ef..a8cb005eb 100644 --- a/tests/models/minimax_h3/test_h3_weight_source.py +++ b/tests/models/minimax_h3/test_h3_weight_source.py @@ -8,27 +8,20 @@ import pytest import torch -from rl_engine.validation.models.h3_report import MLP_NAMES, MLP_SHAPES, h3_params +from rl_engine.validation.models.h3_report import h3_weights_source from rl_engine.validation.models.h3_weights import WEIGHTS_ENV, h3_weights_dir from tools.validation.models.h3_evidence import check_weight_source, measure_case -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def test_without_the_pinned_checkpoint_the_params_are_marked_random(monkeypatch): +def test_without_the_pinned_checkpoint_the_source_is_random(monkeypatch): monkeypatch.delenv(WEIGHTS_ENV, raising=False) - params, source = h3_params(MLP_NAMES, MLP_SHAPES) - assert source == "random" - assert [tuple(p.shape) for p in params] == MLP_SHAPES - assert all(p.is_cuda for p in params) + assert h3_weights_source() == "random" -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def test_with_the_pinned_checkpoint_the_params_are_marked_pinned(): +def test_with_the_pinned_checkpoint_the_source_is_pinned(): if h3_weights_dir() is None: pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py") - params, source = h3_params(MLP_NAMES, MLP_SHAPES) - assert source == "pinned" - assert [tuple(p.shape) for p in params] == MLP_SHAPES + assert h3_weights_source() == "pinned" def test_mlp_evidence_is_refused_without_pinned_weights(): diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py index bb8aee771..ecea9ab1e 100644 --- a/tools/validation/models/plot_h3_evidence.py +++ b/tools/validation/models/plot_h3_evidence.py @@ -221,7 +221,9 @@ def mlp_title(report: dict[str, Any]) -> str: def plot_mlp(report: dict[str, Any]): - fig, (left, right) = _figure(report, mlp_title(report)) + 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") _latency_panel(left, report["perf"], "diffusers (cuBLAS)") acc = report["accuracy"] _error_boxes( From 92175dbe23351b8a2f4190e1a38bc6c69048cdcd Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:31:25 +0000 Subject: [PATCH 14/19] fix(h3): refuse a random-weight MLP evidence run before measuring h3_evidence.py checked the weight source only after the accuracy and every performance case had run, so a run without RL_KERNEL_H3_WEIGHTS and without --allow-random-weights spent the whole GPU time and then wrote nothing. require_pinned_weights() makes the same decision from h3_weights_source() right after the arguments are parsed; the check on the finished report stays. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/test_h3_weight_source.py | 14 +++++++++- tools/validation/models/h3_evidence.py | 27 ++++++++++++++----- 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/tests/models/minimax_h3/test_h3_weight_source.py b/tests/models/minimax_h3/test_h3_weight_source.py index a8cb005eb..ad1cf09ae 100644 --- a/tests/models/minimax_h3/test_h3_weight_source.py +++ b/tests/models/minimax_h3/test_h3_weight_source.py @@ -10,7 +10,11 @@ from rl_engine.validation.models.h3_report import h3_weights_source from rl_engine.validation.models.h3_weights import WEIGHTS_ENV, h3_weights_dir -from tools.validation.models.h3_evidence import check_weight_source, measure_case +from tools.validation.models.h3_evidence import ( + check_weight_source, + measure_case, + require_pinned_weights, +) def test_without_the_pinned_checkpoint_the_source_is_random(monkeypatch): @@ -35,6 +39,14 @@ def test_mlp_evidence_is_refused_without_pinned_weights(): check_weight_source("timestep_sinusoid_h3", {"accuracy": {}}, allow_random=False) +def test_mlp_evidence_is_refused_before_measuring_without_pinned_weights(monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + with pytest.raises(SystemExit, match="pinned time_embedder weights"): + require_pinned_weights("timestep_mlp_fp32", allow_random=False) + require_pinned_weights("timestep_mlp_fp32", allow_random=True) + require_pinned_weights("timestep_sinusoid_h3", allow_random=False) + + def test_the_figure_title_names_the_recorded_weight_source(): pytest.importorskip("matplotlib") from tools.validation.models.plot_h3_evidence import mlp_title diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index e67333b8b..5aa5e4f73 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -25,9 +25,20 @@ from rl_engine.runtime.registry import KernelRegistry # noqa: E402 from rl_engine.validation.models.h3_chain import environment, git_state # noqa: E402 -from rl_engine.validation.models.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 +from rl_engine.validation.models.h3_report import ( # noqa: E402 + ACCURACY, + PERF_CASES, + h3_weights_source, + measure, +) from rl_engine.validation.models.h3_weights import load_h3_manifest # noqa: E402 +_NEEDS_PINNED_MLP = ( + "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " + "RL_KERNEL_H3_WEIGHTS (tools/weights/prepare_h3_weights.py) " + "or pass --allow-random-weights" +) + def measure_case(case: dict) -> dict: """Time one perf case and keep the weight source the case was built with next to the timings.""" @@ -46,11 +57,14 @@ def check_weight_source(op: str, report: dict, allow_random: bool) -> None: and report["accuracy"]["weights_source"] != "pinned" and not allow_random ): - raise SystemExit( - "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " - "RL_KERNEL_H3_WEIGHTS (tools/weights/prepare_h3_weights.py) " - "or pass --allow-random-weights" - ) + raise SystemExit(_NEEDS_PINNED_MLP) + + +def require_pinned_weights(op: str, allow_random: bool) -> None: + """Stop before any measurement when timestep_mlp_fp32 would run on random weights.""" + + if op == "timestep_mlp_fp32" and h3_weights_source() != "pinned" and not allow_random: + raise SystemExit(_NEEDS_PINNED_MLP) def main() -> None: @@ -64,6 +78,7 @@ def main() -> None: "(smoke runs only)", ) args = parser.parse_args() + require_pinned_weights(args.op, args.allow_random_weights) if not torch.cuda.is_available(): raise SystemExit("needs a CUDA device") torch.backends.cuda.matmul.allow_tf32 = False From 0f66d75515b48375a9dcc9f0ec91498ee5d29050 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:52:11 +0000 Subject: [PATCH 15/19] fix(h3): do not consume the weight registers of columns past n_out The forward kernel loads weights only for n0 + c < n_out, but the loop that unpacks them and feeds the FMAs ran for every c < kCols. On the last column tile of an N that is not a multiple of kCols (any odd N for FP32, kCols = 2; N % 4 for BF16) it read raw[u][c] uninitialized. Those columns are never stored, but the read is undefined behaviour. The loop is guarded with n0 + c < n_out like the load. test_h3_det_linear_odd_width.py covers N in {1, 2, 3, 5, 7, 9, 33, 65} for FP32 and BF16: forward, input gradient, weight gradient and bias gradient against an FP64 golden, and bitwise against the same product with zero columns padded on. Run on a fresh build; the other kernels of det_linear.cu index their columns by n < n_out already. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 10 ++- .../test_h3_det_linear_odd_width.py | 86 +++++++++++++++++++ 2 files changed, 93 insertions(+), 3 deletions(-) create mode 100644 tests/models/minimax_h3/test_h3_det_linear_odd_width.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index ddaf199ee..7af6592b4 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -149,10 +149,14 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) XV::load(x + (t0 + r) * k_in + k0, xv); #pragma unroll for (int c = 0; c < kCols; ++c) { - float wv[kVec]; - WV::unpack(raw[u][c], wv); + // raw[u][c] is only loaded for columns inside the output; the last column tile of + // an N that is not a multiple of kCols must not read the other registers. + if (n0 + c < n_out) { + float wv[kVec]; + WV::unpack(raw[u][c], wv); #pragma unroll - for (int j = 0; j < kVec; ++j) acc[c][r] = fmaf(xv[j], wv[j], acc[c][r]); + for (int j = 0; j < kVec; ++j) acc[c][r] = fmaf(xv[j], wv[j], acc[c][r]); + } } } } 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 new file mode 100644 index 000000000..19f95784c --- /dev/null +++ b/tests/models/minimax_h3/test_h3_det_linear_odd_width.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""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). +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, +) + +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_*"), +] + +DTYPES = [torch.float32, torch.bfloat16] +WIDTHS = [1, 2, 3, 5, 7, 9, 33, 65] # odd widths, and tails past a 2- or 4-column tile +K_IN = 64 +ROWS = [1, 3, 5] + + +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() + weight = (torch.randn(n_out, K_IN, generator=g) / K_IN**0.5).to(dtype).cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + return x, weight, bias + + +def _pad(weight, bias, extra): + zeros = torch.zeros(extra, weight.shape[1], dtype=weight.dtype, device=weight.device) + return torch.cat([weight, zeros]), torch.cat([bias, bias.new_zeros(extra)]) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("n_out", WIDTHS) +@pytest.mark.parametrize("rows", ROWS) +def test_forward_matches_golden_and_ignores_the_tail(dtype, n_out, rows): + 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() + golden = x.double() @ weight.double().T + bias.double() + tol = 1e-5 if dtype == torch.float32 else 1e-2 + torch.testing.assert_close(out.double(), golden, atol=tol, rtol=tol) + padded = det_linear_forward(x, *_pad(weight, bias, 3))[0] + assert torch.equal(out, padded[:, :n_out]) + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("n_out", WIDTHS) +def test_backward_matches_golden_and_ignores_the_tail(dtype, n_out): + rows = 3 + x, weight, _ = _inputs(dtype, rows, n_out) + g = torch.Generator(device="cpu").manual_seed(1) + grad = torch.randn(rows, n_out, generator=g).cuda() + tol = 1e-5 if dtype == torch.float32 else 1e-2 + + d_input = det_linear_backward_input(grad, weight, torch.float32) + torch.testing.assert_close( + d_input.double(), grad.double() @ weight.double(), atol=tol, rtol=tol + ) + d_weight, d_bias = det_linear_backward_weight(grad, x, dtype, with_bias=True) + torch.testing.assert_close(d_weight.double(), grad.double().T @ x.double(), atol=tol, rtol=tol) + torch.testing.assert_close(d_bias.double(), grad.double().sum(0), atol=tol, rtol=tol) + + zero = torch.zeros(rows, 3, device="cuda") + padded_weight, _ = _pad(weight, weight[:, 0], 3) + padded_grad = torch.cat([grad, zero], dim=1) + assert torch.equal( + d_input, det_linear_backward_input(padded_grad, padded_weight, torch.float32) + ) + d_weight_p, d_bias_p = det_linear_backward_weight(padded_grad, x, dtype, with_bias=True) + assert torch.equal(d_weight, d_weight_p[:n_out]) and torch.equal(d_bias, d_bias_p[:n_out]) From 0e911bed680d8932407106ef9650b25f63af7316 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:11:36 +0000 Subject: [PATCH 16/19] fix(h3): fail at build time or with a clear error when det_linear launch shapes overflow - det_linear_forward_kernel: static_assert(kCols * kRows <= kWarp). The epilogue stores through lane == c * kRows + r, so a tuned config above a warp would silently drop stores. - h3_det_linear_backward_input: TORCH_CHECK on the chunk count (blockIdx.y <= 65535) and on the fold launch size instead of 'invalid configuration argument'. - test_h3_det_linear_odd_width.py: K that is a vector multiple but not a multiple of the 32-lane stride (BF16 8, 24, 40, 72, 264, 520; FP32 4, 12, 20, 36, 132, 1044) against the FP64 golden. The same shapes, with the backward kernels, ran under compute-sanitizer memcheck with 0 errors. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 9 +++++++++ .../test_h3_det_linear_odd_width.py | 19 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 7af6592b4..13798a0e8 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -35,6 +35,7 @@ #include #include +#include #include namespace { @@ -110,6 +111,8 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) int64_t k_in, int64_t activation) { using WV = Vec16; using XV = Vec16; + static_assert(kCols * kRows <= kWarp, + "the epilogue stores through lane == c * kRows + r, so kCols * kRows must fit a warp"); constexpr int kVec = WV::kN; static_assert(XV::kN == kVec, "x and weight share a dtype"); constexpr int64_t kStride = static_cast(kWarp) * kVec; @@ -405,6 +408,12 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei const int64_t k_in = weight.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); 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. + TORCH_CHECK(chunks <= 65535, "N=", n_out, " needs ", chunks, + " input-gradient chunks; at most 65535 are supported"); + 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(); 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 19f95784c..2b0ee1aac 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 @@ -84,3 +84,22 @@ def test_backward_matches_golden_and_ignores_the_tail(dtype, n_out): ) d_weight_p, d_bias_p = det_linear_backward_weight(padded_grad, x, dtype, with_bias=True) assert torch.equal(d_weight, d_weight_p[:n_out]) and torch.equal(d_bias, d_bias_p[:n_out]) + + +@pytest.mark.parametrize( + "dtype, k_in", + [(torch.bfloat16, k) for k in (8, 24, 40, 72, 264, 520)] + + [(torch.float32, k) for k in (4, 12, 20, 36, 132, 1044)], +) +@pytest.mark.parametrize("n_out", [1, 3, 33]) +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.""" + + 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() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + out = det_linear_forward(x, weight, bias)[0] + tol = 1e-5 if dtype == torch.float32 else 1e-2 + golden = x.double() @ weight.double().T + bias.double() + torch.testing.assert_close(out.double(), golden, atol=tol, rtol=tol) From e1274b0e61b486ffcdb7a3624a27e0973fe6e2e6 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:01:25 +0000 Subject: [PATCH 17/19] docs(h3): state the K and N tail contract of det_linear next to the loads The 16-byte loads of the forward kernels start at multiples of VEC below K and the host requires K % VEC == 0, so k < K implies k + VEC <= K; say so in the contract header and at the load, together with the way N and the row count tails are masked. Also note that db is the dW chain with the multiplier 1 (a plain FP32 add) and that check_matrix admits float32 and bfloat16 only before the dW dispatch. No code change. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 13798a0e8..d9a7d5653 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -26,7 +26,14 @@ // ascending chunk order and cast once. // d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] // * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one -// cross-row reduction; its order is the logical row order. +// cross-row reduction; its order is the logical row order. db is the same +// 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. #include #include @@ -135,6 +142,8 @@ __global__ void __launch_bounds__(kWarpsPerBlock * kWarp) uint4 raw[kLoadAhead][kCols]; #pragma unroll for (int u = 0; u < kLoadAhead; ++u) { + // K % kVec == 0 (host check) and k0 is a multiple of kVec, so k0 < k_in means the whole + // 16-byte chunk lies inside the row. const int64_t k0 = g0 + u * kStride; #pragma unroll for (int c = 0; c < kCols; ++c) { @@ -464,6 +473,7 @@ std::vector h3_det_linear_backward_weight(torch::Tensor grad, tor const int64_t total = n_out * k_in; const unsigned blocks = static_cast(grid_for(total, threads)); const float* g = grad.data_ptr(); + // check_matrix(x) above admits float32 and bfloat16 only, so the last branch is bfloat16. #define H3_DW_LAUNCH(XT, XC, WT, WC) \ det_linear_dweight_kernel<<>>( \ g, cptr(x), mptr(dw), rows, n_out, k_in) From 1adec0431149018537449833c1fd9ac9195ead71 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:22:26 +0000 Subject: [PATCH 18/19] fix(h3): report a w1 that is not a matrix as a dimension error, and say which MLP layer has the SiLU validate_h3_timestep_mlp unpacked (None, None) for a w1 that is not 2-D, so the next checks reported "w1 must be (H, K)" and "b1 must be (None,)". It now rejects the rank first with a message that says so. The operator doc stated that SiLU follows the bias addition of every output; it applies to linear_1 only, linear_2 returns its biased linear output. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-timestep-mlp.md | 5 +++-- rl_engine/reference/minimax_h3/timestep_mlp.py | 4 +++- tests/models/minimax_h3/test_h3_timestep_mlp.py | 7 +++++++ 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index b65b19fca..69e6683f5 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -56,8 +56,9 @@ Forward, for each output `y[t, n]`: - Each lane accumulates its chunks in ascending order, and the elements inside a chunk in ascending order, using `fmaf` into an FP32 accumulator that starts at 0. - The 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1). -- The bias is added once, then SiLU runs in FP32 as `v / (1 + expf(-v))` - (PyTorch's formula). The result is stored as FP32. +- The bias is added once. In `linear_1` SiLU then runs in FP32 as `v / (1 + expf(-v))` + (PyTorch's formula); `linear_2` returns its biased linear output without an activation. + The result is stored as FP32. How many columns a warp owns, the row tile (1, 2 or 4), and how many chunks are loaded ahead change only when loads are issued, never the order in which a column diff --git a/rl_engine/reference/minimax_h3/timestep_mlp.py b/rl_engine/reference/minimax_h3/timestep_mlp.py index ac54d295f..f8348d99c 100644 --- a/rl_engine/reference/minimax_h3/timestep_mlp.py +++ b/rl_engine/reference/minimax_h3/timestep_mlp.py @@ -40,7 +40,9 @@ def validate_h3_timestep_mlp( raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") if x.dim() != 2 or x.shape[0] == 0: raise ValueError(f"x must be a non-empty (T, K) matrix, got {tuple(x.shape)}") - hidden, k_in = w1.shape if w1.dim() == 2 else (None, None) + if w1.dim() != 2: + raise ValueError(f"w1 must be a 2-D (H, {x.shape[1]}) matrix, got {tuple(w1.shape)}") + hidden, k_in = w1.shape if k_in != x.shape[1]: raise ValueError(f"w1 must be (H, {x.shape[1]}), got {tuple(w1.shape)}") if b1.shape != (hidden,): diff --git a/tests/models/minimax_h3/test_h3_timestep_mlp.py b/tests/models/minimax_h3/test_h3_timestep_mlp.py index bc6059125..09d1b2dd4 100644 --- a/tests/models/minimax_h3/test_h3_timestep_mlp.py +++ b/tests/models/minimax_h3/test_h3_timestep_mlp.py @@ -90,6 +90,13 @@ def test_rejects_shape_mismatches(self): with pytest.raises(ValueError): op.forward(x[:0], w1, b1, w2, b2) + @pytest.mark.parametrize("bad_w1", [torch.zeros(256), torch.zeros(2, 5, 256)]) + def test_a_w1_that_is_not_a_matrix_gets_a_dimension_error(self, bad_w1): + x = _features(2, device="cpu") + _, b1, w2, b2 = _synthetic_params(256, 64, 32, device="cpu") + with pytest.raises(ValueError, match="w1 must be a 2-D"): + NativeH3TimestepMLPOp().forward(x, bad_w1, b1, w2, b2) + @requires_cuda class TestCudaRealWeights: From 2ee5abc074f8f27ed889069b344e36eee1cf3f2e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:05:21 +0000 Subject: [PATCH 19/19] fix(h3): skip the MLP bias fold when no bias needs a gradient, and pin the device and K = 0 behaviour of det_linear * `timestep_mlp` backward passes `with_bias=` to `det_linear_backward_weight`, so the bias fold of a layer whose weight alone needs a gradient is not run and not allocated (`adaln_modulation` already did this). A test freezes each input in turn, compares the remaining gradients bitwise with the all-trainable run and records the `with_bias` argument of every call. * The det_linear entry points are registered as plain function pointers; the `CUDAGuard` inside each of them is what puts the launch on the tensors' device. A two-GPU test calls forward and both backwards with the tensors on cuda:1 while cuda:0 is current and compares the bytes with the same call on cuda:0, for FP32 and BF16. * `h3_det_linear_backward_weight` with K = 0 returns an empty dW and the real bias gradient (the column sums of grad); a test pins that, so an early return for K = 0 cannot hand back an uninitialised bias. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../model_specific/minimax_h3/timestep_mlp.py | 11 ++- .../models/minimax_h3/test_h3_timestep_mlp.py | 97 +++++++++++++++++++ 2 files changed, 105 insertions(+), 3 deletions(-) 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 3244d49d6..87763d308 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 @@ -7,7 +7,6 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( ACT_NONE, ACT_SILU, @@ -18,6 +17,7 @@ det_linear_forward, silu_backward_fp32, ) +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.timestep_mlp import validate_h3_timestep_mlp KERNEL_ID = "rl_engine._C.h3_det_linear_forward[silu]+rl_engine._C.h3_det_linear_forward" @@ -45,13 +45,18 @@ def backward(ctx, grad_out): grad_out = grad_out.float().contiguous() dw2 = db2 = dw1 = db1 = dx = None if need_w2 or need_b2: - dw2, db2 = det_linear_backward_weight(grad_out, hidden, torch.float32) + # The bias fold is skipped when only the weight needs a gradient. + dw2, *rest = det_linear_backward_weight( + grad_out, hidden, torch.float32, with_bias=need_b2 + ) + db2 = rest[0] if need_b2 else None d_pre = None if need_x or need_w1 or need_b1: d_hidden = det_linear_backward_input(grad_out, w2, torch.float32) d_pre = silu_backward_fp32(d_hidden, pre).contiguous() if need_w1 or need_b1: - dw1, db1 = det_linear_backward_weight(d_pre, x, torch.float32) + dw1, *rest = det_linear_backward_weight(d_pre, x, torch.float32, with_bias=need_b1) + db1 = rest[0] if need_b1 else None if need_x: dx = det_linear_backward_input(d_pre, w1, torch.float32) record_backward( diff --git a/tests/models/minimax_h3/test_h3_timestep_mlp.py b/tests/models/minimax_h3/test_h3_timestep_mlp.py index 09d1b2dd4..28359f0c2 100644 --- a/tests/models/minimax_h3/test_h3_timestep_mlp.py +++ b/tests/models/minimax_h3/test_h3_timestep_mlp.py @@ -150,6 +150,25 @@ def test_other_sizes(self, k_in, hidden, out): rtol=1e-4, ) + @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 + # it must be computed (not skipped like the empty dW). + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_weight, + ) + + _cuda_op() + grad = torch.randn(5, 7, device="cuda") + x = torch.empty(5, 0, device="cuda") + result = det_linear_backward_weight(grad, x, torch.float32, with_bias=with_bias) + assert len(result) == (2 if with_bias else 1) + assert result[0].shape == (7, 0) + if with_bias: + torch.testing.assert_close( + result[1].double(), grad.double().sum(0), atol=1e-5, rtol=1e-5 + ) + def test_rejects_unaligned_k(self): params = _synthetic_params(6, 8, 8) with pytest.raises(RuntimeError, match="multiple of 4"): @@ -192,6 +211,51 @@ def dx(rows): for i in range(6): assert torch.equal(dx(slice(i, i + 1))[0], full[i]) + @pytest.mark.parametrize( + "frozen", + [("b1",), ("b2",), ("b1", "b2"), ("w1",), ("w2",), ("x",), ("w1", "b1", "w2", "b2")], + ) + def test_frozen_inputs_get_no_gradient_and_do_not_change_the_others(self, frozen, monkeypatch): + from rl_engine.backends.cuda.model_specific.minimax_h3 import timestep_mlp + + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=6) + x = _features(3, seed=6) + grad_out = torch.randn(3, 2688, device="cuda") + + def grads(frozen_names): + leaves = [t.detach().clone() for t in (x, *params)] + for name, leaf in zip(NAMES, leaves, strict=True): + leaf.requires_grad_(name not in frozen_names) + if all(not leaf.requires_grad for leaf in leaves): + return leaves + op(*leaves).backward(grad_out) + return leaves + + full = grads(()) + calls = [] + real = timestep_mlp.det_linear_backward_weight + + def recording(*args, **kwargs): + calls.append(kwargs.get("with_bias", True)) + return real(*args, **kwargs) + + monkeypatch.setattr(timestep_mlp, "det_linear_backward_weight", recording) + part = grads(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 + # The bias fold runs only for a layer whose bias needs a gradient; a layer none of whose + # parameters need one is not folded at all. + expected = [] + if "w2" not in frozen or "b2" not in frozen: + expected.append("b2" not in frozen) + if "w1" not in frozen or "b1" not in frozen: + expected.append("b1" not in frozen) + assert calls == expected, (frozen, calls) + def test_registry_dispatches_cuda(self): from rl_engine.runtime.registry import KernelRegistry @@ -235,3 +299,36 @@ def test_backward_weight_rejects_input_on_another_device(self, dtype): x = torch.randn(2, 8, device="cuda:1", dtype=dtype) with pytest.raises(RuntimeError, match="grad and x must be on the same device"): det_linear_backward_weight(grad, x, dtype) + + def test_ops_run_on_the_device_of_their_tensors_not_the_current_one(self, dtype): + # The entry points are registered as plain function pointers; the device guard inside + # each of them (not the binding) is what puts the launch on the tensors' device. + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + ) + + _cuda_op() + g = torch.Generator("cpu").manual_seed(0) + + def run(device): + x = torch.randn(3, 64, generator=g).to(device=device, dtype=dtype) + weight = torch.randn(40, 64, generator=g).to(device=device, dtype=dtype) + bias = torch.randn(40, generator=g).to(device=device, dtype=dtype) + grad = torch.randn(3, 40, generator=g).to(device) + (out,) = det_linear_forward(x, weight, bias) + dx = det_linear_backward_input(grad, weight, torch.float32) + dw, db = det_linear_backward_weight(grad, x, dtype) + return [t.cpu() for t in (out, dx, dw, db)] + + # Same bytes from the same generator state on both devices, called while device 0 is + # current: nothing may run on (or touch the memory of) the wrong device. + torch.cuda.set_device(0) + g.manual_seed(0) + on_zero = run("cuda:0") + g.manual_seed(0) + on_one = run("cuda:1") + assert torch.cuda.current_device() == 0 + for a, b in zip(on_zero, on_one, strict=True): + assert torch.equal(a, b)