vllm.models.qwen4_exp.nvidia.hyperconnection

HyperConnection (Gated Residual) utilities — NVIDIA model variant.

Implements the HyperConnection residual scheme proposed in "HyperConnections" (https://arxiv.org/abs/2409.19606). This NVIDIA variant delays each HC combine to the following HC mix boundary. HC glue kernels, including fused combine+RMSNorm, live in ops/hc.py; projections remain standard vLLM Linear modules.

Hidden states between layers have shape [..., HC*HS] with HS inner (HC outer, HS inner — checkpoint-native layout).

Typical usage inside a transformer decoder layer::

self.attn_hc = GatedResidual(hc_config)

hidden_states, block_input, injection = self.attn_hc.mix(hidden_states)
attention_output = attention(block_input)
hidden_states, block_input, injection = self.mlp_hc.combine_and_mix(
    hidden_states, attention_output, injection
)

Classes:

• GatedResidual –

  Gated HyperConnection with learnable low-rank mixing and injection.

• HyperConnectionConfig –

  Configuration shared by all HyperConnection variants.

GatedResidual

Bases: Module

Gated HyperConnection with learnable low-rank mixing and injection.

combine_and_mix() runs the pre pipeline (grouped GemmaRMSNorm -> merged low-rank down+inject GEMM -> silu -> up GEMM -> sigmoid -> gated mean over the HC streams). When passed a pending block output, it fuses its residual combine with the RMSNorm. A missing injection selects unit-weight combine. Final mixers use use_combine=False and do not produce a new injection.

Weights: the norm owns the grouped GemmaRMSNorm affine; the projections are vLLM Linear modules (merged replicated linear for down+inject). Eligible down+inject projections use the fused SiLU GEMM; other projections use their Linear module's GEMM dispatch.

Methods:

• combine_and_mix –

  Consume a pending combine, then prepare the next block input.

Source code in vllm/models/qwen4_exp/nvidia/hyperconnection.py

class GatedResidual(nn.Module):
    """Gated HyperConnection with learnable low-rank mixing and injection.

    ``combine_and_mix()`` runs the pre pipeline (grouped GemmaRMSNorm -> merged
    low-rank down+inject GEMM -> silu -> up GEMM -> sigmoid -> gated mean
    over the HC streams). When passed a pending block output, it fuses its
    residual combine with the RMSNorm. A missing injection selects unit-weight
    combine. Final mixers use ``use_combine=False`` and do not produce a new
    injection.

    Weights: the norm owns the grouped GemmaRMSNorm affine; the projections
    are vLLM Linear modules (merged replicated linear for down+inject).
    Eligible down+inject projections use the fused SiLU GEMM; other projections
    use their Linear module's GEMM dispatch.
    """

    def __init__(
        self,
        config: HyperConnectionConfig,
        use_combine: bool = True,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.config = config
        self.lora_rank = config.hc_lowrank
        self.hc_count = config.hc_count
        self.hidden_size = config.hidden_size
        self.use_combine = use_combine

        norm_size = (
            self.hyper_hidden_size if config.hc_per_branch_norm else config.hidden_size
        )
        group_size = config.hidden_size if config.hc_per_branch_norm else None
        # Normalize each H-sized HC stream independently while retaining a
        # separate affine weight for every element of the HC*H layout.
        self.hc_norm = GroupedGemmaRMSNorm(
            norm_size,
            eps=config.rms_norm_eps,
            group_size=group_size,
            dtype=config.params_dtype,
        )

        # -- vLLM Linear weights --------------------------------------------
        # The merged skinny-GEMM shape is physically padded to 16 rows to ensure
        # good alignment and performant implementation chosen by CuBLAS heuristics.
        self.pad_size = (-(self.lora_rank + self.hc_count)) % 16 if use_combine else 0
        if use_combine:
            self.input_mix_weight_down_block_inject = MergedColumnParallelLinear(
                self.hyper_hidden_size,
                [self.lora_rank, self.hc_count]
                + ([self.pad_size] if self.pad_size else []),
                bias=False,
                params_dtype=config.params_dtype,
                quant_config=None,
                prefix=maybe_prefix(prefix, "input_mix_weight_down_block_inject"),
                return_bias=False,
                disable_tp=True,
            )
            weight = self.input_mix_weight_down_block_inject.weight
            self._use_hc_down_silu = (
                weight.shape[1] % 8 == 0
                and weight.dtype == torch.bfloat16
                and current_platform.has_device_capability(90)
            )
        else:
            self.input_mix_weight_down = ReplicatedLinear(
                self.hyper_hidden_size,
                self.lora_rank,
                bias=False,
                params_dtype=config.params_dtype,
                quant_config=None,
                prefix=maybe_prefix(prefix, "input_mix_weight_down"),
                return_bias=False,
            )
        self.input_mix_weight_up = ReplicatedLinear(
            self.lora_rank,
            self.hyper_hidden_size,
            bias=False,
            params_dtype=config.params_dtype,
            quant_config=None,
            prefix=maybe_prefix(prefix, "input_mix_weight_up"),
            return_bias=False,
        )

    def _down_and_inject(
        self, xn: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Down projection + SiLU; also returns injection logits if combined."""
        if not self.use_combine:
            return hc_silu(self.input_mix_weight_down(xn), self.hc_count), None

        use_fused = (
            self._use_hc_down_silu
            and not envs.VLLM_BATCH_INVARIANT
            and 1 <= xn.shape[0] <= MAX_FUSED_M
        )
        if use_fused:
            return hc_down_silu(
                xn,
                self.input_mix_weight_down_block_inject.weight,
                self.lora_rank,
                self.hc_count,
            )
        down_and_injection = self.input_mix_weight_down_block_inject(xn)
        split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
        lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
        return hc_silu(lora, self.hc_count), injection

    def mix(
        self, hidden_states: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
        xn = grouped_gemma_rmsnorm(
            hidden_states,
            self.hc_norm.weight,
            self.config.rms_norm_eps,
            self.hc_count,
        )

        lora, injection = self._down_and_inject(xn)
        gate = self.input_mix_weight_up(lora)  # [M, D]
        block_input = hc_gate_mix(xn, gate, self.hc_count)

        return hidden_states, block_input, injection

    def combine_and_mix(
        self,
        hidden_states: torch.Tensor,
        prev_block_output: torch.Tensor,
        prev_injection: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
        """Consume a pending combine, then prepare the next block input.

        ``hidden_states`` is the multi-stream state from before the pending
        block's mix. Its combine with ``block_output`` is fused with this
        module's input RMSNorm. A missing injection applies the block output
        to every stream with unit weight.
        """
        hidden_states, xn = hc_combine_norm(
            hidden_states,
            prev_block_output,
            prev_injection,
            self.hc_norm.weight,
            self.config.rms_norm_eps,
            self.hc_count,
        )

        lora, injection = self._down_and_inject(xn)
        gate = self.input_mix_weight_up(lora)  # [M, D]
        block_input = hc_gate_mix(xn, gate, self.hc_count)

        return hidden_states, block_input, injection

    def combine(
        self,
        hidden_states: torch.Tensor,
        block_output: torch.Tensor,
        injection: torch.Tensor | None,
    ) -> torch.Tensor:
        return hc_combine(hidden_states, block_output, injection, self.hc_count)

    @property
    def hyper_hidden_size(self) -> int:
        return self.hc_count * self.hidden_size

_down_and_inject(xn)

Down projection + SiLU; also returns injection logits if combined.

Source code in vllm/models/qwen4_exp/nvidia/hyperconnection.py

def _down_and_inject(
    self, xn: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """Down projection + SiLU; also returns injection logits if combined."""
    if not self.use_combine:
        return hc_silu(self.input_mix_weight_down(xn), self.hc_count), None

    use_fused = (
        self._use_hc_down_silu
        and not envs.VLLM_BATCH_INVARIANT
        and 1 <= xn.shape[0] <= MAX_FUSED_M
    )
    if use_fused:
        return hc_down_silu(
            xn,
            self.input_mix_weight_down_block_inject.weight,
            self.lora_rank,
            self.hc_count,
        )
    down_and_injection = self.input_mix_weight_down_block_inject(xn)
    split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
    lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
    return hc_silu(lora, self.hc_count), injection

combine_and_mix(hidden_states, prev_block_output, prev_injection)

Consume a pending combine, then prepare the next block input.

hidden_states is the multi-stream state from before the pending block's mix. Its combine with block_output is fused with this module's input RMSNorm. A missing injection applies the block output to every stream with unit weight.

Source code in vllm/models/qwen4_exp/nvidia/hyperconnection.py

def combine_and_mix(
    self,
    hidden_states: torch.Tensor,
    prev_block_output: torch.Tensor,
    prev_injection: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """Consume a pending combine, then prepare the next block input.

    ``hidden_states`` is the multi-stream state from before the pending
    block's mix. Its combine with ``block_output`` is fused with this
    module's input RMSNorm. A missing injection applies the block output
    to every stream with unit weight.
    """
    hidden_states, xn = hc_combine_norm(
        hidden_states,
        prev_block_output,
        prev_injection,
        self.hc_norm.weight,
        self.config.rms_norm_eps,
        self.hc_count,
    )

    lora, injection = self._down_and_inject(xn)
    gate = self.input_mix_weight_up(lora)  # [M, D]
    block_input = hc_gate_mix(xn, gate, self.hc_count)

    return hidden_states, block_input, injection

HyperConnectionConfig dataclass

Configuration shared by all HyperConnection variants.

Source code in vllm/models/qwen4_exp/common/hyperconnection.py

@dataclass
class HyperConnectionConfig:
    """Configuration shared by all HyperConnection variants."""

    hc_count: int = 4
    hidden_size: int = 64
    params_dtype: torch.dtype = torch.bfloat16
    mtp_hc: bool = False
    hc_lowrank: int = 16
    rms_norm_eps: float = 1e-6
    hc_per_branch_norm: bool = False