Add batch-invariant h_aggregate kernel

[nodeeeeee]

Summary

This PR adds a CUDA implementation of the MHC H Aggregate operation:

residual: [T, 4, H], BF16
pre:      [T, 4],    FP32
output:   [T, H],    BF16

For every token and hidden position, the kernel computes:

output = bf16(
    (pre[0] * residual[0] + pre[1] * residual[1])
  + (pre[2] * residual[2] + pre[3] * residual[3])
)

The change only covers mhc_pre / h_aggregate. It does not modify Sinkhorn,
GEMM, RMSNorm, MHC Post, or residual RMSNorm.

Batch invariance

The 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;
• FP32 __fmul_rn and __fadd_rn operations;
• fixed (0 + 1) + (2 + 3) addition tree;
• BF16 conversion only at the input and output boundaries.

Even hidden sizes use __nv_bfloat162 loads and stores. Odd hidden sizes use a
scalar fallback. The two values in each BF16 pair are still calculated
independently.

Correctness

tests/test_mhc_pre_h_aggregate.py directly calls
rl_engine._C.mhc_pre_h_aggregate and checks:

1. The same input produces byte-identical output across 100 calls.
2. A full T=129 batch and 129 one-token calls produce byte-identical output.
3. The result matches the PyTorch reference within atol=5e-2, rtol=1e-2.

T=129 intentionally exercises different launch configurations: the full
batch uses 512 threads per block, while each one-token call uses 1024 threads.
This verifies that changing the batch partition does not change output bytes.

Additional H100 checks covered (T, H) values from (1, 4096) through
(1024, 7168) and the odd-size scalar fallback (3, 1279). All cases were
byte-identical across 100 calls and matched the PyTorch reference. The maximum
absolute difference from the original TileLang stage was 0.015625.

Performance

Environment: NVIDIA H100 80GB HBM3, Torch 2.13.0+cu129, CUDA toolkit 12.8,
BF16 residual/output, and FP32 pre/accumulator. The baseline is the original
TileLang H Aggregate stage. Timing used CUDA Graph replay, cold L2, and the
FlashInfer timing helper with its CUDA Event fallback because CUPTI was not
available.

T	H	Max diff vs TileLang	TileLang	CUDA	Speedup	CUDA GB/s
1	4096	0	9.376 us	6.288 us	1.491x	6.5
8	4096	0	9.840 us	6.576 us	1.496x	49.8
128	4096	0.015625	11.104 us	8.256 us	1.345x	635.3
1	7168	0	13.056 us	7.280 us	1.793x	9.8
8	7168	0.0009765625	13.744 us	7.632 us	1.801x	75.2
128	7168	0.0078125	15.584 us	9.920 us	1.571x	925.1
1024	7168	0.015625	33.536 us	33.888 us	0.990x	2166.5

Model run

The model-level test uses Hugging Face DeepseekV4ForCausalLM with random BF16
weights. The only architectural size change is num_hidden_layers, from 43 to
1; the retained layer keeps the official hidden size, attention dimensions,
MHC expansion, and MoE dimensions from deepseek-ai/DeepSeek-V4-Flash. The
resulting model has
7,635,547,771 parameters and 14.22 GiB of BF16 weights.

The model and test live in the vLLM workspace rather than this repository:

/home/ubuntu/vllm/models/DeepSeek-V4-One-Layer-MHC
/home/ubuntu/vllm/tests/models/run_deepseek_v4_mhc_h_aggregate.py

The test replaces both MHC collapse sites inside the decoder layer and the
final HyperHead collapse with rl_engine._C.mhc_pre_h_aggregate.

The model test verified:

• all 12 expected H Aggregate calls were observed across the test run;
• repeated model forwards produced byte-identical logits;
• logits were byte-identical to the original PyTorch mHC model path.

The full-batch and per-sample model logits were not byte-identical and had a
maximum absolute difference of 0.03125. The isolated H Aggregate test is
byte-identical under the same partition change, so this model-level drift comes
from other batch-sensitive operations in the complete layer.

This is an integration test, not a model-quality evaluation. The official
DeepSeek V4 Flash checkpoint has hundreds of billions of parameters and cannot
run on the single 80GB H100 test machine.

Commands and results

The extension build completed successfully:

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

The lightweight RL-Kernel operator test passed:

RL_KERNEL_REQUIRE_EXT=1 \
  .venv/bin/python -m pytest tests/test_mhc_pre_h_aggregate.py -v

1 passed in 1.97s

The heavier model run was executed from the vLLM repository:

cd /home/ubuntu/vllm
RL_KERNEL_ROOT=/home/ubuntu/RL-Kernel \
  .venv/bin/python tests/models/run_deepseek_v4_mhc_h_aggregate.py

kernel_calls=12
repeat_byte_equal=True
batch_split_byte_equal=False
batch_split_max_abs=0.03125
pytorch_reference_max_abs=0.0
PASS

[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: Pro Plus

Run ID: bc6a0c1f-4cdf-457e-8855-15577b4c76ae

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!