[Kernel][Inkling] Fuse shared-expert partial addition into the Lamport collective (#50697)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
Canlin Guo committed
Aug 5, 2026 at 02:41 UTC
edbc4969a76b25c531f7b9bb16b79984e65ef023
3 files changed
+48
-6
vllm/models/inkling/nvidia/model.py
+18
-5
@@ -5,7 +5,7 @@
5
from __future__ import annotations
6
7
from collections.abc import Iterable
8
-from typing import Any
8
+from typing import Any, TypeAlias
9
10
import regex as re
11
import torch
@@ -57,6 +57,8 @@ from .ops.norm import add_rmsnorm, embed_rmsnorm
57
from .sconv_swa_attn import _ATTN, _MLP, InklingConvState, InklingSconvMetadata
58
from .short_conv import InklingShortConv
59
60
+InklingDelta: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
61
+
62
63
def _layer_id(name: str) -> int | None:
64
m = re.search(r"\.layers\.(\d+)\.", name)
@@ -64,7 +66,7 @@ def _layer_id(name: str) -> int | None:
66
67
68
def _sconv_add_norm(
67
- delta: torch.Tensor,
69
+ delta: InklingDelta,
70
hidden: torch.Tensor,
71
sconv: InklingShortConv,
72
norm: InklingRMSNorm | None,
@@ -84,6 +86,10 @@ def _sconv_add_norm(
86
off_s, ws = sconv.owner.stream_ranges[sconv.stream_idx]
87
norm_w = norm.weight if norm is not None else None
88
eps = norm.variance_epsilon if norm is not None else 0.0
89
+ if isinstance(delta, tuple):
90
+ delta, shared_delta = delta
91
+ else:
92
+ shared_delta = None
93
94
mm = get_lamport_rs_conv(hidden.shape[-1], sconv.kernel_size)
95
if mm is not None and mm.usable(delta.shape[0]) and m is not None:
@@ -103,9 +109,12 @@ def _sconv_add_norm(
109
off_s,
110
ws,
111
sconv.owner.block_size,
112
+ shared_tensor=shared_delta,
113
)
114
115
# Fallback: NCCL RS -> shard sconv -> AG -> fused add(+rmsnorm).
116
+ if shared_delta is not None:
117
+ delta.add_(shared_delta)
118
shard = tensor_model_parallel_reduce_scatter(delta, dim=-1)
119
shard = sconv(shard.contiguous(), positions)
120
full = tensor_model_parallel_all_gather(shard, dim=-1)
@@ -194,7 +203,7 @@ class InklingDecoderLayer(nn.Module):
203
self,
204
positions: torch.Tensor,
205
hidden_states: torch.Tensor,
197
- pending: tuple[torch.Tensor | None, InklingShortConv] | None = None,
206
+ pending: tuple[InklingDelta, InklingShortConv] | None = None,
207
defer_mlp_add: bool = False,
208
attn_in: torch.Tensor | None = None,
209
log_scaling: torch.Tensor | None = None,
@@ -217,7 +226,11 @@ class InklingDecoderLayer(nn.Module):
226
mlp_in, hidden_states = _sconv_add_norm(
227
attn_output, hidden_states, self.attn_sconv, self.mlp_norm, positions
228
)
220
- mlp_output = self.mlp(mlp_in)
229
+ mlp_output = (
230
+ self.mlp.forward_partials(mlp_in)
231
+ if isinstance(self.mlp, InklingMoE)
232
+ else self.mlp(mlp_in)
233
+ )
234
if defer_mlp_add:
235
# Caller folds mlp_output (pre-reduce, pre-sconv) into the next
236
# fused sconv+add+rmsnorm.
@@ -326,7 +339,7 @@ class InklingModel(nn.Module):
339
self.config.log_scaling_alpha,
340
)
341
329
- pending: tuple[torch.Tensor | None, InklingShortConv] | None = None
342
+ pending: tuple[InklingDelta, InklingShortConv] | None = None
343
for layer in self.layers[self.start_layer : self.end_layer]:
344
hidden_states, pending = layer(
345
positions,
vllm/models/inkling/nvidia/moe.py
+5
-1
@@ -501,7 +501,7 @@ class InklingMoE(nn.Module):
501
weights, ids = self.gate.select_experts(gating_output)
502
return weights[:, :topk].contiguous(), ids[:, :topk].contiguous()
503
504
- def forward(self, x: torch.Tensor) -> torch.Tensor | None:
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
@@ -527,6 +527,10 @@ class InklingMoE(nn.Module):
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 ----------------------------------------------------
vllm/models/inkling/nvidia/ops/lamport.py
+25
@@ -88,15 +88,18 @@ def _wait_pairs(ptr, offsets, mask):
88
@triton.jit
89
def _publish_input_kernel(
90
stage_ptr,
91
+ shared_ptr,
92
peer_ptrs,
93
peer_offset_u32,
94
stride_stage_t,
95
+ stride_shared_t,
96
C: tl.constexpr,
97
CS: tl.constexpr,
98
CS_P2: tl.constexpr,
99
SPLITS: tl.constexpr,
100
RANK: tl.constexpr,
101
WORLD: tl.constexpr,
102
+ HAS_SHARED: tl.constexpr,
103
USE_PDL: tl.constexpr,
104
launch_pdl: tl.constexpr,
105
):
@@ -120,6 +123,13 @@ def _publish_input_kernel(
123
mask=elem_mask,
124
other=0.0,
125
)
126
+ if HAS_SHARED:
127
+ shared = tl.load(
128
+ shared_ptr + token * stride_shared_t + owner * CS + split * CSS + elem,
129
+ mask=elem_mask,
130
+ other=0.0,
131
+ )
132
+ values = (values.to(tl.float32) + shared.to(tl.float32)).to(tl.bfloat16)
133
packed = _pack_bf16_pairs(values)
134
base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32))
135
dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair
@@ -595,6 +605,7 @@ class LamportRSConv:
605
off_s: int,
606
ws: int,
607
block_size: int,
608
+ shared_tensor: torch.Tensor | None = None,
609
) -> tuple[torch.Tensor | None, torch.Tensor]:
610
"""Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``."""
611
tokens, hidden_size = residual.shape
@@ -608,6 +619,13 @@ class LamportRSConv:
619
or input_tensor.stride(1) != 1
620
):
621
raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]")
622
+ if shared_tensor is not None and (
623
+ shared_tensor.shape != input_tensor.shape
624
+ or shared_tensor.dtype != input_tensor.dtype
625
+ or shared_tensor.device != input_tensor.device
626
+ or shared_tensor.stride(1) != 1
627
+ ):
628
+ raise ValueError("shared_tensor must match input_tensor")
629
shard_size = hidden_size // self.world_size
630
if conv_weight.shape != (shard_size, self.window_size):
631
raise ValueError(
@@ -648,15 +666,22 @@ class LamportRSConv:
666
gather_warps = 4 if tokens >= 256 else 8
667
_publish_input_kernel[phase_grid](
668
input_tensor,
669
+ shared_tensor if shared_tensor is not None else input_tensor,
670
self.input_peer_ptrs,
671
input_offset,
672
input_tensor.stride(0),
673
+ (
674
+ shared_tensor.stride(0)
675
+ if shared_tensor is not None
676
+ else input_tensor.stride(0)
677
+ ),
678
C=hidden_size,
679
CS=shard_size,
680
CS_P2=phase_tile_p2,
681
SPLITS=phase_splits,
682
RANK=self.rank,
683
WORLD=self.world_size,
684
+ HAS_SHARED=shared_tensor is not None,
685
USE_PDL=self.use_pdl,
686
launch_pdl=self.use_pdl,
687
num_warps=phase_warps,