vllm.model_executor.layers.hpc.hpc_ihc ¶

class HpcIHCHead(HpcModule):
    """Fused iHC head block: merge the hc channels back into one hidden state.

    Same structure as HpcIHCPre but the projection emits only hc_mult gate
    logits and there is no H_post output. Called once per forward, after the
    last decoder layer.
    """

    def __init__(
        self,
        hc_mult: int,
        hidden_size: int,
        hc_eps: float,
        norm_eps: float,
        fallback_op: torch.nn.Module,
    ) -> None:
        super().__init__()
        self.hc_mult = hc_mult
        self.hidden_size = hidden_size
        self.hc_eps = hc_eps
        self.norm_eps = norm_eps
        # See HpcIHCPre: keep the owner out of the module tree.
        object.__setattr__(self, "_fallback_op", fallback_op)
        self.weight = None
        self.hc_scale = None
        self.hc_base = None

    @classmethod
    def support(cls, hc_mult: int, hidden_size: int) -> bool:
        return _ihc_supported(hc_mult, hidden_size)

    def process_weights_after_loading(self, model=None):
        # hc_head_fn weight is [hc_mult, hc_mult * d], float32.
        owner = self._fallback_op
        self.weight = owner.hc_head_fn.weight
        self.hc_scale = owner.hc_head_scale
        self.hc_base = owner.hc_head_base

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        import hpc

        return hpc.fuse_ihc_head(
            x,
            self.weight,
            self.hc_scale,
            self.hc_base,
            self.norm_eps,
            self.hc_eps,
        )

class HpcIHCPost(HpcModule):
    """Fused iHC post block: H_post gating plus multi-channel residual add.

    Unlike mHC there is no comb matrix, so each output channel only needs its
    own residual channel and the whole thing is one fused multiply-add per
    element.
    """

    def __init__(self, hc_mult: int, hidden_size: int) -> None:
        super().__init__()
        self.hc_mult = hc_mult
        self.hidden_size = hidden_size

    @classmethod
    def support(cls, hc_mult: int, hidden_size: int) -> bool:
        return _ihc_supported(hc_mult, hidden_size)

    def forward(
        self, x: torch.Tensor, residual: torch.Tensor, H_post: torch.Tensor
    ) -> torch.Tensor:
        import hpc

        return hpc.fuse_ihc_post(x, residual, H_post)

class HpcIHCPre(HpcModule):
    """Fused iHC pre block.

    Computes, in one kernel:
        x_flat = x.flatten(1)
        r      = rsqrt(x_flat.square().mean(-1) + rms_norm_eps)
        mixes  = (x_flat @ w.T) * r
        H_pre  = sigmoid(mixes[:, :hc] * hc_scale[0] + hc_base[:hc]) + hc_eps
        H_post = magnitude * sigmoid(
                     mixes[:, hc:] * hc_scale[1] + hc_base[hc:]) + hc_eps
        y      = sum_i H_pre[:, i] * x[:, i, :]

    Args:
        hc_mult: HC expand ratio.
        hidden_size: Model hidden dimension.
        magnitude: H_post multiplier (config.hc_magnitude).
        hc_eps: Epsilon added to both gates (config.hc_eps).
        norm_eps: Epsilon inside the rsqrt (config.rms_norm_eps).
        fallback_op: The eager HYV4HCPreLayer to source weights from.
        norm_owner: Optional RMSNorm that immediately follows this pre block.
            When given, its weight/eps are folded into the kernel so a single
            launch covers pre + RMSNorm (the caller then skips the separate
            layernorm). None keeps the original pre-only behaviour.
    """

    def __init__(
        self,
        hc_mult: int,
        hidden_size: int,
        magnitude: float,
        hc_eps: float,
        norm_eps: float,
        fallback_op: torch.nn.Module,
        norm_owner: torch.nn.Module | None = None,
    ) -> None:
        super().__init__()
        self.hc_mult = hc_mult
        self.hidden_size = hidden_size
        self.magnitude = magnitude
        self.hc_eps = hc_eps
        self.norm_eps = norm_eps
        # Stash the owner outside nn.Module's attribute machinery: callers pass
        # their own `self` here, and registering that as a submodule would make
        # the module tree cyclic (state_dict() then hits RecursionError).
        object.__setattr__(self, "_fallback_op", fallback_op)
        # Same anti-cycle idiom for the fused RMSNorm owner (a sibling module).
        object.__setattr__(self, "_norm_owner", norm_owner)
        self.weight = None
        self.hc_scale = None
        self.hc_base = None
        self.rms_weight = None
        self.rms_eps = 0.0

    @classmethod
    def support(cls, hc_mult: int, hidden_size: int) -> bool:
        return _ihc_supported(hc_mult, hidden_size)

    def process_weights_after_loading(self, model=None):
        # hc_fn is a ReplicatedLinear whose weight is [2 * hc_mult, hc_mult * d].
        # It stays float32 -- the HPC kernel consumes fp32 weights directly, so
        # there is no conversion (and no precision loss) here.
        owner = self._fallback_op
        self.weight = owner.hc_fn.weight
        self.hc_scale = owner.hc_scale
        self.hc_base = owner.hc_base
        if self._norm_owner is not None:
            self.rms_weight = self._norm_owner.weight
            self.rms_eps = self._norm_owner.variance_epsilon

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        import hpc

        return hpc.fuse_ihc_pre(
            x,
            self.weight,
            self.hc_scale,
            self.hc_base,
            self.norm_eps,
            self.hc_eps,
            self.magnitude,
            self.rms_weight,
            self.rms_eps,
        )