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