From edbc4969a76b25c531f7b9bb16b79984e65ef023 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Wed, 5 Aug 2026 02:41:35 +0800 Subject: [PATCH] [Kernel][Inkling] Fuse shared-expert partial addition into the Lamport collective (#50697) Signed-off-by: Canlin Guo --- vllm/models/inkling/nvidia/model.py | 23 ++++++++++++++++----- vllm/models/inkling/nvidia/moe.py | 6 +++++- vllm/models/inkling/nvidia/ops/lamport.py | 25 +++++++++++++++++++++++ 3 files changed, 48 insertions(+), 6 deletions(-) diff --git a/vllm/models/inkling/nvidia/model.py b/vllm/models/inkling/nvidia/model.py index 04ba0ee085..a7df6243e7 100644 --- a/vllm/models/inkling/nvidia/model.py +++ b/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, diff --git a/vllm/models/inkling/nvidia/moe.py b/vllm/models/inkling/nvidia/moe.py index 36ffd07b92..a8ea968076 100644 --- a/vllm/models/inkling/nvidia/moe.py +++ b/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 ---------------------------------------------------- diff --git a/vllm/models/inkling/nvidia/ops/lamport.py b/vllm/models/inkling/nvidia/ops/lamport.py index c02290dc48..f863d92c38 100644 --- a/vllm/models/inkling/nvidia/ops/lamport.py +++ b/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,