Module hc

Module hc 

Source
Expand description

DeepSeek-V4.1 block scaffold: hyper-connections (hc_mixes with the Sinkhorn split, hc_pre, hc_post) and RMSNorm, CPU reference against the golden.

The residual stream is hc_mult parallel copies [tokens, hc, dim]. A sub-layer sits between hc_pre (collapse the copies to one input, weighted by a pre mix) and hc_post (expand its output back out, mixing the residual in through a doubly-stochastic comb). hc_mixes derives all three coefficient sets from the stream itself, one projection of the flattened hc * dim stream normalised by its RMS, then split by hc_split_sinkhorn: pre = sigmoid(.) + eps, post = 2 sigmoid(.), comb = row softmax + eps, then alternating column / row normalisations (sinkhorn_iters in total).

The mixes are DELAYED: a block’s attention uses the pre its predecessor’s FFN produced, and its FFN uses the pre its own attention produced (Block.forward). The tests reproduce that order exactly and check every intermediate the golden captured.

Orientation of comb, the one place this is easy to get backwards: with comb[j][k], the output copy is k and the residual copies j are summed, y[k] = post[k] * x + sum_j comb[j][k] * residual[j] (the reference’s sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2)).

Functions§

hc_mixes
Block.hc_mixes. x: [tokens, hc, dim] (bf16 values); hc_fn: [(2+hc)*hc, hc*dim] f32. Returns (pre[tokens, hc], post[tokens, hc], comb[tokens, hc, hc]).
hc_post
Block.hc_post: y[k][d] = post[k] * x[d] + sum_j comb[j][k] * residual[j][d], to bf16.
hc_pre
Block.hc_pre: collapse the copies, y[d] = sum_c pre[c] * x[c][d], back to bf16.
rms_norm
RMSNorm.forward: w * (x * rsqrt(mean(x^2) + eps)) in f32, back to bf16.
split_sinkhorn
ds41_ref_shims.hc_split_sinkhorn for one token. mixes: [(2 + hc) * hc]. Returns (pre[hc], post[hc], comb[hc * hc] row-major [j][k]), all f32.