| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """iHC (independent Hyper-Connections) layers for HY V4 (NVIDIA). |
| 4 | |
| 5 | iHC replaces the single residual stream of a standard transformer with |
| 6 | ``hc_mult`` parallel residual channels. Each decoder sub-block reduces the |
| 7 | channels to one hidden state (``HYV4HCPreLayer``), runs the sub-block, then |
| 8 | scatters the result back over the channels (``HYV4HCPostLayer``). The final |
| 9 | ``HYV4HCHeadLayer`` merges the channels before the model's output norm. |
| 10 | |
| 11 | NOTE: Each of the three steps has an optional single-kernel HPC replacement |
| 12 | (``HpcIHCPre`` / ``HpcIHCPost`` / ``HpcIHCHead``). Pre and post fall back to |
| 13 | in-tree Triton kernels on CUDA when HPC is unavailable, then to the eager path. |
| 14 | TODO: port the cross-layer post+pre fusion (``HpcIHCPostPre``) as well; it |
| 15 | requires restructuring the decoder-layer forward scheduling. |
| 16 | """ |
| 17 | |
| 18 | import torch |
| 19 | from torch import nn |
| 20 | from transformers import PreTrainedConfig |
| 21 | |
| 22 | from vllm.model_executor.layers.hpc import HpcIHCHead, HpcIHCPost, HpcIHCPre |
| 23 | from vllm.model_executor.layers.linear import ReplicatedLinear |
| 24 | from vllm.models.hy_v4.nvidia.triton_ihc import ( |
| 25 | triton_ihc_post, |
| 26 | triton_ihc_pre, |
| 27 | triton_ihc_supported, |
| 28 | ) |
| 29 | |
| 30 | |
| 31 | class HYV4HCPreLayer(nn.Module): |
| 32 | """iHC pre-processing layer (2D-activation adaptation). |
| 33 | |
| 34 | Steps: |
| 35 | 1. RMS-normalize the flattened ``[num_tokens, hc * d]`` input. |
| 36 | 2. Project to the pre/post gating logits. |
| 37 | 3. Turn the logits into sigmoid gates. |
| 38 | 4. Reduce over the channel dim with the pre gates. |
| 39 | """ |
| 40 | |
| 41 | def __init__( |
| 42 | self, |
| 43 | config: PreTrainedConfig, |
| 44 | hidden_dim: int, |
| 45 | hc_mult: int = 4, |
| 46 | magnitude: float = 2.0, |
| 47 | init_std: float = 6e-3, |
| 48 | base_noise_std: float = 0.0, |
| 49 | hc_eps: float = 1e-6, |
| 50 | layernorm_epsilon: float = 1e-5, |
| 51 | prefix: str = "", |
| 52 | ): |
| 53 | super().__init__() |
| 54 | self.config = config |
| 55 | self.hidden_dim = hidden_dim |
| 56 | self.hc_mult = hc_mult |
| 57 | self.magnitude = magnitude |
| 58 | self.hc_eps = hc_eps |
| 59 | self.layernorm_epsilon = layernorm_epsilon |
| 60 | hc_in_dim = hc_mult * hidden_dim |
| 61 | mix_hc = 2 * hc_mult # pre + post only, no comb |
| 62 | |
| 63 | self.hc_fn = ReplicatedLinear( |
| 64 | input_size=hc_in_dim, |
| 65 | output_size=mix_hc, |
| 66 | params_dtype=torch.float32, |
| 67 | bias=False, |
| 68 | prefix=f"{prefix}.hc_fn", |
| 69 | ) |
| 70 | self.hc_scale = nn.Parameter(torch.empty(2, dtype=torch.float32)) |
| 71 | self.hc_base = nn.Parameter(torch.empty(mix_hc, dtype=torch.float32)) |
| 72 | |
| 73 | self.reset_parameters(init_std, base_noise_std) |
| 74 | |
| 75 | # Optional single-kernel HPC replacement for the whole forward below. |
| 76 | # ``norm_owner`` is left unset: the RMSNorm that follows the pre block |
| 77 | # lives in the decoder layer and is not folded in yet. |
| 78 | self.hpc_op: HpcIHCPre | None = None |
| 79 | if HpcIHCPre.support(hc_mult, hidden_dim): |
| 80 | self.hpc_op = HpcIHCPre( |
| 81 | hc_mult=hc_mult, |
| 82 | hidden_size=hidden_dim, |
| 83 | magnitude=magnitude, |
| 84 | hc_eps=hc_eps, |
| 85 | norm_eps=layernorm_epsilon, |
| 86 | fallback_op=self, |
| 87 | ) |
| 88 | |
| 89 | def reset_parameters(self, init_std: float, base_noise_std: float = 0.0) -> None: |
| 90 | """Initialize the gate scale and per-channel gate bias.""" |
| 91 | del init_std # hc_fn is initialized by ReplicatedLinear |
| 92 | nn.init.constant_(self.hc_scale, 0.01) |
| 93 | with torch.no_grad(): |
| 94 | self.hc_base[: self.hc_mult].fill_( |
| 95 | -torch.log(torch.tensor(self.hc_mult - 1.0, dtype=self.hc_base.dtype)) |
| 96 | ) |
| 97 | self.hc_base[self.hc_mult : 2 * self.hc_mult].fill_(0.0) |
| 98 | if base_noise_std > 0.0: |
| 99 | self.hc_base.add_(torch.randn_like(self.hc_base) * base_noise_std) |
| 100 | |
| 101 | def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 102 | """Reduce the iHC channels and emit the post gates. |
| 103 | |
| 104 | Args: |
| 105 | x: Input of shape ``[num_tokens, hc, d]``. |
| 106 | |
| 107 | Returns: |
| 108 | A tuple of the pre-gated reduction ``[num_tokens, d]`` and the post |
| 109 | gates ``[num_tokens, hc]`` consumed by `HYV4HCPostLayer`. |
| 110 | |
| 111 | """ |
| 112 | if self.hpc_op is not None: |
| 113 | return self.hpc_op(x) |
| 114 | if triton_ihc_supported(x): |
| 115 | return triton_ihc_pre( |
| 116 | x, |
| 117 | self.hc_fn.weight, |
| 118 | self.hc_scale, |
| 119 | self.hc_base, |
| 120 | self.magnitude, |
| 121 | self.hc_eps, |
| 122 | self.layernorm_epsilon, |
| 123 | ) |
| 124 | |
| 125 | shape = x.size() # [num_tokens, hc, d] |
| 126 | hc = self.hc_mult |
| 127 | hc_eps = self.hc_eps |
| 128 | |
| 129 | x_flat = x.flatten(1).float() # [num_tokens, hc*d] |
| 130 | |
| 131 | rsqrt = torch.rsqrt( |
| 132 | x_flat.square().mean(-1, keepdim=True) + self.layernorm_epsilon |
| 133 | ) |
| 134 | mixes = self.hc_fn(x_flat)[0] * rsqrt # [num_tokens, 2*hc] |
| 135 | |
| 136 | pre_raw = mixes[..., :hc] |
| 137 | post_raw = mixes[..., hc : 2 * hc] |
| 138 | |
| 139 | pre = ( |
| 140 | torch.sigmoid( |
| 141 | pre_raw * self.hc_scale[0].float() + self.hc_base[:hc].float() |
| 142 | ) |
| 143 | + hc_eps |
| 144 | ) |
| 145 | post = ( |
| 146 | self.magnitude |
| 147 | * torch.sigmoid( |
| 148 | post_raw * self.hc_scale[1].float() + self.hc_base[hc : 2 * hc].float() |
| 149 | ) |
| 150 | + hc_eps |
| 151 | ) |
| 152 | |
| 153 | y = torch.sum(pre.unsqueeze(-1) * x.reshape(shape), dim=1) # [num_tokens, d] |
| 154 | return y.to(x.dtype), post |
| 155 | |
| 156 | |
| 157 | class HYV4HCPostLayer(nn.Module): |
| 158 | """iHC post-processing layer (2D-activation adaptation). |
| 159 | |
| 160 | Applies post-gating to the sub-block output and adds the multi-channel |
| 161 | residual (no comb mixing):: |
| 162 | |
| 163 | y[n, i, d] = post[n, i] * x[n, d] + residual[n, i, d] |
| 164 | """ |
| 165 | |
| 166 | def __init__(self, config: PreTrainedConfig): |
| 167 | super().__init__() |
| 168 | self.config = config |
| 169 | |
| 170 | # Optional single-kernel HPC replacement for the whole forward below. |
| 171 | # Only constructed under enable_ihc (see HYV4HCLayer), so hc_mult is |
| 172 | # present; getattr keeps this robust if that ever changes. |
| 173 | self.hpc_op: HpcIHCPost | None = None |
| 174 | hc_mult = getattr(config, "hc_mult", 0) |
| 175 | if HpcIHCPost.support(hc_mult, config.hidden_size): |
| 176 | self.hpc_op = HpcIHCPost(hc_mult=hc_mult, hidden_size=config.hidden_size) |
| 177 | |
| 178 | def forward( |
| 179 | self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor |
| 180 | ) -> torch.Tensor: |
| 181 | """Scatter the sub-block output back onto the iHC channels. |
| 182 | |
| 183 | Args: |
| 184 | x: Attention/MLP output of shape ``[num_tokens, d]``. |
| 185 | residual: Multi-channel residual ``[num_tokens, hc, d]``. |
| 186 | post: Post gates ``[num_tokens, hc]`` from `HYV4HCPreLayer`. |
| 187 | |
| 188 | Returns: |
| 189 | The updated residual channels ``[num_tokens, hc, d]``. |
| 190 | |
| 191 | """ |
| 192 | if self.hpc_op is not None: |
| 193 | return self.hpc_op(x, residual, post) |
| 194 | if triton_ihc_supported(x): |
| 195 | return triton_ihc_post(x, residual, post) |
| 196 | |
| 197 | dtype = x.dtype |
| 198 | x = x.float() |
| 199 | residual = residual.float() |
| 200 | post = post.float() |
| 201 | |
| 202 | post_gated = post.unsqueeze(-1) * x.unsqueeze(-2) # [num_tokens, hc, d] |
| 203 | y = post_gated + residual |
| 204 | return y.to(dtype) |
| 205 | |
| 206 | |
| 207 | class HYV4HCHeadLayer(nn.Module): |
| 208 | """iHC head layer (2D-activation adaptation). |
| 209 | |
| 210 | Merges the iHC channels back into a single hidden state before the final |
| 211 | layer norm, using an RMS-normed projection plus sigmoid-gated reduction. |
| 212 | """ |
| 213 | |
| 214 | def __init__( |
| 215 | self, |
| 216 | config: PreTrainedConfig, |
| 217 | hidden_size: int, |
| 218 | hc_mult: int = 4, |
| 219 | hc_eps: float = 1e-6, |
| 220 | init_std: float = 6e-3, |
| 221 | base_noise_std: float = 0.0, |
| 222 | prefix: str = "", |
| 223 | ): |
| 224 | super().__init__() |
| 225 | self.config = config |
| 226 | self.hidden_size = hidden_size |
| 227 | self.hc_mult = hc_mult |
| 228 | self.hc_eps = hc_eps |
| 229 | self.hc_head_fn = ReplicatedLinear( |
| 230 | input_size=hc_mult * hidden_size, |
| 231 | output_size=hc_mult, |
| 232 | params_dtype=torch.float32, |
| 233 | bias=False, |
| 234 | prefix=f"{prefix}.hc_head_fn", |
| 235 | ) |
| 236 | self.hc_head_base = nn.Parameter(torch.empty(hc_mult, dtype=torch.float32)) |
| 237 | self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32)) |
| 238 | |
| 239 | self.reset_parameters(init_std, base_noise_std) |
| 240 | |
| 241 | # Optional single-kernel HPC replacement for the whole forward below. |
| 242 | self.hpc_op: HpcIHCHead | None = None |
| 243 | if HpcIHCHead.support(hc_mult, hidden_size): |
| 244 | self.hpc_op = HpcIHCHead( |
| 245 | hc_mult=hc_mult, |
| 246 | hidden_size=hidden_size, |
| 247 | hc_eps=hc_eps, |
| 248 | norm_eps=config.rms_norm_eps, |
| 249 | fallback_op=self, |
| 250 | ) |
| 251 | |
| 252 | def reset_parameters( |
| 253 | self, init_std: float = 6e-3, base_noise_std: float = 0.0 |
| 254 | ) -> None: |
| 255 | """Initialize the head gate scale and per-channel gate bias.""" |
| 256 | del init_std # hc_head_fn is initialized by ReplicatedLinear |
| 257 | nn.init.constant_(self.hc_head_scale, 0.01) |
| 258 | with torch.no_grad(): |
| 259 | self.hc_head_base.fill_( |
| 260 | -torch.log( |
| 261 | torch.tensor(self.hc_mult - 1.0, dtype=self.hc_head_base.dtype) |
| 262 | ) |
| 263 | ) |
| 264 | if base_noise_std > 0.0: |
| 265 | self.hc_head_base.add_( |
| 266 | torch.randn_like(self.hc_head_base) * base_noise_std |
| 267 | ) |
| 268 | |
| 269 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 270 | """Merge the iHC channels into a single hidden state. |
| 271 | |
| 272 | Args: |
| 273 | x: Input of shape ``[num_tokens, hc, d]``. |
| 274 | |
| 275 | Returns: |
| 276 | The merged hidden state ``[num_tokens, d]``. |
| 277 | |
| 278 | """ |
| 279 | if self.hpc_op is not None: |
| 280 | return self.hpc_op(x) |
| 281 | |
| 282 | shape, x_dtype = x.size(), x.dtype |
| 283 | |
| 284 | x = x.flatten(1).float() # [num_tokens, hc*d] |
| 285 | rsqrt = torch.rsqrt( |
| 286 | x.square().mean(-1, keepdim=True) + self.config.rms_norm_eps |
| 287 | ) |
| 288 | mixes = self.hc_head_fn(x)[0] * rsqrt # [num_tokens, hc] |
| 289 | pre = ( |
| 290 | torch.sigmoid( |
| 291 | mixes * self.hc_head_scale.float() + self.hc_head_base.float() |
| 292 | ) |
| 293 | + self.hc_eps |
| 294 | ) |
| 295 | y = torch.sum(pre.unsqueeze(-1) * x.reshape(shape), dim=1) # [num_tokens, d] |
| 296 | return y.to(x_dtype) |
| 297 | |
| 298 | |
| 299 | class HYV4HCLayer(nn.Module): |
| 300 | """Wrapper owning one iHC boundary (pre + post) of a decoder sub-block.""" |
| 301 | |
| 302 | def __init__( |
| 303 | self, |
| 304 | config: PreTrainedConfig, |
| 305 | layer_idx: int, |
| 306 | init_std: float = 6e-3, |
| 307 | base_noise_std: float = 0.0, |
| 308 | prefix: str = "", |
| 309 | ): |
| 310 | super().__init__() |
| 311 | self.config = config |
| 312 | self.layer_idx = layer_idx |
| 313 | self.enable_ihc = getattr(config, "enable_ihc", False) |
| 314 | if self.enable_ihc: |
| 315 | self.hc_pre = HYV4HCPreLayer( |
| 316 | config, |
| 317 | config.hidden_size, |
| 318 | config.hc_mult, |
| 319 | config.hc_magnitude, |
| 320 | init_std, |
| 321 | base_noise_std, |
| 322 | config.hc_eps, |
| 323 | config.rms_norm_eps, |
| 324 | prefix=f"{prefix}.hc_pre", |
| 325 | ) |
| 326 | self.hc_post = HYV4HCPostLayer(config) |
| 327 | |
| 328 | def prepare_input(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 329 | """Normalize the sub-block input to 3D when iHC is enabled.""" |
| 330 | if not self.enable_ihc: |
| 331 | return hidden_states |
| 332 | return self._prepare_input_to_3d(hidden_states) |
| 333 | |
| 334 | def _prepare_input_to_3d(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 335 | """Reshape the iHC input to ``[num_tokens, hc, h]``. |
| 336 | |
| 337 | Accepted inputs are ``[num_tokens, hc, h]`` (no-op), ``[num_tokens, h]`` |
| 338 | (broadcast over the channels) and ``[num_tokens, hc * h]`` (reshape). |
| 339 | """ |
| 340 | if hidden_states.dim() == 3: |
| 341 | return hidden_states |
| 342 | if hidden_states.dim() != 2: |
| 343 | raise RuntimeError( |
| 344 | f"HC expects a 2D/3D tensor, got shape={tuple(hidden_states.shape)}" |
| 345 | ) |
| 346 | |
| 347 | d0, d1 = hidden_states.shape |
| 348 | h = self.config.hidden_size |
| 349 | hc = self.config.hc_mult |
| 350 | |
| 351 | if d1 == h: |
| 352 | return hidden_states.unsqueeze(1).repeat(1, hc, 1) |
| 353 | |
| 354 | expected = hc * h |
| 355 | if d1 == expected: |
| 356 | return hidden_states.reshape(d0, hc, h) |
| 357 | |
| 358 | raise RuntimeError( |
| 359 | f"HC expects last dim to be hidden_size ({h})" |
| 360 | f"or hc_mult*hidden_size ({expected}), got {d1}. " |
| 361 | f"hc_mult={hc}, hidden_size={h}." |
| 362 | ) |
| 363 | |
| 364 | def pre( |
| 365 | self, hidden_states: torch.Tensor |
| 366 | ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: |
| 367 | """Reduce the iHC channels and produce the post gates. |
| 368 | |
| 369 | Returns: |
| 370 | A tuple of the reduced hidden states ``[num_tokens, d]``, the post |
| 371 | gates ``[num_tokens, hc]`` (``None`` when iHC is disabled) and the |
| 372 | residual (the untouched input). |
| 373 | |
| 374 | """ |
| 375 | if not self.enable_ihc: |
| 376 | return hidden_states, None, hidden_states |
| 377 | reduced, post_gates = self.hc_pre(hidden_states) |
| 378 | return reduced, post_gates, hidden_states |
| 379 | |
| 380 | def post( |
| 381 | self, |
| 382 | output_with_bias: torch.Tensor, |
| 383 | residual: torch.Tensor, |
| 384 | post_gates: torch.Tensor | None, |
| 385 | ) -> torch.Tensor: |
| 386 | """Apply post-gating and add the residual.""" |
| 387 | if not self.enable_ihc: |
| 388 | return output_with_bias + residual |
| 389 | assert post_gates is not None |
| 390 | return self.hc_post(output_with_bias, residual, post_gates) |