[WS1] Add batch-invariant h_aggregate kernel

[nodeeeeee]

[CUDA][MHC] Add batch-invariant H Aggregate forward and FP32 backward

Summary

This PR adds a deterministic CUDA implementation of the unfused
mhc_pre / h_aggregate boundary, including an explicit FP32 backward:

h_aggregate_fwd(
    residual: BF16[T, 4, 4096],
    pre:      FP32[T, 4],
) -> H:       BF16[T, 4096]

h_aggregate_bwd(
    dH:       BF16[T, 4096],
    residual: BF16[T, 4, 4096],
    pre:      FP32[T, 4],
) -> (
    dR_from_aggregate: FP32[T, 4, 4096],
    dPRE:              FP32[T, 4],
)

The change only covers H Aggregate. It does not implement or modify Sinkhorn,
the controller GEMM/RMSNorm path, MHC Post, or residual RMSNorm.

The backward interface deliberately keeps dR_from_aggregate in FP32 so that
a future full mhc_pre_bwd can merge it with dR_controller in FP32 before
the next permitted downcast boundary.

Fixed numerical paths

For every token and hidden position, forward uses four separately rounded FP32
products and the fixed (0 + 1) + (2 + 3) addition tree:

s01 = pre[0] * residual[0] + pre[1] * residual[1]
s23 = pre[2] * residual[2] + pre[3] * residual[3]
H   = BF16(s01 + s23)

There is one FP32-to-BF16 conversion at the final H store. There are no BF16
multiply or accumulation intermediates.

Backward computes:

dR_from_aggregate[i, d] = FP32(dH[d]) * pre[i]
dPRE[i] = fixed_tree_sum_d(FP32(dH[d]) * FP32(residual[i, d]))

One block owns one token. For dPRE, 256 threads accumulate fixed strided
subsets of the 4096 hidden elements, reduce within eight warps, and then use
warp zero to reduce the eight warp sums. Compile-time assertions enforce
complete warps, CUDA's block-size limit, the one-warp collection limit, and a
fixed number of hidden elements per thread.

No Split-K, Stream-K, atomics, runtime reduction-tree selection, or fast-math
is permitted. The extension build fails closed if
KERNEL_ALIGN_USE_FAST_MATH=1 is requested while this kernel is included.

Batch invariance

The forward arithmetic path is fixed for every output element:

• one block handles one token;
• each output element has one writer;
• no atomics or cross-thread reduction;
• explicit round-to-nearest FP32 multiply and add operations;
• fixed (0 + 1) + (2 + 3) addition tree;
• one BF16 conversion at the final output boundary.

The backward path is also fixed:

• one block owns all four dPRE values for one token;
• each thread visits the same 16 hidden positions;
• the two-level FP32 reduction tree is independent of batch partitioning;
• each dR_from_aggregate element has one writer;
• both backward outputs remain FP32.

Supported profile and fail-closed behavior

This boundary supports only:

• hc_mult = 4;
• hidden size D = 4096;
• contiguous BF16 residual and dH;
• contiguous FP32 pre;
• all tensors on the same CUDA device;
• compute capability 8.0 or newer;
• the unfused, aggregate-only boundary.

Unsupported shapes, dtypes, strides, or devices raise an error. The Python
wrapper does not silently call .contiguous() and does not register an
autograd fallback. Training integration must invoke the explicit
backward_fp32 boundary.

Fusion modes, trainability modes, TP/PP placement metadata, and controller
gradient merge are not accepted by this aggregate-only interface. They belong
to the future full mhc_pre composite boundary.

Forward experiments

Isolated operator correctness

The isolated CUDA operator was checked for:

1. byte-identical output across repeated calls;
2. byte-identical full-batch and token-by-token execution;
3. byte-identical logical rows with right and left padding;
4. exact byte equality with the explicitly associated FP32 oracle after the
   single final BF16 conversion;
5. invariance across the T=128 launch threshold, where the forward block size
   changes from 1024 to 512 threads.

Forward performance

