Repository navigation
[WS1][CUDA][MiniMax-H3] adaln_gate_residual: gated residual with in-kernel gate gather (RFC #420 stack 6/10) - #488
fusheng-ji wants to merge 258 commits into
Conversation
RFC RL-Align#420 step 4, row `h3_rmsnorm` (stacked on adaln_row_gather). - CUDA kernel replaying PyTorch's own RMSNorm reduction order (vectorized_layer_norm_kernel: (32, 4) block per row, 4-element vectors, shuffle-down and cross-warp trees, rsqrtf, w * (rstd * x)): bitwise equal to nn.RMSNorm on all four pinned H3 norm weights and batch/position invariant. - Fused AdaLN modulation n * (1 + scale[i]) + shift[i] with the rows gathered in-kernel from strided table views and each step rounded where the eager expression rounds: bitwise equal to diffusers for the block (adaln_indices) and norm_out (timestep_indices), without materialising the (S, H) gathered tensors; 2.5x faster forward at S = 131072. - Deterministic backward: row-local dx, dweight over fixed 256-row tiles, dshift/dscale as sorted segment sums (diffusers: non-deterministic BF16 atomics, ~15x larger error). - Manifest pins the block/refiner-final/norm_out norm weights and norm_out.linear (all in shard 1); chain gains a per-mode history and the block norm1 + modulation stage (covered by the end-to-end test). - Registry, gtest, report/figure hooks, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Produced by scripts/h3_evidence.py at f36aa60 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>
…ernel gate gather RFC RL-Align#420 step 6, row `adaln_gate_residual` (stacked on h3_rmsnorm). - residual + gate[adaln_index] * sublayer_output in one pass, the gate row read from the AdaLN table view in the kernel; the two roundings sit where the eager expression rounds (__fmul_rn/__fadd_rn, no FMA contraction), so the forward is bitwise equal to diffusers in bf16/fp16/fp32. 1.7x faster forward at S = 131072. - Backward: d_residual = grad; d_sublayer bitwise equal to the eager VJP (exact product, one rounding); d_gate a deterministic FP32 segment sum (diffusers: non-deterministic BF16 atomics, 13-49x larger error, varying run to run). The PyTorch reference gets an equally deterministic gate gradient. - Chain stage residual + gate_msa * (stand-in attention output), covered by the end-to-end test; registry, gtest, report/figure hooks, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Produced by scripts/h3_evidence.py at a26b41a on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Reference contracts and numerical behavior rl_engine/reference/minimax_h3/*, rl_engine/validation/models/h3_provider.py, rl_engine/validation/models/h3_cases.py |
Adds reference implementations, validation rules, H3 layout constants, provider replay helpers, and fixed-order FP32 reductions. |
CUDA kernels and extension bindings csrc/cuda/h3/*, csrc/bindings/ops.cpp, build_tools/extensions.py, rl_engine/_C.pyi |
Adds CUDA forward and backward implementations for sinusoidal features, deterministic linear operations, AdaLN row gathering, RMSNorm, and gated residuals. Registers native entry points and controls H3 sources for precise-math CUDA builds. |
CUDA wrappers and runtime dispatch rl_engine/backends/cuda/model_specific/minimax_h3/*, rl_engine/runtime/registry.py, rl_engine/validation/operators/* |
Adds CUDA autograd wrappers, availability checks, input construction, operator specifications, and CUDA-first dispatch with PyTorch fallback on other backends. |
Conditioning chain and replay rl_engine/validation/models/h3_chain.py, tests/models/minimax_h3/test_h3_conditioning_e2e.py, tools/validation/models/h3_chain_replay.py |
Adds staged candidate, provider, and golden forward replay plus repeated backward comparisons for the H3 conditioning chain. |
Weights, measurements, and evidence rl_engine/validation/models/h3_manifest.json, rl_engine/validation/models/h3_weights.py, rl_engine/validation/models/h3_report.py, benchmarks/models/*, tools/validation/models/*, tools/weights/* |
Adds pinned-weight metadata and verification, benchmark and evidence CLIs, prior-art and bitwise comparison runners, and report plotting. |
Operator documentation and migration docs/operators/*, docs/.nav.yml, docs/operators/README.md, tools/migration/* |
Adds six H3 operator guides, navigation entries, and migration mappings for H3 implementations, reports, utilities, and tests. |
Operator and infrastructure tests tests/build/*, tests/models/minimax_h3/*, tests/runtime/test_dispatch.py, tests/validation/operators/* |
Adds coverage for native validation, numerical behavior, deterministic gradients, dispatch, timing, weight integrity, CLI behavior, and end-to-end chain replay. |
Priority:
Estimated code review effort: 5 (Critical) | ~90 minutes
Change: Feature
Sequence Diagram(s)
sequenceDiagram
participant Caller
participant Registry as KernelRegistry
participant CUDA as H3 CUDA operators
participant Reference as H3 PyTorch references
Caller->>Registry: request H3 conditioning operator
Registry->>CUDA: try CUDA backend on CUDA
CUDA-->>Registry: return output or report unavailable
Registry->>Reference: use PyTorch fallback when required
Reference-->>Caller: return reference output
Possibly related issues
- RL-Align/RL-Kernel#420 — This issue describes the H3 WS1 roadmap; the pull request adds the gated-residual operator and CUDA validation covered by that roadmap.
Merge Risk: 
60e77
-
This PR adds the gated-residual operator and the other H3 conditioning operators with the reference implementations, tests and tooling that go with them. No production defect remains open. Some tests still compare tuples without checking that their lengths match, which could hide a missing output in a future regression and also trips the repository's lint rule. This is a small follow-up and does not block the merge.
Security Architecture Review
-
Security architecture risk: 🟡 Moderate · up to
6ba92The new operators enforce inputs differently across their callable interfaces. Malformed lower-level inputs can reach invalid GPU memory accesses. Default high-level calls validate inputs, which substantially limits exposure; no remotely reachable attack path has been established.
Retained concerns
- Medium · security · observed: The newly bound native interfaces do not enforce the same input-safety contract as their wrappers. Forward row values can reach unchecked pointer offsets; gate and RMSNorm index-device ownership is not enforced natively; backward reduction metadata lacks consistent device, cardinality and value-bound checks. Malformed direct calls can therefore cause out-of-bounds reads or cross-device failures. Default wrappers validate inputs and generate metadata internally, so this concern is limited to lower-level callers and unchecked fast paths rather than ordinary validated dispatch.
Security review details
-
Security Blast Radius
- inferred — The demonstrated attack surface requires invoking the bound native tensor functions or supplying malformed indices to an integration that disables range checks. Potential exposure is other memory reachable in the calling process's CUDA address space and failure of its GPU execution context. Cross-process, cross-tenant, host-memory or privilege-escalation reachability has not been established.
Security Findings and Attack Paths
- inferred — A lower-level caller can provide an out-of-range gate row or malformed sorted-position and tile boundaries. Native validation accepts some such tensors based on layout and type, then uses their values in raw-pointer reads. If the resulting address is mapped, unrelated GPU data may influence returned outputs; otherwise execution may fault. These outcomes follow from the access path and were not demonstrated at runtime.
Trust Boundaries and Controls
- observed — Default high-level validators require matching devices and valid row ranges. Wrapper-generated backward metadata comes from stable sorting and bounded segment construction. Native code also preserves activation-layout checks, alignment-aware vectorization, device guards and current-stream launches. These controls protect normal wrapper use but do not make direct native inputs trustworthy.
Resilience and Maintainability Implications
- observed — Weight preparation rejects source-hash and tensor-structure mismatches before saving the extracted file. Repository code writes directly to the fixed target without an explicit temporary-file publication or locking protocol. Atomicity supplied by the serialization dependency is unresolved; this is a bounded tooling recovery gap, not an established security bypass.
Hardening Proposals
- proposed — Establish one explicit enforcement boundary for callable native operators: validate device ownership, row bounds and reduction-metadata bounds at exposed interfaces, or separate checked entrypoints from internal prevalidated helpers with documented caller obligations.
Pre-merge checks: 
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly identifies the gated-residual CUDA operator and its in-kernel gate gather, which matches the stated pull request objective. Docstring Coverage Docstring coverage is inconclusive: files from earlier reviews still need to be rechecked within the check's analysis limits. Not blocking. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request.
Finishing Touches: 💡 1
-
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
fusheng-ji
changed the title
Flink-ddd
added
platform: cuda
|
@coderabbitai full review |
|
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> # Conflicts: # csrc/ops.cpp
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> # Conflicts: # scripts/plot_h3_evidence.py
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…once, and align the reference validators rmsnorm_modulate.cu: make_modulation rejects a shift/scale table whose row stride is smaller than N (an expanded or overlapping view is not a table of rows; a view of a wider table, such as table.chunk(2, dim=-1) with row stride 2N, is still read by its row pitch and a test pins that it is bitwise the copy). The modulated backward runs the shared tile-metadata pre-pass (tile_metadata.cuh, RL-Align#486) before the segment kernels: sorted positions, tiles and segment boundaries that do not describe the rows end in a device-side assert instead of an out-of-bounds read, and the fold launch is skipped for zero segments. N is checked against the int and grid.y limits of the segment kernels. The column indices of the segment kernels were already int64 in the tile term, so nothing changes there. Tests run corrupted metadata in a child process. rmsnorm.py (CUDA op): a table that is not made of dense rows (stride(0) < H) is copied before the call, like mixed layouts; a wider-table view goes through as it is. Reference: the modulation backward takes the positions of all rows from positions_by_row (one sort and one host read instead of two torch.nonzero scans per row; a test compares the gradient bitwise with the loop it replaces). forward_modulated_fp32 takes check_range like forward_modulated. A device mismatch of shift/scale is a ValueError, as in the other H3 validators; a dtype mismatch stays a TypeError. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Brings in the shared tile-metadata pre-pass, the rmsnorm hardening and the fixes of the lower PRs. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> # Conflicts: # csrc/cuda/h3/rmsnorm_modulate.cu # tests/models/minimax_h3/test_h3_report_timing.py
…n the reference's rows once gate_residual.cu: the backward runs the shared tile-metadata pre-pass (tile_metadata.cuh, RL-Align#486) before the segment kernels, so sorted positions, tiles and segment boundaries that do not describe the rows end in a device-side assert instead of an out-of-bounds read, including the case of no tiles with a non-trivial seg_first_tile (the partial buffer then holds one uninitialized row). Host checks: sorted_pos has M entries, tile_begin and tile_end have the same length, seg_first_tile has at least one entry, and the column blocks fit grid.y (65535); the fold launch is skipped for zero segments. Native gate backward, idle B200, min of 400: 86.6 -> 90.7 us at S = 4097, 362.7 -> 367.5 us at S = 32768, 1315.0 -> 1319.7 us at S = 131072 (one extra launch of about 4 us). Tests run corrupted metadata in a child process and the host checks. gate_residual.py: dy is viewed as the shape of y (validate_h3_gate_residual makes it the residual's shape, so nothing changes for valid calls) and the saved shapes are named for the tensor they belong to. Reference: the gate gradient takes the positions of all rows from positions_by_row (one sort and one host read instead of one torch.nonzero scan and sync per row); a test compares the gradient bitwise with the loop it replaces. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
The B200 report and figure are regenerated from a clean tree at c416a10 on an otherwise idle GPU (utilization 0% before and after) with the backward cases built by candidate_backward_setup / provider_backward_setup. The backward columns of the latency table are now backward-only: 1.01 vs 1.30 ms at S = 4097, 1.76 vs 8.00 ms at S = 32768, 5.59 vs 32.2 ms at S = 131072 (CUDA vs index_select); the peak-memory sentence describes the backward-only increment. The H100 report and chain_replay.json are not regenerated and still carry the earlier timing. Links point at the pr_assets commit that holds the new files. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Brings in the regenerated row-gather evidence and its doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> # Conflicts: # docs/operators/h3-adaln-row-gather.md
Brings in the row-gather doc merge of the previous PR. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| record_backward( | ||
| "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" | ||
| ) | ||
| return grad, dy.view(ctx.y_shape), dgate, None |
There was a problem hiding this comment.
The residual VJP is returned as the raw incoming grad without reshaping, but the forward flattens the inputs: out = _C.h3_gate_residual_forward(res2, y2, ...) then return out.view(residual.shape). validate_h3_gate_residual only requires residual to be (..., S, H) with dim() >= 2, so a batched input ((B, S, H), which the reference explicitly supports and where autograd will pass a (B, S, H) grad) returns a 2-D-shaped residual gradient; and the eager reference (residual + gathered * y) returns the plain first VJP, so d_y/d_gate semantics must come from the same flattened graph. At minimum, return grad.view(ctx.residual_shape) (the y branch already does dy.view(ctx.y_shape)), otherwise a 3-D residual gets a mis-shaped/incorrectly-strided gradient from this custom Function.
Suggestion:
| return grad, dy.view(ctx.y_shape), dgate, None | |
| return grad.view(ctx.residual_shape), dy.view(ctx.y_shape), dgate, None |
There was a problem hiding this comment.
Not a defect, made explicit in ab407b3. forward returns out.view(residual.shape), so the cotangent autograd hands to backward already has the residual's shape (3-D for a (B, S, H) input), and the residual VJP of residual + gate * y is that cotangent itself. The existing test_backward runs a (2, 4097, 5376) case and requires d_residual bitwise equal to the eager VJP, which cannot hold with a different shape. ctx.residual_shape was saved and never used, which is what read like a missing reshape; the return is now grad.reshape(ctx.residual_shape) (a no-op for the same shape) and test_residual_gradient_keeps_the_residual_shape_for_batched_and_strided_gradients feeds a 3-D input with a non-contiguous incoming gradient. view (the suggestion in the other thread) would raise on that gradient, so reshape.
…els out, and keep the sinusoid sources self-contained * NVCC_*_FLAGS can switch fast math on through `--options-file` / `-optf` (a comma-separated list, as `--options-file=a,b` or `--options-file a,b`): nvcc expands those files itself, while it does not expand `@file`. The predicate read only the `@file` spelling, so `NVCC_APPEND_FLAGS="--options-file flags.txt"` with `--use_fast_math` inside built the H3 kernels with approximated math. Both spellings are read now; a file that cannot be opened, or an option nvcc rejects for lack of a file name, counts as relaxed (the build would fail anyway). * The build warning names the variable and the reason (`NVCC_APPEND_FLAGS contains --ftz true`, `... cannot read options file ...`, `... cannot be parsed`), so an unreadable value is not mistaken for a relaxed one. * A test uses nvcc as the oracle: `nvcc --dryrun` prints the cicc flags it would use, and the predicate must agree with them for every spelling in the list (plain, split, dashes, options files, nested files, PREPEND and APPEND). Another test checks the nvcc list handed to the compiler: a relaxed flag in it and the H3 define never appear together. * The dispatch fallback test removes the symbol from the extension (instead of replacing the probe) and checks that the CUDA backend was rejected. * `timestep_sinusoid.cu` includes <algorithm> for std::min; `frequencies_fp32` rejects an odd or non-positive channel count like the validator; the H3 block of `_C.pyi` says the symbols exist only in builds that compile the H3 sources. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…n 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>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> # Conflicts: # tests/models/minimax_h3/test_h3_timestep_mlp.py
…jection bias fold, reject a weight without rows * `det_linear_device_supported` returns False where torch raises (an index outside this host's devices, a device string it cannot parse, a CUDA runtime that fails to start), so a constructor that uses it falls back to the PyTorch reference instead of failing; integer indices keep working (`torch.device(0)` is `cuda:0`). Tests cover each case. * `adaln_projection` backward passes `with_bias=` to `det_linear_backward_weight`, like the fused modulation op, so the bias fold does not run when only the weight needs a gradient. A test freezes each input in turn, compares the remaining gradients bitwise with the all-trainable run and records the argument of every call. * `adaln_hidden_size` rejects a (0, D) weight: it divided by 18 and gave H = 0 and six empty outputs. * `det_linear.cu`: the two FP32 launch-shape knobs (`H3_FWD_F32_COLS`, `H3_FWD_F32_AHEAD`) must be overridden together; a build that sets one fails in the preprocessor with that message instead of with an undeclared identifier. The test runs the preprocessor on stand-ins for the torch headers. * `h3_bitwise_reuse._compare_case` records a different number of outputs as a difference (a bare zip stopped at the shorter list and could call the case bitwise equal) and zips strictly otherwise. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
…d show the device asserts survive a release build * The tile-metadata pre-pass also requires `seg_first_tile[0] == 0` and `seg_first_tile[R] == tiles`. The old bound (`<= tiles`) kept every read inside `partial`, but a tile that no segment owns, or a first segment that starts after tile 0, would drop contributions from a row's sum without an error. Two more corrupted-metadata cases run in the child-process test. The launch uses at least one block, so an empty list is a no-op instead of an invalid configuration. * `CUDA_KERNEL_ASSERT` is not removed by NDEBUG on CUDA (c10 expands it to `if (!cond) __assert_fail(...)`; only Android, Apple and FreeBSD compile it out). The header says so, and `test_h3_assert_in_release_build.py` compiles a one-kernel program with `-O3 -DNDEBUG` against the torch headers in use and requires the "device-side assert triggered" stop (and a clean run for a valid value) so the claim is checked for the torch this runs against. * `test_backward_writes_the_requested_dtype_whatever_the_gradient_dtype` pins that the gather backward allocates the table gradient in `out_dtype` for FP32, BF16 and FP16 gradients (exact integer data, torch.equal against the FP64 sum). * `_C.pyi`: contracts for `h3_adaln_row_gather_forward` / `_backward`; the modulation tests zip strictly (Ruff B905). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> # Conflicts: # rl_engine/_C.pyi
…odulation index once before the dweight loop * `check_xw`, called by every rmsnorm entry point, rejects N or M above INT_MAX (the kernels take N as `int` and launch one block per row), and the modulated backward bounds the table rows and tiles it launches as blocks. A test allocates a (1, 2**31) FP16 row and checks forward and backward. * The per-row `CUDA_KERNEL_ASSERT` on the modulation index inside `rmsnorm_dweight_partial_kernel` cost about 11% of the modulated backward (A/B, interleaved min of 300, B200: 1770 -> 1586 us at S = 32768, 6346 -> 5694 us at S = 131072, 378 -> 371 us at S = 4097 once removed). The loop now reads the index unchecked and the entry point runs `h3_tiles::check_index_range` over the S index values first, a 4 us launch like the tile-metadata pre-pass. The forward and dx kernels keep their once-per-row assert. A child-process test requires the device-side assert for negative, past-the-table and far indices in both forward and backward, with a valid-index control. * The backward header states the exact rounding chain of dx (coef = ((rstd*rstd)*rstd)*dot/N, one cast, FMA contraction allowed with fixed flags) and that the backward is measured against FP64, not claimed bitwise to autograd; the empty `Modulation` initialiser says why `seq` is 1. * `_C.pyi`: contracts for `h3_rmsnorm_forward` / `h3_rmsnorm_backward`. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> # Conflicts: # rl_engine/_C.pyi
…l gradient shape * `make_gate` refuses a table whose row stride is below N (an expanded table would make every row alias the same memory), as the modulation tables already did; `H3GateResidualCudaOp` copies such a table to dense rows first, so an expanded gate still works and gives the bytes of its dense copy. Views of a wider AdaLN table (unit column stride, larger row stride) go to the kernel unchanged. * The residual's VJP is the incoming cotangent, which already has the residual's shape because forward returns `out.view(residual.shape)`; `ctx.residual_shape` was saved and unused, which read like a missing reshape. The return is now `grad.reshape(ctx.residual_shape)` (a no-op for equal shapes, and it accepts a non-contiguous gradient where `view` would not), and a test feeds a (B, S, H) residual with a transposed cotangent and compares every gradient with the eager VJP bitwise. * `_C.pyi`: contracts for `h3_gate_residual_forward` / `_backward`. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
| temb, act, weight, timestep_indices, token_tags = ctx.saved_tensors | ||
| need_temb, need_weight, need_bias = ctx.needs_input_grad[:3] |
There was a problem hiding this comment.
_H3AdaLNModulationCuda.backward only saves temb, act, weight, timestep_indices, token_tags, but ctx.needs_input_grad[:3] refers to the first three inputs (temb, weight, bias) — that part is fine. The real problem: backward unconditionally computes _C.h3_adaln_row_gather_backward(...) and record_backward(...) even when no input requires a gradient, which is the normal case for inference (all inputs are parameters under torch.no_grad() there is no backward at all, but any partial requires_grad combination still pays the tile build for all 3T rows). More importantly, record_backward is emitted before the early-out checks, so a pure-inference eval graph with a single differentiable leaf still records a fabricated backward implementation id. Add an early return when not any(ctx.needs_input_grad[:3]).
There was a problem hiding this comment.
Not changed. Autograd runs backward only when some input needs a gradient (the outputs require grad iff an input does, and the index tensors are integers), and the gather backward is the first stage of every gradient here: dW, db and d_temb all start from d_table. The input-specific work is already gated by needs_input_grad (the bias fold, the weight gradient, the input gradient), and record_backward runs after the work that ran, so it records what executed.
| if rows.stride(1) != 1: | ||
| rows = rows.contiguous() |
There was a problem hiding this comment.
H3AdaLNRowGatherCudaOp.forward copies rows with .contiguous() when rows.stride(1) != 1, but _H3AdaLNRowGatherCuda.forward records ctx.rows_meta = (rows.shape[0], rows.dtype) from the tensor it receives. That is consistent because the contiguous copy has the same shape. However the native h3_adaln_row_gather_forward requires rows to be 2-D with unit column stride and rows.size(0) % H3_MODALITY_NUM == 0; rows.size(0) is the number of table rows (3T), correct. The gap: gather_chunks() concatenates six (3T, H) tensors from diffusers into (3T, 6H), which is correct. The one real issue is that forward() rebinds rows locally to a contiguous copy, so the returned packed is detached from the caller's rows graph node — for the rows.stride(1) != 1 branch, the autograd edge is to the copy, which still carries the grad back to the original, so it is fine. No defect confirmed beyond the missing rows.size(0) % H3_MODALITY_NUM Python pre-check (native raises).
There was a problem hiding this comment.
Not changed; the finding concludes "no defect confirmed". The rows.size(0) % H3_MODALITY_NUM pre-check it mentions is made by the native entry point, which raises a TORCH_CHECK with the numbers.
| def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: | ||
| validate_h3_rmsnorm(x, weight, eps) |
There was a problem hiding this comment.
The unmodulated path is reachable but wrong at the dtype/device boundary is not the issue here; what is: _C.h3_rmsnorm_forward is called with the optional shift, scale, index passed through positionally as None, which is fine, but forward() validates with validate_h3_rmsnorm (which accepts fp16) whereas the native check_xw requires x.size(1) % kVec == 0 and rejects at::kDouble. The Python validator accepts any hidden size, so an H = 3 or similarly unvectorized width passes validation and then raises a TORCH_CHECK from the extension instead of a clear Python error; the documented "plain norm is bitwise equal to nn.RMSNorm" claim also only holds for widths on PyTorch's vectorized path. Either also pre-check the width in the Python wrapper or document/reject it. This is a contract/validity gap triggered by any non-multiple-of-kVec hidden size.
There was a problem hiding this comment.
Not changed. A width that is not a multiple of 4 is refused with a clear message from the native entry point (RuntimeError: N=3 must be a multiple of 4 (PyTorch's vectorized RMSNorm path), checked by running it). That restriction is the reason the "bitwise equal to nn.RMSNorm" claim holds: the kernel replays PyTorch's vectorized reduction, and other widths are refused rather than computed on a different path. H3's hidden size is 5376. A Python pre-check would repeat the same message.
…unknown flag values, bound the options-file read The relaxed-math check split every shlex word again on whitespace and ';', so --options-file "my opts.txt" was read as the file 'my' (and a decoy named 'my' hid the real flags). Quoted words now stay whole for file names and are also scanned piece by piece for relaxed flags. A value nvcc does not take (--ftz=yes, --prec-div=no) is reported instead of read as the exact spelling. Options files are read only if they are regular files of at most 1 MiB, so /dev/zero or a FIFO cannot hang the configure step. The unused _external_relaxed_math helper is removed (get_extensions uses _relaxed_math_reason). The nvcc --dryrun oracle test now also asserts that the three cicc options are present and which flag sets relax. H3_ADALN_CHUNKS is the length of H3_ADALN_CHUNK_NAMES, and the sinusoid docstring states the divisor as half. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
check_tile_metadata bounded every tile and cut the tile list exactly by segment, but a gap between tiles (positions dropped from every row's sum) or an overlap (positions counted twice) passed all of it. Each tile must now start where the one before it ends, the first at 0 and the last at the number of positions, and a tile list that is empty over a non-empty sequence is refused on the host. Mutation tests cover a gap, an overlap, a first tile not at 0, a last tile that stops short, and the empty list with and without matching segment bounds. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
validate_h3_modulation read index.dtype first, so forward_modulated(..., index=None) failed with an AttributeError. All three tensors are required for the modulated form; a missing one now raises a TypeError that names it. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
The gate residual backward handed back the cotangent as the residual's gradient whether or not the residual required one. It now follows needs_input_grad[0], like the other ports of this group. The tile-metadata mutation test gains a gap, an overlap and a short last tile, and expects the empty tile list to be refused on the host. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
| const int64_t blocks = grid_for((n_out + kMmaCols - 1) / kMmaCols, kWarpsPerBlock); | ||
| det_linear_forward_bf16_mma_kernel<<<static_cast<unsigned>(blocks), kWarpsPerBlock * kWarp, | ||
| 0, stream>>>( |
There was a problem hiding this comment.
The BF16 MMA forward path has no k_in % vec == 0 / 16-byte-alignment guarantee for the row start. The MMA kernel does 16-byte uint4 loads at w_lo + k and x_row + k with k = kb + 32*u + 8*q, and bx[u] is an unguarded *reinterpret_cast<const uint4*> load (unlike the weight loads, which check k_ok). When k_in is a multiple of 8 but not of 16 (e.g. K=8, K=24), 8*q can be 8, so a lane reads 16 bytes starting at a non-16-byte-offset pointer: misaligned access (undefined behavior / wrong results) and a read past the end of the last row. The FP32 branch is protected by the k_in % vec check, but with x.scalar_type() != at::kFloat that check uses vec = 16/element_size() = 8 for BF16, which is exactly the insufficient bound. Note each lane loads x_row + 8q + {0..7} BF16 elements = 16 bytes, so the k_ok guard must also require k + 16 <= k_in, not merely k < k_in.
Suggestion:
| const int64_t blocks = grid_for((n_out + kMmaCols - 1) / kMmaCols, kWarpsPerBlock); | |
| det_linear_forward_bf16_mma_kernel<<<static_cast<unsigned>(blocks), kWarpsPerBlock * kWarp, | |
| 0, stream>>>( | |
| 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"); | |
| // the BF16 mma path loads 16 bytes at k = kb + 32u + 8q, so require full 16-byte chunks | |
| TORCH_CHECK(x.scalar_type() != at::kBFloat16 || k_in % (16 / x.element_size() * 2) == 0, ...); |
There was a problem hiding this comment.
Not changed (answered in the previous rounds). Same point as the K-tail thread on #486: the load is 8 BF16 elements (16 bytes) at a k that is a multiple of 8, and the host requires K % 8 == 0, so k < K already means k + 8 <= K; the uint4 read of x_row + k is 16-byte aligned because the row stride is a multiple of 8 elements. The x load is guarded by k_ok like the weight loads (bx[u] = (k_ok && x_ok) ? ... : zero). Requiring k + 16 <= K would zero real data for K = 8, 24, ... and change the result.
| index = h3_adaln_indices(timestep_indices, token_tags) | ||
| tiles = _segment_tiles(index, num_timesteps * H3_MODALITY_NUM) | ||
| d_rows = _C.h3_adaln_row_gather_backward(grad_packed.contiguous(), *tiles, torch.float32) | ||
| d_table = d_rows.view(num_timesteps, width) |
There was a problem hiding this comment.
d_rows comes back from _C.h3_adaln_row_gather_backward with shape (num_rows, chunk_width) where chunk_width = H3_ADALN_CHUNKS * hidden, but the table viewed here is (num_timesteps, width) with width = table.shape[1] = H3_ADALN_CHUNKS * hidden * H3_MODALITY_NUM. When H3_MODALITY_NUM > 1 these are different shapes and this view raises (num_rows = num_timesteps * H3_MODALITY_NUM, so the element count matches but the requested shape (num_timesteps, width) does not). It only works by accident when num_rows * chunk_width == num_timesteps * width, i.e. never for the 3-modality table. Use d_matrix = d_rows.reshape(num_timesteps * H3_MODALITY_NUM, -1) and drive det_linear_backward_weight/det_linear_backward_input off that, or reshape with the total element count.
There was a problem hiding this comment.
Not changed; the shapes agree. rows = table.view(-1, 6H) has T * 3 rows, so h3_adaln_row_gather_backward returns d_rows of shape (T * 3, 6H), and T * 3 * 6H == T * 18H == num_timesteps * width, which is the shape view(num_timesteps, width) asks for (the gather chunk width is 6H, not 18H: the three modalities are separate table rows). The forward/backward tests for all three modalities compare every gradient with an FP64 golden, and view raises on any mismatch.
| declared = F.silu(temb.float()).to(weight.dtype).double() | ||
| act = act + (declared - act).detach() |
There was a problem hiding this comment.
act = act + (declared - act).detach() makes the golden's forward value exactly declared, i.e. FP32 F.silu(temb) rounded to the weight dtype, computed by the same provider path the golden is meant to validate. If F.silu or the cast in the provider is wrong, the golden reproduces it and the check passes. The straight-through trick is justified for the gradient (identity VJP), but the forward should keep the independent FP64 value (e.g. declared + (act - declared).detach(), or a proper custom_fwd/identity-VJP cast).
There was a problem hiding this comment.
Not changed, by design; the module and method docstrings state it. The cast of the activation to the weight dtype is part of the declared model (the provider rounds the FP32 SiLU to the weight dtype before the projection), so the golden rounds it too: an FP64 SiLU can round to the other BF16 neighbour near a tie and shift a whole output column, which would flag a correct kernel. What the golden keeps independent is everything after the cast: the projection and the reductions run in FP64 with an FP64 gradient end to end. A wrong FP32 SiLU in the provider is caught by the SiLU/MLP tests, which compare against an FP64 autograd reference with no cast at all (_fp64_grads in test_h3_timestep_mlp.py).
| n = n + (n.to(x.dtype).double() - n).detach() | ||
| one_plus_scale = 1.0 + scale.double().index_select(0, index) | ||
| one_plus_scale = ( | ||
| one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() | ||
| ) |
There was a problem hiding this comment.
Same coupling as the AdaLN projection golden: the straight-through terms make the golden's forward output use the x.dtype-rounded n and 1 + scale, i.e. values produced by the provider rounding path, so the golden cannot detect a wrong rounding in forward/forward_modulated. Keep the FP64 values in the forward and apply straight-through only on the gradient path (st = rounded + (exact - rounded).detach() with the exact value as the returned term, e.g. n_exact + (n_rounded - n_exact).detach() replaced by returning n_exact and letting only the jacobian be identity).
There was a problem hiding this comment.
Not changed, same reasoning as the projection golden. The norm is an x.dtype module and diffusers forms 1 + scale in x.dtype, so those roundings are model semantics, and the comment above the straight-through lines says why: 1 + scale is shared by every position of a row, so leaving it unrounded skews d_norm_weight systematically with S, and the rounding of n enters d_scale. The straight-through term only changes the gradient path (identity VJP); the FP64 reduction (mean, rsqrt, the products) is independent of the provider. A wrong rounding in the provider's forward_modulated would differ from this golden in the cast point, which the tests compare at the declared tolerance.
What
residual + gate[adaln_index] * sublayer_outputin one pass. The gate row (gate_msa/gate_mlp) is read from the AdaLN table view in the kernel. The two roundings fall exactly where the eager expression rounds, with__fmul_rn/__fadd_rnso there is no FMA contraction. The result is bitwise equal to diffusers in bf16, fp16 and fp32.Backward:
d_residualis the incoming gradient;d_sublayeris bitwise equal to the eager VJP, because the product of two 16-bit values is exact and then rounded once;d_gateis a deterministic FP32 segment sum.The PyTorch reference gets the same deterministic gate gradient.
Chain: the stage
hidden + gate_msa * (stand-in attention output), which the e2e test covers.Prior art & reuse decision
Measured with
tools/validation/models/h3_prior_art.py, committed in this stack.gate_residual.jsonwas written from a clean tree atbe7c6a99on an otherwise idle B200, and the operator doc's "Existing implementations" section has the figure and the command. Batch invariance is bitwise, for the forward, the per-row gradients and repeatability of every table and parameter gradient. It has three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows over three seeds; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. Errors are max|err| / max|ref| against FP64; latency is the median forward + backward time. Libraries checked: diffusers, SGLang 0.5.21, Liger 0.8.4, Transformer Engine 2.20.2, vLLM 0.30.0 and Megatron-LM52fbcbc(the training path of Miles and Vime); verl, Miles and Vime have no H3 diffusion operators.The forward and the residual / sublayer gradients are batch-invariant in both, and this kernel's forward is bitwise equal to diffusers'. The gate-table gradient differs:
index_selectaccumulates it with atomics, so it changes from run to run and loses accuracy (0.14 at S = 32768, against 2.0e-3 here). This kernel is slower at S = 4097 (1.2 ms vs 0.64 ms forward + backward) and faster at S = 32768 (1.6 ms vs 2.5 ms). SGLang 0.5.21's fused residual kernels accept per-token gates but define no backward; Liger 0.8.4 and FLA 0.5.2 have no gated-residual op with a per-timestep table. The gate-table gradient in diffusers is not repeatable, so this PR writes its own backward.Direct bitwise comparison (2026-10-13). There is no reuse candidate beyond diffusers'
residual + gate.index_select(...) * y. The forward,d_residualandd_sublayerare bitwise identical to it at every S from 1 to 32768. The gate-table gradient matches at S = 1 and differs from S = 257 on (atomics vs a fixed FP32 order).Direct bitwise comparison
Measured with
tools/validation/models/h3_bitwise_reuse.py(committed in #483, merged down the stack), which runs every candidate thath3_prior_art.candidates()defines on identical inputs, with each global batch-invariant mode in its own process, and compares every forward output and leaf gradient byte for byte with diffusers' path and with this kernel.gate_residual.jsonwas written from a clean tree ata746bff0(python tools/validation/models/h3_bitwise_reuse.py --op gate_residual); the operator doc's "Direct bitwise comparison" section has the figure and the plot command.Results (B200, B = 1, H = 5376, BF16)
d_sublayervs diffusersd_gate: repeat-bitwise / error vs FP64At S = 4097 the backward is 1.02 ms, against diffusers' 0.55 ms, because of the fixed sort and tile setup. The data is in
reports/experiments/h3-adaln-gate-residual-b200/report.json, written from a clean tree at239ef808on an otherwise idle B200; backward timings exclude the leaf creation, the forward pass and the graph construction.Status
tools/validation/models/h3_prior_art.py; report and figure for this op inreports/experiments/h3-prior-art-b200). Sincea05952fthe branch gained only that runner, its evidence and the merges carrying them through the stack (no change outsidedocs/and those scripts). See Prior art.a05952f), carrying the review fixes of [WS1][CUDA][MiniMax-H3] adaln_projection_3mod: three-modality AdaLN projection (RFC #420 stack 3/10) #485–[WS1][CUDA][MiniMax-H3] h3_rmsnorm: RMSNorm with fused AdaLN modulation (RFC #420 stack 5/10) #487 and [WS1][CUDA][MiniMax-H3] h3_rmsnorm: RMSNorm with fused AdaLN modulation (RFC #420 stack 5/10) #487's modulated-golden fix. The native RMSNorm input checks now combine both branches' checks; the empty row-index and tile-metadata errors use one wording across [WS1][CUDA][MiniMax-H3] h3_rmsnorm: RMSNorm with fused AdaLN modulation (RFC #420 stack 5/10) #487 and this PR, and the tests were updated to match.tests/models/minimax_h3ata05952fwith two B200s visible: 399 passed, 1 skipped (the plotting test, because matplotlib is not installed in the test environment).History
h3_report.py). The case now prepares the graph before the timed call; the report is rerun from a clean tree at239ef808with the alternating execution order. Backward: 1.02 / 1.08 / 2.21 ms for S = 4097 / 32768 / 131072 (was 1.19 / 1.37 / 3.29 ms), diffusers 0.55 / 2.24 / 8.92 ms (was 0.68 / 2.71 / 10.6 ms); the speedup at S = 131072 is 4.0x (was 3.2x). Accuracy is unchanged (thed_gateerror of diffusers varies run to run, 4.2e-2 now). The figure andreport.jsonmoved to thepr_assetscommit named in the links above. H100 evidence was rerun ata2d0674b(see Results (H100 / Hopper)).pr_assets(see the evidence note at the top), and all commits got aSigned-off-bytrailer. Commit hashes in this body, including the clean-tree commit of the measurement, are the rewritten ones; the ZIPs and JSON files onpr_assetsstill carry the pre-rewrite hashes.KERNEL_ALIGN_USE_FAST_MATH=1adds--use_fast_mathto the whole_Cbuild, which approximatesexpf/sinf/cosf/rsqrtf/division and flushes denormals. Before the fix a fast-math build of the top of the stack failed 33 H3 tests (thetimestep_sinusoid_h3andh3_rmsnormbitwise tests and the chain/block tests built on them). nvcc 13 defines no fast-math macro, so the guard is in the build: everycsrc/cuda/h3source is listed inH3_BITWISE_SOURCESand left out of fast-math builds, the bindings are gated onKERNEL_ALIGN_WITH_H3_KERNELS, dispatch falls back to the PyTorch references, and tests that assert CUDA dispatch skip throughrequire_h3_kernels. On this branch: mergef75ff235(addsgate_residual.cu). Results for both builds are under Tests.test-h3@16b81175874e(notice): H3 CUDA code moved torl_engine/backends/cuda/model_specific/minimax_h3, references torl_engine/reference/minimax_h3, the harness torl_engine/validation/models, with registration, build, tests, tools, docs andtools/migration/layout.jsonupdated. Native arithmetic is unchanged; pinned evidence keeps its original commit metadata.Pre-refactor validation (historical, before 2026-10-10)
python -m pytest tests/models/minimax_h3/test_h3_adaln_gate_residual.py tests/models/minimax_h3/test_h3_conditioning_e2e.py -v python tools/validation/operators/check_operator.py --op adaln_gate_residual --candidate cuda --device cuda --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad python tools/validation/models/h3_evidence.py --op adaln_gate_residual --out report.json && python tools/validation/models/plot_h3_evidence.py report.jsontests/models/minimax_h3+ dispatch + gtest-GPUtest-h3--check-grad, CUDA and PyTorchThe repository's CI workflows only run on PRs to
mainortest, so ontest-h3only the DCO and CodeRabbit checks run. The results above come from local runs on one B200.Tests
Measured at commit
9b95c4a7, an earlier head of this PR (later commits are listed under History and these runs were not repeated for the current head), 2× B200, pinnedMiniMax-H3@42ed227weights; each build compiled from a clean checkout of that commit:_CKERNEL_ALIGN_USE_FAST_MATH=1Default-build skips (environment only): 1 × could not import 'matplotlib'; 1 × could not import 'sglang.kernels.ops.diffusion.modulate.timestep_embedding_jit'; 1 × extension enforcement is CI-only (orchestrator sets RL_KERNEL_REQUIRE_EXT=1).
Fast-math regression tests:
tests/build/test_build_platform_collectives.py::test_cuda_build_compiles_h3_kernels_without_fast_math,::test_fast_math_build_excludes_bitwise_h3_kernels,tests/runtime/test_dispatch.py::test_h3_timestep_falls_back_when_fast_math_build_omits_kernel,::test_h3_ops_prefer_cuda_and_fall_back_to_pytorch_reference.On this PR the GitHub checks are DCO and CodeRabbit only (the CPU lint/unit/docs workflow for
test-h3arrives with #586), so the GPU results above are the test evidence.Scope / Known limitations
CUDA only. The ROCm column stays open. Until the attention and FFN rows land, the chain feeds a seeded stand-in for the sublayer output.
.coderabbit.yamlis not part of this operator. The branch carries thetest-h3CodeRabbit auto-review config, inherited through the stack from commitsd113891and5c79216. The file is byte-identical to #490's, so it drops out of this diff once #490 merges intotest-h3.Summary by CodeRabbit
Results (H100 / Hopper)
Measured on 4 x NVIDIA H100 94GB (sm90, 95830 MiB reported), driver 580.173.02, CUDA 13.0, torch 2.13.0+cu130, pinned
MiniMax-H3@42ed227, clean tree ata2d0674b. Timing evidence was taken on an otherwise idle GPU (idle_before_evidence=yes). B200 and H100 are independent measurement sets (RFC #420); timings and raw bytes are not compared across GPUs. Evidence onpr_assets@5f14a845underpr488/reports/experiments/*-head-a2d0674b-h100/.Tests (
pytest tests/models/minimax_h3 tests/build tests/runtime/test_dispatch.py)KERNEL_ALIGN_USE_FAST_MATH=1build: 165 passed, 411 skipped (H3 CUDA sources omitted for bitwise parity)Bitwise vs diffusers: every shared entry matches the B200 outcome (same equal / not-equal on diffusers vs rl_kernel).
Backward (H100): gradients on identical inputs and upstream gradients, each run twice, sizes 1, 257, 4097, 32768. The kernel gradient is bitwise repeatable for every gradient at every size; native diffusers is not repeatable for
gate(atomicindex_selectbackward). Errors are max|err| / max|ref| against FP64 autograd at size 4097; "det diffusers" usestorch.use_deterministic_algorithms(True).resygateData:
gate_residual-head-a2d0674b/.Repeat check with 30 replays per size (
native-repeat30-head-a2d0674b/gate_residual/): native diffusers gives 30 distinct bitwise results forgate(max relative spread 0.16) at 32768 rows. Deterministic diffusers and the kernel give 1 distinct result for every gradient and size; all other gradients are repeatable in all three.Diffusers-side fix (report @
1a50f2a1): stable-sort plus FP32 segmented sum (RowGather) — validation harness only, unchanged from prior H100 publication.For the gradients that go through the gather the kernel is about 10 to 30x more accurate than diffusers (native or deterministic): diffusers accumulates BF16 values in a fixed order, the kernel reduces in FP32 in a fixed tree. Element-wise and non-gather gradients are at the same rounding level as diffusers.
Figures
