perf(dsv4): compute the hyper-connection pre-norm in one kernel

[gdevenyi]

The problem

Block.hc_pre and Transformer.hc_head compute the inverse RMS of the hidden state:

xf    = x.flatten(2).float()
rsqrt = torch.rsqrt(xf.square().mean(-1, keepdim=True) + eps)

square() writes a second full fp32 copy of the hidden state. mean() then reads that copy back and reduces it to one value per row. At DeepSeek-V4-Flash shapes (hc_dim 16384, T=8192) the pair moves 1.6 GB to produce a [T, 1] result.

The change

This PR adds inv_rms to kernel/triton/dsv4/norm.py, beside the existing rms_norm. It reduces each row in one pass and returns the scalar.

rms_norm cannot serve here. It applies the scale to x and writes a full-size output, but the hyper-connection pre-norm multiplies the scalar into a [.., mix_hc] tensor instead.

The kernel reads the bf16 source rather than the fp32 upcast. The upcast is exact, so the squares do not change, and this halves the bytes again.

Measurements

One RTX 6000 Ada, hc_dim 16384, measured on main.

rows	before	after	speedup
64	0.020 ms	0.017 ms	1.2x
512	0.030 ms	0.016 ms	1.9x
4096	0.972 ms	0.153 ms	6.3x
8192	1.944 ms	0.304 ms	6.4x

The whole pre-norm block, which also includes the upcast and the narrow fp32 linear that follows:

T	before	after	speedup
1024	0.307 ms	0.108 ms	2.83x
4096	1.819 ms	0.966 ms	1.88x
8192	3.628 ms	1.969 ms	1.84x

At 8192 rows the kernel moves 268 MB in 0.304 ms. That is 880 GB/s, against a 960 GB/s peak on this card.

Effect on a running server

Two nsys captures of main, one without the change and one with it. DeepSeek-V4-Flash, --moe-backend offload, TP=1, 16 concurrent requests, prompts of about 1700 tokens. Both runs completed 33 requests in the 45-second window, so the work is comparable.

kernel	baseline	with inv_rms
vectorized_elementwise_kernel (the square)	633 launches, 181.8 ms	0
reduce_kernel<512,1,ReduceOp<float>> (the mean)	633 launches, 91.9 ms	0
_inv_rms_kernel	absent	629 launches, 48.9 ms
total GPU kernel time	5.830 s	5.587 s

273.7 ms of kernel time becomes 48.9 ms. Total GPU kernel time falls 4.2%.

That figure is small in wall-clock terms on the machine it was measured on, because --moe-backend offload puts the routed experts on the CPU and leaves the GPU 13% busy. A GPU-resident deployment spends a larger share of its time here.

Accuracy

The result is not bit-identical to square().mean(), and it cannot be. ATen chooses its reduction order per shape. A tree-of-two fold reproduces it exactly at M=4, and no candidate order reproduces it at M=64, because the block and grid split changes with M. Any fused form therefore moves the last bit.

What matters is that accuracy does not get worse. Measured against an fp64 reference:

rows	ATen mean relative error	inv_rms
64	3.00e-08	3.24e-08
512	3.47e-08	3.50e-08
4096	3.63e-08	3.60e-08
8192	3.60e-08	3.65e-08

inv_rms is ahead at 4096 rows and behind at the other three, always within 8% of ATen's own distance from the truth. A test pins this.

Two faster-looking forms were rejected on that test. linalg.vector_norm matches the speed but is consistently about 23% worse, at 4.47e-08, because taking a square root and then squaring it rounds twice. einsum over bf16 accumulates in bf16 and reaches 1.9e-03.

Bit-identical forms exist, but none is faster. vecdot(xf, xf) lowers to the same path and takes 1.944 ms. vecdot with fp32 accumulation over bf16 takes 4.557 ms.

Effect on generated text

The server is deterministic. Two baseline runs produced byte-identical greedy output. Moving the last bit therefore changes greedy continuations. Reviewers who depend on exact-output tests should treat this as a deliberate trade, not an invisible one.

Testing

tests/dsv4/test_hc_inv_rms.py covers shape, dtype, leading dimensions, magnitude range, and the accuracy floor.

On main with this change: tests/dsv4 tests/models tests/kernels gives 265 passed, 1 skipped, 1 failed. The failure is test_e4m3_compat.py::test_forced_emu_matches_native, which also fails on main without this change. #85 fixes it.

🤖 Generated with Claude Code

https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun