@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 280 lines 9.58 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 from functools import partial
6 from typing import ClassVar
7
8 import torch
9 import torch.distributed as dist
10
11 from vllm.distributed import get_tp_group
12 from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
13 from vllm.model_executor.warmup.cutedsl_warmup import (
14 CuTeDSLCompileUnit,
15 register_cutedsl_warmup_provider,
16 )
17
18 _MAX_NUM_TOKENS = 128
19 _SKINNY_MAX_NUM_TOKENS = 5
20 _MMA_TILER_MN = (64, 32)
21 _GEMM_CLUSTER_MN = (1, 8)
22 _B_PRIME_STAGES = 2
23 _COLLECTIVE_TOKEN_CTAS = 32
24 _LAMPORT_COPY_CTAS = 32
25 _LAMPORT_COPY_THREADS = 128
26 _SUPPORTED_TP_SIZES = (8, 16)
27
28
29 @dataclass(frozen=True)
30 class KimiK3LatentMoETailContract:
31 tp_group_id: int
32 tp_size: int
33 device: torch.device
34 dtype: torch.dtype
35 hidden_size: int
36 latent_size: int
37 max_num_tokens: int
38 rms_eps: float
39 experts_per_token: int
40
41
42 class KimiK3LatentMoETailOp:
43 """Process-wide cached K3 latent-MoE tail implementation."""
44
45 _instances: ClassVar[
46 dict[KimiK3LatentMoETailContract, "KimiK3LatentMoETailOp"]
47 ] = {}
48
49 @classmethod
50 def _contract_and_group(
51 cls,
52 *,
53 hidden_size: int,
54 latent_size: int,
55 dtype: torch.dtype,
56 device: torch.device,
57 rms_eps: float,
58 experts_per_token: int = 0,
59 ) -> tuple[KimiK3LatentMoETailContract, dist.ProcessGroup]:
60 tp = get_tp_group()
61 group = tp.device_group
62 device = torch.device(device)
63 tp_device = torch.device(tp.device)
64 if device != tp_device:
65 raise ValueError(
66 f"Input device {device} does not match TP device {tp_device}."
67 )
68 return (
69 KimiK3LatentMoETailContract(
70 tp_group_id=id(group),
71 tp_size=dist.get_world_size(group),
72 device=device,
73 dtype=dtype,
74 hidden_size=hidden_size,
75 latent_size=latent_size,
76 max_num_tokens=_MAX_NUM_TOKENS,
77 rms_eps=float(rms_eps),
78 experts_per_token=experts_per_token,
79 ),
80 group,
81 )
82
83 @classmethod
84 def initialize(
85 cls,
86 *,
87 hidden_size: int,
88 latent_size: int,
89 dtype: torch.dtype,
90 device: torch.device,
91 rms_eps: float,
92 experts_per_token: int = 0,
93 ) -> "KimiK3LatentMoETailOp":
94 contract, group = cls._contract_and_group(
95 hidden_size=hidden_size,
96 latent_size=latent_size,
97 experts_per_token=experts_per_token,
98 dtype=dtype,
99 device=device,
100 rms_eps=rms_eps,
101 )
102 op = cls._instances.get(contract)
103 if op is None:
104 op = cls(contract, group)
105 cls._instances[contract] = op
106 return op
107
108 def __init__(
109 self,
110 contract: KimiK3LatentMoETailContract,
111 group: dist.ProcessGroup,
112 ) -> None:
113 if contract.tp_size not in _SUPPORTED_TP_SIZES:
114 raise ValueError(
115 "K3 latent-MoE tail fusion requires TP in "
116 f"{_SUPPORTED_TP_SIZES}, got {contract.tp_size}."
117 )
118 if contract.device.type != "cuda":
119 raise ValueError("K3 latent-MoE tail fusion requires a CUDA device.")
120 if torch.cuda.get_device_capability(contract.device)[0] != 10:
121 raise ValueError("K3 latent-MoE tail fusion requires SM100.")
122 if contract.dtype != torch.bfloat16:
123 raise ValueError("K3 latent-MoE tail fusion requires bfloat16.")
124 if (contract.hidden_size, contract.latent_size) != (7168, 3584):
125 raise ValueError(
126 "K3 latent-MoE tail fusion requires hidden_size=7168 and "
127 "latent_size=3584."
128 )
129
130 self.contract = contract
131 self.rank = dist.get_rank(group)
132 from .cute_dsl.latent_moe_tail import (
133 AdaptiveUpProjectionKernel,
134 CollectiveKernel,
135 LamportCopyKernel,
136 )
137
138 with torch.accelerator.device_index(contract.device.index):
139 self._collective = CollectiveKernel(
140 group=group,
141 rank=self.rank,
142 tp_size=contract.tp_size,
143 latent_dim=contract.latent_size,
144 hidden_dim=contract.hidden_size,
145 max_m=contract.max_num_tokens,
146 max_token_ctas=_COLLECTIVE_TOKEN_CTAS,
147 rms_eps=contract.rms_eps,
148 fp32_internal=False,
149 top_k=contract.experts_per_token,
150 )
151 self._up_projection = AdaptiveUpProjectionKernel(
152 group=group,
153 rank=self.rank,
154 tp_size=contract.tp_size,
155 latent_dim=contract.latent_size,
156 hidden_dim=contract.hidden_size,
157 max_m=contract.max_num_tokens,
158 skinny_max_m=_SKINNY_MAX_NUM_TOKENS,
159 mma_tiler_mn=_MMA_TILER_MN,
160 cluster_shape_mn=_GEMM_CLUSTER_MN,
161 b_prime_stages=_B_PRIME_STAGES,
162 )
163 self._lamport_copy = LamportCopyKernel(
164 hidden_dim=contract.hidden_size,
165 max_m=contract.max_num_tokens,
166 ctas=_LAMPORT_COPY_CTAS,
167 threads=_LAMPORT_COPY_THREADS,
168 )
169 register_cutedsl_warmup_provider(self)
170
171 def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
172 contract = self.contract
173 skinny_units = tuple(
174 CuTeDSLCompileUnit(
175 name="K3 latent MoE tail Skinny up-projection",
176 key=(
177 "k3-latent-moe-tail-skinny-up-projection",
178 contract,
179 m,
180 ),
181 compile=partial(self._up_projection.compile_skinny, m),
182 )
183 for m in range(1, self._up_projection.skinny_max_m + 1)
184 )
185 return skinny_units + (
186 CuTeDSLCompileUnit(
187 name="K3 latent MoE tail dynamic up-projection",
188 key=(
189 "k3-latent-moe-tail-dynamic-up-projection",
190 contract,
191 _MMA_TILER_MN,
192 _GEMM_CLUSTER_MN,
193 _B_PRIME_STAGES,
194 ),
195 compile=self._up_projection.compile_dynamic,
196 ),
197 )
198
199 def __call__(
200 self,
201 routed_output: torch.Tensor | UnfinalizedMoEOutput,
202 shared_output: torch.Tensor,
203 rms_weight: torch.Tensor,
204 up_weight: torch.Tensor,
205 ) -> torch.Tensor:
206 self._validate_inputs(
207 routed_output,
208 shared_output,
209 rms_weight,
210 up_weight,
211 )
212 num_tokens = (
213 routed_output.expanded_idx_to_permuted_idx.shape[0]
214 if isinstance(routed_output, UnfinalizedMoEOutput)
215 else routed_output.shape[0]
216 )
217 self._up_projection.ensure_compiled(num_tokens)
218 latent, shared_shard = self._collective(
219 routed_output,
220 shared_output,
221 rms_weight,
222 )
223 local_hidden_size = self.contract.hidden_size // self.contract.tp_size
224 local_up_weight = up_weight.narrow(
225 0,
226 self.rank * local_hidden_size,
227 local_hidden_size,
228 )
229 mailbox = self._up_projection(
230 latent,
231 local_up_weight,
232 shared_shard,
233 )
234 return self._lamport_copy(
235 mailbox,
236 m=num_tokens,
237 ).squeeze(0)
238
239 def _validate_inputs(
240 self,
241 routed_output: torch.Tensor | UnfinalizedMoEOutput,
242 shared_output: torch.Tensor,
243 rms_weight: torch.Tensor,
244 up_weight: torch.Tensor,
245 ) -> None:
246 contract = self.contract
247 if isinstance(routed_output, UnfinalizedMoEOutput):
248 num_tokens = routed_output.expanded_idx_to_permuted_idx.shape[0]
249 else:
250 if routed_output.ndim != 2:
251 raise ValueError("routed_output must be a 2D tensor.")
252 num_tokens = routed_output.shape[0]
253 if routed_output.shape != (num_tokens, contract.latent_size):
254 raise ValueError(
255 f"routed_output must have shape [M, {contract.latent_size}]."
256 )
257 if shared_output.shape != (num_tokens, contract.hidden_size):
258 raise ValueError(
259 f"shared_output must have shape [M, {contract.hidden_size}]."
260 )
261 if rms_weight.shape != (contract.latent_size,):
262 raise ValueError(f"rms_weight must have shape [{contract.latent_size}].")
263 if up_weight.shape != (contract.hidden_size, contract.latent_size):
264 raise ValueError(
265 "up_weight must have shape "
266 f"[{contract.hidden_size}, {contract.latent_size}]."
267 )
268 if not 1 <= num_tokens <= contract.max_num_tokens:
269 raise ValueError(
270 "K3 latent-MoE tail fusion requires between 1 and "
271 f"{contract.max_num_tokens} tokens."
272 )
273
274 tensors = (shared_output, rms_weight, up_weight)
275 if any(tensor.device != contract.device for tensor in tensors):
276 raise ValueError("All inputs must be on the contract device.")
277 if any(tensor.dtype != contract.dtype for tensor in tensors):
278 raise ValueError("All inputs must use the contract dtype.")
279 if any(not tensor.is_contiguous() for tensor in tensors):
280 raise ValueError("All inputs must be contiguous.")