[Model] Add Inkling LoRA support [4/N] (#48884)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Woosuk Kwon committed
Jul 16, 2026 at 18:52 UTC
f3e9497e921a16741401c5e93af0c2c29ea74907
12 files changed
+338
-48
tests/models/inkling/test_moe_weight_layout.py
+25
@@ -6,7 +6,9 @@ from types import SimpleNamespace
6
import pytest
7
import torch
8
9
+from vllm.lora.utils import get_supported_lora_modules
10
from vllm.models.inkling.nvidia import moe
11
+from vllm.models.inkling.nvidia.model import _TmlForCausalLMBase
12
from vllm.platforms import current_platform
13
14
@@ -62,6 +64,29 @@ def test_gate_uses_ll_bf16_gemm_through_token_limit(
64
assert logits.dtype == torch.float32
65
66
67
+def test_gate_is_not_a_lora_target() -> None:
68
+ model = torch.nn.Module()
69
+ model.gate = moe.InklingGate(
70
+ d_model=4,
71
+ n_routed_experts=5,
72
+ n_shared_experts=2,
73
+ experts_per_token=2,
74
+ route_scale=1.0,
75
+ )
76
+
77
+ assert "gate" not in get_supported_lora_modules(model)
78
+
79
+
80
+def test_custom_embedding_is_not_a_lora_target() -> None:
81
+ model = torch.nn.Module()
82
+ model.embedding_modules = _TmlForCausalLMBase.embedding_modules
83
+
84
+ supported = get_supported_lora_modules(model)
85
+
86
+ assert "embed_tokens" not in supported
87
+ assert "lm_head" in supported
88
+
89
+
90
@pytest.mark.parametrize(("projection", "amax"), [("w13", 4.375), ("w2", 2960.0)])
91
def test_moe_loads_calibrated_input_scale(projection: str, amax: float) -> None:
92
experts = SimpleNamespace(
vllm/config/lora.py
+19
-1
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import os
5
from typing import TYPE_CHECKING, Any, Literal
6
7
import torch
@@ -75,8 +76,14 @@ class LoRAConfig:
76
"""If True, force the engine to use the universal 2D MoE LoRA wrapper
77
(`FusedMoEWithLoRA`) regardless of the model's `is_3d_moe_weight` flag, so
78
that 2D-format and 3D-format MoE LoRA adapters can be served in the same
78
- deployment. Only meaningful forMoE models; ignored otherwise. Default False
79
+ deployment. Only meaningful for MoE models; ignored otherwise. Default False
80
keeps the existing model-driven behavior."""
81
+ enable_moe_shared_loras: bool = False
82
+ """If True, load MoE expert adapters in the "shared-outer" layout, where the
83
+ gate/up (`w1`/`w3`) lora_A and the down (`w2`) lora_B are shared across all
84
+ experts (stored once with expert-dim 1) instead of per-expert. The shared
85
+ factors are broadcast to the expert count at kernel time. Only meaningful for
86
+ MoE models whose adapters use this layout; ignored otherwise."""
87
88
def compute_hash(self) -> str:
89
"""
@@ -97,6 +104,7 @@ class LoRAConfig:
104
factors.append(self.lora_dtype)
105
factors.append(self.enable_tower_connector_lora)
106
factors.append(self.enable_mixed_moe_lora_format)
107
+ factors.append(self.enable_moe_shared_loras)
108
# target_modules affects which modules get LoRA applied
109
factors.append(
110
tuple(sorted(self.target_modules)) if self.target_modules else None
@@ -129,3 +137,13 @@ class LoRAConfig:
137
self.lora_dtype = model_config.dtype
138
elif isinstance(self.lora_dtype, str):
139
self.lora_dtype = getattr(torch, self.lora_dtype)
140
+
141
+ architectures = getattr(model_config, "architectures", None) or []
142
+ is_inkling = any("Inkling" in arch for arch in architectures)
143
+ if is_inkling and os.environ.get("INKLING_MULTIMEM_AR", "1") != "0":
144
+ raise ValueError(
145
+ "Inkling LoRA requires INKLING_MULTIMEM_AR=0: the Lamport "
146
+ "fused-collective path bypasses the LoRA-wrapped wo_ud and dense "
147
+ "down_proj layers on decode-sized batches, silently dropping "
148
+ "their LoRA. Set INKLING_MULTIMEM_AR=0."
149
+ )
vllm/engine/arg_utils.py
+6
@@ -602,6 +602,7 @@ class EngineArgs:
602
enable_tower_connector_lora: bool = LoRAConfig.enable_tower_connector_lora
603
specialize_active_lora: bool = LoRAConfig.specialize_active_lora
604
enable_mixed_moe_lora_format: bool = LoRAConfig.enable_mixed_moe_lora_format
605
+ enable_moe_shared_loras: bool = LoRAConfig.enable_moe_shared_loras
606
607
ray_workers_use_nsight: bool = ParallelConfig.ray_workers_use_nsight
608
num_gpu_blocks_override: int | None = CacheConfig.num_gpu_blocks_override
@@ -1351,6 +1352,10 @@ class EngineArgs:
1352
"--enable-mixed-moe-lora-format",
1353
**lora_kwargs["enable_mixed_moe_lora_format"],
1354
)
1355
+ lora_group.add_argument(
1356
+ "--enable-moe-shared-loras",
1357
+ **lora_kwargs["enable_moe_shared_loras"],
1358
+ )
1359
1360
# Observability arguments
1361
observability_kwargs = get_kwargs(ObservabilityConfig)
@@ -2213,6 +2218,7 @@ class EngineArgs:
2218
enable_tower_connector_lora=self.enable_tower_connector_lora,
2219
specialize_active_lora=self.specialize_active_lora,
2220
enable_mixed_moe_lora_format=self.enable_mixed_moe_lora_format,
2221
+ enable_moe_shared_loras=self.enable_moe_shared_loras,
2222
max_cpu_loras=self.max_cpu_loras
2223
if self.max_cpu_loras and self.max_cpu_loras > 0
2224
else None,
vllm/lora/layers/fused_moe.py
+51
-15
@@ -57,6 +57,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
57
# For non-gated MoE (is_act_and_mul=False), only 1 slice is needed
58
# since there's only up_proj (w1), not gate_proj + up_proj (w1 + w3)
59
self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1
60
+ # Set from lora_config.enable_moe_shared_loras in create_lora_weights.
61
+ # When True, w13 lora_A and w2 lora_B are stored once for all experts.
62
+ self.enable_moe_shared_loras = False
63
# Mirrors per-(lora_id) layout of `self.lora_a_stacked` (built in
64
# `create_lora_weights`) so `create_dummy_lora`'s n_slices fallback
65
# matches `lora_a_stacked` length under EP.
@@ -150,10 +153,21 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
153
local_num_experts=self.local_num_experts,
154
punica_wrapper=self.punica_wrapper,
155
use_tuned_config=bool(envs.VLLM_TUNED_CONFIG_FOLDER),
156
+ enable_moe_shared_loras=self.enable_moe_shared_loras,
157
aux_stream=self._lora_stream if use_dual_stream else None,
158
events=self._events if use_dual_stream else None,
159
)
160
161
+ @property
162
+ def _w13_a_num_experts(self) -> int:
163
+ """Expert-dim of the w13 lora_A buffer: 1 when shared."""
164
+ return 1 if self.enable_moe_shared_loras else self.local_num_experts
165
+
166
+ @property
167
+ def _w2_b_num_experts(self) -> int:
168
+ """Expert-dim of the w2 lora_B buffer: 1 when shared."""
169
+ return 1 if self.enable_moe_shared_loras else self.local_num_experts
170
+
171
def _create_lora_a_weights(
172
self,
173
max_loras: int,
@@ -163,7 +177,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
177
torch.zeros(
178
(
179
max_loras,
166
- self.local_num_experts,
180
+ self._w13_a_num_experts,
181
lora_config.max_lora_rank
182
if not self.fully_sharded
183
else divide(lora_config.max_lora_rank, self.tp_size),
@@ -205,7 +219,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
219
torch.zeros(
220
(
221
max_loras,
208
- self.local_num_experts,
222
+ self._w2_b_num_experts,
223
self.hidden_size
224
if not self.fully_sharded
225
else divide(self.hidden_size, self.tp_size),
@@ -247,6 +261,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
261
self._verify_ep_fs(lora_config)
262
self.max_loras = lora_config.max_loras
263
self.fully_sharded = lora_config.fully_sharded_loras
264
+ self.enable_moe_shared_loras = lora_config.enable_moe_shared_loras
265
266
self.adapter_enabled = torch.tensor(
267
[0] * (max_loras + 1), dtype=torch.int, device=self.device
@@ -263,8 +278,11 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
278
for experts_id in range(self.local_num_experts):
279
# For gated MoE: gate_proj (w1), down_proj (w2), up_proj (w3)
280
# For non-gated MoE: up_proj (w1), down_proj (w2)
281
+ # Shared factors are collapsed (expert-dim 1): index 0.
282
+ w13_a_eid = 0 if self.enable_moe_shared_loras else experts_id
283
+ w2_b_eid = 0 if self.enable_moe_shared_loras else experts_id
284
self.lora_a_stacked.append(
267
- self.w13_lora_a_stacked[0][lora_id][experts_id]
285
+ self.w13_lora_a_stacked[0][lora_id][w13_a_eid]
286
)
287
self.lora_a_stacked.append(
288
self.w2_lora_a_stacked[0][lora_id][experts_id]
@@ -273,14 +291,12 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
291
self.lora_b_stacked.append(
292
self.w13_lora_b_stacked[0][lora_id][experts_id]
293
)
276
- self.lora_b_stacked.append(
277
- self.w2_lora_b_stacked[0][lora_id][experts_id]
278
- )
294
+ self.lora_b_stacked.append(self.w2_lora_b_stacked[0][lora_id][w2_b_eid])
295
296
# Only add w3 (up_proj) for gated MoE (_w13_slices == 2)
297
if self._w13_slices == 2:
298
self.lora_a_stacked.append(
283
- self.w13_lora_a_stacked[1][lora_id][experts_id]
299
+ self.w13_lora_a_stacked[1][lora_id][w13_a_eid]
300
)
301
self.lora_b_stacked.append(
302
self.w13_lora_b_stacked[1][lora_id][experts_id]
@@ -340,6 +356,22 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
356
357
return w2_lora_b[:, start_idx:end_idx, :]
358
359
+ def _match_expert_dim(
360
+ self, src: torch.Tensor, buffer: torch.Tensor
361
+ ) -> torch.Tensor:
362
+ """Align a LoRA factor's expert-dim (dim 0) to its stacked buffer.
363
+
364
+ Equal dims pass through. When the buffer is collapsed to expert-dim 1
365
+ (a shared factor) but the source carries per-expert copies, keep
366
+ the first — every expert shares the same factor, so the copies are
367
+ identical for a real adapter and irrelevant (zeros) for the dummy.
368
+ """
369
+ tgt = buffer.shape[1]
370
+ if src.shape[0] == tgt:
371
+ return src
372
+ assert tgt == 1, f"expert-dim mismatch: source {src.shape[0]} vs buffer {tgt}"
373
+ return src[:1]
374
+
375
def reset_lora(self, index: int):
376
"""Resets the lora weights at index back to 0."""
377
for pos in range(self._w13_slices):
@@ -366,20 +398,24 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
398
self.reset_lora(index)
399
self.adapter_enabled[index] = 1
400
369
- num_experts = self.w13_lora_a_stacked[0].shape[1]
370
-
401
w1_lora_a, w2_lora_a, w3_lora_a = lora_a
402
w1_lora_b, w2_lora_b, w3_lora_b = lora_b
403
404
# EP slicing is done once at add time in
405
# LoRAModelManager._slice_moe_lora_ep, so by here the cached
406
# tensors already match the local-expert dim of the stacked buffers.
377
- assert (
378
- num_experts
379
- == w1_lora_a.shape[0]
380
- == w2_lora_a.shape[0]
381
- == w3_lora_a.shape[0]
382
- )
407
+ # For shared-loras adapters the w13 lora_A and w2 lora_B buffers are
408
+ # collapsed to expert-dim 1: match each factor to the expert-dim of
409
+ # the buffer it is copied into. A real shared factor already has
410
+ # expert-dim 1 (no-op); a per-expert source (e.g. the warmup dummy
411
+ # LoRA, which always stacks to num_experts) is collapsed to the shared
412
+ # slot since every expert shares the same factor.
413
+ w1_lora_a = self._match_expert_dim(w1_lora_a, self.w13_lora_a_stacked[0])
414
+ w3_lora_a = self._match_expert_dim(w3_lora_a, self.w13_lora_a_stacked[-1])
415
+ w2_lora_a = self._match_expert_dim(w2_lora_a, self.w2_lora_a_stacked[0])
416
+ w1_lora_b = self._match_expert_dim(w1_lora_b, self.w13_lora_b_stacked[0])
417
+ w3_lora_b = self._match_expert_dim(w3_lora_b, self.w13_lora_b_stacked[-1])
418
+ w2_lora_b = self._match_expert_dim(w2_lora_b, self.w2_lora_b_stacked[0])
419
420
slliced_w1_lora_a = self._slice_w13_a(w1_lora_a)
421
slliced_w1_lora_b = self._slice_w13_b(w1_lora_b)
vllm/lora/lora_weights.py
+33
@@ -227,6 +227,39 @@ class PackedLoRALayerWeights(LoRALayerWeights):
227
)
228
return obj
229
230
+ @classmethod
231
+ def pack_moe_stacked(
232
+ cls,
233
+ loras: GenericSequence["LoRALayerWeights | None"],
234
+ module_name: str,
235
+ ) -> "PackedLoRALayerWeights":
236
+ """Pack pre-stacked (3D) w1/w2/w3 expert LoRAs into a single LoRA.
237
+
238
+ Unlike :meth:`pack_moe`, which stacks one 2D tensor per expert, this
239
+ expects each of the three input LoRAs to already carry the expert
240
+ dimension (``experts.w{1,2,3}``). It is used for "shared-outer"
241
+ adapters where some factors are shared across experts (expert-dim 1)
242
+ and cannot be reconstructed by stacking per-expert tensors. The
243
+ produced ``lora_a``/``lora_b`` are the ``[w1, w2, w3]`` lists that
244
+ ``FusedMoEWithLoRA.set_lora`` consumes directly.
245
+ """
246
+ assert len(loras) == 3, (
247
+ "shared-outer expert LoRA expects exactly w1/w2/w3 stacked tensors"
248
+ )
249
+ w1_lora, w2_lora, w3_lora = loras
250
+ assert w1_lora is not None and w2_lora is not None and w3_lora is not None
251
+ rank = w1_lora.rank
252
+ lora_alpha = w1_lora.lora_alpha
253
+ scaling = lora_alpha / rank
254
+ return cls(
255
+ module_name,
256
+ rank,
257
+ [lora_alpha, lora_alpha, lora_alpha],
258
+ [w1_lora.lora_a, w2_lora.lora_a, w3_lora.lora_a],
259
+ [w1_lora.lora_b, w2_lora.lora_b, w3_lora.lora_b],
260
+ scaling=[scaling, scaling, scaling],
261
+ )
262
+
263
def optimize(self) -> "PackedLoRALayerWeights":
264
"""Optimize the LoRA by merging the scaling into lora_b."""
265
for i in range(len(self.lora_b)):
vllm/lora/model_manager.py
+27
-25
@@ -132,8 +132,13 @@ class LoRAModelManager:
132
and self.model.is_3d_moe_weight
133
and not self._enable_mixed_moe_lora_format
134
)
135
+ # Shared MoE adapters: w13 lora_A / w2 lora_B shared across experts,
136
+ # stored as pre-stacked experts.w{1,2,3} tensors (startup opt-in).
137
+ self._enable_moe_shared_loras = is_moe and lora_config.enable_moe_shared_loras
138
self.packed_modules_mapping = process_packed_modules_mapping(
136
- self.model, force_2d_moe=self._enable_mixed_moe_lora_format
139
+ self.model,
140
+ force_2d_moe=self._enable_mixed_moe_lora_format,
141
+ enable_moe_shared_loras=self._enable_moe_shared_loras,
142
)
143
self._is_non_gated_moe = is_moe and self.model.is_non_gated_moe
144
self._use_ep = bool(
@@ -761,6 +766,16 @@ class LoRAModelManager:
766
if lora_model.check_lora_name(replaced_module_name):
767
module_name = replaced_module_name
768
if module_name.endswith(".experts"):
769
+ if self._enable_moe_shared_loras:
770
+ lora_model.loras[module_name] = (
771
+ PackedLoRALayerWeights.pack_moe_stacked(
772
+ replacement_loras,
773
+ module_name,
774
+ )
775
+ )
776
+ for module in packed_module_names:
777
+ lora_model.loras.pop(module, None)
778
+ continue
779
if self._is_non_gated_moe and len(replacement_loras) > 0:
780
replacement_loras = self._pad_lora_pairs_to_triplets(
781
replacement_loras
@@ -1019,20 +1034,10 @@ class LoRAModelManager:
1034
module: FusedMoEWithLoRA,
1035
module_name: str,
1036
) -> None:
1022
- """Slice the cached LoRA tensors down to this rank's local experts.
1023
-
1024
- The 2D MoE checkpoint enters as a list of per-(w1/w2/w3) tensors of
1025
- shape (num_experts, rank, in) / (num_experts, out, rank). When EP
1026
- is active each rank only owns local_num_experts; without this slice
1027
- the CPU LoRAModel keeps the full global weight and set_lora has to
1028
- re-slice on every activation.
1029
-
1030
- With the load-time / pack-time slicing in
1031
- ``_restrict_to_local_experts``, the stacked tensors already match
1032
- ``local_num_experts`` and the inner branch becomes a no-op. The
1033
- guard remains so checkpoints that bypassed the pre-slicing (e.g.
1034
- ``.bin``/``.pt`` adapters with weights mappers we don't recognize)
1035
- still get sliced here.
1037
+ """Slice cached LoRA tensors down to this rank's local experts.
1038
+
1039
+ Shared factors have expert dimension one and remain available on every
1040
+ rank. Per-expert factors are narrowed to the local expert block.
1041
"""
1042
if not module.use_ep:
1043
return
@@ -1046,16 +1051,13 @@ class LoRAModelManager:
1051
expert_start = ep_rank * local_num_experts
1052
expert_end = expert_start + local_num_experts
1053
1049
- new_lora_a: list[torch.Tensor | None] = []
1050
- new_lora_b: list[torch.Tensor | None] = []
1051
- for a, b in zip(module_lora.lora_a, module_lora.lora_b):
1052
- if a is not None and b is not None and a.shape[0] == global_num_experts:
1053
- a = a[expert_start:expert_end].contiguous()
1054
- b = b[expert_start:expert_end].contiguous()
1055
- new_lora_a.append(a)
1056
- new_lora_b.append(b)
1057
- module_lora.lora_a = new_lora_a
1058
- module_lora.lora_b = new_lora_b
1054
+ def _slice_local(t: torch.Tensor | None) -> torch.Tensor | None:
1055
+ if t is not None and t.shape[0] == global_num_experts:
1056
+ return t[expert_start:expert_end].contiguous()
1057
+ return t
1058
+
1059
+ module_lora.lora_a = [_slice_local(a) for a in module_lora.lora_a]
1060
+ module_lora.lora_b = [_slice_local(b) for b in module_lora.lora_b]
1061
1062
def _restrict_to_local_experts(
1063
self, module_name: str, new_module_names: list[str]
vllm/lora/utils.py
+14
-2
@@ -369,7 +369,9 @@ def get_adapter_absolute_path(lora_path: str) -> str:
369
370
371
def process_packed_modules_mapping(
372
- model: nn.Module, force_2d_moe: bool = False
372
+ model: nn.Module,
373
+ force_2d_moe: bool = False,
374
+ enable_moe_shared_loras: bool = False,
375
) -> dict[str, list[str]]:
376
if is_moe_model(model):
377
# This method generates and returns a dictionary mapping packed module
@@ -382,7 +384,17 @@ def process_packed_modules_mapping(
384
# the engine forces the universal 2D wrapper via
385
# enable_mixed_moe_lora_format (so 3D models can also load 2D
386
# adapters through FusedMoEWithLoRA).
385
- if (not model.is_3d_moe_weight) or force_2d_moe:
387
+ if enable_moe_shared_loras:
388
+ # Shared MoE adapters store one pre-stacked tensor per
389
+ # expert-projection (experts.w1/w2/w3) rather than a per-expert
390
+ # tensor list, so the packed mapping references the three stack
391
+ # names directly (drives expected_lora_modules and the packing).
392
+ packed_modules_mapping["experts"] = [
393
+ "experts.w1",
394
+ "experts.w2",
395
+ "experts.w3",
396
+ ]
397
+ elif (not model.is_3d_moe_weight) or force_2d_moe:
398
# Filter out malformed entries: non-gated MoE has empty
399
# ckpt_up_proj_name which results in weight_name containing ".."
400
# (e.g., "experts.0.." instead of "experts.0.layer_name.")
vllm/model_executor/layers/fused_moe/experts/lora_context.py
+5
@@ -43,6 +43,11 @@ class MoELoRAContext:
43
# try_get_optimal_moe_lora_config for Triton kernel tile configs.
44
use_tuned_config: bool
45
46
+ # Shared MoE LoRA: w13 lora_A and w2 lora_B are shared across all experts
47
+ # (stored collapsed with expert-dim 1). When True, LoRAExpertsMixin expands
48
+ # them to local_num_experts via a stride-0 view before the kernel.
49
+ enable_moe_shared_loras: bool = False
50
+
51
# Optional dual-stream support for overlapping each (base GEMM, LoRA)
52
# pair. When aux_stream is None, the experts.apply() path runs the
53
# original sequential schedule. When set, base GEMM runs on the default
vllm/model_executor/layers/fused_moe/experts/lora_experts_mixin.py
+19
-2
@@ -52,10 +52,19 @@ class LoRAExpertsMixin:
52
torch.Tensor | None,
53
torch.Tensor | None,
54
]:
55
+ w13_lora_a_stacked = lora_context.w13_lora_a_stacked
56
+ if lora_context.enable_moe_shared_loras:
57
+ # w13 lora_A is shared across experts (collapsed expert-dim 1);
58
+ # broadcast to local_num_experts via a stride-0 view. The kernel
59
+ # derives num_experts from shape[1] and indexes via stride(1)==0.
60
+ w13_lora_a_stacked = tuple(
61
+ a.expand(-1, lora_context.local_num_experts, -1, -1)
62
+ for a in w13_lora_a_stacked
63
+ )
64
return lora_context.punica_wrapper.add_lora_w13(
65
y,
66
x,
58
- lora_context.w13_lora_a_stacked,
67
+ w13_lora_a_stacked,
68
lora_context.w13_lora_b_stacked,
69
topk_ids,
70
topk_weights,
@@ -92,11 +101,19 @@ class LoRAExpertsMixin:
101
top_k_num: int,
102
add_inputs: bool = True,
103
) -> None:
104
+ w2_lora_b_stacked = lora_context.w2_lora_b_stacked
105
+ if lora_context.enable_moe_shared_loras:
106
+ # w2 lora_B is shared across experts (collapsed expert-dim 1);
107
+ # broadcast to local_num_experts via a stride-0 view.
108
+ w2_lora_b_stacked = tuple(
109
+ b.expand(-1, lora_context.local_num_experts, -1, -1)
110
+ for b in w2_lora_b_stacked
111
+ )
112
lora_context.punica_wrapper.add_lora_w2(
113
y,
114
x,
115
lora_context.w2_lora_a_stacked,
99
- lora_context.w2_lora_b_stacked,
116
+ w2_lora_b_stacked,
117
topk_weights,
118
sorted_token_ids_lora,
119
expert_ids_lora,
vllm/models/inkling/nvidia/logits_processor.py
+57
-1
@@ -1,6 +1,22 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Inkling logits processor with muP output scaling."""
3
+"""Inkling logits processor (muP + LoRA aware).
4
+
5
+Inkling divides the final logits by a muP width multiplier
6
+(``logits_mup_width_multiplier``). This applies it two ways, depending on
7
+whether an lm_head LoRA is attached:
8
+
9
+* No LoRA: fold ``1/mup`` into the lm_head GEMM alpha (fp32 epilogue) -- no
10
+ separate elementwise kernel, no extra rounding, no weight mutation.
11
+* LoRA attached: the LoRA manager wraps this layer in
12
+ ``LogitsProcessorWithLoRA``, whose ``forward`` calls
13
+ ``type(base_layer).forward(self=wrapper)`` -- so this ``forward`` runs with
14
+ ``self`` bound to the wrapper. We detect that via ``base_layer`` and take the
15
+ LoRA path: run the wrapper's ``_get_logits`` (base logits + the lm_head LoRA
16
+ delta), then divide the full logits by the multiplier so the delta is scaled
17
+ too. muP thus composes with the LoRA delta, with the dispatch as the only
18
+ model-side branching.
19
+"""
20
21
from __future__ import annotations
22
@@ -47,6 +63,46 @@ class InklingLogitsProcessor(LogitsProcessor):
63
lm_head: VocabParallelEmbedding,
64
hidden_states: torch.Tensor,
65
embedding_bias: torch.Tensor | None = None,
66
+ ) -> torch.Tensor | None:
67
+ # ``base_layer`` exists only on the LogitsProcessorWithLoRA wrapper,
68
+ # which calls this forward with ``self`` bound to the wrapper. The
69
+ # wrapper is not an ``InklingLogitsProcessor`` instance, so dispatch
70
+ # ``_lora_forward`` explicitly through the base_layer's class (it
71
+ # provides ``_get_logits``/``logits_as_input``; only ``_lora_forward``
72
+ # lives on this class).
73
+ if hasattr(self, "base_layer"):
74
+ return type(self.base_layer)._lora_forward(
75
+ self, lm_head, hidden_states, embedding_bias
76
+ )
77
+ return self._base_forward(lm_head, hidden_states, embedding_bias)
78
+
79
+ def _lora_forward(
80
+ self,
81
+ lm_head: VocabParallelEmbedding,
82
+ hidden_states: torch.Tensor,
83
+ embedding_bias: torch.Tensor | None = None,
84
+ ) -> torch.Tensor | None:
85
+ # ``self`` is the LogitsProcessorWithLoRA wrapper here: ``_get_logits``
86
+ # returns the base logits plus the lm_head LoRA delta. Apply the muP
87
+ # divisor on the full logits so the LoRA delta is scaled too.
88
+ mup_multiplier = self.base_layer.logits_mup_width_multiplier
89
+ mup = 1.0 / mup_multiplier if mup_multiplier else None
90
+ if self.logits_as_input:
91
+ logits = hidden_states
92
+ else:
93
+ logits = self._get_logits(hidden_states, lm_head, embedding_bias)
94
+ # TODO: fuse this multiplication
95
+ if logits is not None and mup:
96
+ assert self.base_layer.soft_cap is None
97
+ assert self.base_layer.scale == 1.0
98
+ logits = logits * mup
99
+ return logits
100
+
101
+ def _base_forward(
102
+ self,
103
+ lm_head: VocabParallelEmbedding,
104
+ hidden_states: torch.Tensor,
105
+ embedding_bias: torch.Tensor | None = None,
106
) -> torch.Tensor | None:
107
mup = self.logits_mup_width_multiplier
108
if not mup:
vllm/models/inkling/nvidia/model.py
+11
-1
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
24
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
25
from vllm.model_executor.models.interfaces import (
26
MultiModalEmbeddings,
27
+ SupportsLoRA,
28
SupportsMultiModal,
29
SupportsPP,
30
)
@@ -360,7 +361,7 @@ class InklingModel(nn.Module):
361
return self.norm(hidden_states)
362
363
363
-class _TmlForCausalLMBase(nn.Module, SupportsPP):
364
+class _TmlForCausalLMBase(nn.Module, SupportsPP, SupportsLoRA):
365
"""Shared text-backbone causal-LM scaffolding for both entry classes."""
366
367
hf_to_vllm_mapper = WeightsMapper(
@@ -380,6 +381,8 @@ class _TmlForCausalLMBase(nn.Module, SupportsPP):
381
"model.llm.embed": "model.embed_tokens",
382
"model.llm.norm": "model.norm",
383
"model.llm.unembed": "lm_head",
384
+ "language_model.layers.": "model.layers.",
385
+ "language_model.lm_head.": "lm_head.",
386
},
387
orig_to_new_suffix={
388
# NVFP4 scale
@@ -389,6 +392,13 @@ class _TmlForCausalLMBase(nn.Module, SupportsPP):
392
".w2_weight.scale2": ".w2_weight_scale_2",
393
},
394
)
395
+ packed_modules_mapping = {
396
+ "qkvr": ["wq_du", "wk_dv", "wv_dv", "wr_du"],
397
+ "w13": ["w1", "w3"],
398
+ }
399
+ embedding_modules = {
400
+ "lm_head": "output_embeddings",
401
+ }
402
403
def _build(
404
self,
vllm/models/inkling/nvidia/moe.py
+71
-1
@@ -343,6 +343,71 @@ class InklingSinkExperts(nn.Module):
343
return h @ self.w2_weight.T # (T, D)
344
345
346
+class InklingSinkExpertsLinear(nn.Module):
347
+ """LoRA-capable implementation of the Inkling sink experts."""
348
+
349
+ def __init__(
350
+ self,
351
+ n_experts: int,
352
+ d_model: int,
353
+ d_mlp: int,
354
+ *,
355
+ prefix: str = "",
356
+ ) -> None:
357
+ super().__init__()
358
+ from vllm.model_executor.layers.linear import (
359
+ MergedColumnParallelLinear,
360
+ RowParallelLinear,
361
+ )
362
+
363
+ self.n_experts = n_experts
364
+ self.d_mlp = d_mlp
365
+ total = n_experts * d_mlp
366
+ self.w13 = MergedColumnParallelLinear(
367
+ input_size=d_model,
368
+ output_sizes=[total, total],
369
+ bias=False,
370
+ prefix=f"{prefix}.w13",
371
+ )
372
+ self.w2 = RowParallelLinear(
373
+ input_size=total,
374
+ output_size=d_model,
375
+ bias=False,
376
+ reduce_results=False,
377
+ prefix=f"{prefix}.w2",
378
+ )
379
+ self._w2_input_pp = self.w2.input_size_per_partition
380
+ self._col_expert: torch.Tensor | None = None
381
+
382
+ def _gamma_expand(self, gammas: torch.Tensor) -> torch.Tensor:
383
+ if self._col_expert is None or self._col_expert.device != gammas.device:
384
+ local = self._w2_input_pp
385
+ start = get_tensor_model_parallel_rank() * local
386
+ cols = torch.arange(start, start + local, device=gammas.device)
387
+ self._col_expert = (cols // self.d_mlp).long()
388
+ return gammas[:, self._col_expert]
389
+
390
+ def load_weight(self, key: str, weight: torch.Tensor) -> list[str]:
391
+ if key == "w13_weight":
392
+ d_model = weight.shape[-1]
393
+ gate = weight[:, 0::2, :].reshape(-1, d_model).contiguous()
394
+ up = weight[:, 1::2, :].reshape(-1, d_model).contiguous()
395
+ self.w13.weight_loader(self.w13.weight, gate, 0)
396
+ self.w13.weight_loader(self.w13.weight, up, 1)
397
+ return ["w13.weight"]
398
+ w = weight.permute(1, 0, 2).reshape(weight.shape[1], -1).contiguous()
399
+ self.w2.weight_loader(self.w2.weight, w)
400
+ return ["w2.weight"]
401
+
402
+ def forward(self, x: torch.Tensor, gammas: torch.Tensor) -> torch.Tensor:
403
+ gate_up, _ = self.w13(x)
404
+ gate, up = gate_up.chunk(2, dim=-1)
405
+ hidden_states = torch.nn.functional.silu(gate) * up
406
+ hidden_states = (hidden_states * self._gamma_expand(gammas)).to(x.dtype)
407
+ output, _ = self.w2(hidden_states)
408
+ return output
409
+
410
+
411
class InklingMoE(nn.Module):
412
def __init__(
413
self,
@@ -417,7 +482,12 @@ class InklingMoE(nn.Module):
482
layer_id
483
), f"layer {layer_id}: NVFP4 shared experts are not supported"
484
420
- self.sink_experts = InklingSinkExperts(
485
+ sink_experts_cls = (
486
+ InklingSinkExpertsLinear
487
+ if get_current_vllm_config().lora_config is not None
488
+ else InklingSinkExperts
489
+ )
490
+ self.sink_experts = sink_experts_cls(
491
n_experts=n_shared,
492
d_model=config.hidden_size,
493
d_mlp=config.intermediate_size,