| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inkling mixture-of-experts on vLLM's FusedMoE abstraction. |
| 4 | |
| 5 | Overfit to the served checkpoint: sigmoid gate (+ selection bias) top-k over |
| 6 | the routed experts, log-sigmoid renormalization over the k routed + S shared |
| 7 | "sink" logits, scaled by route_scale * global_scale. The routed top-k goes |
| 8 | through vLLM's FusedMoE (which handles TP/EP); the sink experts run in |
| 9 | :class:`InklingSinkExperts` -- replicated across EP ranks (every token |
| 10 | activates every sink) and always bf16 (the checkpoint excludes every |
| 11 | ``shared_experts`` from quantization). |
| 12 | |
| 13 | NVFP4 routed experts reuse vLLM's fused-MoE methods; excluded (bf16) layers |
| 14 | fall back to the unquantized method. Checkpoint fused stacked tensors are |
| 15 | translated to the standard per-expert loads in |
| 16 | :meth:`InklingMoE.load_expert_weight`. |
| 17 | """ |
| 18 | |
| 19 | from __future__ import annotations |
| 20 | |
| 21 | import math |
| 22 | from typing import TYPE_CHECKING |
| 23 | |
| 24 | import torch |
| 25 | from torch import nn |
| 26 | from torch.nn.parameter import Parameter |
| 27 | |
| 28 | import vllm.envs as envs |
| 29 | from vllm.config import get_current_vllm_config |
| 30 | from vllm.distributed import ( |
| 31 | get_dp_group, |
| 32 | get_pcp_group, |
| 33 | get_tensor_model_parallel_rank, |
| 34 | get_tensor_model_parallel_world_size, |
| 35 | ) |
| 36 | from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16 |
| 37 | from vllm.model_executor.layers.fused_moe import FusedMoEFactory |
| 38 | from vllm.model_executor.utils import set_weight_attrs |
| 39 | from vllm.platforms import current_platform |
| 40 | from vllm.triton_utils import tl, tldevice, triton |
| 41 | from vllm.utils.multi_stream_utils import maybe_execute_in_parallel |
| 42 | from vllm.utils.torch_utils import aux_stream |
| 43 | |
| 44 | from ..configs import InklingModelConfig |
| 45 | |
| 46 | if TYPE_CHECKING: |
| 47 | from vllm.model_executor.layers.fused_moe.routed_experts import ( |
| 48 | RoutedExperts, |
| 49 | ) |
| 50 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 51 | |
| 52 | # --------------------------------------------------------------------------- |
| 53 | # Gate / expert selection |
| 54 | # --------------------------------------------------------------------------- |
| 55 | |
| 56 | _INKLING_LL_BF16_MAX_TOKENS = 64 |
| 57 | _NVFP4_INPUT_SCALE_DENOMINATOR = torch.finfo(torch.float8_e4m3fn).max * 6.0 |
| 58 | |
| 59 | |
| 60 | def _linear_with_fp32_out(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: |
| 61 | leading = list(x.shape[:-1]) |
| 62 | flat = x.flatten(0, -2) |
| 63 | if ( |
| 64 | flat.shape[0] <= _INKLING_LL_BF16_MAX_TOKENS |
| 65 | and flat.dtype == torch.bfloat16 |
| 66 | and weight.dtype == torch.bfloat16 |
| 67 | and flat.is_cuda |
| 68 | and flat.is_contiguous() |
| 69 | and weight.is_contiguous() |
| 70 | and flat.shape[1] % 8 == 0 |
| 71 | and current_platform.has_device_capability(90) |
| 72 | and ll_bf16.is_available() |
| 73 | ): |
| 74 | out = ll_bf16.ll_bf16_gemm(flat, weight) |
| 75 | else: |
| 76 | out = torch.mm(flat, weight.T, out_dtype=torch.float32) |
| 77 | return out.view(*leading, weight.shape[0]) |
| 78 | |
| 79 | |
| 80 | @triton.jit(do_not_specialize=["T", "route_scale"]) |
| 81 | def _inkling_gate_select_kernel( |
| 82 | logits_ptr, # [T, G] fp32 gate logits (stride_logits_0 may include pad) |
| 83 | bias_ptr, # [R] fp32 selection bias (or 0 ptr if HAS_BIAS=False) |
| 84 | global_scale_ptr, # [1] fp32 (or unused if HAS_GSCALE=False) |
| 85 | ids_ptr, # [T, K + S] int32 out: selected expert ids |
| 86 | weights_ptr, # [T, K + S] fp32 out: renormalized weights |
| 87 | route_scale, |
| 88 | T, |
| 89 | G: tl.constexpr, # total gate experts (routed + shared) |
| 90 | stride_logits_0, |
| 91 | R: tl.constexpr, # routed experts |
| 92 | K: tl.constexpr, # top-k routed |
| 93 | S: tl.constexpr, # shared (sink) experts |
| 94 | HAS_BIAS: tl.constexpr, |
| 95 | HAS_GSCALE: tl.constexpr, |
| 96 | BLOCK_G: tl.constexpr, |
| 97 | ): |
| 98 | pid = tl.program_id(0).to(tl.int64) |
| 99 | if pid >= T: |
| 100 | return |
| 101 | offs = tl.arange(0, BLOCK_G) |
| 102 | mask_r = offs < R |
| 103 | logits = tl.load( |
| 104 | logits_ptr + pid * stride_logits_0 + offs, |
| 105 | mask=offs < G, |
| 106 | other=float("-inf"), |
| 107 | ).to(tl.float32) |
| 108 | |
| 109 | # Selection scores: sigmoid(routed logits) (+ bias), non-routed lanes -inf. |
| 110 | sel = tl.where(mask_r, tl.sigmoid(logits), float("-inf")) |
| 111 | if HAS_BIAS: |
| 112 | bias = tl.load(bias_ptr + offs, mask=mask_r, other=0.0).to(tl.float32) |
| 113 | sel = tl.where(mask_r, sel + bias, float("-inf")) |
| 114 | |
| 115 | scale = route_scale |
| 116 | if HAS_GSCALE: |
| 117 | scale = scale * tl.load(global_scale_ptr).to(tl.float32) |
| 118 | |
| 119 | # Iterative top-K (K is small); argmax tie-breaks to the lowest index |
| 120 | # (stable ordering). |
| 121 | A: tl.constexpr = K + S |
| 122 | offs_a = tl.arange(0, A) |
| 123 | top_ids = tl.zeros([A], dtype=tl.int32) |
| 124 | active = tl.zeros([A], dtype=tl.float32) |
| 125 | for kk in tl.static_range(K): |
| 126 | idx = tl.argmax(sel, axis=0).to(tl.int32) |
| 127 | raw = tl.max(tl.where(offs == idx, logits, float("-inf")), axis=0) |
| 128 | top_ids = tl.where(offs_a == kk, idx, top_ids) |
| 129 | active = tl.where(offs_a == kk, raw, active) |
| 130 | sel = tl.where(offs == idx, float("-inf"), sel) |
| 131 | if S > 0: |
| 132 | # Shared sink logits sit at the tail of the gate output; their expert |
| 133 | # ids continue after the routed range (R + j). |
| 134 | for jj in tl.static_range(S): |
| 135 | raw = tl.max(tl.where(offs == R + jj, logits, float("-inf")), axis=0) |
| 136 | top_ids = tl.where(offs_a == K + jj, tl.full([], R + jj, tl.int32), top_ids) |
| 137 | active = tl.where(offs_a == K + jj, raw, active) |
| 138 | |
| 139 | # Log-sigmoid renormalization over the K + S active logits. |
| 140 | abs_l = tl.abs(active) |
| 141 | min_l = tl.minimum(active, 0.0) |
| 142 | log_probs = min_l - tldevice.log1p(tldevice.exp(-abs_l)) |
| 143 | max_lp = tl.max(log_probs, axis=0) |
| 144 | exp_shifted = tldevice.exp(log_probs - max_lp) |
| 145 | sum_exp = tl.sum(exp_shifted, axis=0) |
| 146 | weights = exp_shifted / sum_exp * scale |
| 147 | |
| 148 | tl.store(ids_ptr + pid * A + offs_a, top_ids) |
| 149 | tl.store(weights_ptr + pid * A + offs_a, weights) |
| 150 | |
| 151 | |
| 152 | def inkling_gate_select( |
| 153 | logits: torch.Tensor, # [T, >=G] fp32 (rows may carry GEMM padding) |
| 154 | n_gate_experts: int, |
| 155 | n_routed_experts: int, |
| 156 | topk: int, |
| 157 | n_shared_experts: int, |
| 158 | bias: torch.Tensor | None, |
| 159 | route_scale: float, |
| 160 | global_scale: torch.Tensor | None, |
| 161 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 162 | """Sigmoid + bias + top-k + log-sigmoid renorm; returns (weights, ids).""" |
| 163 | assert logits.dtype == torch.float32 |
| 164 | tokens = logits.shape[0] |
| 165 | active = topk + n_shared_experts |
| 166 | topk_ids = torch.empty((tokens, active), dtype=torch.int32, device=logits.device) |
| 167 | topk_weights = torch.empty( |
| 168 | (tokens, active), dtype=torch.float32, device=logits.device |
| 169 | ) |
| 170 | if tokens == 0: |
| 171 | return topk_weights, topk_ids |
| 172 | _inkling_gate_select_kernel[(tokens,)]( |
| 173 | logits, |
| 174 | bias if bias is not None else logits, |
| 175 | global_scale if global_scale is not None else logits, |
| 176 | topk_ids, |
| 177 | topk_weights, |
| 178 | route_scale, |
| 179 | tokens, |
| 180 | n_gate_experts, |
| 181 | logits.stride(0), |
| 182 | n_routed_experts, |
| 183 | topk, |
| 184 | n_shared_experts, |
| 185 | HAS_BIAS=bias is not None, |
| 186 | HAS_GSCALE=global_scale is not None, |
| 187 | BLOCK_G=triton.next_power_of_2(n_gate_experts), |
| 188 | ) |
| 189 | return topk_weights, topk_ids |
| 190 | |
| 191 | |
| 192 | class InklingGate(nn.Module): |
| 193 | """Sigmoid gate with selection bias, log-sigmoid renorm after top-k, and |
| 194 | global scale (the served checkpoint's only configuration).""" |
| 195 | |
| 196 | def __init__( |
| 197 | self, |
| 198 | d_model: int, |
| 199 | n_routed_experts: int, |
| 200 | n_shared_experts: int, |
| 201 | experts_per_token: int, |
| 202 | route_scale: float, |
| 203 | *, |
| 204 | use_global_scale: bool = False, |
| 205 | use_gate_bias: bool = False, |
| 206 | ) -> None: |
| 207 | super().__init__() |
| 208 | self.n_routed_experts = n_routed_experts |
| 209 | self.n_shared_experts = n_shared_experts |
| 210 | self.n_total_experts = n_routed_experts + n_shared_experts |
| 211 | self.topk = experts_per_token |
| 212 | self.route_scale = route_scale |
| 213 | |
| 214 | padded_experts = self.n_total_experts + (-self.n_total_experts) % 8 |
| 215 | self.weight = Parameter( |
| 216 | torch.empty(padded_experts, d_model), requires_grad=False |
| 217 | ) |
| 218 | set_weight_attrs(self.weight, {"weight_loader": self._load_weight}) |
| 219 | if use_global_scale: |
| 220 | self.global_scale = Parameter( |
| 221 | torch.empty(1, dtype=torch.float32), requires_grad=False |
| 222 | ) |
| 223 | else: |
| 224 | self.global_scale = None |
| 225 | if use_gate_bias: |
| 226 | self.bias = Parameter( |
| 227 | torch.empty(n_routed_experts, dtype=torch.float32), |
| 228 | requires_grad=False, |
| 229 | ) |
| 230 | else: |
| 231 | self.bias = None |
| 232 | |
| 233 | @staticmethod |
| 234 | def _load_weight(param: Parameter, loaded_weight: torch.Tensor) -> None: |
| 235 | param.data.zero_() |
| 236 | param.data[: loaded_weight.shape[0]].copy_(loaded_weight) |
| 237 | |
| 238 | def compute_logits(self, x: torch.Tensor) -> torch.Tensor: |
| 239 | """fp32 gate logits [T, n_total_experts + pad] (pad columns are junk).""" |
| 240 | return _linear_with_fp32_out(x, self.weight) |
| 241 | |
| 242 | def select_experts( |
| 243 | self, gating_output: torch.Tensor |
| 244 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 245 | """Full selection: (weights, ids) of [T, K + S]. The first K entries |
| 246 | are the routed top-k; the S trailing entries are the sink gammas.""" |
| 247 | return inkling_gate_select( |
| 248 | gating_output, |
| 249 | self.n_total_experts, |
| 250 | self.n_routed_experts, |
| 251 | self.topk, |
| 252 | self.n_shared_experts, |
| 253 | self.bias, |
| 254 | self.route_scale, |
| 255 | self.global_scale, |
| 256 | ) |
| 257 | |
| 258 | |
| 259 | # --------------------------------------------------------------------------- |
| 260 | # MoE layer |
| 261 | # --------------------------------------------------------------------------- |
| 262 | |
| 263 | |
| 264 | def _inkling_moe_ep_size() -> int: |
| 265 | """EP size the FusedMoE layer will run with (mirrors |
| 266 | FusedMoEParallelConfig.make: experts shard over tp * dp * pcp when |
| 267 | expert parallelism is enabled).""" |
| 268 | parallel_config = get_current_vllm_config().parallel_config |
| 269 | if not parallel_config.enable_expert_parallel: |
| 270 | return 1 |
| 271 | world = ( |
| 272 | get_tensor_model_parallel_world_size() |
| 273 | * get_dp_group().world_size |
| 274 | * get_pcp_group().world_size |
| 275 | ) |
| 276 | return world if world > 1 else 1 |
| 277 | |
| 278 | |
| 279 | class InklingSinkExperts(nn.Module): |
| 280 | """Shared "sink" experts with per-token gammas, in bf16. |
| 281 | |
| 282 | Replicated across EP ranks (every token activates every sink, so |
| 283 | EP-sharding them would hotspot the owning rank) and TP-sharded on the |
| 284 | intermediate dim so the output remains a TP-partial sum like the routed |
| 285 | output. The sinks are always bf16 (the checkpoint excludes every |
| 286 | ``shared_experts`` from quantization): the experts concatenate into two |
| 287 | plain dense GEMMs with the fused sink epilogue between them. |
| 288 | """ |
| 289 | |
| 290 | def __init__( |
| 291 | self, n_experts: int, d_model: int, d_mlp: int, *, prefix: str = "" |
| 292 | ) -> None: |
| 293 | super().__init__() |
| 294 | self.n_experts = n_experts |
| 295 | tp_size = get_tensor_model_parallel_world_size() |
| 296 | self.tp_rank = get_tensor_model_parallel_rank() |
| 297 | intermediate_pp = d_mlp // tp_size |
| 298 | self.w13_weight = Parameter( |
| 299 | torch.empty(n_experts, 2 * intermediate_pp, d_model), |
| 300 | requires_grad=False, |
| 301 | ) |
| 302 | self.w2_weight = Parameter( |
| 303 | torch.empty(d_model, n_experts * intermediate_pp), |
| 304 | requires_grad=False, |
| 305 | ) |
| 306 | self._unit: torch.Tensor | None = None |
| 307 | |
| 308 | def load_weight(self, key: str, weight: torch.Tensor) -> list[str]: |
| 309 | """Load one checkpoint sink tensor (stacked over the S experts).""" |
| 310 | if key == "w13_weight": |
| 311 | if weight.shape != self.w13_weight.shape: |
| 312 | shard = self.w13_weight.shape[1] |
| 313 | weight = weight.narrow(1, self.tp_rank * shard, shard) |
| 314 | self.w13_weight.data.copy_(weight) |
| 315 | return [key] |
| 316 | |
| 317 | assert key == "w2_weight" |
| 318 | shard = self.w2_weight.shape[1] // self.n_experts |
| 319 | shard_start = 0 if weight.shape[2] == shard else self.tp_rank * shard |
| 320 | for expert_idx, expert_weight in enumerate(weight): |
| 321 | local_weight = expert_weight.narrow(1, shard_start, shard) |
| 322 | start = expert_idx * shard |
| 323 | self.w2_weight.data[:, start : start + shard].copy_(local_weight) |
| 324 | return [key] |
| 325 | |
| 326 | def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor: |
| 327 | """``sum_e gammas[:, e] * MLP_e(x)`` (TP-partial along d_mlp).""" |
| 328 | from .ops import sink_silu_mul_epilogue |
| 329 | |
| 330 | # One GEMM over the experts' stacked w13 (a view), fused epilogue, |
| 331 | # then one GEMM whose K-reduction over the K-concatenated w2 performs |
| 332 | # the expert sum. |
| 333 | if self._unit is None or self._unit.device != x.device: |
| 334 | self._unit = torch.ones( |
| 335 | self.n_experts, dtype=torch.float32, device=x.device |
| 336 | ) |
| 337 | raw = x @ self.w13_weight.view(-1, x.shape[-1]).T # (T, S*2F) |
| 338 | h = sink_silu_mul_epilogue( |
| 339 | raw, self._unit, gammas, self._unit, self.n_experts, x.dtype |
| 340 | ) |
| 341 | return h @ self.w2_weight.T # (T, D) |
| 342 | |
| 343 | |
| 344 | class InklingSinkExpertsLinear(nn.Module): |
| 345 | """LoRA-capable implementation of the Inkling sink experts.""" |
| 346 | |
| 347 | def __init__( |
| 348 | self, |
| 349 | n_experts: int, |
| 350 | d_model: int, |
| 351 | d_mlp: int, |
| 352 | *, |
| 353 | prefix: str = "", |
| 354 | ) -> None: |
| 355 | super().__init__() |
| 356 | from vllm.model_executor.layers.linear import ( |
| 357 | MergedColumnParallelLinear, |
| 358 | RowParallelLinear, |
| 359 | ) |
| 360 | |
| 361 | self.n_experts = n_experts |
| 362 | self.d_mlp = d_mlp |
| 363 | total = n_experts * d_mlp |
| 364 | self.w13 = MergedColumnParallelLinear( |
| 365 | input_size=d_model, |
| 366 | output_sizes=[total, total], |
| 367 | bias=False, |
| 368 | prefix=f"{prefix}.w13", |
| 369 | ) |
| 370 | self.w2 = RowParallelLinear( |
| 371 | input_size=total, |
| 372 | output_size=d_model, |
| 373 | bias=False, |
| 374 | reduce_results=False, |
| 375 | prefix=f"{prefix}.w2", |
| 376 | ) |
| 377 | self._w2_input_pp = self.w2.input_size_per_partition |
| 378 | self._col_expert: torch.Tensor | None = None |
| 379 | |
| 380 | def _gamma_expand(self, gammas: torch.Tensor) -> torch.Tensor: |
| 381 | if self._col_expert is None or self._col_expert.device != gammas.device: |
| 382 | local = self._w2_input_pp |
| 383 | start = get_tensor_model_parallel_rank() * local |
| 384 | cols = torch.arange(start, start + local, device=gammas.device) |
| 385 | self._col_expert = (cols // self.d_mlp).long() |
| 386 | return gammas[:, self._col_expert] |
| 387 | |
| 388 | def load_weight(self, key: str, weight: torch.Tensor) -> list[str]: |
| 389 | if key == "w13_weight": |
| 390 | d_model = weight.shape[-1] |
| 391 | gate = weight[:, 0::2, :].reshape(-1, d_model).contiguous() |
| 392 | up = weight[:, 1::2, :].reshape(-1, d_model).contiguous() |
| 393 | self.w13.weight_loader(self.w13.weight, gate, 0) |
| 394 | self.w13.weight_loader(self.w13.weight, up, 1) |
| 395 | return ["w13.weight"] |
| 396 | w = weight.permute(1, 0, 2).reshape(weight.shape[1], -1).contiguous() |
| 397 | self.w2.weight_loader(self.w2.weight, w) |
| 398 | return ["w2.weight"] |
| 399 | |
| 400 | def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor: |
| 401 | gate_up, _ = self.w13(x) |
| 402 | gate, up = gate_up.chunk(2, dim=-1) |
| 403 | hidden_states = torch.nn.functional.silu(gate) * up |
| 404 | hidden_states = (hidden_states * self._gamma_expand(gammas)).to(x.dtype) |
| 405 | output, _ = self.w2(hidden_states) |
| 406 | return output |
| 407 | |
| 408 | |
| 409 | class InklingMoE(nn.Module): |
| 410 | def __init__( |
| 411 | self, |
| 412 | config: InklingModelConfig, |
| 413 | *, |
| 414 | prefix: str = "", |
| 415 | quant_config: QuantizationConfig | None = None, |
| 416 | ) -> None: |
| 417 | super().__init__() |
| 418 | # Overfit to the served checkpoint: sigmoid gate renormalized after |
| 419 | # top-k, shared sink experts, interleaved gate/up checkpoint rows. |
| 420 | assert config.gate_activation == "sigmoid" and config.norm_after_topk |
| 421 | assert config.n_shared_experts > 0 and config.shared_expert_sink |
| 422 | assert config.inference_moe_w13_interleaved |
| 423 | n_routed = config.n_routed_experts |
| 424 | n_shared = config.n_shared_experts |
| 425 | self.n_routed_experts = n_routed |
| 426 | self.gate = InklingGate( |
| 427 | d_model=config.hidden_size, |
| 428 | n_routed_experts=n_routed, |
| 429 | n_shared_experts=n_shared, |
| 430 | experts_per_token=config.num_experts_per_tok, |
| 431 | route_scale=config.route_scale, |
| 432 | use_global_scale=config.use_global_scale, |
| 433 | use_gate_bias=config.use_gate_bias, |
| 434 | ) |
| 435 | |
| 436 | # TRTLLM MoE kernels assume equal, contiguous per-rank expert slabs |
| 437 | # (local_expert_offset = ep_rank * local_num_experts), so pad the |
| 438 | # expert count to a multiple of the EP size. A no-op for the usual |
| 439 | # power-of-two EP sizes (n_routed is a power of two). |
| 440 | num_experts = n_routed + (-n_routed) % _inkling_moe_ep_size() |
| 441 | |
| 442 | self.experts = FusedMoEFactory( |
| 443 | num_experts=num_experts, |
| 444 | top_k=config.num_experts_per_tok, |
| 445 | hidden_size=config.hidden_size, |
| 446 | intermediate_size=config.intermediate_size, |
| 447 | renormalize=False, |
| 448 | quant_config=quant_config, |
| 449 | prefix=f"{prefix}.experts", |
| 450 | custom_routing_function=self._select_routed, |
| 451 | router_logits_dtype=torch.float32, |
| 452 | activation="silu", |
| 453 | ) |
| 454 | # The decoder layer reduce-scatters the MoE delta into the sconv |
| 455 | # stream itself (RS -> shard sconv -> AG); the runner must return the |
| 456 | # per-rank partial sum instead of all-reducing. |
| 457 | self.experts.moe_config.skip_final_all_reduce = True |
| 458 | |
| 459 | self._routed_sel: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None |
| 460 | |
| 461 | sink_experts_cls = ( |
| 462 | InklingSinkExpertsLinear |
| 463 | if get_current_vllm_config().lora_config is not None |
| 464 | else InklingSinkExperts |
| 465 | ) |
| 466 | self.sink_experts = sink_experts_cls( |
| 467 | n_experts=n_shared, |
| 468 | d_model=config.hidden_size, |
| 469 | d_mlp=config.intermediate_size, |
| 470 | prefix=f"{prefix}.shared_experts", |
| 471 | ) |
| 472 | |
| 473 | # Sink chain overlaps the routed MoE call on the aux stream for |
| 474 | # decode-sized batches (same pattern as the runner's SharedExperts |
| 475 | # multi-stream overlap). The routed GEMM runs on the default stream and |
| 476 | # the sink chain on the aux stream, joined via these two events by |
| 477 | # ``maybe_execute_in_parallel``. |
| 478 | self._sink_stream: torch.cuda.Stream | None = aux_stream() |
| 479 | self._sink_events = (torch.cuda.Event(), torch.cuda.Event()) |
| 480 | |
| 481 | def _select_routed( |
| 482 | self, |
| 483 | hidden_states: torch.Tensor, |
| 484 | gating_output: torch.Tensor, |
| 485 | topk: int, |
| 486 | renormalize: bool, |
| 487 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 488 | """FusedMoE ``custom_routing_function``: the routed top-k slice of the |
| 489 | full (routed + sink) selection. |
| 490 | |
| 491 | forward() stashes its selection (keyed by logits identity) so the |
| 492 | gate select runs once per layer; the fallback covers paths where the |
| 493 | runner re-derives the logits (e.g. naive DP dispatch). |
| 494 | """ |
| 495 | del hidden_states, renormalize |
| 496 | assert topk == self.gate.topk |
| 497 | cached = self._routed_sel |
| 498 | self._routed_sel = None |
| 499 | if cached is not None and cached[0] is gating_output: |
| 500 | return cached[1], cached[2] |
| 501 | weights, ids = self.gate.select_experts(gating_output) |
| 502 | return weights[:, :topk].contiguous(), ids[:, :topk].contiguous() |
| 503 | |
| 504 | def forward_partials(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| 505 | router_logits = self.gate.compute_logits(x) |
| 506 | num_tokens = x.shape[0] |
| 507 | # One gate select per layer: the routed slice is stashed for the |
| 508 | # routing function inside the FusedMoE op; the sink gammas are the |
| 509 | # trailing columns. |
| 510 | k = self.gate.topk |
| 511 | weights, ids = self.gate.select_experts(router_logits) |
| 512 | self._routed_sel = ( |
| 513 | router_logits, |
| 514 | weights[:, :k].contiguous(), |
| 515 | ids[:, :k].contiguous(), |
| 516 | ) |
| 517 | gammas = weights[:, k:] |
| 518 | |
| 519 | out, sink_out = maybe_execute_in_parallel( |
| 520 | lambda: self.experts(hidden_states=x, router_logits=router_logits), |
| 521 | lambda: self.sink_experts(x, gammas), |
| 522 | self._sink_events[0], |
| 523 | self._sink_events[1], |
| 524 | self._sink_stream |
| 525 | if num_tokens <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD |
| 526 | else None, |
| 527 | ) |
| 528 | self._routed_sel = None |
| 529 | |
| 530 | return out, sink_out |
| 531 | |
| 532 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 533 | out, sink_out = self.forward_partials(x) |
| 534 | return out.add_(sink_out) |
| 535 | |
| 536 | # -- weight loading ---------------------------------------------------- |
| 537 | |
| 538 | def _local_expert_slots(self) -> dict[int, int]: |
| 539 | """Global expert id -> local slot for this rank's expert partition.""" |
| 540 | manager = self.experts.routed_experts.expert_map_manager |
| 541 | if manager.expert_map is None: |
| 542 | return {g: g for g in range(manager.global_num_experts)} |
| 543 | emap = manager.expert_map.tolist() |
| 544 | return {g: slot for g, slot in enumerate(emap) if slot >= 0} |
| 545 | |
| 546 | def load_expert_weight(self, name: str, weight: torch.Tensor) -> list[str]: |
| 547 | """Load one checkpoint expert tensor. |
| 548 | |
| 549 | ``name`` is relative to the mlp module: ``experts.<t>`` (routed |
| 550 | stack) or ``shared_experts.shared_<t>`` (sink experts). Returns the |
| 551 | loaded param names (relative to this module). |
| 552 | """ |
| 553 | if name.startswith("shared_experts."): |
| 554 | key = name.split(".", 1)[1].replace("shared_", "", 1) |
| 555 | return [ |
| 556 | f"sink_experts.{p}" for p in self.sink_experts.load_weight(key, weight) |
| 557 | ] |
| 558 | |
| 559 | experts: RoutedExperts = self.experts.routed_experts |
| 560 | key = name.split(".", 1)[1] |
| 561 | |
| 562 | # original_shape is unused by the vLLM serving layout. |
| 563 | if key.endswith(".original_shape"): |
| 564 | return [] |
| 565 | if key.endswith(".input_amax"): |
| 566 | projection = "w13" if key.startswith("w13") else "w2" |
| 567 | amax = float(weight.max()) |
| 568 | assert math.isfinite(amax) and amax > 0, ( |
| 569 | f"bad {projection} input_amax: {amax}" |
| 570 | ) |
| 571 | input_scale = getattr(experts, f"{projection}_input_scale") |
| 572 | input_scale.data.fill_(amax / _NVFP4_INPUT_SCALE_DENOMINATOR) |
| 573 | return [f"experts.routed_experts.{projection}_input_scale"] |
| 574 | |
| 575 | param = getattr(experts, key) |
| 576 | slots = self._local_expert_slots() |
| 577 | gids = sorted(slots) |
| 578 | lids = [slots[g] for g in gids] |
| 579 | tp_rank = experts.moe_config.moe_parallel_config.tp_rank |
| 580 | |
| 581 | if key.endswith(("_scale_2", "_global_scale")): |
| 582 | # Per-expert scalars, vectorized over the local experts. The |
| 583 | # fused w13 params carry one slot per gate/up half. A single |
| 584 | # checkpoint value is shared by both halves. |
| 585 | vals = weight[gids].float().reshape(len(gids), -1) |
| 586 | target_width = math.prod(param.shape[1:]) |
| 587 | if vals.shape[1] == 1: |
| 588 | vals = vals.expand(-1, target_width) |
| 589 | elif vals.shape[1] != target_width: |
| 590 | raise ValueError( |
| 591 | f"cannot load {tuple(weight.shape)} into {tuple(param.shape)}" |
| 592 | ) |
| 593 | param.data[lids] = vals.reshape(len(gids), *param.shape[1:]).to( |
| 594 | param.device |
| 595 | ) |
| 596 | elif key == "w2_weight_scale" and weight.shape[-1] == 1: |
| 597 | # Per-output-channel scales are replicated across TP ranks. |
| 598 | param.data[lids] = weight[gids].to(device=param.device, dtype=param.dtype) |
| 599 | elif key.startswith("w13"): |
| 600 | # Checkpoint w13 rows are interleaved [g0, u0, g1, u1, ...]; the |
| 601 | # fused param layout is [w1(gate); w3(up)]. The TP-local rows form |
| 602 | # one contiguous slab of the interleaved tensor, so upload just |
| 603 | # that slab (a single bounded synchronous H2D; pre-uploading whole |
| 604 | # untrimmed tensors pins the mmap pages of the entire checkpoint |
| 605 | # and OOMs the host) and de-interleave on device. |
| 606 | half = param.shape[1] // 2 |
| 607 | for gid, lid in slots.items(): |
| 608 | slab = weight[gid].narrow(0, tp_rank * 2 * half, 2 * half) |
| 609 | slab = slab.to(param.device) |
| 610 | param.data[lid, :half].copy_(slab[0::2]) |
| 611 | param.data[lid, half:].copy_(slab[1::2]) |
| 612 | else: |
| 613 | # w2: shard the packed intermediate (last) dim. |
| 614 | shard = param.shape[2] |
| 615 | for gid, lid in slots.items(): |
| 616 | param.data[lid].copy_(weight[gid].narrow(1, tp_rank * shard, shard)) |
| 617 | return [f"experts.routed_experts.{key}"] |
| 618 | |
| 619 | def finalize_load(self) -> list[str]: |
| 620 | """Post-load fixups for zeroed padding experts.""" |
| 621 | experts = self.experts.routed_experts |
| 622 | out: list[str] = [] |
| 623 | # Zero the EP-alignment padding experts (if any) so their |
| 624 | # (never-routed) slots hold defined values. |
| 625 | slots = self._local_expert_slots() |
| 626 | for gid in range(self.n_routed_experts, experts.global_num_experts): |
| 627 | lid = slots.get(gid) |
| 628 | if lid is None: |
| 629 | continue |
| 630 | for pname in ( |
| 631 | "w13_weight", |
| 632 | "w2_weight", |
| 633 | "w13_weight_scale", |
| 634 | "w2_weight_scale", |
| 635 | "w13_weight_scale_2", |
| 636 | "w2_weight_scale_2", |
| 637 | ): |
| 638 | p = getattr(experts, pname, None) |
| 639 | if p is not None: |
| 640 | p.data[lid].zero_() |
| 641 | return out |