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_sinkhornfor one token.mixes:[(2 + hc) * hc]. Returns (pre[hc],post[hc],comb[hc * hc]row-major[j][k]), all f32.