[Kernel][Inkling] Fuse shared-expert partial addition into the Lamport collective (#50697)

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>

vllm/models/inkling/nvidia/model.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, TypeAlias
 
 import regex as re
 import torch
@@ -57,6 +57,8 @@ from .ops.norm import add_rmsnorm, embed_rmsnorm
 from .sconv_swa_attn import _ATTN, _MLP, InklingConvState, InklingSconvMetadata
 from .short_conv import InklingShortConv
 
+InklingDelta: TypeAlias = torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+
 
 def _layer_id(name: str) -> int | None:
     m = re.search(r"\.layers\.(\d+)\.", name)
@@ -64,7 +66,7 @@ def _layer_id(name: str) -> int | None:
 
 
 def _sconv_add_norm(
-    delta: torch.Tensor,
+    delta: InklingDelta,
     hidden: torch.Tensor,
     sconv: InklingShortConv,
     norm: InklingRMSNorm | None,
@@ -84,6 +86,10 @@ def _sconv_add_norm(
     off_s, ws = sconv.owner.stream_ranges[sconv.stream_idx]
     norm_w = norm.weight if norm is not None else None
     eps = norm.variance_epsilon if norm is not None else 0.0
+    if isinstance(delta, tuple):
+        delta, shared_delta = delta
+    else:
+        shared_delta = None
 
     mm = get_lamport_rs_conv(hidden.shape[-1], sconv.kernel_size)
     if mm is not None and mm.usable(delta.shape[0]) and m is not None:
@@ -103,9 +109,12 @@ def _sconv_add_norm(
             off_s,
             ws,
             sconv.owner.block_size,
+            shared_tensor=shared_delta,
         )
 
     # Fallback: NCCL RS -> shard sconv -> AG -> fused add(+rmsnorm).
+    if shared_delta is not None:
+        delta.add_(shared_delta)
     shard = tensor_model_parallel_reduce_scatter(delta, dim=-1)
     shard = sconv(shard.contiguous(), positions)
     full = tensor_model_parallel_all_gather(shard, dim=-1)
@@ -194,7 +203,7 @@ class InklingDecoderLayer(nn.Module):
         self,
         positions: torch.Tensor,
         hidden_states: torch.Tensor,
-        pending: tuple[torch.Tensor | None, InklingShortConv] | None = None,
+        pending: tuple[InklingDelta, InklingShortConv] | None = None,
         defer_mlp_add: bool = False,
         attn_in: torch.Tensor | None = None,
         log_scaling: torch.Tensor | None = None,
@@ -217,7 +226,11 @@ class InklingDecoderLayer(nn.Module):
         mlp_in, hidden_states = _sconv_add_norm(
             attn_output, hidden_states, self.attn_sconv, self.mlp_norm, positions
         )
-        mlp_output = self.mlp(mlp_in)
+        mlp_output = (
+            self.mlp.forward_partials(mlp_in)
+            if isinstance(self.mlp, InklingMoE)
+            else self.mlp(mlp_in)
+        )
         if defer_mlp_add:
             # Caller folds mlp_output (pre-reduce, pre-sconv) into the next
             # fused sconv+add+rmsnorm.
@@ -326,7 +339,7 @@ class InklingModel(nn.Module):
                 self.config.log_scaling_alpha,
             )
 
-        pending: tuple[torch.Tensor | None, InklingShortConv] | None = None
+        pending: tuple[InklingDelta, InklingShortConv] | None = None
         for layer in self.layers[self.start_layer : self.end_layer]:
             hidden_states, pending = layer(
                 positions,

vllm/models/inkling/nvidia/moe.py

@@ -501,7 +501,7 @@ class InklingMoE(nn.Module):
         weights, ids = self.gate.select_experts(gating_output)
         return weights[:, :topk].contiguous(), ids[:, :topk].contiguous()
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor | None:
+    def forward_partials(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         router_logits = self.gate.compute_logits(x)
         num_tokens = x.shape[0]
         # One gate select per layer: the routed slice is stashed for the
@@ -527,6 +527,10 @@ class InklingMoE(nn.Module):
         )
         self._routed_sel = None
 
+        return out, sink_out
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out, sink_out = self.forward_partials(x)
         return out.add_(sink_out)
 
     # -- weight loading ----------------------------------------------------

vllm/models/inkling/nvidia/ops/lamport.py

@@ -88,15 +88,18 @@ def _wait_pairs(ptr, offsets, mask):
 @triton.jit
 def _publish_input_kernel(
     stage_ptr,
+    shared_ptr,
     peer_ptrs,
     peer_offset_u32,
     stride_stage_t,
+    stride_shared_t,
     C: tl.constexpr,
     CS: tl.constexpr,
     CS_P2: tl.constexpr,
     SPLITS: tl.constexpr,
     RANK: tl.constexpr,
     WORLD: tl.constexpr,
+    HAS_SHARED: tl.constexpr,
     USE_PDL: tl.constexpr,
     launch_pdl: tl.constexpr,
 ):
@@ -120,6 +123,13 @@ def _publish_input_kernel(
             mask=elem_mask,
             other=0.0,
         )
+        if HAS_SHARED:
+            shared = tl.load(
+                shared_ptr + token * stride_shared_t + owner * CS + split * CSS + elem,
+                mask=elem_mask,
+                other=0.0,
+            )
+            values = (values.to(tl.float32) + shared.to(tl.float32)).to(tl.bfloat16)
         packed = _pack_bf16_pairs(values)
         base = tl.load(ptrs + owner).to(tl.pointer_type(tl.uint32))
         dst = (token * WORLD + RANK) * (CS // 2) + split * (CSS // 2) + pair
@@ -595,6 +605,7 @@ class LamportRSConv:
         off_s: int,
         ws: int,
         block_size: int,
+        shared_tensor: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor | None, torch.Tensor]:
         """Return ``(normed | None, new_residual)``, both shaped ``[T, 6144]``."""
         tokens, hidden_size = residual.shape
@@ -608,6 +619,13 @@ class LamportRSConv:
             or input_tensor.stride(1) != 1
         ):
             raise ValueError("input_tensor must be channel-contiguous bf16 [T, 6144]")
+        if shared_tensor is not None and (
+            shared_tensor.shape != input_tensor.shape
+            or shared_tensor.dtype != input_tensor.dtype
+            or shared_tensor.device != input_tensor.device
+            or shared_tensor.stride(1) != 1
+        ):
+            raise ValueError("shared_tensor must match input_tensor")
         shard_size = hidden_size // self.world_size
         if conv_weight.shape != (shard_size, self.window_size):
             raise ValueError(
@@ -648,15 +666,22 @@ class LamportRSConv:
         gather_warps = 4 if tokens >= 256 else 8
         _publish_input_kernel[phase_grid](
             input_tensor,
+            shared_tensor if shared_tensor is not None else input_tensor,
             self.input_peer_ptrs,
             input_offset,
             input_tensor.stride(0),
+            (
+                shared_tensor.stride(0)
+                if shared_tensor is not None
+                else input_tensor.stride(0)
+            ),
             C=hidden_size,
             CS=shard_size,
             CS_P2=phase_tile_p2,
             SPLITS=phase_splits,
             RANK=self.rank,
             WORLD=self.world_size,
+            HAS_SHARED=shared_tensor is not None,
             USE_PDL=self.use_pdl,
             launch_pdl=self.use_pdl,
             num_warps=phase_warps,