For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

HyperConnection

HyperConnection

class max.nn.legacy.HyperConnection(hidden_size, hc_mult, device, dtype=float32, *, hc_sinkhorn_iters=20, hc_eps=1e-06, rms_norm_eps=1e-06)

source

Bases: Module

Mixes hc_mult parallel residual streams around a transformer sublayer.

Manifold-Constrained Hyper-Connections (mHC) generalize the residual connection: instead of one stream carrying the residual between sublayers, hc_mult streams run in parallel, and a learned gate decides how they collapse into the sublayer’s input and how the sublayer’s output is written back. Reference: Xie et al. 2026, section 2.2 equation 8.

A single projection of the normalized streams produces the three gates:

• pre: how the parallel streams collapse into the single sequence the sublayer consumes. Applied here, producing collapsed.
• post: where the sublayer’s output lands across the streams, in [0, 2]. Returned for the caller to apply.
• comb: an hc_mult x hc_mult mixer over the streams, projected onto the doubly-stochastic manifold by Sinkhorn-Knopp. Returned for the caller to apply.

A decoder layer holds two of these, one for the attention site and one for the feed-forward site. The caller drives the surrounding residual update:

post, comb, collapsed, pre = hc(hidden_streams)
sublayer_out = sublayer(collapsed)
hidden_streams = comb @ hidden_streams + post.unsqueeze(-1) * sublayer_out

A model that carries the collapse weights from one sublayer to the next passes pre to the next call as its pre_mix.

The gate math always runs in float32, regardless of dtype, matching the reference implementation. collapsed is cast back to the input’s dtype. On gfx950, a call without pre_mix on bfloat16 streams runs as one fused op (mhc_pre()) where the shape allows it.

from max.dtype import DType
from max.graph import DeviceRef, Graph, TensorType
from max.nn.legacy import HyperConnection

device_ref = DeviceRef.GPU()
hc = HyperConnection(
    hidden_size=4096, hc_mult=4, device=device_ref
)

input_type = TensorType(
    DType.bfloat16, ["total_seq_len", 4, 4096], device=device_ref
)
with Graph("hyper_connection", input_types=[input_type]) as graph:
    post, comb, collapsed, pre = hc(graph.inputs[0].tensor)
    graph.output(post, comb, collapsed, pre)

Parameters:

• hidden_size (int) – Width of a single residual stream.
• hc_mult (int) – Number of parallel residual streams. Must be a power of two whose square fits in one warp.
• device (DeviceRef) – Device the weights and the gate computation live on.
• dtype (DType) – Dtype the three weights are stored in. They are cast to float32 for the gate computation.
• hc_sinkhorn_iters (int) – Sinkhorn-Knopp iterations used to project comb.
• hc_eps (float) – Epsilon guarding the Sinkhorn divisions.
• rms_norm_eps (float) – Epsilon of the unweighted RMSNorm applied to the streams before the projection.

hc_base

hc_base: Weight

source

Per-output bias of shape [2 * hc_mult + hc_mult ** 2], in pre, post, comb order.

hc_fn

hc_fn: Weight

source

Stream projection of shape [2 * hc_mult + hc_mult ** 2, hc_mult * hidden_size].

hc_scale

hc_scale: Weight

source

Per-output scale of shape [3], in pre, post, comb order.