vllm.models.qwen4_exp.amd.hyperconnection

HyperConnection (Gated Residual) utilities for the AMD model variant.

Implements the HyperConnection residual scheme proposed in "HyperConnections" (https://arxiv.org/abs/2409.19606). This AMD 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 and an injection, it fuses their residual combine with the RMSNorm. 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), so GEMM dispatch (e.g. the low-latency skinny GEMM) applies through the standard quant_method mechanism.

Methods:

• combine_and_mix –

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

Source code in vllm/models/qwen4_exp/amd/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 and an injection,
    it fuses their residual combine with the RMSNorm. 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), so
    GEMM dispatch (e.g. the low-latency skinny GEMM) applies through the
    standard quant_method mechanism.
    """

    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 for
        # alignment and efficient backend dispatch.
        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,
            )
        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 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,
        )

        if self.use_combine:
            # produce injection logits for combine
            split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
            down_and_injection = self.input_mix_weight_down_block_inject(xn)
            lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
        else:
            lora = self.input_mix_weight_down(xn)
            injection = None

        lora = hc_silu(lora, self.hc_count)
        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,
    ) -> 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.
        """
        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,
        )

        if self.use_combine:
            # produce injection logits for combine
            split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
            down_and_injection = self.input_mix_weight_down_block_inject(xn)
            lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
        else:
            lora = self.input_mix_weight_down(xn)
            injection = None

        lora = hc_silu(lora, self.hc_count)
        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,
    ) -> 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

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.

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

def combine_and_mix(
    self,
    hidden_states: torch.Tensor,
    prev_block_output: torch.Tensor,
    prev_injection: torch.Tensor,
) -> 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.
    """
    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,
    )

    if self.use_combine:
        # produce injection logits for combine
        split_sizes = [self.lora_rank, self.hc_count, self.pad_size]
        down_and_injection = self.input_mix_weight_down_block_inject(xn)
        lora, injection, _ = down_and_injection.split(split_sizes, dim=-1)
    else:
        lora = self.input_mix_weight_down(xn)
        injection = None

    lora = hc_silu(lora, self.hc_count)
    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