@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 237 lines 8.15 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """HPC fused iHC (independent Hyper-Connections) kernels for HY V4.
4
5 Replaces the eager HYV4HCPreLayer / HYV4HCPostLayer / HYV4HCHeadLayer bodies
6 with single-kernel HPC implementations:
7
8 pre : rms square sum + w projection + sigmoid gates + hc-dim weighted sum
9 post : y[n, i, :] = H_post[n, i] * x[n, :] + residual[n, i, :]
10 head : same as pre but hc_mult projection rows and no H_post output
11
12 The eager path issues 20 / 5 / 15 kernels for pre / post / head; each HPC op is
13 a single kernel.
14
15 Constraints:
16 - Requires VLLM_ENABLE_HPC_OPS=1
17 - Requires the hpc package (.so) built for the current arch
18 - Only sm100 / sm103 (compute capability 100, 103)
19 - Only hc_mult == 4 and hidden_size in {4096, 6144}
20 - x / residual must be bfloat16; hc_fn weights stay float32 (the checkpoint
21 keeps them out of fp8 quantization via modules_to_not_convert)
22
23 The fused modules own no parameters: their forwards read the weights of the
24 eager layer they replace directly via the captured owner reference.
25
26 NOTE: The reference implementation also offers a cross-layer post+pre fusion
27 (``HpcIHCPostPre``) that folds one segment's post into the next segment's pre.
28 It requires reworking the decoder-layer forward scheduling and is not ported
29 here. TODO: add it once the decoder dataflow is restructured for it.
30 """
31
32 import torch
33
34 from vllm import envs
35 from vllm.logger import init_logger
36 from vllm.model_executor.layers.hpc.hpc_module import HpcModule
37 from vllm.utils.hpc import has_hpc
38
39 logger = init_logger(__name__)
40
41 # Shapes the HPC kernels are instantiated for. Other shapes launch nothing,
42 # so gate on this rather than relying on a silent no-op.
43 _SUPPORTED_HIDDEN_SIZES: frozenset[int] = frozenset({4096, 6144})
44 _SUPPORTED_HC_MULTS: frozenset[int] = frozenset({4})
45 _SUPPORTED_CAPABILITIES: frozenset[int] = frozenset({100, 103})
46
47
48 def _ihc_supported(hc_mult: int, hidden_size: int) -> bool:
49 """Shared gate for all three iHC ops."""
50 if not envs.VLLM_ENABLE_HPC_OPS:
51 return False
52
53 if not has_hpc():
54 logger.warning_once(
55 "HPC iHC disabled: 'hpc' package is not installed. "
56 "Please install the HPC library (.so) to enable fused kernels."
57 )
58 return False
59
60 from vllm.platforms import current_platform
61
62 if not current_platform.is_cuda():
63 logger.warning_once("HPC iHC disabled: only CUDA is supported.")
64 return False
65
66 capability = current_platform.get_device_capability()
67 if capability is None or capability.to_int() not in _SUPPORTED_CAPABILITIES:
68 logger.warning_once(
69 "HPC iHC disabled: compute capability %s not in %s.",
70 capability,
71 _SUPPORTED_CAPABILITIES,
72 )
73 return False
74
75 if hc_mult not in _SUPPORTED_HC_MULTS:
76 logger.warning_once(
77 "HPC iHC disabled: hc_mult=%d not in %s.", hc_mult, _SUPPORTED_HC_MULTS
78 )
79 return False
80
81 if hidden_size not in _SUPPORTED_HIDDEN_SIZES:
82 logger.warning_once(
83 "HPC iHC disabled: hidden_size=%d not in %s.",
84 hidden_size,
85 _SUPPORTED_HIDDEN_SIZES,
86 )
87 return False
88
89 if envs.VLLM_BATCH_INVARIANT:
90 # Batch-invariant mode relies on torch's deterministic reductions; the
91 # HPC kernel reduces in a different order, so keep the two regimes
92 # apart.
93 logger.warning_once(
94 "HPC iHC disabled: not supported under VLLM_BATCH_INVARIANT."
95 )
96 return False
97
98 logger.info_once("HPC iHC enabled by set VLLM_ENABLE_HPC_OPS.")
99 return True
100
101
102 class HpcIHCPre(HpcModule):
103 """Fused iHC pre block.
104
105 Computes, in one kernel:
106 x_flat = x.flatten(1)
107 r = rsqrt(x_flat.square().mean(-1) + rms_norm_eps)
108 mixes = (x_flat @ w.T) * r
109 H_pre = sigmoid(mixes[:, :hc] * hc_scale[0] + hc_base[:hc]) + hc_eps
110 H_post = magnitude * sigmoid(
111 mixes[:, hc:] * hc_scale[1] + hc_base[hc:]) + hc_eps
112 y = sum_i H_pre[:, i] * x[:, i, :]
113
114 Args:
115 hc_mult: HC expand ratio.
116 hidden_size: Model hidden dimension.
117 magnitude: H_post multiplier (config.hc_magnitude).
118 hc_eps: Epsilon added to both gates (config.hc_eps).
119 norm_eps: Epsilon inside the rsqrt (config.rms_norm_eps).
120 fallback_op: The eager HYV4HCPreLayer to source weights from.
121 norm_owner: Optional RMSNorm that immediately follows this pre block.
122 When given, its weight/eps are folded into the kernel so a single
123 launch covers pre + RMSNorm (the caller then skips the separate
124 layernorm). None keeps the original pre-only behaviour.
125
126 """
127
128 def __init__(
129 self,
130 hc_mult: int,
131 hidden_size: int,
132 magnitude: float,
133 hc_eps: float,
134 norm_eps: float,
135 fallback_op: torch.nn.Module,
136 norm_owner: torch.nn.Module | None = None,
137 ) -> None:
138 super().__init__()
139 self.hc_mult = hc_mult
140 self.hidden_size = hidden_size
141 self.magnitude = magnitude
142 self.hc_eps = hc_eps
143 self.norm_eps = norm_eps
144 # Stash the owner outside nn.Module's attribute machinery: callers pass
145 # their own `self` here, and registering that as a submodule would make
146 # the module tree cyclic (state_dict() then hits RecursionError).
147 object.__setattr__(self, "_fallback_op", fallback_op)
148 # Same anti-cycle idiom for the fused RMSNorm owner (a sibling module).
149 object.__setattr__(self, "_norm_owner", norm_owner)
150
151 @classmethod
152 def support(cls, hc_mult: int, hidden_size: int) -> bool:
153 return _ihc_supported(hc_mult, hidden_size)
154
155 def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
156 import hpc
157
158 owner = self._fallback_op
159 norm_owner = self._norm_owner
160 return hpc.fuse_ihc_pre(
161 x,
162 owner.hc_fn.weight,
163 owner.hc_scale,
164 owner.hc_base,
165 self.norm_eps,
166 self.hc_eps,
167 self.magnitude,
168 norm_owner.weight if norm_owner is not None else None,
169 norm_owner.variance_epsilon if norm_owner is not None else 0.0,
170 )
171
172
173 class HpcIHCPost(HpcModule):
174 """Fused iHC post block: H_post gating plus multi-channel residual add.
175
176 Unlike mHC there is no comb matrix, so each output channel only needs its
177 own residual channel and the whole thing is one fused multiply-add per
178 element.
179 """
180
181 def __init__(self, hc_mult: int, hidden_size: int) -> None:
182 super().__init__()
183 self.hc_mult = hc_mult
184 self.hidden_size = hidden_size
185
186 @classmethod
187 def support(cls, hc_mult: int, hidden_size: int) -> bool:
188 return _ihc_supported(hc_mult, hidden_size)
189
190 def forward(
191 self, x: torch.Tensor, residual: torch.Tensor, H_post: torch.Tensor
192 ) -> torch.Tensor:
193 import hpc
194
195 return hpc.fuse_ihc_post(x, residual, H_post)
196
197
198 class HpcIHCHead(HpcModule):
199 """Fused iHC head block: merge the hc channels back into one hidden state.
200
201 Same structure as HpcIHCPre but the projection emits only hc_mult gate
202 logits and there is no H_post output. Called once per forward, after the
203 last decoder layer.
204 """
205
206 def __init__(
207 self,
208 hc_mult: int,
209 hidden_size: int,
210 hc_eps: float,
211 norm_eps: float,
212 fallback_op: torch.nn.Module,
213 ) -> None:
214 super().__init__()
215 self.hc_mult = hc_mult
216 self.hidden_size = hidden_size
217 self.hc_eps = hc_eps
218 self.norm_eps = norm_eps
219 # See HpcIHCPre: keep the owner out of the module tree.
220 object.__setattr__(self, "_fallback_op", fallback_op)
221
222 @classmethod
223 def support(cls, hc_mult: int, hidden_size: int) -> bool:
224 return _ihc_supported(hc_mult, hidden_size)
225
226 def forward(self, x: torch.Tensor) -> torch.Tensor:
227 import hpc
228
229 owner = self._fallback_op
230 return hpc.fuse_ihc_head(
231 x,
232 owner.hc_head_fn.weight,
233 owner.hc_head_scale,
234 owner.hc_head_base,
235 self.norm_eps,
236 self.hc_eps,
237 )