The performance experiment was rerun after the FP32 numerical contract was
finalized, using the kernel source at PR head 38e3f82. Environment: NVIDIA
H100 80GB HBM3, Torch 2.13.0+cu129 with CUDA 12.9 runtime, CUDA 12.8 toolkit,
BF16 residual/output, and FP32 pre/multiply/accumulation.

The baseline is the original TileLang H Aggregate stage. Both paths were timed
as kernel-only CUDA Graph replays with cold L2. Each result is the median of
five interleaved measurements. The FlashInfer timing helper used its CUDA Event
fallback because CUPTI was unavailable.

T	H	Max diff vs TileLang	Byte equal to fixed FP32 oracle	TileLang	Current CUDA	Speedup	CUDA GB/s
1	4096	0	yes	9.408 us	6.368 us	1.477x	6.4
8	4096	0.0009765625	yes	9.856 us	6.464 us	1.525x	50.7
128	4096	0.03125	yes	11.200 us	8.096 us	1.383x	647.8

Model integration experiment

A one-layer DeepseekV4ForCausalLM integration experiment replaced both MHC
collapse sites in the decoder layer and the final HyperHead collapse with the
CUDA H Aggregate operator. The retained layer kept the official hidden size,
attention dimensions, MHC expansion, and MoE dimensions.

The experiment observed all 12 expected H Aggregate calls. Repeated model
forwards produced byte-identical logits, and the logits were byte-identical to
the original PyTorch MHC path. Full-batch and per-sample model logits differed
by at most 0.03125; the isolated H Aggregate operator remained byte-identical
under the same partition change, so that model-level drift originates in other
batch-sensitive operations in the complete layer.

This was an integration test, not a model-quality evaluation.

Backward experiments

The explicit CUDA backward was validated on H100 against an oracle that does
not use PyTorch autograd. The experiments check:

1. dR_from_aggregate and dPRE are FP32 at the public boundary;
2. both outputs are byte-identical to the explicit fixed-tree FP32 oracle;
3. repeated backward calls and recomputation are byte-identical;
4. full-batch and token-by-token backward execution are byte-identical;
5. logical gradients remain byte-identical under right padding, left padding,
   sample permutation, and chunked execution;
6. unsupported shapes, dtypes, and non-contiguous layouts fail closed;
7. gradient-enabled forward calls fail closed instead of silently using an
   autograd or reference fallback.

These experiments validate the raw aggregate gradients. The future composite
boundary must separately validate the FP32
dR_from_aggregate + dR_controller merge.

Reference and test coverage

NativeMHCPreHAggregateOp provides an explicitly associated FP32 forward
oracle and an explicit backward_fp32 oracle. The backward oracle uses the
same 256-thread/two-level fixed reduction tree as CUDA and does not call
PyTorch autograd.

tests/test_mhc_pre_h_aggregate.py, tests/test_operator_inputs.py, and the
registered gradient-invariance tests cover the fixed mixed-precision profile,
accuracy, determinism, batch/chunk/padding/permutation invariance, and
fail-closed behavior.

H100 validation:

MAX_JOBS=8 TORCH_CUDA_ARCH_LIST=9.0 \
  .venv/bin/python setup.py build_ext --inplace

RL_KERNEL_REQUIRE_EXT=1 \
  .venv/bin/python -m pytest \
    tests/test_mhc_pre_h_aggregate.py \
    tests/test_operator_inputs.py \
    tests/test_gradient_invariance.py -q

53 passed, 6 skipped, 0 failed

The skipped cases are conditionally inapplicable cases, not test failures.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Team

Run ID: 1849f379-a0d7-487e-9a9e-1bafc77ce21d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[zhangj1an]

Thanks for your contribution!

Upon a quick scan, please add:

1. backward feature for h_aggregate, just so that training loss can converge, and
2. register this op into gtest. (gtest is a rl-kernel tool to benchmark the following 4 items:

 forward_accuracy   forward_invariance
  gradient_accuracy  gradient_invariance

after that i will review the new files by each line.
Thanks again for your incredible speed!