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.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 = hc(hidden_streams)
sublayer_out = sublayer(collapsed)
hidden_streams = comb @ hidden_streams + post.unsqueeze(-1) * sublayer_out

The gate math always runs in float32, regardless of dtype, matching the reference implementation. collapsed is cast back to the input’s dtype.

from max.dtype import DType
from max.graph import DeviceRef, Graph, TensorType
from max.nn 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 = hc(graph.inputs[0].tensor)
    graph.output(post, comb, collapsed)

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.