dsv4.1: mHC computation and compensated projections

[hnyls2002]

Summary

• Batch-invariant mHC mixing statistics for DeepSeek-V4.1's hyper-connection layers, in mhc.py: hc_mix_stats computes the per-row mixing projections and sum of squares over a fixed number of K slices with tf32x3 dots, so a row's result is bitwise identical whether computed alone or in a batch; hc_mix_stats_sinkhorn fuses the slice reduction with the sinkhorn normalisation that hc_split_sinkhorn performs.
• Two compensated variants of the same projection for large batches, keeping the fp32 weight exact on the tensor cores: hc_mix_stats_sinkhorn_bf16x3 (Triton, three bf16 weight components, one activation read) and hc_mix_stats_sinkhorn_deepgemm (tf32 high part plus fp32 residual through DeepGEMM), with split_bf16_hc_weight / split_tf32_hc_weight preparing the weight parts. Both accumulate over _HC_MIX_COMPENSATED_SLICES and share the sinkhorn reduce kernel; sizes come from the input shapes and hc_mult.
• hc_split_sinkhorn answers an empty batch directly instead of launching a zero-sized grid, which the TileLang backend rejects.

Changes to existing kernels

• hc_split_sinkhorn: the empty-batch return only. Its existing caller already returns before reaching it with no tokens, so results on existing models are unchanged; the guard is for callers that do not pre-check.

Verification

• Sinkhorn arithmetic (eps placement, iteration count, RMS scaling order) matches the existing Triton hc_split_sinkhorn port and the reference implementation.
• The two compensated variants produce bitwise identical outputs to their previous standalone-module versions for M in {4096, 5000, 8192} (bf16x3) and {128, 1000, 4096} (DeepGEMM) at K = 20480.
• test_mhc_kernels.py and test_hc_combine.py cover the untouched kernels of the module.

---

CI States

Latest PR Test (Base): ❌ Run #35165253652
Latest PR Test (Extra): ❌ Run #35165253438
Latest PR Test (AMD ROCm 10): ❌ Run #35165253625

[hnyls2002]

/rerun-test test/registered/kernels/ops/layernorm/test_mhc_kernels.py test/registered/kernels/ops/layernorm/test_hc_combine.py

[github-actions]

Results for /rerun-test test/registered/kernels/ops/layernorm/test_mhc_kernels.py test/registered/kernels/ops/layernorm/test_hc_combine.py:

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/layernorm/test_mhc_kernels.py
cd test/ && python3 registered/kernels/ops/layernorm/test_hc_combine.py