minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

qwen 3.8 rebase (#35758)

Browse Source 
Co-authored-by: cherichy <cherichy@outlook.com>
Co-authored-by: guangyunh-nv <guangyunh@nvidia.com>
Co-authored-by: jiahanc <jiahanc@nvidia.com>
Co-authored-by: jinyangyuan-nvidia <joyuan@nvidia.com>
Co-authored-by: Cheng Hang <chang@nvidia.com>
Co-authored-by: Yicheng Qiang <yqiang@nvidia.com>
Co-authored-by: Sam Li <lsam@nvidia.com>
Co-authored-by: Tom-Zheng <tizheng@nvidia.com>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: xiaoweiw-nv <xiaoweiw@nvidia.com>
Co-authored-by: Zheng Li <lizheng.cs@zju.edu.cn>
Co-authored-by: yizhang2077 <1109276519@qq.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Zijie Xia <zijie.xia@radixark.ai>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Qiaolin Yu
2026-08-28 20:41:34 -07:00
committed by GitHub
co-authored by cherichy guangyunh-nv jiahanc jinyangyuan-nvidia Cheng Hang Yicheng Qiang Sam Li Tom-Zheng Yangmin Li xiaoweiw-nv Zheng Li yizhang2077 Ke Bao Xinyuan Tong Yuhao Yang Zijie Xia github-actions[bot]
parent ca8cc101b8
commit 5f216fc33f
97 changed files with 14483 additions and 373 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/minimax/per_token_quant_ue8m0.cuh
+1 -2
View File
@@ -151,11 +151,10 @@ __global__ __launch_bounds__(1024, 2) void //
// Read this token's kTopK destinations once (fully unrolled).
const auto* src2dst_row = params.src2dst + static_cast<uint64_t>(token_id) * kTopK;
const auto* topk_ids_row = params.topk_ids + static_cast<uint64_t>(token_id) * kTopK;
int32_t dst_rows[kTopK];
#pragma unroll
for (uint32_t i = 0; i < kTopK; ++i) {
dst_rows[i] = (topk_ids_row[i] >= 0) ? src2dst_row[i] : -1;
dst_rows[i] = src2dst_row[i];
}
const uint32_t group_id = tid / kThreadsPerGroup;
python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py
+306
View File
@@ -330,6 +330,312 @@ def fused_qkvzba_split_reshape_cat_contiguous(
return mixed_qkv, z, b, a
# Fusion begins after the quantized GEMMs: qkvz=[q|k|v|z], ba=[b|a].
# This tail unpacks both projections and updates the causal Conv1D state.
@triton.jit
def _fused_qkvzba_causal_conv1d_update_contiguous_kernel(
mixed_qkv,
z,
b,
a,
mixed_qkvz,
mixed_ba,
conv_state,
conv_weight,
conv_bias,
conv_state_indices,
stride_qkvz_batch: tl.constexpr,
stride_qkvz_dim: tl.constexpr,
stride_ba_batch: tl.constexpr,
stride_ba_dim: tl.constexpr,
stride_state_batch: tl.constexpr,
stride_state_dim: tl.constexpr,
stride_state_pos: tl.constexpr,
stride_weight_dim: tl.constexpr,
stride_weight_width: tl.constexpr,
stride_state_indices: tl.constexpr,
QKV_DIM: tl.constexpr,
V_DIM: tl.constexpr,
NUM_V_HEADS: tl.constexpr,
NUM_STATE_SLOTS: tl.constexpr,
STATE_LEN: tl.constexpr,
KERNEL_WIDTH: tl.constexpr,
HAS_BIAS: tl.constexpr,
SILU_ACTIVATION: tl.constexpr,
PAD_SLOT_ID: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
batch_idx = tl.program_id(0)
dim_idx = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
qkv_mask = dim_idx < QKV_DIM
x = tl.load(
mixed_qkvz + batch_idx * stride_qkvz_batch + dim_idx * stride_qkvz_dim,
mask=qkv_mask,
other=0.0,
)
state_slot = tl.load(conv_state_indices + batch_idx * stride_state_indices).to(
tl.int64
)
# Treat every out-of-range index as padding so stale replay metadata cannot
# turn an indexed state update into an OOB access.
valid_slot = (
(state_slot != PAD_SLOT_ID) & (state_slot >= 0) & (state_slot < NUM_STATE_SLOTS)
)
state_base = (
conv_state + state_slot * stride_state_batch + dim_idx * stride_state_dim
)
acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
if HAS_BIAS:
acc += tl.load(conv_bias + dim_idx, mask=qkv_mask, other=0.0).to(tl.float32)
# Match the deployed direct-Triton update exactly. Its effective decode
# state length is width-1 even when the physical cache tensor is wider.
for pos in tl.static_range(KERNEL_WIDTH - 1):
state_value = tl.load(
state_base + pos * stride_state_pos,
mask=qkv_mask & valid_slot,
other=0.0,
)
weight_value = tl.load(
conv_weight + dim_idx * stride_weight_dim + pos * stride_weight_width,
mask=qkv_mask,
other=0.0,
)
# Do not force an FP32 multiply here. This expression deliberately
# retains the operand types/order of causal_conv1d_triton.py.
acc += state_value * weight_value
last_weight = tl.load(
conv_weight
+ dim_idx * stride_weight_dim
+ (KERNEL_WIDTH - 1) * stride_weight_width,
mask=qkv_mask,
other=0.0,
)
acc += x * last_weight
if SILU_ACTIVATION:
conv_out = acc / (1.0 + tl.exp(-acc))
else:
conv_out = acc
# The legacy kernel leaves padded rows' input unchanged.
conv_out = tl.where(valid_slot, conv_out, x)
tl.store(
mixed_qkv + batch_idx * QKV_DIM + dim_idx,
conv_out,
mask=qkv_mask,
)
# The direct-Triton wrapper sets effective state_len=width-1 for decode.
for pos in tl.static_range(KERNEL_WIDTH - 2):
next_value = tl.load(
state_base + (pos + 1) * stride_state_pos,
mask=qkv_mask & valid_slot,
other=0.0,
)
tl.store(
state_base + pos * stride_state_pos,
next_value,
mask=qkv_mask & valid_slot,
)
tl.store(
state_base + (KERNEL_WIDTH - 2) * stride_state_pos,
x,
mask=qkv_mask & valid_slot,
)
# The first feature lanes also materialize the smaller downstream tensors.
z_mask = dim_idx < V_DIM
z_value = tl.load(
mixed_qkvz
+ batch_idx * stride_qkvz_batch
+ (QKV_DIM + dim_idx) * stride_qkvz_dim,
mask=z_mask,
other=0.0,
)
tl.store(z + batch_idx * V_DIM + dim_idx, z_value, mask=z_mask)
gate_mask = dim_idx < NUM_V_HEADS
b_value = tl.load(
mixed_ba + batch_idx * stride_ba_batch + dim_idx * stride_ba_dim,
mask=gate_mask,
other=0.0,
)
a_value = tl.load(
mixed_ba
+ batch_idx * stride_ba_batch
+ (NUM_V_HEADS + dim_idx) * stride_ba_dim,
mask=gate_mask,
other=0.0,
)
tl.store(b + batch_idx * NUM_V_HEADS + dim_idx, b_value, mask=gate_mask)
tl.store(a + batch_idx * NUM_V_HEADS + dim_idx, a_value, mask=gate_mask)
def can_use_fused_qkvzba_causal_conv1d_update_contiguous(
mixed_qkvz: torch.Tensor,
mixed_ba: torch.Tensor,
conv_state: torch.Tensor,
conv_weight: torch.Tensor,
conv_bias: torch.Tensor | None,
conv_state_indices: torch.Tensor,
*,
qkv_dim: int,
v_dim: int,
num_v_heads: int,
activation: str | None,
) -> tuple[bool, str]:
"""Return an explicit eligibility decision for the decode fusion."""
tensors = (mixed_qkvz, mixed_ba, conv_state, conv_weight, conv_state_indices)
if not all(isinstance(tensor, torch.Tensor) for tensor in tensors):
return False, "all inputs must be torch.Tensor instances"
if not all(tensor.is_cuda for tensor in tensors):
return False, "CUDA tensors are required"
if mixed_qkvz.ndim != 2 or mixed_ba.ndim != 2:
return False, "projection outputs must be rank-2"
if conv_state.ndim != 3 or conv_weight.ndim != 2:
return False, "Conv1D state/weight ranks must be 3/2"
if conv_state_indices.ndim != 1:
return False, "conv_state_indices must be rank-1"
batch = mixed_qkvz.shape[0]
if mixed_ba.shape[0] != batch or conv_state_indices.shape[0] != batch:
return False, "batch dimensions must match"
if qkv_dim <= 0 or v_dim <= 0 or num_v_heads <= 0:
return False, "TP-local dimensions must be positive"
if mixed_qkvz.shape[1] != qkv_dim + v_dim:
return False, "qkvz layout is not contiguous [Q|K|V|Z]"
if mixed_ba.shape[1] != 2 * num_v_heads:
return False, "ba layout is not contiguous [B|A]"
if conv_state.shape[1] != qkv_dim or conv_weight.shape[0] != qkv_dim:
return False, "Conv1D feature dimension does not match packed QKV"
width = conv_weight.shape[1]
if width < 2 or width > 4:
return False, "only Conv1D widths 2 through 4 are supported"
if conv_state.shape[2] < width - 1:
return False, "Conv1D state is shorter than width - 1"
supported_dtypes = (torch.float16, torch.bfloat16, torch.float32)
if mixed_qkvz.dtype not in supported_dtypes:
return False, "QKVZ activation dtype must be FP16, BF16, or FP32"
if conv_state.dtype != mixed_qkvz.dtype or conv_weight.dtype != mixed_qkvz.dtype:
return False, "QKVZ, Conv1D state, and weight dtypes must match"
if mixed_ba.dtype not in supported_dtypes:
return False, "BA activation dtype must be FP16, BF16, or FP32"
if conv_bias is not None:
if (
not isinstance(conv_bias, torch.Tensor)
or not conv_bias.is_cuda
or conv_bias.ndim != 1
or conv_bias.shape[0] != qkv_dim
or conv_bias.dtype != mixed_qkvz.dtype
):
return False, "Conv1D bias contract is incompatible"
if activation not in (None, "silu", "swish"):
return False, "activation must be None, silu, or swish"
if mixed_qkvz.stride(1) != 1 or mixed_ba.stride(1) != 1:
return False, "projection feature dimensions must be contiguous"
if conv_weight.stride(1) != 1:
return False, "Conv1D weight width dimension must be contiguous"
if conv_state_indices.dtype not in (torch.int32, torch.int64):
return False, "conv_state_indices must be int32 or int64"
return True, "eligible"
def fused_qkvzba_causal_conv1d_update_contiguous(
mixed_qkvz: torch.Tensor,
mixed_ba: torch.Tensor,
conv_state: torch.Tensor,
conv_weight: torch.Tensor,
conv_bias: torch.Tensor | None,
conv_state_indices: torch.Tensor,
*,
qkv_dim: int,
v_dim: int,
num_v_heads: int,
head_v_dim: int,
activation: str | None,
pad_slot_id: int = -1,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Decode-only fused Qwen3.5 projection unpack and Conv1D state update."""
eligible, reason = can_use_fused_qkvzba_causal_conv1d_update_contiguous(
mixed_qkvz,
mixed_ba,
conv_state,
conv_weight,
conv_bias,
conv_state_indices,
qkv_dim=qkv_dim,
v_dim=v_dim,
num_v_heads=num_v_heads,
activation=activation,
)
if not eligible:
raise ValueError(f"Ineligible fused GDN decode projection/Conv1D: {reason}")
if v_dim != num_v_heads * head_v_dim:
raise ValueError(
"Ineligible fused GDN decode projection/Conv1D: "
"v_dim must equal num_v_heads * head_v_dim"
)
batch = mixed_qkvz.shape[0]
mixed_qkv = torch.empty(
(batch, qkv_dim), dtype=mixed_qkvz.dtype, device=mixed_qkvz.device
)
z = torch.empty(
(batch, num_v_heads, head_v_dim),
dtype=mixed_qkvz.dtype,
device=mixed_qkvz.device,
)
b = torch.empty(
(batch, num_v_heads),
dtype=mixed_ba.dtype,
device=mixed_ba.device,
)
a = torch.empty_like(b)
block_size = 256
grid = (batch, triton.cdiv(qkv_dim, block_size))
_fused_qkvzba_causal_conv1d_update_contiguous_kernel[grid](
mixed_qkv,
z,
b,
a,
mixed_qkvz,
mixed_ba,
conv_state,
conv_weight,
conv_bias,
conv_state_indices,
mixed_qkvz.stride(0),
mixed_qkvz.stride(1),
mixed_ba.stride(0),
mixed_ba.stride(1),
conv_state.stride(0),
conv_state.stride(1),
conv_state.stride(2),
conv_weight.stride(0),
conv_weight.stride(1),
conv_state_indices.stride(0),
QKV_DIM=qkv_dim,
V_DIM=v_dim,
NUM_V_HEADS=num_v_heads,
NUM_STATE_SLOTS=conv_state.shape[0],
STATE_LEN=conv_state.shape[2],
KERNEL_WIDTH=conv_weight.shape[1],
HAS_BIAS=conv_bias is not None,
SILU_ACTIVATION=activation in ("silu", "swish"),
PAD_SLOT_ID=pad_slot_id,
BLOCK_SIZE=block_size,
num_warps=8,
num_stages=2,
)
return mixed_qkv, z, b, a
@triton.jit
def fused_qkv_split_gdn_prefill_kernel(
q,
python/sglang/kernels/ops/communication/mnnvl_cutedsl/__init__.py
+52
View File
@@ -0,0 +1,52 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""MNNVL CuTe DSL AllReduce fusion backend internals."""
from importlib import import_module
from .config import (
KernelTarget,
MNNVLCuteDSLConfig,
MRangeDispatch,
ProtocolKind,
StaticProfile,
)
def __getattr__(name: str):
if name in {
"BT_ONLY_CONFIG",
"DEFAULT_CONFIG",
"HT_ONLY_CONFIG",
"LL_ONLY_CONFIG",
}:
presets = import_module(f"{__name__}.presets")
value = getattr(presets, name)
globals()[name] = value
return value
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
__all__ = [
"BT_ONLY_CONFIG",
"DEFAULT_CONFIG",
"HT_ONLY_CONFIG",
"LL_ONLY_CONFIG",
"KernelTarget",
"MNNVLCuteDSLConfig",
"MRangeDispatch",
"ProtocolKind",
"StaticProfile",
]
python/sglang/kernels/ops/communication/mnnvl_cutedsl/config.py
+231
View File
@@ -0,0 +1,231 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Configuration and token-count routing for the MNNVL CuTe DSL backend."""
from __future__ import annotations
from bisect import bisect_left
from dataclasses import dataclass
from enum import Enum
from typing import Generic, TypeVar
import torch
__all__ = [
"KernelTarget",
"MNNVLCuteDSLConfig",
"MRangeDispatch",
"ProtocolKind",
"StaticProfile",
]
class ProtocolKind(Enum):
LL = "ll"
BT = "bt"
HT = "ht"
PresetT = TypeVar("PresetT")
@dataclass(frozen=True, slots=True)
class KernelTarget(Generic[PresetT]):
protocol: ProtocolKind
preset: PresetT
TargetT = TypeVar("TargetT")
@dataclass(frozen=True, slots=True)
class MRangeDispatch(Generic[TargetT]):
"""Map contiguous positive token-count ranges to kernel targets."""
upper_bounds: tuple[int | None, ...]
targets: tuple[TargetT, ...]
def __post_init__(self) -> None:
if not self.upper_bounds:
raise ValueError("M range dispatch must contain at least one range")
if len(self.upper_bounds) != len(self.targets):
raise ValueError("M range upper bounds and targets must have equal length")
previous = 0
for index, upper_bound in enumerate(self.upper_bounds):
if upper_bound is None:
if index != len(self.upper_bounds) - 1:
raise ValueError("An unbounded M range must be the final range")
continue
if upper_bound <= previous:
raise ValueError("M range upper bounds must be strictly increasing")
previous = upper_bound
@property
def is_unbounded(self) -> bool:
return self.upper_bounds[-1] is None
@property
def finite_upper_bound(self) -> int | None:
return None if self.is_unbounded else self.upper_bounds[-1]
def supports(self, m: int) -> bool:
if m <= 0:
return False
upper_bound = self.finite_upper_bound
return upper_bound is None or m <= upper_bound
def select(self, m: int) -> TargetT:
if not self.supports(m):
raise ValueError(f"No kernel route supports M={m}")
finite_bounds = tuple(
upper_bound for upper_bound in self.upper_bounds if upper_bound is not None
)
index = bisect_left(finite_bounds, m)
return self.targets[index]
def referenced_protocols(self) -> frozenset[ProtocolKind]:
protocols = {
target.protocol
for target in self.targets
if isinstance(target, KernelTarget)
}
return frozenset(protocols)
def targets_for_capacity(self, capacity_m: int) -> tuple[TargetT, ...]:
if capacity_m <= 0:
return ()
selected = []
lower_bound = 1
for upper_bound, target in zip(self.upper_bounds, self.targets, strict=True):
if lower_bound > capacity_m:
break
selected.append(target)
if upper_bound is None:
break
lower_bound = upper_bound + 1
return tuple(selected)
def max_m_for_protocol(
self, protocol: ProtocolKind, *, capacity_m: int
) -> int | None:
lower_bound = 1
maximum = None
for upper_bound, target in zip(self.upper_bounds, self.targets, strict=True):
effective_upper_bound = capacity_m if upper_bound is None else upper_bound
if (
isinstance(target, KernelTarget)
and target.protocol is protocol
and lower_bound <= capacity_m
):
maximum = min(effective_upper_bound, capacity_m)
lower_bound = effective_upper_bound + 1
return maximum
@dataclass(frozen=True, slots=True)
class StaticProfile:
tp_size: int
hidden_size: int
top_k: int
dtype: torch.dtype
finalize_routes: MRangeDispatch[KernelTarget[object]]
all_reduce_routes: MRangeDispatch[KernelTarget[object]]
def __post_init__(self) -> None:
if self.hidden_size <= 0 or self.hidden_size % 8:
raise ValueError("hidden_size must be a positive multiple of 8")
def matches(
self,
*,
tp_size: int,
hidden_size: int,
top_k: int,
dtype: torch.dtype,
) -> bool:
return (
self.tp_size == tp_size
and self.hidden_size == hidden_size
and self.top_k == top_k
and self.dtype == dtype
)
def validate_capacity(self, capacity_m: int) -> None:
if capacity_m <= 0:
raise ValueError("capacity_m must be positive")
if not self.finalize_routes.supports(capacity_m):
raise ValueError(
"Finalize routes do not cover the requested workspace capacity"
)
if not self.all_reduce_routes.supports(capacity_m):
raise ValueError(
"AllReduce routes do not cover the requested workspace capacity"
)
@property
def referenced_protocols(self) -> frozenset[ProtocolKind]:
return (
self.finalize_routes.referenced_protocols()
| self.all_reduce_routes.referenced_protocols()
)
def protocol_capacity(
self, protocol: ProtocolKind, *, capacity_m: int
) -> int | None:
maxima = (
self.finalize_routes.max_m_for_protocol(protocol, capacity_m=capacity_m),
self.all_reduce_routes.max_m_for_protocol(protocol, capacity_m=capacity_m),
)
present = tuple(value for value in maxima if value is not None)
return max(present) if present else None
@dataclass(frozen=True, slots=True)
class MNNVLCuteDSLConfig:
"""Static profiles and routing policy for one backend configuration."""
profiles: tuple[StaticProfile, ...]
def __post_init__(self) -> None:
keys = [
(profile.tp_size, profile.hidden_size, profile.top_k, profile.dtype)
for profile in self.profiles
]
if not keys:
raise ValueError("A backend config must contain at least one profile")
if len(keys) != len(set(keys)):
raise ValueError("Backend config profiles must have unique static shapes")
def resolve(
self,
*,
tp_size: int,
hidden_size: int,
top_k: int,
dtype: torch.dtype,
capacity_m: int,
) -> StaticProfile:
for profile in self.profiles:
if profile.matches(
tp_size=tp_size,
hidden_size=hidden_size,
top_k=top_k,
dtype=dtype,
):
profile.validate_capacity(capacity_m)
return profile
raise ValueError("No MNNVL CuTe DSL profile supports this static shape")
python/sglang/kernels/ops/communication/mnnvl_cutedsl/cute_dsl_primitives.py
+876
View File
@@ -0,0 +1,876 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Small standalone CuTe DSL and PTX primitives shared by Kernel backends."""
import cutlass
import cutlass.cute as cute
from cutlass import BFloat16, Float32, Int32, Int64, Uint16, Uint32
from cutlass._mlir import ir
from cutlass._mlir.dialects import llvm, vector
from cutlass.cutlass_dsl import T, dsl_user_op
WARP_SIZE = 32
VEC_BF16 = 8
QUAD_BF16 = 4
NEGATIVE_ZERO_BF16_BITS = 0x8000
NEGATIVE_ZERO_BF16_PAIR = 0x80008000
# CUTLASS cute::TMA::CacheHintSm100::EVICT_FIRST policy descriptor.
L2_EVICT_FIRST = 0x12F0000000000000
@dsl_user_op
def load_global_u32x4(
pointer: cute.Pointer,
*,
volatile: cutlass.Constexpr[bool] = False,
loc=None,
ip=None,
):
address = pointer.toint(loc=loc, ip=ip)
if volatile:
opcode = "ld.volatile.global.v4.u32"
else:
opcode = "ld.global.v4.u32"
loaded = llvm.inline_asm(
llvm.StructType.get_literal([T.i32()] * 4),
[address.ir_value(loc=loc, ip=ip)],
f"{opcode} {{$0, $1, $2, $3}}, [$4];",
"=r,=r,=r,=r,l",
has_side_effects=volatile,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
packed = vector.from_elements(
ir.VectorType.get([4], T.i32(), loc=loc),
[
llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
for index in range(4)
],
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 4, Uint32)
@dsl_user_op
def load_global_u32x4_predicated(
pointer: cute.Pointer,
predicate: Int32,
*,
loc=None,
ip=None,
):
address = pointer.toint(loc=loc, ip=ip)
loaded = llvm.inline_asm(
llvm.StructType.get_literal([T.i32()] * 4),
[
address.ir_value(loc=loc, ip=ip),
Int32(predicate).ir_value(loc=loc, ip=ip),
],
(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.s32 p, $5, 0;\n\t"
"@!p mov.u32 $0, 0;\n\t"
"@!p mov.u32 $1, 0;\n\t"
"@!p mov.u32 $2, 0;\n\t"
"@!p mov.u32 $3, 0;\n\t"
"@p ld.global.v4.u32 {$0, $1, $2, $3}, [$4];\n\t"
"}"
),
"=r,=r,=r,=r,l,r",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
packed = vector.from_elements(
ir.VectorType.get([4], T.i32(), loc=loc),
[
llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
for index in range(4)
],
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 4, Uint32)
@dsl_user_op
def load_global_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32:
address = pointer.toint(loc=loc, ip=ip)
return Uint32(
llvm.inline_asm(
T.i32(),
[address.ir_value(loc=loc, ip=ip)],
"ld.global.u32 $0, [$1];",
"=r,l",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def load_global_u32_predicated(
pointer: cute.Pointer,
predicate: Int32,
*,
loc=None,
ip=None,
) -> Uint32:
address = pointer.toint(loc=loc, ip=ip)
return Uint32(
llvm.inline_asm(
T.i32(),
[
address.ir_value(loc=loc, ip=ip),
Int32(predicate).ir_value(loc=loc, ip=ip),
],
(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.s32 p, $2, 0;\n\t"
"@!p mov.u32 $0, 0;\n\t"
"@p ld.global.u32 $0, [$1];\n\t"
"}"
),
"=r,l,r",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def load_global_u32x2(pointer: cute.Pointer, *, loc=None, ip=None):
address = pointer.toint(loc=loc, ip=ip)
loaded = llvm.inline_asm(
llvm.StructType.get_literal([T.i32()] * 2),
[address.ir_value(loc=loc, ip=ip)],
"ld.global.v2.u32 {$0, $1}, [$2];",
"=r,=r,l",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
packed = vector.from_elements(
ir.VectorType.get([2], T.i32(), loc=loc),
[
llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
for index in range(2)
],
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 2, Uint32)
@dsl_user_op
def load_global_u32x2_predicated(
pointer: cute.Pointer,
predicate: Int32,
*,
loc=None,
ip=None,
):
address = pointer.toint(loc=loc, ip=ip)
loaded = llvm.inline_asm(
llvm.StructType.get_literal([T.i32()] * 2),
[
address.ir_value(loc=loc, ip=ip),
Int32(predicate).ir_value(loc=loc, ip=ip),
],
(
"{\n\t"
".reg .pred p;\n\t"
"setp.ne.s32 p, $3, 0;\n\t"
"@!p mov.u32 $0, 0;\n\t"
"@!p mov.u32 $1, 0;\n\t"
"@p ld.global.v2.u32 {$0, $1}, [$2];\n\t"
"}"
),
"=r,=r,l,r",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
packed = vector.from_elements(
ir.VectorType.get([2], T.i32(), loc=loc),
[
llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
for index in range(2)
],
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 2, Uint32)
@dsl_user_op
def store_global_u32x4(address: Int64, packed, *, loc=None, ip=None) -> None:
words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)]
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip), *words],
"st.global.v4.u32 [$0], {$1, $2, $3, $4};",
"l,r,r,r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def store_global_u32_address(
address: Int64,
value: Uint32,
*,
loc=None,
ip=None,
) -> None:
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
"st.global.u32 [$0], $1;",
"l,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def store_global_u32x2(address: Int64, packed, *, loc=None, ip=None) -> None:
words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(2)]
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip), *words],
"st.global.v2.u32 [$0], {$1, $2};",
"l,r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def store_global_u16_bits(
address: Int64,
value: Uint32,
*,
loc=None,
ip=None,
) -> None:
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
(
"{\n\t"
".reg .b16 bits;\n\t"
"cvt.u16.u32 bits, $1;\n\t"
"st.global.u16 [$0], bits;\n\t"
"}"
),
"l,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def store_lamport_sentinel_u32x4(
address: Int64,
*,
loc=None,
ip=None,
) -> None:
sentinel = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip)
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip), sentinel, sentinel, sentinel, sentinel],
"st.global.v4.u32 [$0], {$1, $2, $3, $4};",
"l,r,r,r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def load_global_bf16_as_f32(
address: Int64,
*,
loc=None,
ip=None,
) -> Float32:
return Float32(
llvm.inline_asm(
T.f32(),
[address.ir_value(loc=loc, ip=ip)],
(
"{\n\t"
".reg .b16 bits;\n\t"
"ld.global.b16 bits, [$1];\n\t"
"cvt.f32.bf16 $0, bits;\n\t"
"}"
),
"=f,l",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def load_global_bf16_as_f32_predicated(
address: Int64,
predicate: Int32,
*,
loc=None,
ip=None,
) -> Float32:
return Float32(
llvm.inline_asm(
T.f32(),
[
address.ir_value(loc=loc, ip=ip),
Int32(predicate).ir_value(loc=loc, ip=ip),
],
(
"{\n\t"
".reg .pred p;\n\t"
".reg .b16 bits;\n\t"
"setp.ne.s32 p, $2, 0;\n\t"
"@!p mov.b16 bits, 0;\n\t"
"@p ld.global.b16 bits, [$1];\n\t"
"cvt.f32.bf16 $0, bits;\n\t"
"}"
),
"=f,l,r",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def f32_to_bf16_bits(value: Float32, *, loc=None, ip=None) -> Uint32:
return Uint32(
llvm.inline_asm(
T.i32(),
[value.ir_value(loc=loc, ip=ip)],
(
"{\n\t"
".reg .b16 bits;\n\t"
"cvt.rn.bf16.f32 bits, $1;\n\t"
"cvt.u32.u16 $0, bits;\n\t"
"}"
),
"=r,f",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def shuffle_sync_idx_u32(
value: Uint32,
source_lane: Int32,
*,
loc=None,
ip=None,
) -> Uint32:
return Uint32(
llvm.inline_asm(
T.i32(),
[
value.ir_value(loc=loc, ip=ip),
source_lane.ir_value(loc=loc, ip=ip),
],
"shfl.sync.idx.b32 $0, $1, $2, 0x1f, 0xffffffff;",
"=r,r,r",
# Preserve full-warp execution across later divergent consumers.
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def load_volatile_u32(pointer: cute.Pointer, *, loc=None, ip=None) -> Uint32:
address = pointer.toint(loc=loc, ip=ip)
return Uint32(
llvm.inline_asm(
T.i32(),
[address.ir_value(loc=loc, ip=ip)],
"ld.volatile.global.u32 $0, [$1];",
"=r,l",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def store_global_u32(
pointer: cute.Pointer,
value: Uint32,
*,
loc=None,
ip=None,
) -> None:
address = pointer.toint(loc=loc, ip=ip)
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip), value.ir_value(loc=loc, ip=ip)],
"st.global.u32 [$0], $1;",
"l,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None):
values = llvm.bitcast(
ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc),
packed.ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
return cute.TensorSSA(values, VEC_BF16, BFloat16)
@dsl_user_op
def packed_u32_to_bf16x2(packed: Uint32, *, loc=None, ip=None):
values = llvm.bitcast(
ir.VectorType.get([2], BFloat16.mlir_type, loc=loc),
packed.ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
return cute.TensorSSA(values, 2, BFloat16)
@dsl_user_op
def packed_u32x2_to_bf16x4(packed, *, loc=None, ip=None):
values = llvm.bitcast(
ir.VectorType.get([QUAD_BF16], BFloat16.mlir_type, loc=loc),
packed.ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
return cute.TensorSSA(values, QUAD_BF16, BFloat16)
@dsl_user_op
def bf16x8_to_packed_u32x4(values, *, loc=None, ip=None):
packed = llvm.bitcast(
ir.VectorType.get([4], T.i32(), loc=loc),
values.ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 4, Uint32)
@dsl_user_op
def bf16x2_to_packed_u32(values, *, loc=None, ip=None) -> Uint32:
return Uint32(
llvm.bitcast(
T.i32(),
values.ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
)
@dsl_user_op
def bf16x4_to_packed_u32x2(values, *, loc=None, ip=None):
packed = llvm.bitcast(
ir.VectorType.get([2], T.i32(), loc=loc),
values.ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 2, Uint32)
@cute.jit
def sanitize_negative_zero_u32x4(packed):
sanitized = cute.make_rmem_tensor(cute.make_layout((4,)), Uint32)
for index in cutlass.range_constexpr(4):
sanitized[index] = sanitize_negative_zero_u32(packed[index])
return sanitized.load()
@cute.jit
def sanitize_negative_zero_u32(word: Uint32) -> Uint32:
low = Uint16(word & Uint32(0xFFFF))
high = Uint16(word >> Uint32(16))
if low == Uint16(NEGATIVE_ZERO_BF16_BITS):
word = word & Uint32(0xFFFF0000)
if high == Uint16(NEGATIVE_ZERO_BF16_BITS):
word = word & Uint32(0x0000FFFF)
return word
@cute.jit
def sanitize_negative_zero_u32x2(packed):
sanitized = cute.make_rmem_tensor(cute.make_layout((2,)), Uint32)
for index in cutlass.range_constexpr(2):
sanitized[index] = sanitize_negative_zero_u32(packed[index])
return sanitized.load()
@cute.jit
def fragment_has_negative_zero(packed):
dirty = False
for index in cutlass.range_constexpr(4):
word = packed[index]
dirty = (
dirty
| (Uint16(word & Uint32(0xFFFF)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
| (Uint16(word >> Uint32(16)) == Uint16(NEGATIVE_ZERO_BF16_BITS))
)
return dirty
@dsl_user_op
def map_shared_to_peer(
smem_pointer: cute.Pointer,
peer_rank: Int32,
*,
loc=None,
ip=None,
) -> Int32:
address = smem_pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
return Int32(
llvm.inline_asm(
T.i32(),
[address, peer_rank.ir_value(loc=loc, ip=ip)],
"mapa.shared::cluster.u32 $0, $1, $2;",
"=r,r,r",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def store_shared_cluster_f32(
remote_address: Int32,
value: Float32,
*,
loc=None,
ip=None,
) -> None:
llvm.inline_asm(
None,
[
remote_address.ir_value(loc=loc, ip=ip),
value.ir_value(loc=loc, ip=ip),
],
"st.shared::cluster.f32 [$0], $1;",
"r,f",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def load_shared_u32x4(pointer: cute.Pointer, *, loc=None, ip=None):
address = pointer.toint(loc=loc, ip=ip)
# Prevent motion across the named-barrier pipeline protocol.
loaded = llvm.inline_asm(
llvm.StructType.get_literal([T.i32()] * 4),
[Int32(address).ir_value(loc=loc, ip=ip)],
"ld.shared.v4.u32 {$0, $1, $2, $3}, [$4];",
"=r,=r,=r,=r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
packed = vector.from_elements(
ir.VectorType.get([4], T.i32(), loc=loc),
[
llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
for index in range(4)
],
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 4, Uint32)
@dsl_user_op
def store_shared_u32x4(
pointer: cute.Pointer,
packed,
*,
loc=None,
ip=None,
) -> None:
address = pointer.toint(loc=loc, ip=ip)
words = [packed[index].ir_value(loc=loc, ip=ip) for index in range(4)]
llvm.inline_asm(
None,
[Int32(address).ir_value(loc=loc, ip=ip), *words],
"st.shared.v4.u32 [$0], {$1, $2, $3, $4};",
"r,r,r,r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def load_global_u32x4_address(
address: Int64,
*,
volatile: cutlass.Constexpr[bool] = False,
loc=None,
ip=None,
):
opcode = "ld.volatile.global.v4.u32" if volatile else "ld.global.v4.u32"
loaded = llvm.inline_asm(
llvm.StructType.get_literal([T.i32()] * 4),
[address.ir_value(loc=loc, ip=ip)],
f"{opcode} {{$0, $1, $2, $3}}, [$4];",
"=r,=r,=r,=r,l",
has_side_effects=volatile,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
packed = vector.from_elements(
ir.VectorType.get([4], T.i32(), loc=loc),
[
llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
for index in range(4)
],
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 4, Uint32)
@dsl_user_op
def packed_negative_zero_bf16x8(*, loc=None, ip=None):
word = Uint32(NEGATIVE_ZERO_BF16_PAIR).ir_value(loc=loc, ip=ip)
packed = vector.from_elements(
ir.VectorType.get([4], T.i32(), loc=loc),
[word, word, word, word],
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 4, Uint32)
@dsl_user_op
def cpasync_bulk_g2s(
gmem_ptr: cute.Pointer,
smem_ptr: cute.Pointer,
barrier_ptr: cute.Pointer,
size_bytes: Int32,
*,
loc=None,
ip=None,
) -> None:
operands = [
gmem_ptr.toint(loc=loc, ip=ip).ir_value(),
smem_ptr.toint(loc=loc, ip=ip).ir_value(),
barrier_ptr.toint(loc=loc, ip=ip).ir_value(),
size_bytes.ir_value(loc=loc, ip=ip),
Int64(L2_EVICT_FIRST).ir_value(),
]
llvm.inline_asm(
None,
operands,
(
"cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes"
".L2::cache_hint [$1], [$0], $3, [$2], $4;"
),
"l,r,r,r,l",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def fence_proxy_async_shared_cta(*, loc=None, ip=None) -> None:
llvm.inline_asm(
None,
[],
"fence.proxy.async.shared::cta;",
"",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def remote_release_add1_u32(address: Int64, *, loc=None, ip=None) -> None:
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip)],
"red.release.sys.global.add.u32 [$0], 1;",
"l",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def ldmc_bf16x8(address: Int64, *, loc=None, ip=None):
loaded = llvm.inline_asm(
llvm.StructType.get_literal([T.i32()] * 4),
[address.ir_value(loc=loc, ip=ip)],
"multimem.ld_reduce.relaxed.sys.global.add.acc::f32.v4.bf16x2 {$0, $1, $2, $3}, [$4];",
"=r,=r,=r,=r,l",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
packed = vector.from_elements(
ir.VectorType.get([4], T.i32(), loc=loc),
[
llvm.extractvalue(T.i32(), loaded, [index], loc=loc, ip=ip)
for index in range(4)
],
loc=loc,
ip=ip,
)
return cute.TensorSSA(packed, 4, Uint32)
@dsl_user_op
def stmc_bf16x2(
address: Int64,
packed: Uint32,
*,
loc=None,
ip=None,
) -> None:
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip), packed.ir_value(loc=loc, ip=ip)],
"multimem.st.relaxed.sys.global.bf16x2 [$0], $1;",
"l,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def stmc_bf16x4(address: Int64, values, *, loc=None, ip=None) -> None:
words = [values[index].ir_value(loc=loc, ip=ip) for index in range(2)]
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip), *words],
"multimem.st.relaxed.sys.global.v2.bf16x2 [$0], {$1, $2};",
"l,r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def stmc_bf16x8(address: Int64, values, *, loc=None, ip=None) -> None:
words = [values[index].ir_value(loc=loc, ip=ip) for index in range(4)]
llvm.inline_asm(
None,
[address.ir_value(loc=loc, ip=ip), *words],
"multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};",
"l,r,r,r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
python/sglang/kernels/ops/communication/mnnvl_cutedsl/kernel_bt/__init__.py
+43
View File
@@ -0,0 +1,43 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Balanced MNNVL protocol."""
from .protocol import (
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0,
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1,
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0,
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1,
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0,
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1,
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0,
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1,
BTAllReduceTuning,
BTCollectiveTuning,
BTFinalizeTuning,
)
__all__ = [
"BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0",
"BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1",
"BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0",
"BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1",
"BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0",
"BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1",
"BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0",
"BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1",
"BTAllReduceTuning",
"BTCollectiveTuning",
"BTFinalizeTuning",
]
python/sglang/kernels/ops/communication/mnnvl_cutedsl/kernel_bt/device_kernels.py
+1285
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/ops/communication/mnnvl_cutedsl/kernel_bt/protocol.py
+511
View File
@@ -0,0 +1,511 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Balanced MNNVL protocol and its two operation paths."""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Any, TypedDict, cast
import cutlass.cute as cute
import torch
import torch.distributed as dist
from cutlass import BFloat16, Int32, Int64
from cutlass.cute.runtime import make_fake_compact_tensor
from ..cute_dsl_primitives import VEC_BF16
from ..runtime import (
current_cu_stream,
make_fake_dynamic_compact_tensor,
to_cute,
to_cute_dynamic,
)
from ..symmetric_buffer import SymmetricBuffer
from .device_kernels import (
LAMPORT_GENERATIONS,
_MaterializeRMSNormDeviceKernel,
_NarrowVectorFinalizeUnicastDeviceKernel,
_OwnerReduceMulticastDeviceKernel,
_ScalarFinalizeUnicastDeviceKernel,
_SharedOnlyPublishDeviceKernel,
_VectorFinalizeUnicastDeviceKernel,
)
@dataclass(frozen=True, slots=True)
class BTCollectiveTuning:
reduction_threads: int = 128
rms_threads: int = 1024
enable_pdl: bool = True
@dataclass(frozen=True, slots=True)
class BTFinalizeTuning:
elements_per_thread: int = VEC_BF16
threads: int = 128
prefetch_group: int = 1
load_shared_expert_before_pdl: bool = False
collective: BTCollectiveTuning = BTCollectiveTuning()
@dataclass(frozen=True, slots=True)
class BTAllReduceTuning:
publish_threads: int = 128
publish_vectors_per_thread: int = 1
collective: BTCollectiveTuning = BTCollectiveTuning(reduction_threads=32)
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0 = BTFinalizeTuning(
elements_per_thread=2, threads=256
)
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1 = BTFinalizeTuning()
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0 = BTFinalizeTuning(
elements_per_thread=2, threads=256
)
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1 = BTFinalizeTuning()
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0 = BTAllReduceTuning()
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1 = BTAllReduceTuning(
collective=BTCollectiveTuning(reduction_threads=320)
)
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0 = BTAllReduceTuning()
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1 = BTAllReduceTuning(
collective=BTCollectiveTuning(reduction_threads=320)
)
@dataclass(slots=True)
class BTProtocolState:
contribution_mailbox: SymmetricBuffer
prenorm_mailbox: SymmetricBuffer
stage_state: torch.Tensor
@dataclass(frozen=True, slots=True)
class _CompiledTail:
reduce: Any
rms_norm: Any
@dataclass(frozen=True, slots=True)
class _CompiledFinalize:
publish: Any
tail: _CompiledTail
@dataclass(frozen=True, slots=True)
class _CompiledAllReduce:
publish: Any
tail: _CompiledTail
class _PathKwargs(TypedDict):
hidden_size: int
top_k: int
capacity_m: int
write_residual_output: bool
class _BTPath:
def __init__(
self,
*,
hidden_size: int,
top_k: int,
capacity_m: int,
write_residual_output: bool,
) -> None:
self.hidden_size = hidden_size
self.top_k = top_k
self.capacity_m = capacity_m
self.write_residual_output = write_residual_output
def _outputs(
self,
m: int,
norm_output: torch.Tensor | None,
residual_output: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
shape = (m, self.hidden_size)
device = torch.device("cuda", torch.cuda.current_device())
if norm_output is None:
norm_output = torch.empty(shape, dtype=torch.bfloat16, device=device)
if self.write_residual_output and residual_output is None:
residual_output = torch.empty(shape, dtype=torch.bfloat16, device=device)
return norm_output, residual_output
def _validate_state(self, state: BTProtocolState, m: int) -> None:
if not 1 <= m <= self.capacity_m:
raise ValueError(f"m must be in [1, {self.capacity_m}]")
if state.contribution_mailbox.peer_addresses is None:
raise ValueError("BT contribution mailbox requires peer addresses")
address = state.prenorm_mailbox.multicast_address
if address is None or address % 16:
raise ValueError(
"BT prenorm mailbox requires a 16-byte-aligned multicast address"
)
def _launch_tail(
self,
tail: _CompiledTail,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
state: BTProtocolState,
norm_output: torch.Tensor,
residual_output: torch.Tensor | None,
m: int,
) -> None:
residual_arg = residual_source if residual_source is not None else norm_output
residual_output_arg = (
residual_output if residual_output is not None else norm_output
)
stream = current_cu_stream()
tail.reduce(
to_cute(state.contribution_mailbox.tensor.flatten(), 16),
to_cute_dynamic(residual_arg.flatten(), 16, divisibility=self.hidden_size),
to_cute(state.stage_state, 4),
Int64(cast(int, state.prenorm_mailbox.multicast_address)),
Int32(m),
stream,
)
tail.rms_norm(
to_cute(state.prenorm_mailbox.tensor.flatten(), 16),
to_cute_dynamic(
residual_output_arg.flatten(),
16,
divisibility=self.hidden_size,
),
to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size),
to_cute(gamma, 16),
to_cute(state.stage_state, 4),
Int32(m),
stream,
)
class FinalizeAllReduceRMSNormBTKernel(_BTPath):
def __init__(self, *, compiled: _CompiledFinalize, **kwargs) -> None:
super().__init__(**kwargs)
self._compiled = compiled
def __call__(
self,
routed_output: torch.Tensor,
expert_weights: torch.Tensor,
permuted_indices: torch.Tensor,
shared_output: torch.Tensor | None,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
m: int,
*,
state: BTProtocolState,
norm_output: torch.Tensor | None = None,
residual_output: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
self._validate_state(state, m)
norm_output, residual_output = self._outputs(m, norm_output, residual_output)
shared_arg = shared_output if shared_output is not None else norm_output
peers = cast(torch.Tensor, state.contribution_mailbox.peer_addresses)
self._compiled.publish(
to_cute_dynamic(routed_output.flatten(), 16, divisibility=self.hidden_size),
to_cute_dynamic(expert_weights.flatten(), 2, divisibility=self.top_k),
to_cute_dynamic(permuted_indices.flatten(), 4, divisibility=self.top_k),
to_cute_dynamic(shared_arg.flatten(), 16, divisibility=self.hidden_size),
to_cute(state.stage_state, 4),
to_cute(peers, 8),
Int32(m),
current_cu_stream(),
)
self._launch_tail(
self._compiled.tail,
residual_source,
gamma,
state,
norm_output,
residual_output,
m,
)
return norm_output, residual_output
class AllReduceRMSNormBTKernel(_BTPath):
def __init__(self, *, compiled: _CompiledAllReduce, **kwargs) -> None:
super().__init__(**kwargs)
self._compiled = compiled
def __call__(
self,
local_contribution: torch.Tensor,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
m: int,
*,
state: BTProtocolState,
norm_output: torch.Tensor | None = None,
residual_output: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
self._validate_state(state, m)
norm_output, residual_output = self._outputs(m, norm_output, residual_output)
peers = cast(torch.Tensor, state.contribution_mailbox.peer_addresses)
self._compiled.publish(
to_cute_dynamic(
local_contribution.flatten(),
16,
divisibility=self.hidden_size,
),
to_cute(state.stage_state, 4),
to_cute(peers, 8),
Int32(m),
current_cu_stream(),
)
self._launch_tail(
self._compiled.tail,
residual_source,
gamma,
state,
norm_output,
residual_output,
m,
)
return norm_output, residual_output
class BTProtocol:
"""Own BT State and protocol-local compiled variants for both paths."""
def __init__(
self,
hidden_size: int,
top_k: int,
tp_size: int,
rank: int,
capacity_m: int,
rms_epsilon: float,
routed_scaling_factor: float,
weight_bias: float,
*,
include_shared_expert: bool,
add_residual: bool,
write_residual_output: bool,
finalize_tunings: tuple[BTFinalizeTuning, ...],
all_reduce_tunings: tuple[BTAllReduceTuning, ...],
group: dist.ProcessGroup,
) -> None:
self.hidden_size = hidden_size
self.top_k = top_k
self.tp_size = tp_size
self.rank = rank
self.capacity_m = capacity_m
self.local_capacity = math.ceil(capacity_m / tp_size)
self.rms_epsilon = rms_epsilon
self.routed_scaling_factor = routed_scaling_factor
self.weight_bias = weight_bias
self.include_shared_expert = include_shared_expert
self.add_residual = add_residual
self.write_residual_output = write_residual_output
tail_cache = {
tuning: self._compile_tail(tuning)
for tuning in {
*(item.collective for item in finalize_tunings),
*(item.collective for item in all_reduce_tunings),
}
}
self.finalize_kernels = {
tuning: FinalizeAllReduceRMSNormBTKernel(
compiled=_CompiledFinalize(
publish=self._compile_finalize(tuning),
tail=tail_cache[tuning.collective],
),
**self._path_kwargs(),
)
for tuning in dict.fromkeys(finalize_tunings)
}
self.all_reduce_kernels = {
tuning: AllReduceRMSNormBTKernel(
compiled=_CompiledAllReduce(
publish=self._compile_all_reduce_publish(tuning),
tail=tail_cache[tuning.collective],
),
**self._path_kwargs(),
)
for tuning in dict.fromkeys(all_reduce_tunings)
}
self.state = self._create_state(group)
def _path_kwargs(self) -> _PathKwargs:
return {
"hidden_size": self.hidden_size,
"top_k": self.top_k,
"capacity_m": self.capacity_m,
"write_residual_output": self.write_residual_output,
}
def _compile_finalize(self, tuning: BTFinalizeTuning):
if tuning.elements_per_thread not in (1, 2, 4, VEC_BF16):
raise ValueError("BT finalize elements_per_thread must be 1, 2, 4, or 8")
kwargs: dict[str, Any] = {
"hidden_size": self.hidden_size,
"top_k": self.top_k,
"tp_size": self.tp_size,
"rank": self.rank,
"local_capacity": self.local_capacity,
"threads": tuning.threads,
"routed_scaling_factor": self.routed_scaling_factor,
"include_shared_expert": self.include_shared_expert,
"load_shared_expert_before_pdl": tuning.load_shared_expert_before_pdl,
"enable_pdl": tuning.collective.enable_pdl,
"prefetch_group": tuning.prefetch_group,
}
device_kernel: Any
if tuning.elements_per_thread == 1:
device_kernel = _ScalarFinalizeUnicastDeviceKernel(**kwargs)
elif tuning.elements_per_thread == VEC_BF16:
device_kernel = _VectorFinalizeUnicastDeviceKernel(**kwargs)
else:
device_kernel = _NarrowVectorFinalizeUnicastDeviceKernel(
**kwargs, elements_per_thread=tuning.elements_per_thread
)
return cute.compile(
device_kernel,
*self._publish_compile_args(include_routed=True),
)
def _compile_all_reduce_publish(self, tuning: BTAllReduceTuning):
device_kernel = _SharedOnlyPublishDeviceKernel(
hidden_size=self.hidden_size,
tp_size=self.tp_size,
rank=self.rank,
local_capacity=self.local_capacity,
threads=tuning.publish_threads,
vectors_per_thread=tuning.publish_vectors_per_thread,
enable_pdl=tuning.collective.enable_pdl,
)
return cute.compile(
device_kernel,
*self._publish_compile_args(include_routed=False),
)
def _publish_compile_args(self, *, include_routed: bool) -> tuple:
activation = make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
)
common = (
make_fake_compact_tensor(Int32, (2,), assumed_align=4),
make_fake_compact_tensor(Int64, (self.tp_size,), assumed_align=8),
Int32(self.capacity_m),
current_cu_stream(),
)
if not include_routed:
return (activation, *common)
return (
activation,
make_fake_dynamic_compact_tensor(
BFloat16, alignment=2, divisibility=self.top_k
),
make_fake_dynamic_compact_tensor(
Int32, alignment=4, divisibility=self.top_k
),
activation,
*common,
)
def _compile_tail(self, tuning: BTCollectiveTuning) -> _CompiledTail:
reduce_kernel = _OwnerReduceMulticastDeviceKernel(
hidden_size=self.hidden_size,
tp_size=self.tp_size,
rank=self.rank,
capacity_m=self.capacity_m,
local_capacity=self.local_capacity,
threads=tuning.reduction_threads,
add_residual=self.add_residual,
enable_pdl=tuning.enable_pdl,
)
reduce_elements = (
LAMPORT_GENERATIONS * self.tp_size * self.local_capacity * self.hidden_size
)
reduce = cute.compile(
reduce_kernel,
make_fake_compact_tensor(BFloat16, (reduce_elements,), assumed_align=16),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_compact_tensor(Int32, (2,), assumed_align=4),
Int64(0),
Int32(self.capacity_m),
current_cu_stream(),
)
rms_kernel = _MaterializeRMSNormDeviceKernel(
hidden_size=self.hidden_size,
capacity_m=self.capacity_m,
threads=tuning.rms_threads,
rms_epsilon=self.rms_epsilon,
weight_bias=self.weight_bias,
write_residual_output=self.write_residual_output,
enable_pdl=tuning.enable_pdl,
)
prenorm_elements = LAMPORT_GENERATIONS * self.capacity_m * self.hidden_size
rms_norm = cute.compile(
rms_kernel,
make_fake_compact_tensor(BFloat16, (prenorm_elements,), assumed_align=16),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_compact_tensor(BFloat16, (self.hidden_size,), assumed_align=16),
make_fake_compact_tensor(Int32, (2,), assumed_align=4),
Int32(self.capacity_m),
current_cu_stream(),
)
return _CompiledTail(reduce=reduce, rms_norm=rms_norm)
def _create_state(self, group: dist.ProcessGroup) -> BTProtocolState:
if dist.get_world_size(group) != self.tp_size:
raise ValueError("ProcessGroup size does not match tp_size")
if dist.get_rank(group) != self.rank:
raise ValueError("ProcessGroup rank does not match rank")
device = torch.device("cuda", torch.cuda.current_device())
contribution = SymmetricBuffer.allocate(
(
LAMPORT_GENERATIONS,
self.tp_size,
self.local_capacity,
self.hidden_size,
),
torch.bfloat16,
device,
group,
materialize_peer_addresses=True,
)
contribution.tensor.view(torch.int16).fill_(-32768)
prenorm = SymmetricBuffer.allocate(
(
LAMPORT_GENERATIONS,
self.capacity_m,
self.hidden_size,
),
torch.bfloat16,
device,
group,
require_multicast=True,
)
prenorm.tensor.view(torch.int16).fill_(-32768)
return BTProtocolState(
contribution_mailbox=contribution,
prenorm_mailbox=prenorm,
stage_state=torch.zeros((2,), dtype=torch.int32, device=device),
)
python/sglang/kernels/ops/communication/mnnvl_cutedsl/kernel_ht/__init__.py
+37
View File
@@ -0,0 +1,37 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""High-throughput MNNVL protocol."""
from .protocol import (
HT_ALL_REDUCE_GB300_TP8_H8192,
HT_ALL_REDUCE_GB300_TP16_H8192,
HT_FINALIZE_GB300_TP8_H8192_K10,
HT_FINALIZE_GB300_TP8_H8192_K10_M_GE_2049,
HT_FINALIZE_GB300_TP8_H8192_K10_M_LE_2048,
HT_FINALIZE_GB300_TP16_H8192_K10,
HTAllReduceTuning,
HTFinalizeTuning,
)
__all__ = [
"HT_ALL_REDUCE_GB300_TP8_H8192",
"HT_ALL_REDUCE_GB300_TP16_H8192",
"HT_FINALIZE_GB300_TP8_H8192_K10",
"HT_FINALIZE_GB300_TP8_H8192_K10_M_GE_2049",
"HT_FINALIZE_GB300_TP8_H8192_K10_M_LE_2048",
"HT_FINALIZE_GB300_TP16_H8192_K10",
"HTAllReduceTuning",
"HTFinalizeTuning",
]
python/sglang/kernels/ops/communication/mnnvl_cutedsl/kernel_ht/device_kernel.py
+978
View File
@@ -0,0 +1,978 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Persistent BF16 MoE finalize, TP reduction, and RMSNorm for SM100."""
from __future__ import annotations
import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import cutlass.pipeline as pipeline
import cutlass.utils as utils
from cutlass import BFloat16, Float32, Int32, Int64, Uint32
from ..cute_dsl_primitives import (
VEC_BF16,
WARP_SIZE,
bf16x8_to_packed_u32x4,
cpasync_bulk_g2s,
fence_proxy_async_shared_cta,
fragment_has_negative_zero,
ldmc_bf16x8,
load_global_bf16_as_f32,
load_global_u32x4_address,
load_shared_u32x4,
packed_negative_zero_bf16x8,
packed_u32x4_to_bf16x8,
remote_release_add1_u32,
sanitize_negative_zero_u32x4,
stmc_bf16x8,
store_global_u32x4,
store_shared_u32x4,
)
SMEM_ALIGNMENT = 1024
class _MoeFinalizeAllReduceRMSNormHTDeviceKernel:
def __init__(
self,
*,
hidden: int,
top_k: int,
tp: int,
rank: int,
active_ctas: int,
stages: int,
consumer_threads: int,
vectors_per_thread: int,
reduction_warps: int,
reduction_cta_groups: int | None,
rms_token_groups: int,
rms_pipeline_stages: int,
rms_shard_major: bool,
rms_epsilon: float,
routed_scaling_factor: float,
weight_bias: float,
include_shared_expert: bool,
add_residual: bool,
write_residual_output: bool,
enable_pdl: bool,
) -> None:
if tp not in (2, 4, 8, 16):
raise ValueError("tp must be 2, 4, 8, or 16")
if rank < 0 or rank >= tp:
raise ValueError("rank must be in [0, tp)")
if hidden <= 0 or hidden % VEC_BF16:
raise ValueError("hidden must be a positive multiple of 8")
if top_k < 0:
raise ValueError("top_k must be nonnegative")
if active_ctas <= 0 or active_ctas % tp:
raise ValueError("active_ctas must be positive and divisible by tp")
if stages < 2:
raise ValueError("stages must be at least 2")
if consumer_threads <= 0 or consumer_threads % WARP_SIZE:
raise ValueError("consumer_threads must be a positive warp multiple")
if vectors_per_thread <= 0:
raise ValueError("vectors_per_thread must be positive")
if reduction_warps not in (1, 2, 4, 8):
raise ValueError("reduction_warps must be 1, 2, 4, or 8")
if rms_token_groups not in (1, 2, 4):
raise ValueError("rms_token_groups must be 1, 2, or 4")
if consumer_threads % rms_token_groups:
raise ValueError("consumer threads must divide across RMS token groups")
if rms_pipeline_stages not in (1, 2, 3):
raise ValueError("rms_pipeline_stages must be 1, 2, or 3")
block_threads = consumer_threads + (2 + reduction_warps) * WARP_SIZE
if block_threads > 1024:
raise ValueError("warp roles exceed the CUDA block limit")
shard_elements = consumer_threads * VEC_BF16 * vectors_per_thread
if hidden <= 0 or hidden % shard_elements:
raise ValueError(f"hidden must be divisible by {shard_elements}")
cta_groups = active_ctas // tp
if reduction_cta_groups is None:
reduction_cta_groups = active_ctas // tp
if reduction_cta_groups <= 0 or reduction_cta_groups * tp > active_ctas:
raise ValueError("reduction CTA groups and shards must fit the grid")
contributions = top_k + int(include_shared_expert)
if contributions <= 0:
raise ValueError("at least one local contribution is required")
self.hidden = hidden
self.top_k = top_k
self.tp = tp
self.rank = rank
self.active_ctas = active_ctas
self.stages = stages
self.vectors_per_thread = vectors_per_thread
self.consumer_threads = consumer_threads
self.reduction_warps = reduction_warps
self.reduction_cta_groups = reduction_cta_groups
self.reduction_ctas = reduction_cta_groups * tp
self.rms_token_groups = rms_token_groups
self.rms_pipeline_stages = rms_pipeline_stages
self.rms_shard_major = rms_shard_major
self.rms_epsilon = rms_epsilon
self.routed_scaling_factor = routed_scaling_factor
self.weight_bias = weight_bias
self.include_shared_expert = include_shared_expert
self.add_residual = add_residual
self.write_residual_output = write_residual_output
self.enable_pdl = enable_pdl
self.metadata_chunks = (top_k + WARP_SIZE - 1) // WARP_SIZE
self.metadata_slots = max(top_k, 1)
self.consumer_warps = consumer_threads // WARP_SIZE
self.rms_threads_per_token = consumer_threads // rms_token_groups
self.rms_warps_per_token = self.rms_threads_per_token // WARP_SIZE
self.rms_stage_slots = rms_token_groups * rms_pipeline_stages
self.rms_warp_sum_slots = self.rms_stage_slots * self.rms_warps_per_token
self.publisher_warp = 1 + self.consumer_warps
self.reduction_warp_begin = self.publisher_warp + 1
self.reduction_threads = reduction_warps * WARP_SIZE
self.block_threads = block_threads
self.shard_elements = shard_elements
self.shard_bytes = shard_elements * 2
self.hidden_shards = hidden // shard_elements
self.contributions = contributions
if (
rms_pipeline_stages > 1
and self.rms_stage_slots * hidden > self.shard_elements * stages
):
raise ValueError("finalize stage storage cannot hold the RMS pipeline")
self.cta_groups = cta_groups
self.packs_per_token = hidden // VEC_BF16
if self.packs_per_token % tp:
raise ValueError("hidden vector count must be divisible by tp")
if self.packs_per_token % consumer_threads:
raise ValueError("token vectors must divide evenly across consumers")
self.clear_vectors_per_thread = self.packs_per_token // consumer_threads
self.copy_threads = self.rms_threads_per_token
if self.packs_per_token % self.copy_threads:
raise ValueError("token vectors must divide evenly across finalize threads")
self.rms_vectors_per_thread = self.packs_per_token // self.copy_threads
self.packs_per_reduction_shard = self.packs_per_token // tp
if rms_shard_major:
if tp < self.rms_warps_per_token or tp % self.rms_warps_per_token:
raise ValueError(
"shard-major RMS requires an integer number of reduction "
"shards per RMS warp"
)
self.reduction_shards_per_rms_warp = tp // self.rms_warps_per_token
if (
self.rms_vectors_per_thread * WARP_SIZE
!= self.packs_per_reduction_shard * self.reduction_shards_per_rms_warp
):
raise ValueError(
"shard-major RMS warp coverage must match its reduction shards"
)
else:
self.reduction_shards_per_rms_warp = 0
if self.packs_per_reduction_shard % self.reduction_threads:
raise ValueError("the reduction shard must divide evenly across threads")
self.reduction_vectors_per_thread = (
self.packs_per_reduction_shard // self.reduction_threads
)
@cute.jit
def _rms_arrive_and_wait(self, rms_group: Int32) -> None:
barrier_0 = pipeline.NamedBarrier(
barrier_id=2, num_threads=self.rms_threads_per_token
)
if cutlass.const_expr(self.rms_token_groups > 1):
barrier_1 = pipeline.NamedBarrier(
barrier_id=3, num_threads=self.rms_threads_per_token
)
if cutlass.const_expr(self.rms_token_groups == 4):
barrier_2 = pipeline.NamedBarrier(
barrier_id=4, num_threads=self.rms_threads_per_token
)
barrier_3 = pipeline.NamedBarrier(
barrier_id=5, num_threads=self.rms_threads_per_token
)
if rms_group == 0:
barrier_0.arrive_and_wait()
elif cutlass.const_expr(self.rms_token_groups == 2): # noqa: SIM114
barrier_1.arrive_and_wait()
elif rms_group == 1:
barrier_1.arrive_and_wait()
elif rms_group == 2:
barrier_2.arrive_and_wait()
else:
barrier_3.arrive_and_wait()
else:
barrier_0.arrive_and_wait()
@cute.jit
def __call__(
self,
routed_output: cute.Tensor,
expert_weights: cute.Tensor,
permuted_indices: cute.Tensor,
shared_output: cute.Tensor,
residual_source: cute.Tensor,
gamma: cute.Tensor,
local_contributions: cute.Tensor,
prenorm_mailbox: cute.Tensor,
residual_output: cute.Tensor,
norm_output: cute.Tensor,
ready_counter_peer_addresses: cute.Tensor,
ready_counters: cute.Tensor,
processed_counters: cute.Tensor,
local_contributions_multicast_address: Int64,
prenorm_mailbox_multicast_address: Int64,
m: Int32,
stream: cuda.CUstream,
) -> None:
smem_layout = cute.make_layout(
(self.shard_elements * self.stages,), stride=(1,)
)
@cute.struct
class SharedStorage:
barriers: cute.struct.MemRange[Int64, 2 * self.stages]
stage_probs: cute.struct.MemRange[Float32, self.stages]
cached_rows: cute.struct.MemRange[Int32, self.metadata_slots]
cached_probs: cute.struct.MemRange[Float32, self.metadata_slots]
consumer_progress: cute.struct.MemRange[Int32, self.consumer_warps]
norm_warp_sums: cute.struct.MemRange[Float32, self.rms_warp_sum_slots]
norm_inv_rms: cute.struct.MemRange[Float32, self.rms_stage_slots]
rows: cute.struct.Align[
cute.struct.MemRange[BFloat16, cute.cosize(smem_layout)],
SMEM_ALIGNMENT,
]
self.shared_storage: type[cute.struct.Struct] = SharedStorage
self.kernel(
routed_output,
shared_output,
residual_source,
gamma,
expert_weights,
permuted_indices,
local_contributions,
prenorm_mailbox,
residual_output,
norm_output,
ready_counter_peer_addresses,
ready_counters,
processed_counters,
local_contributions_multicast_address,
prenorm_mailbox_multicast_address,
m,
smem_layout,
).launch(
grid=(self.active_ctas, 1, 1),
block=(self.block_threads, 1, 1),
min_blocks_per_mp=1,
use_pdl=self.enable_pdl,
stream=stream,
)
@cute.kernel
def kernel(
self,
routed_source: cute.Tensor,
shared_source: cute.Tensor,
residual_source: cute.Tensor,
gamma: cute.Tensor,
expert_weights: cute.Tensor,
permuted_indices: cute.Tensor,
local_contributions: cute.Tensor,
prenorm_mailbox: cute.Tensor,
residual_output: cute.Tensor,
norm_output: cute.Tensor,
ready_counter_peer_addresses: cute.Tensor,
ready_counters: cute.Tensor,
processed_counters: cute.Tensor,
local_contributions_multicast_address: Int64,
prenorm_mailbox_multicast_address: Int64,
m: Int32,
smem_layout: cute.Layout,
) -> None:
block = cute.arch.block_idx()[0]
tidx = cute.arch.thread_idx()[0]
warp = cute.arch.make_warp_uniform(cute.arch.warp_idx())
lane = cute.arch.lane_idx()
cta_group = block // self.tp
cta_slot = block % self.tp
wave = Int64(cta_group)
token = wave * self.tp + cta_slot
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
rows = storage.rows.get_tensor(smem_layout)
barrier_storage = storage.barriers.data_ptr()
stage_probs = storage.stage_probs.data_ptr()
cached_rows = storage.cached_rows.data_ptr()
cached_probs = storage.cached_probs.data_ptr()
consumer_progress = storage.consumer_progress.data_ptr()
norm_warp_sums = storage.norm_warp_sums.data_ptr()
norm_inv_rms = storage.norm_inv_rms.data_ptr()
if tidx < self.consumer_warps:
cute.arch.store((consumer_progress + tidx).llvm_ptr, Int32(0))
cute.arch.sync_threads()
if cutlass.const_expr(self.enable_pdl):
cute.arch.griddepcontrol_wait()
load_pipeline = pipeline.PipelineTmaAsync.create(
barrier_storage=barrier_storage,
num_stages=self.stages,
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 1),
consumer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread, self.consumer_warps
),
tx_count=self.shard_bytes,
)
if warp == 0:
producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.stages
)
peek_empty = cutlass.Boolean(1)
if token < Int64(m):
peek_empty = load_pipeline.producer_try_acquire(producer_state)
while token < Int64(m):
for metadata_chunk in cutlass.range_constexpr(self.metadata_chunks):
metadata_slot = metadata_chunk * WARP_SIZE + lane
if metadata_slot < self.top_k:
item = Int64(token) * self.top_k + metadata_slot
row = cute.arch.load(
(permuted_indices.iterator + item).llvm_ptr, Int32
)
prob = load_global_bf16_as_f32(
Int64((expert_weights.iterator + item).toint())
)
if cutlass.const_expr(self.routed_scaling_factor != 1.0):
prob = prob * Float32(self.routed_scaling_factor)
if row == Int32(-1):
row = Int32(0)
prob = Float32(0.0)
cute.arch.store((cached_rows + metadata_slot).llvm_ptr, row)
cute.arch.store((cached_probs + metadata_slot).llvm_ptr, prob)
cute.arch.sync_warp()
for shard in cutlass.range_constexpr(self.hidden_shards):
for contribution in cutlass.range_constexpr(self.contributions):
load_pipeline.producer_acquire(producer_state, peek_empty)
if lane == 0:
if cutlass.const_expr(contribution < self.top_k):
prob = cute.arch.load(
(cached_probs + contribution).llvm_ptr,
Float32,
)
row = cute.arch.load(
(cached_rows + contribution).llvm_ptr,
Int32,
)
source_element = (
Int64(row) * self.hidden
+ shard * self.shard_elements
)
source = routed_source.iterator + source_element
else:
prob = Float32(1.0)
source_element = (
Int64(token) * self.hidden
+ shard * self.shard_elements
)
source = shared_source.iterator + source_element
cute.arch.store(
(stage_probs + producer_state.index).llvm_ptr,
prob,
)
fence_proxy_async_shared_cta()
cpasync_bulk_g2s(
source,
rows.iterator
+ producer_state.index * self.shard_elements,
load_pipeline.producer_get_barrier(producer_state),
Int32(self.shard_bytes),
)
producer_state.advance()
peek_empty = load_pipeline.producer_try_acquire(producer_state)
wave += self.cta_groups
token = wave * self.tp + cta_slot
load_pipeline.producer_tail(producer_state)
elif warp > 0 and warp <= self.consumer_warps:
finalize_join = pipeline.NamedBarrier(
barrier_id=6, num_threads=self.consumer_threads
)
consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.stages
)
consumer_tid = tidx - WARP_SIZE
consumer_wave = Int64(cta_group)
token = consumer_wave * self.tp + cta_slot
consumer_token_progress = Int32(0)
while token < Int64(m):
for shard in cutlass.range_constexpr(self.hidden_shards):
if cutlass.const_expr(self.top_k == 0):
load_pipeline.consumer_wait(consumer_state)
for trip in cutlass.range_constexpr(self.vectors_per_thread):
output_element = (
Int64(token) * self.hidden
+ shard * self.shard_elements
+ trip * self.consumer_threads * VEC_BF16
+ consumer_tid * VEC_BF16
)
store_global_u32x4(
Int64(
(
local_contributions.iterator + output_element
).toint()
),
load_shared_u32x4(
rows.iterator
+ consumer_state.index * self.shard_elements
+ trip * self.consumer_threads * VEC_BF16
+ consumer_tid * VEC_BF16
),
)
load_pipeline.consumer_release(consumer_state)
consumer_state.advance()
accum = cute.make_rmem_tensor(
cute.make_layout(
(self.vectors_per_thread, VEC_BF16),
stride=(VEC_BF16, 1),
),
Float32,
)
accum.fill(Float32(0.0))
for _ in cutlass.range_constexpr(
self.contributions if self.top_k > 0 else 0
):
load_pipeline.consumer_wait(consumer_state)
prob = cute.arch.load(
(stage_probs + consumer_state.index).llvm_ptr,
Float32,
)
if prob != Float32(0.0):
for trip in cutlass.range_constexpr(
self.vectors_per_thread
):
stage_ptr = (
rows.iterator
+ consumer_state.index * self.shard_elements
+ trip * self.consumer_threads * VEC_BF16
+ consumer_tid * VEC_BF16
)
values = packed_u32x4_to_bf16x8(
load_shared_u32x4(stage_ptr)
).to(Float32)
accum[trip, None].store(
accum[trip, None].load() + values * prob
)
load_pipeline.consumer_release(consumer_state)
consumer_state.advance()
for trip in cutlass.range_constexpr(
self.vectors_per_thread if self.top_k > 0 else 0
):
output_element = (
Int64(token) * self.hidden
+ shard * self.shard_elements
+ trip * self.consumer_threads * VEC_BF16
+ consumer_tid * VEC_BF16
)
store_global_u32x4(
Int64(
(local_contributions.iterator + output_element).toint()
),
bf16x8_to_packed_u32x4(
accum[trip, None].load().to(BFloat16)
),
)
clear_value = packed_negative_zero_bf16x8()
token_pack = token * self.packs_per_token
for clear_item in cutlass.range_constexpr(
self.clear_vectors_per_thread
):
clear_pack = consumer_tid + clear_item * self.consumer_threads
clear_element = (token_pack + clear_pack) * VEC_BF16
store_global_u32x4(
Int64((prenorm_mailbox.iterator + clear_element).toint()),
clear_value,
)
cute.arch.sync_warp()
consumer_token_progress += 1
if lane == 0:
cute.arch.store(
(consumer_progress + warp - 1).llvm_ptr,
consumer_token_progress,
sem="release",
scope="cta",
)
consumer_wave += self.cta_groups
token = consumer_wave * self.tp + cta_slot
finalize_join.arrive_and_wait()
if cutlass.const_expr(self.rms_token_groups > 1):
rms_group = (warp - 1) // self.rms_warps_per_token
rms_group_warp = warp - 1 - rms_group * self.rms_warps_per_token
copy_tid = rms_group_warp * WARP_SIZE + lane
else:
rms_group = Int32(0)
rms_group_warp = warp - 1
copy_tid = consumer_tid
rms_pack_base = copy_tid
rms_pack_stride = self.copy_threads
if cutlass.const_expr(self.rms_shard_major):
rms_pack_base = (
rms_group_warp
* self.reduction_shards_per_rms_warp
* self.packs_per_reduction_shard
+ lane
)
rms_pack_stride = WARP_SIZE
copy_wave = Int64(cta_group) + Int64(rms_group) * self.cta_groups
copy_token = copy_wave * self.tp + cta_slot
if cutlass.const_expr(self.rms_pipeline_stages > 1):
rms_wave_stride = self.cta_groups * self.rms_token_groups
while copy_token < Int64(m):
for rms_stage in cutlass.range_constexpr(self.rms_pipeline_stages):
stage_wave = copy_wave + rms_stage * rms_wave_stride
stage_token = stage_wave * self.tp + cta_slot
if stage_token < Int64(m):
stage_slot = (
rms_group * self.rms_pipeline_stages + rms_stage
)
token_pack = stage_token * self.packs_per_token
copy_fragments = cute.make_rmem_tensor(
cute.make_layout(
(self.rms_vectors_per_thread, 4),
stride=(4, 1),
),
Uint32,
)
all_ready = cutlass.Boolean(0)
while not all_ready:
all_ready = cutlass.Boolean(1)
for item in cutlass.range_constexpr(
self.rms_vectors_per_thread
):
pack = rms_pack_base + item * rms_pack_stride
linear_pack = token_pack + pack
packed = load_global_u32x4_address(
Int64(
(
prenorm_mailbox.iterator
+ linear_pack * VEC_BF16
).toint()
),
volatile=True,
)
copy_fragments[item, None].store(packed)
all_ready = all_ready and (
not fragment_has_negative_zero(packed)
)
thread_sum = Float32(0.0)
if cutlass.const_expr(self.write_residual_output):
prenorm_packed = []
for item in cutlass.range_constexpr(
self.rms_vectors_per_thread
):
pack = rms_pack_base + item * rms_pack_stride
prenorm = packed_u32x4_to_bf16x8(
copy_fragments[item, None].load()
)
packed_prenorm = bf16x8_to_packed_u32x4(prenorm)
if cutlass.const_expr(self.write_residual_output):
prenorm_packed.append(packed_prenorm)
# Finalize has drained `rows`; __init__ verifies the RMS layout fits.
store_shared_u32x4(
rows.iterator
+ stage_slot * self.hidden
+ pack * VEC_BF16,
packed_prenorm,
)
prenorm_f32 = prenorm.to(Float32)
thread_sum = thread_sum + (
prenorm_f32 * prenorm_f32
).reduce(
cute.ReductionOp.ADD,
init_val=Float32(0.0),
reduction_profile=0,
)
if cutlass.const_expr(self.write_residual_output):
for item in cutlass.range_constexpr(
self.rms_vectors_per_thread
):
pack = rms_pack_base + item * rms_pack_stride
linear_pack = token_pack + pack
residual_address = Int64(
(
residual_output.iterator
+ linear_pack * VEC_BF16
).toint()
)
store_global_u32x4(
residual_address, prenorm_packed[item]
)
warp_sum = cute.arch.warp_reduction_sum(thread_sum)
if lane == 0:
cute.arch.store(
(
norm_warp_sums
+ stage_slot * self.rms_warps_per_token
+ rms_group_warp
).llvm_ptr,
warp_sum,
)
self._rms_arrive_and_wait(rms_group)
if warp == 1 + rms_group * self.rms_warps_per_token:
for rms_stage in cutlass.range_constexpr(
self.rms_pipeline_stages
):
stage_wave = copy_wave + rms_stage * rms_wave_stride
stage_token = stage_wave * self.tp + cta_slot
if stage_token < Int64(m):
stage_slot = (
rms_group * self.rms_pipeline_stages + rms_stage
)
cta_sum = Float32(0.0)
if lane < self.rms_warps_per_token:
cta_sum = cute.arch.load(
(
norm_warp_sums
+ stage_slot * self.rms_warps_per_token
+ lane
).llvm_ptr,
Float32,
)
cta_sum = cute.arch.warp_reduction_sum(cta_sum)
if lane == 0:
inv_rms = cute.math.rsqrt(
cta_sum / Float32(self.hidden)
+ Float32(self.rms_epsilon),
fastmath=True,
)
cute.arch.store(
(norm_inv_rms + stage_slot).llvm_ptr,
inv_rms,
)
self._rms_arrive_and_wait(rms_group)
gamma_values = []
for item in cutlass.range_constexpr(self.rms_vectors_per_thread):
pack = rms_pack_base + item * rms_pack_stride
gamma_value = packed_u32x4_to_bf16x8(
load_global_u32x4_address(
Int64((gamma.iterator + pack * VEC_BF16).toint())
)
).to(Float32)
if cutlass.const_expr(self.weight_bias != 0.0):
gamma_value = gamma_value + Float32(self.weight_bias)
gamma_values.append(gamma_value)
for rms_stage in cutlass.range_constexpr(self.rms_pipeline_stages):
stage_wave = copy_wave + rms_stage * rms_wave_stride
stage_token = stage_wave * self.tp + cta_slot
if stage_token < Int64(m):
stage_slot = (
rms_group * self.rms_pipeline_stages + rms_stage
)
token_pack = stage_token * self.packs_per_token
inv_rms = cute.arch.load(
(norm_inv_rms + stage_slot).llvm_ptr, Float32
)
norm_packed = []
for item in cutlass.range_constexpr(
self.rms_vectors_per_thread
):
pack = rms_pack_base + item * rms_pack_stride
prenorm = packed_u32x4_to_bf16x8(
load_shared_u32x4(
rows.iterator
+ stage_slot * self.hidden
+ pack * VEC_BF16
)
).to(Float32)
result = (prenorm * inv_rms * gamma_values[item]).to(
BFloat16
)
norm_packed.append(bf16x8_to_packed_u32x4(result))
for item in cutlass.range_constexpr(
self.rms_vectors_per_thread
):
pack = rms_pack_base + item * rms_pack_stride
linear_pack = token_pack + pack
store_global_u32x4(
Int64(
(
norm_output.iterator
+ linear_pack * VEC_BF16
).toint()
),
norm_packed[item],
)
copy_wave += rms_wave_stride * self.rms_pipeline_stages
copy_token = copy_wave * self.tp + cta_slot
if cutlass.const_expr(self.rms_pipeline_stages == 1):
while copy_token < Int64(m):
token_pack = copy_token * self.packs_per_token
copy_values = []
copy_sources = []
if cutlass.const_expr(self.write_residual_output):
copy_destinations = []
for item in cutlass.range_constexpr(self.rms_vectors_per_thread):
pack = rms_pack_base + item * rms_pack_stride
linear_pack = token_pack + pack
source_address = Int64(
(prenorm_mailbox.iterator + linear_pack * VEC_BF16).toint()
)
copy_sources.append(source_address)
if cutlass.const_expr(self.write_residual_output):
copy_destinations.append(
Int64(
(
residual_output.iterator
+ linear_pack * VEC_BF16
).toint()
)
)
copy_fragments = cute.make_rmem_tensor(
cute.make_layout(
(self.rms_vectors_per_thread, 4),
stride=(4, 1),
),
Uint32,
)
all_ready = cutlass.Boolean(0)
while not all_ready:
all_ready = cutlass.Boolean(1)
for item in cutlass.range_constexpr(
self.rms_vectors_per_thread
):
packed = load_global_u32x4_address(
copy_sources[item],
volatile=True,
)
copy_fragments[item, None].store(packed)
all_ready = all_ready and (
not fragment_has_negative_zero(packed)
)
for item in cutlass.range_constexpr(self.rms_vectors_per_thread):
copy_values.append(copy_fragments[item, None].load())
prenorm_fragments = cute.make_rmem_tensor(
cute.make_layout(
(self.rms_vectors_per_thread, VEC_BF16),
stride=(VEC_BF16, 1),
),
BFloat16,
)
thread_sum = Float32(0.0)
if cutlass.const_expr(self.write_residual_output):
prenorm_packed = []
for item in cutlass.range_constexpr(self.rms_vectors_per_thread):
pack = rms_pack_base + item * rms_pack_stride
prenorm = packed_u32x4_to_bf16x8(copy_values[item])
prenorm_fragments[item, None].store(prenorm)
packed_prenorm = bf16x8_to_packed_u32x4(prenorm)
if cutlass.const_expr(self.write_residual_output):
prenorm_packed.append(packed_prenorm)
prenorm_f32 = prenorm.to(Float32)
thread_sum = thread_sum + (prenorm_f32 * prenorm_f32).reduce(
cute.ReductionOp.ADD,
init_val=Float32(0.0),
reduction_profile=0,
)
if cutlass.const_expr(self.write_residual_output):
for item in cutlass.range_constexpr(
self.rms_vectors_per_thread
):
store_global_u32x4(
copy_destinations[item], prenorm_packed[item]
)
warp_sum = cute.arch.warp_reduction_sum(thread_sum)
if lane == 0:
cute.arch.store((norm_warp_sums + warp - 1).llvm_ptr, warp_sum)
self._rms_arrive_and_wait(rms_group)
if warp == 1 + rms_group * self.rms_warps_per_token:
cta_sum = Float32(0.0)
if lane < self.rms_warps_per_token:
cta_sum = cute.arch.load(
(
norm_warp_sums
+ rms_group * self.rms_warps_per_token
+ lane
).llvm_ptr,
Float32,
)
cta_sum = cute.arch.warp_reduction_sum(cta_sum)
if lane == 0:
inv_rms = cute.math.rsqrt(
cta_sum / Float32(self.hidden)
+ Float32(self.rms_epsilon),
fastmath=True,
)
cute.arch.store(
(norm_inv_rms + rms_group).llvm_ptr, inv_rms
)
self._rms_arrive_and_wait(rms_group)
inv_rms = cute.arch.load(
(norm_inv_rms + rms_group).llvm_ptr, Float32
)
gamma_values = []
for item in cutlass.range_constexpr(self.rms_vectors_per_thread):
pack = rms_pack_base + item * rms_pack_stride
gamma_value = packed_u32x4_to_bf16x8(
load_global_u32x4_address(
Int64((gamma.iterator + pack * VEC_BF16).toint())
)
).to(Float32)
if cutlass.const_expr(self.weight_bias != 0.0):
gamma_value = gamma_value + Float32(self.weight_bias)
gamma_values.append(gamma_value)
norm_packed = []
for item in cutlass.range_constexpr(self.rms_vectors_per_thread):
pack = rms_pack_base + item * rms_pack_stride
linear_pack = token_pack + pack
prenorm_for_norm = (
prenorm_fragments[item, None].load().to(Float32)
)
result = (prenorm_for_norm * inv_rms * gamma_values[item]).to(
BFloat16
)
norm_packed.append(bf16x8_to_packed_u32x4(result))
for item in cutlass.range_constexpr(self.rms_vectors_per_thread):
pack = rms_pack_base + item * rms_pack_stride
linear_pack = token_pack + pack
store_global_u32x4(
Int64(
(norm_output.iterator + linear_pack * VEC_BF16).toint()
),
norm_packed[item],
)
copy_wave += self.cta_groups * self.rms_token_groups
copy_token = copy_wave * self.tp + cta_slot
elif warp == self.publisher_warp:
owner_ready_address = cute.arch.load(
(ready_counter_peer_addresses.iterator + cta_slot).llvm_ptr,
Int64,
)
first_token = Int64(cta_group) * self.tp + cta_slot
token_count = Int32(0)
if first_token < Int64(m):
token_count = Int32(
(Int64(m) + self.active_ctas - 1 - first_token) // self.active_ctas
)
published = Int32(0)
while published < token_count:
observed = token_count
if lane < self.consumer_warps:
observed = cute.arch.load(
(consumer_progress + lane).llvm_ptr,
Int32,
sem="relaxed",
scope="cta",
)
frontier = cute.arch.warp_reduction(
observed, lambda x, y: cutlass.min(x, y)
)
if frontier > published:
acquired = token_count
if lane < self.consumer_warps:
acquired = cute.arch.load(
(consumer_progress + lane).llvm_ptr,
Int32,
sem="acquire",
scope="cta",
)
frontier = cute.arch.warp_reduction(
acquired, lambda x, y: cutlass.min(x, y)
)
cute.arch.sync_warp()
batch = cutlass.min(frontier - published, Int32(WARP_SIZE))
if lane < batch:
sequence = Int64(published + lane)
publish_token = (
Int64(cta_group) + sequence * self.cta_groups
) * self.tp + cta_slot
owner_token = publish_token // self.tp
remote_release_add1_u32(owner_ready_address + owner_token * 4)
published += batch
elif (
block < self.reduction_ctas
and warp >= self.reduction_warp_begin
and (warp < self.reduction_warp_begin + self.reduction_warps)
):
reduction_warp = warp - self.reduction_warp_begin
reduction_tid = reduction_warp * WARP_SIZE + lane
reduction_barrier = pipeline.NamedBarrier(
barrier_id=1, num_threads=self.reduction_threads
)
reduction_shard = block % self.tp
local_token = Int64(block // self.tp)
token = local_token * self.tp + self.rank
while token < Int64(m):
processed_index = local_token * self.tp + reduction_shard
target = Uint32(0)
if reduction_tid == 0:
ready_counter_address = (
ready_counters.iterator + local_token
).llvm_ptr
processed_counter_address = (
processed_counters.iterator + processed_index
).llvm_ptr
target = cute.arch.load(processed_counter_address, Uint32) + Uint32(
self.tp
)
observed = Uint32(0)
while observed != target:
observed = cute.arch.load(
ready_counter_address,
Uint32,
sem="relaxed",
scope="sys",
)
cute.arch.load(
ready_counter_address,
Uint32,
sem="acquire",
scope="sys",
)
reduction_barrier.arrive_and_wait()
values = []
addresses = []
token_pack = token * self.packs_per_token
shard_pack = reduction_shard * self.packs_per_reduction_shard
for item in cutlass.range_constexpr(self.reduction_vectors_per_thread):
pack = shard_pack + reduction_tid + item * self.reduction_threads
input_address = (
local_contributions_multicast_address + (token_pack + pack) * 16
)
output_address = (
prenorm_mailbox_multicast_address + (token_pack + pack) * 16
)
reduced_packed = ldmc_bf16x8(input_address)
reduced_values = packed_u32x4_to_bf16x8(reduced_packed).to(Float32)
if cutlass.const_expr(self.add_residual):
residual_values = packed_u32x4_to_bf16x8(
load_global_u32x4_address(
Int64(
(
residual_source.iterator
+ (token_pack + pack) * VEC_BF16
).toint()
)
)
).to(Float32)
reduced_values = reduced_values + residual_values
reduced_packed = bf16x8_to_packed_u32x4(reduced_values.to(BFloat16))
values.append(sanitize_negative_zero_u32x4(reduced_packed))
addresses.append(output_address)
for item in cutlass.range_constexpr(self.reduction_vectors_per_thread):
stmc_bf16x8(addresses[item], values[item])
if reduction_tid == 0:
cute.arch.store(
(processed_counters.iterator + processed_index).llvm_ptr,
target,
)
local_token += self.reduction_cta_groups
token = local_token * self.tp + self.rank
cute.arch.sync_threads()
if cutlass.const_expr(self.enable_pdl):
cute.arch.griddepcontrol_launch_dependents()
python/sglang/kernels/ops/communication/mnnvl_cutedsl/kernel_ht/protocol.py
+469
View File
@@ -0,0 +1,469 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""High-throughput MNNVL protocol and its two operation paths."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, TypedDict, cast
import cutlass.cute as cute
import torch
import torch.distributed as dist
from cutlass import BFloat16, Int32, Int64, Uint32
from cutlass.cute.runtime import make_fake_compact_tensor
from ..runtime import (
current_cu_stream,
make_fake_dynamic_compact_tensor,
to_cute,
to_cute_dynamic,
)
from ..symmetric_buffer import SymmetricBuffer
from .device_kernel import _MoeFinalizeAllReduceRMSNormHTDeviceKernel
@dataclass(frozen=True, slots=True)
class HTFinalizeTuning:
persistent_ctas: int | None = None
consumer_threads: int = 512
vectors_per_thread: int = 2
stages: int = 6
reduction_warps: int = 1
reduction_cta_groups: int | None = None
rms_token_groups: int = 2
rms_pipeline_stages: int = 2
rms_shard_major: bool = False
enable_pdl: bool = True
@dataclass(frozen=True, slots=True)
class HTAllReduceTuning:
persistent_ctas: int | None = None
consumer_threads: int = 512
vectors_per_thread: int = 2
stages: int = 2
reduction_warps: int = 2
reduction_cta_groups: int | None = None
rms_token_groups: int = 2
rms_pipeline_stages: int = 1
rms_shard_major: bool = False
enable_pdl: bool = True
HT_FINALIZE_GB300_TP8_H8192_K10_M_LE_2048 = HTFinalizeTuning(
stages=7,
reduction_warps=2,
rms_pipeline_stages=3,
rms_shard_major=True,
)
HT_FINALIZE_GB300_TP8_H8192_K10_M_GE_2049 = HTFinalizeTuning()
HT_FINALIZE_GB300_TP8_H8192_K10 = HT_FINALIZE_GB300_TP8_H8192_K10_M_GE_2049
HT_FINALIZE_GB300_TP16_H8192_K10 = HTFinalizeTuning(
stages=7,
reduction_warps=2,
rms_pipeline_stages=3,
rms_shard_major=True,
)
HT_ALL_REDUCE_GB300_TP8_H8192 = HTAllReduceTuning()
HT_ALL_REDUCE_GB300_TP16_H8192 = HTAllReduceTuning()
@dataclass(slots=True)
class HTProtocolState:
local_contributions: SymmetricBuffer
prenorm_mailbox: SymmetricBuffer
routed_ready_counters: SymmetricBuffer
routed_processed_counters: torch.Tensor
all_reduce_ready_counters: SymmetricBuffer
all_reduce_processed_counters: torch.Tensor
class _PathKwargs(TypedDict):
hidden_size: int
top_k: int
capacity_m: int
write_residual_output: bool
class _HTPath:
def __init__(
self,
*,
compiled: Any,
hidden_size: int,
top_k: int,
capacity_m: int,
write_residual_output: bool,
) -> None:
self._compiled = compiled
self.hidden_size = hidden_size
self.top_k = top_k
self.capacity_m = capacity_m
self.write_residual_output = write_residual_output
def _outputs(
self,
m: int,
norm_output: torch.Tensor | None,
residual_output: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
shape = (m, self.hidden_size)
device = torch.device("cuda", torch.cuda.current_device())
if norm_output is None:
norm_output = torch.empty(shape, dtype=torch.bfloat16, device=device)
if self.write_residual_output and residual_output is None:
residual_output = torch.empty(shape, dtype=torch.bfloat16, device=device)
return norm_output, residual_output
def _state_buffers(
self, state: HTProtocolState
) -> tuple[SymmetricBuffer, SymmetricBuffer]:
return state.local_contributions, state.prenorm_mailbox
def _validate_m(self, m: int) -> None:
if not 1 <= m <= self.capacity_m:
raise ValueError(f"m must be in [1, {self.capacity_m}]")
class FinalizeAllReduceRMSNormHTKernel(_HTPath):
def __call__(
self,
routed_output: torch.Tensor,
expert_weights: torch.Tensor,
permuted_indices: torch.Tensor,
shared_output: torch.Tensor | None,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
m: int,
*,
state: HTProtocolState,
norm_output: torch.Tensor | None = None,
residual_output: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
self._validate_m(m)
norm_output, residual_output = self._outputs(m, norm_output, residual_output)
local, prenorm = self._state_buffers(state)
peers = cast(torch.Tensor, state.routed_ready_counters.peer_addresses)
shared_arg = shared_output if shared_output is not None else norm_output
residual_arg = residual_source if residual_source is not None else norm_output
residual_output_arg = (
residual_output if residual_output is not None else norm_output
)
self._compiled(
to_cute_dynamic(routed_output.flatten(), 16, divisibility=self.hidden_size),
to_cute_dynamic(expert_weights.flatten(), 2, divisibility=self.top_k),
to_cute_dynamic(permuted_indices.flatten(), 4, divisibility=self.top_k),
to_cute_dynamic(shared_arg.flatten(), 16, divisibility=self.hidden_size),
to_cute_dynamic(residual_arg.flatten(), 16, divisibility=self.hidden_size),
to_cute(gamma, 16),
to_cute(local.tensor.flatten(), 16),
to_cute(prenorm.tensor.flatten(), 16),
to_cute_dynamic(
residual_output_arg.flatten(),
16,
divisibility=self.hidden_size,
),
to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size),
to_cute(peers, 8),
to_cute(state.routed_ready_counters.tensor, 4),
to_cute(state.routed_processed_counters.flatten(), 4),
Int64(cast(int, local.multicast_address)),
Int64(cast(int, prenorm.multicast_address)),
Int32(m),
current_cu_stream(),
)
return norm_output, residual_output
class AllReduceRMSNormHTKernel(_HTPath):
def __call__(
self,
local_contribution: torch.Tensor,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
m: int,
*,
state: HTProtocolState,
norm_output: torch.Tensor | None = None,
residual_output: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
self._validate_m(m)
norm_output, residual_output = self._outputs(m, norm_output, residual_output)
local, prenorm = self._state_buffers(state)
peers = cast(torch.Tensor, state.all_reduce_ready_counters.peer_addresses)
residual_arg = residual_source if residual_source is not None else norm_output
residual_output_arg = (
residual_output if residual_output is not None else norm_output
)
index_arg = state.all_reduce_processed_counters.view(torch.int32)
# top_k=0 disables metadata reads, so the aliased placeholders stay unused.
self._compiled(
to_cute_dynamic(
local_contribution.flatten(),
16,
divisibility=self.hidden_size,
),
to_cute_dynamic(local_contribution.flatten(), 2, divisibility=1),
to_cute_dynamic(index_arg.flatten(), 4, divisibility=1),
to_cute_dynamic(
local_contribution.flatten(),
16,
divisibility=self.hidden_size,
),
to_cute_dynamic(residual_arg.flatten(), 16, divisibility=self.hidden_size),
to_cute(gamma, 16),
to_cute(local.tensor.flatten(), 16),
to_cute(prenorm.tensor.flatten(), 16),
to_cute_dynamic(
residual_output_arg.flatten(),
16,
divisibility=self.hidden_size,
),
to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size),
to_cute(peers, 8),
to_cute(state.all_reduce_ready_counters.tensor, 4),
to_cute(state.all_reduce_processed_counters.flatten(), 4),
Int64(cast(int, local.multicast_address)),
Int64(cast(int, prenorm.multicast_address)),
Int32(m),
current_cu_stream(),
)
return norm_output, residual_output
class HTProtocol:
"""Own tuning-independent HT State and both persistent path variants."""
def __init__(
self,
hidden_size: int,
top_k: int,
tp_size: int,
rank: int,
capacity_m: int,
rms_epsilon: float,
routed_scaling_factor: float,
weight_bias: float,
*,
include_shared_expert: bool,
add_residual: bool,
write_residual_output: bool,
finalize_tunings: tuple[HTFinalizeTuning, ...],
all_reduce_tunings: tuple[HTAllReduceTuning, ...],
group: dist.ProcessGroup,
) -> None:
self.hidden_size = hidden_size
self.top_k = top_k
self.tp_size = tp_size
self.rank = rank
self.capacity_m = capacity_m
self.rms_epsilon = rms_epsilon
self.routed_scaling_factor = routed_scaling_factor
self.weight_bias = weight_bias
self.include_shared_expert = include_shared_expert
self.add_residual = add_residual
self.write_residual_output = write_residual_output
self.finalize_kernels = {
tuning: FinalizeAllReduceRMSNormHTKernel(
compiled=self._compile_finalize(tuning),
**self._path_kwargs(),
)
for tuning in dict.fromkeys(finalize_tunings)
}
self.all_reduce_kernels = {
tuning: AllReduceRMSNormHTKernel(
compiled=self._compile_all_reduce(tuning),
**self._path_kwargs(),
)
for tuning in dict.fromkeys(all_reduce_tunings)
}
self.state = self._create_state(group)
def _path_kwargs(self) -> _PathKwargs:
return {
"hidden_size": self.hidden_size,
"top_k": self.top_k,
"capacity_m": self.capacity_m,
"write_residual_output": self.write_residual_output,
}
def _resolve_ctas(self, persistent_ctas: int | None) -> int:
sm_count = torch.cuda.get_device_properties(
torch.cuda.current_device()
).multi_processor_count
# min_blocks_per_mp=1 guarantees one resident CTA per SM for this kernel.
resident_ctas = (sm_count // self.tp_size) * self.tp_size
if resident_ctas == 0:
raise ValueError("tp_size exceeds the available SM count")
if persistent_ctas is None:
return resident_ctas
if persistent_ctas <= 0 or persistent_ctas % self.tp_size:
raise ValueError(
"persistent_ctas must be positive and divisible by tp_size"
)
return min(persistent_ctas, resident_ctas)
def _compile_finalize(self, tuning: HTFinalizeTuning):
active_ctas = self._resolve_ctas(tuning.persistent_ctas)
groups = tuning.reduction_cta_groups or active_ctas // self.tp_size
kernel = _MoeFinalizeAllReduceRMSNormHTDeviceKernel(
hidden=self.hidden_size,
top_k=self.top_k,
tp=self.tp_size,
rank=self.rank,
active_ctas=active_ctas,
stages=tuning.stages,
consumer_threads=tuning.consumer_threads,
vectors_per_thread=tuning.vectors_per_thread,
reduction_warps=tuning.reduction_warps,
reduction_cta_groups=groups,
rms_token_groups=tuning.rms_token_groups,
rms_pipeline_stages=tuning.rms_pipeline_stages,
rms_shard_major=tuning.rms_shard_major,
rms_epsilon=self.rms_epsilon,
routed_scaling_factor=self.routed_scaling_factor,
weight_bias=self.weight_bias,
include_shared_expert=self.include_shared_expert,
add_residual=self.add_residual,
write_residual_output=self.write_residual_output,
enable_pdl=tuning.enable_pdl,
)
return self._compile(kernel, top_k=self.top_k)
def _compile_all_reduce(self, tuning: HTAllReduceTuning):
active_ctas = self._resolve_ctas(tuning.persistent_ctas)
groups = tuning.reduction_cta_groups or active_ctas // self.tp_size
kernel = _MoeFinalizeAllReduceRMSNormHTDeviceKernel(
hidden=self.hidden_size,
top_k=0,
tp=self.tp_size,
rank=self.rank,
active_ctas=active_ctas,
stages=tuning.stages,
consumer_threads=tuning.consumer_threads,
vectors_per_thread=tuning.vectors_per_thread,
reduction_warps=tuning.reduction_warps,
reduction_cta_groups=groups,
rms_token_groups=tuning.rms_token_groups,
rms_pipeline_stages=tuning.rms_pipeline_stages,
rms_shard_major=tuning.rms_shard_major,
rms_epsilon=self.rms_epsilon,
routed_scaling_factor=1.0,
weight_bias=self.weight_bias,
include_shared_expert=True,
add_residual=self.add_residual,
write_residual_output=self.write_residual_output,
enable_pdl=tuning.enable_pdl,
)
return self._compile(kernel, top_k=0)
def _compile(
self,
kernel: _MoeFinalizeAllReduceRMSNormHTDeviceKernel,
*,
top_k: int,
):
activation = self.capacity_m * self.hidden_size
token_slots = (self.capacity_m + self.tp_size - 1) // self.tp_size
args = (
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=2, divisibility=max(top_k, 1)
),
make_fake_dynamic_compact_tensor(
Int32, alignment=4, divisibility=max(top_k, 1)
),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_compact_tensor(BFloat16, (self.hidden_size,), assumed_align=16),
make_fake_compact_tensor(BFloat16, (activation,), assumed_align=16),
make_fake_compact_tensor(BFloat16, (activation,), assumed_align=16),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_compact_tensor(Int64, (self.tp_size,), assumed_align=8),
make_fake_compact_tensor(Uint32, (token_slots,), assumed_align=4),
make_fake_compact_tensor(
Uint32, (token_slots * self.tp_size,), assumed_align=4
),
Int64(0),
Int64(0),
Int32(self.capacity_m),
current_cu_stream(),
)
return cute.compile(kernel, *args)
def _allocate_large_buffers(
self, group: dist.ProcessGroup
) -> tuple[SymmetricBuffer, SymmetricBuffer]:
shape = (self.capacity_m, self.hidden_size)
device = torch.device("cuda", torch.cuda.current_device())
return (
SymmetricBuffer.allocate(
shape,
torch.bfloat16,
device,
group,
require_multicast=True,
),
SymmetricBuffer.allocate(
shape,
torch.bfloat16,
device,
group,
require_multicast=True,
),
)
def _create_state(self, group: dist.ProcessGroup) -> HTProtocolState:
device = torch.device("cuda", torch.cuda.current_device())
token_slots = (self.capacity_m + self.tp_size - 1) // self.tp_size
def counters() -> tuple[SymmetricBuffer, torch.Tensor]:
ready = SymmetricBuffer.allocate(
(token_slots,),
torch.uint32,
device,
group,
materialize_peer_addresses=True,
)
ready.tensor.zero_()
processed = torch.zeros(
(token_slots, self.tp_size), dtype=torch.uint32, device=device
)
return ready, processed
routed_ready, routed_processed = counters()
all_reduce_ready, all_reduce_processed = counters()
local, prenorm = self._allocate_large_buffers(group)
return HTProtocolState(
local_contributions=local,
prenorm_mailbox=prenorm,
routed_ready_counters=routed_ready,
routed_processed_counters=routed_processed,
all_reduce_ready_counters=all_reduce_ready,
all_reduce_processed_counters=all_reduce_processed,
)
python/sglang/kernels/ops/communication/mnnvl_cutedsl/kernel_ll/__init__.py
+47
View File
@@ -0,0 +1,47 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Low-latency MNNVL protocol."""
from .protocol import (
LL_ALL_REDUCE_GB300_TP8_H8192,
LL_ALL_REDUCE_GB300_TP8_H8192_M_GE_5,
LL_ALL_REDUCE_GB300_TP8_H8192_M_LE_4,
LL_ALL_REDUCE_GB300_TP16_H8192,
LL_ALL_REDUCE_GB300_TP16_H8192_M_11_TO_17,
LL_ALL_REDUCE_GB300_TP16_H8192_M_GE_18,
LL_ALL_REDUCE_GB300_TP16_H8192_M_LE_10,
LL_FINALIZE_GB300_TP8_H8192_K10,
LL_FINALIZE_GB300_TP8_H8192_K10_M_GE_20,
LL_FINALIZE_GB300_TP16_H8192_K10,
LLAllReduceTuning,
LLCollectiveTuning,
LLFinalizeTuning,
)
__all__ = [
"LL_ALL_REDUCE_GB300_TP8_H8192",
"LL_ALL_REDUCE_GB300_TP8_H8192_M_GE_5",
"LL_ALL_REDUCE_GB300_TP8_H8192_M_LE_4",
"LL_ALL_REDUCE_GB300_TP16_H8192",
"LL_ALL_REDUCE_GB300_TP16_H8192_M_11_TO_17",
"LL_ALL_REDUCE_GB300_TP16_H8192_M_GE_18",
"LL_ALL_REDUCE_GB300_TP16_H8192_M_LE_10",
"LL_FINALIZE_GB300_TP8_H8192_K10",
"LL_FINALIZE_GB300_TP8_H8192_K10_M_GE_20",
"LL_FINALIZE_GB300_TP16_H8192_K10",
"LLAllReduceTuning",
"LLCollectiveTuning",
"LLFinalizeTuning",
]
python/sglang/kernels/ops/communication/mnnvl_cutedsl/kernel_ll/device_kernels.py
+1003
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/ops/communication/mnnvl_cutedsl/kernel_ll/protocol.py
+461
View File
@@ -0,0 +1,461 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Low-latency MNNVL protocol and its two operation paths."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, TypedDict, cast
import cutlass.cute as cute
import torch
import torch.distributed as dist
from cutlass import BFloat16, Int32, Int64
from cutlass.cute.runtime import make_fake_compact_tensor
from ..cute_dsl_primitives import QUAD_BF16
from ..runtime import (
current_cu_stream,
make_fake_dynamic_compact_tensor,
to_cute,
to_cute_dynamic,
)
from ..symmetric_buffer import SymmetricBuffer
from .device_kernels import (
LAMPORT_GENERATIONS,
_LamportResidualRMSNormDeviceKernel,
_QuadFinalizePublishDeviceKernel,
_ScalarFinalizePublishDeviceKernel,
_SharedOnlyPublishDeviceKernel,
)
@dataclass(frozen=True, slots=True)
class LLCollectiveTuning:
cluster_size: int = 8
rank_lanes: int = 1
threads: int = 128
enable_pdl: bool = True
@dataclass(frozen=True, slots=True)
class LLFinalizeTuning:
elements_per_thread: int = 4
threads: int = 128
prefetch_group: int = 10
load_shared_expert_before_pdl: bool = False
collective: LLCollectiveTuning = LLCollectiveTuning()
@dataclass(frozen=True, slots=True)
class LLAllReduceTuning:
publish_elements_per_thread: int = 8
publish_threads: int = 128
publish_release_before_store: bool = False
collective: LLCollectiveTuning = LLCollectiveTuning()
LL_FINALIZE_GB300_TP8_H8192_K10 = LLFinalizeTuning()
LL_FINALIZE_GB300_TP8_H8192_K10_M_GE_20 = LLFinalizeTuning(
collective=LLCollectiveTuning(cluster_size=16, rank_lanes=2)
)
LL_FINALIZE_GB300_TP16_H8192_K10 = LLFinalizeTuning(
collective=LLCollectiveTuning(cluster_size=16, rank_lanes=2)
)
LL_ALL_REDUCE_GB300_TP8_H8192_M_LE_4 = LLAllReduceTuning(
collective=LLCollectiveTuning(cluster_size=16, threads=64)
)
LL_ALL_REDUCE_GB300_TP8_H8192_M_GE_5 = LLAllReduceTuning()
LL_ALL_REDUCE_GB300_TP16_H8192_M_LE_10 = LLAllReduceTuning(
collective=LLCollectiveTuning(cluster_size=16, threads=64)
)
LL_ALL_REDUCE_GB300_TP16_H8192_M_11_TO_17 = LLAllReduceTuning()
LL_ALL_REDUCE_GB300_TP16_H8192_M_GE_18 = LLAllReduceTuning(
collective=LLCollectiveTuning(cluster_size=16, rank_lanes=2)
)
LL_ALL_REDUCE_GB300_TP8_H8192 = LL_ALL_REDUCE_GB300_TP8_H8192_M_GE_5
LL_ALL_REDUCE_GB300_TP16_H8192 = LL_ALL_REDUCE_GB300_TP16_H8192_M_LE_10
@dataclass(slots=True)
class LLProtocolState:
contribution_mailbox: SymmetricBuffer
stage_state: torch.Tensor
@dataclass(frozen=True, slots=True)
class _CompiledFinalize:
publish: Any
collective: Any
@dataclass(frozen=True, slots=True)
class _CompiledAllReduce:
publish: Any
collective: Any
class _PathKwargs(TypedDict):
hidden_size: int
top_k: int
capacity_m: int
write_residual_output: bool
class _LLPath:
def __init__(
self,
*,
hidden_size: int,
top_k: int,
capacity_m: int,
write_residual_output: bool,
) -> None:
self.hidden_size = hidden_size
self.top_k = top_k
self.capacity_m = capacity_m
self.write_residual_output = write_residual_output
def _outputs(
self,
m: int,
norm_output: torch.Tensor | None,
residual_output: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
shape = (m, self.hidden_size)
device = torch.device("cuda", torch.cuda.current_device())
if norm_output is None:
norm_output = torch.empty(shape, dtype=torch.bfloat16, device=device)
if self.write_residual_output and residual_output is None:
residual_output = torch.empty(shape, dtype=torch.bfloat16, device=device)
return norm_output, residual_output
def _validate_state(self, state: LLProtocolState, m: int) -> None:
if not 1 <= m <= self.capacity_m:
raise ValueError(f"m must be in [1, {self.capacity_m}]")
address = state.contribution_mailbox.multicast_address
if address is None or address % 16:
raise ValueError(
"LL contribution mailbox requires a 16-byte-aligned multicast address"
)
def _launch_collective(
self,
collective,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
state: LLProtocolState,
norm_output: torch.Tensor,
residual_output: torch.Tensor | None,
m: int,
) -> None:
residual_arg = residual_source if residual_source is not None else norm_output
residual_output_arg = (
residual_output if residual_output is not None else norm_output
)
collective(
to_cute(state.contribution_mailbox.tensor.flatten(), 16),
to_cute_dynamic(residual_arg.flatten(), 16, divisibility=self.hidden_size),
to_cute(gamma, 16),
to_cute_dynamic(
residual_output_arg.flatten(),
16,
divisibility=self.hidden_size,
),
to_cute_dynamic(norm_output.flatten(), 16, divisibility=self.hidden_size),
to_cute(state.stage_state, 4),
Int32(m),
current_cu_stream(),
)
class FinalizeAllReduceRMSNormLLKernel(_LLPath):
def __init__(self, *, compiled: _CompiledFinalize, **kwargs) -> None:
super().__init__(**kwargs)
self._compiled = compiled
def __call__(
self,
routed_output: torch.Tensor,
expert_weights: torch.Tensor,
permuted_indices: torch.Tensor,
shared_output: torch.Tensor | None,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
m: int,
*,
state: LLProtocolState,
norm_output: torch.Tensor | None = None,
residual_output: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
self._validate_state(state, m)
norm_output, residual_output = self._outputs(m, norm_output, residual_output)
shared_arg = shared_output if shared_output is not None else norm_output
self._compiled.publish(
to_cute_dynamic(routed_output.flatten(), 16, divisibility=self.hidden_size),
to_cute_dynamic(expert_weights.flatten(), 2, divisibility=self.top_k),
to_cute_dynamic(permuted_indices.flatten(), 4, divisibility=self.top_k),
to_cute_dynamic(shared_arg.flatten(), 16, divisibility=self.hidden_size),
to_cute(state.stage_state, 4),
Int64(cast(int, state.contribution_mailbox.multicast_address)),
Int32(m),
current_cu_stream(),
)
self._launch_collective(
self._compiled.collective,
residual_source,
gamma,
state,
norm_output,
residual_output,
m,
)
return norm_output, residual_output
class AllReduceRMSNormLLKernel(_LLPath):
def __init__(self, *, compiled: _CompiledAllReduce, **kwargs) -> None:
super().__init__(**kwargs)
self._compiled = compiled
def __call__(
self,
local_contribution: torch.Tensor,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
m: int,
*,
state: LLProtocolState,
norm_output: torch.Tensor | None = None,
residual_output: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
self._validate_state(state, m)
norm_output, residual_output = self._outputs(m, norm_output, residual_output)
self._compiled.publish(
to_cute_dynamic(
local_contribution.flatten(),
16,
divisibility=self.hidden_size,
),
to_cute(state.stage_state, 4),
Int64(cast(int, state.contribution_mailbox.multicast_address)),
Int32(m),
current_cu_stream(),
)
self._launch_collective(
self._compiled.collective,
residual_source,
gamma,
state,
norm_output,
residual_output,
m,
)
return norm_output, residual_output
class LLProtocol:
"""Own LL State and protocol-local compiled variants for both paths."""
def __init__(
self,
hidden_size: int,
top_k: int,
tp_size: int,
rank: int,
capacity_m: int,
rms_epsilon: float,
routed_scaling_factor: float,
weight_bias: float,
*,
include_shared_expert: bool,
add_residual: bool,
write_residual_output: bool,
finalize_tunings: tuple[LLFinalizeTuning, ...],
all_reduce_tunings: tuple[LLAllReduceTuning, ...],
group: dist.ProcessGroup,
) -> None:
self.hidden_size = hidden_size
self.top_k = top_k
self.tp_size = tp_size
self.rank = rank
self.capacity_m = capacity_m
self.rms_epsilon = rms_epsilon
self.routed_scaling_factor = routed_scaling_factor
self.weight_bias = weight_bias
self.include_shared_expert = include_shared_expert
self.add_residual = add_residual
self.write_residual_output = write_residual_output
collective_cache = {
tuning: self._compile_collective(tuning)
for tuning in {
*(item.collective for item in finalize_tunings),
*(item.collective for item in all_reduce_tunings),
}
}
self.finalize_kernels = {
tuning: FinalizeAllReduceRMSNormLLKernel(
compiled=_CompiledFinalize(
publish=self._compile_finalize(tuning),
collective=collective_cache[tuning.collective],
),
**self._path_kwargs(),
)
for tuning in dict.fromkeys(finalize_tunings)
}
self.all_reduce_kernels = {
tuning: AllReduceRMSNormLLKernel(
compiled=_CompiledAllReduce(
publish=self._compile_all_reduce_publish(tuning),
collective=collective_cache[tuning.collective],
),
**self._path_kwargs(),
)
for tuning in dict.fromkeys(all_reduce_tunings)
}
self.state = self._create_state(group)
def _path_kwargs(self) -> _PathKwargs:
return {
"hidden_size": self.hidden_size,
"top_k": self.top_k,
"capacity_m": self.capacity_m,
"write_residual_output": self.write_residual_output,
}
def _compile_finalize(self, tuning: LLFinalizeTuning):
if tuning.elements_per_thread not in (1, QUAD_BF16):
raise ValueError("LL finalize elements_per_thread must be 1 or 4")
kwargs: dict[str, Any] = {
"hidden": self.hidden_size,
"top_k": self.top_k,
"tp": self.tp_size,
"rank": self.rank,
"capacity_m": self.capacity_m,
"threads": tuning.threads,
"routed_scaling_factor": self.routed_scaling_factor,
"include_shared_expert": self.include_shared_expert,
"load_shared_expert_before_pdl": tuning.load_shared_expert_before_pdl,
"enable_pdl": tuning.collective.enable_pdl,
"prefetch_group": tuning.prefetch_group,
}
device_kernel = (
_ScalarFinalizePublishDeviceKernel(**kwargs)
if tuning.elements_per_thread == 1
else _QuadFinalizePublishDeviceKernel(**kwargs)
)
args = (
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=2, divisibility=self.top_k
),
make_fake_dynamic_compact_tensor(
Int32, alignment=4, divisibility=self.top_k
),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_compact_tensor(Int32, (2,), assumed_align=4),
Int64(0),
Int32(self.capacity_m),
current_cu_stream(),
)
return cute.compile(device_kernel, *args)
def _compile_all_reduce_publish(self, tuning: LLAllReduceTuning):
device_kernel = _SharedOnlyPublishDeviceKernel(
hidden=self.hidden_size,
tp=self.tp_size,
rank=self.rank,
capacity_m=self.capacity_m,
elements_per_thread=tuning.publish_elements_per_thread,
threads=tuning.publish_threads,
release_before_store=tuning.publish_release_before_store,
enable_pdl=tuning.collective.enable_pdl,
)
args = (
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_compact_tensor(Int32, (2,), assumed_align=4),
Int64(0),
Int32(self.capacity_m),
current_cu_stream(),
)
return cute.compile(device_kernel, *args)
def _compile_collective(self, tuning: LLCollectiveTuning):
device_kernel = _LamportResidualRMSNormDeviceKernel(
hidden=self.hidden_size,
tp=self.tp_size,
capacity_m=self.capacity_m,
cluster_size=tuning.cluster_size,
rank_lanes=tuning.rank_lanes,
threads=tuning.threads,
rms_epsilon=self.rms_epsilon,
weight_bias=self.weight_bias,
add_residual=self.add_residual,
write_residual_output=self.write_residual_output,
enable_pdl=tuning.enable_pdl,
)
activation = self.capacity_m * self.hidden_size
args = (
make_fake_compact_tensor(
BFloat16,
(LAMPORT_GENERATIONS * self.tp_size * activation,),
assumed_align=16,
),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_compact_tensor(BFloat16, (self.hidden_size,), assumed_align=16),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_dynamic_compact_tensor(
BFloat16, alignment=16, divisibility=self.hidden_size
),
make_fake_compact_tensor(Int32, (2,), assumed_align=4),
Int32(self.capacity_m),
current_cu_stream(),
)
return cute.compile(device_kernel, *args)
def _create_state(self, group: dist.ProcessGroup) -> LLProtocolState:
if dist.get_world_size(group) != self.tp_size:
raise ValueError("ProcessGroup size does not match tp_size")
if dist.get_rank(group) != self.rank:
raise ValueError("ProcessGroup rank does not match rank")
device = torch.device("cuda", torch.cuda.current_device())
mailbox = SymmetricBuffer.allocate(
(
LAMPORT_GENERATIONS,
self.tp_size,
self.capacity_m,
self.hidden_size,
),
torch.bfloat16,
device,
group,
require_multicast=True,
)
mailbox.tensor.view(torch.int16).fill_(-32768)
return LLProtocolState(
contribution_mailbox=mailbox,
stage_state=torch.zeros((2,), dtype=torch.int32, device=device),
)
python/sglang/kernels/ops/communication/mnnvl_cutedsl/presets.py
+338
View File
@@ -0,0 +1,338 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Built-in routing configurations for the MNNVL CuTe DSL backend."""
import torch
from .config import (
KernelTarget,
MNNVLCuteDSLConfig,
MRangeDispatch,
ProtocolKind,
StaticProfile,
)
from .kernel_bt import (
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0,
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1,
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0,
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1,
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0,
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1,
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0,
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1,
)
from .kernel_ht import (
HT_ALL_REDUCE_GB300_TP8_H8192,
HT_ALL_REDUCE_GB300_TP16_H8192,
HT_FINALIZE_GB300_TP8_H8192_K10,
HT_FINALIZE_GB300_TP16_H8192_K10,
)
from .kernel_ll import (
LL_ALL_REDUCE_GB300_TP8_H8192,
LL_ALL_REDUCE_GB300_TP16_H8192,
LL_FINALIZE_GB300_TP8_H8192_K10,
LL_FINALIZE_GB300_TP16_H8192_K10,
)
__all__ = [
"BT_ONLY_CONFIG",
"DEFAULT_CONFIG",
"HT_ONLY_CONFIG",
"LL_ONLY_CONFIG",
]
def _target(protocol: ProtocolKind, preset: object) -> KernelTarget[object]:
return KernelTarget(protocol=protocol, preset=preset)
LL_ONLY_CONFIG = MNNVLCuteDSLConfig(
profiles=(
StaticProfile(
tp_size=8,
hidden_size=8192,
top_k=10,
dtype=torch.bfloat16,
finalize_routes=MRangeDispatch(
upper_bounds=(None,),
targets=(
_target(
ProtocolKind.LL,
LL_FINALIZE_GB300_TP8_H8192_K10,
),
),
),
all_reduce_routes=MRangeDispatch(
upper_bounds=(None,),
targets=(
_target(
ProtocolKind.LL,
LL_ALL_REDUCE_GB300_TP8_H8192,
),
),
),
),
StaticProfile(
tp_size=16,
hidden_size=8192,
top_k=10,
dtype=torch.bfloat16,
finalize_routes=MRangeDispatch(
upper_bounds=(None,),
targets=(
_target(
ProtocolKind.LL,
LL_FINALIZE_GB300_TP16_H8192_K10,
),
),
),
all_reduce_routes=MRangeDispatch(
upper_bounds=(None,),
targets=(
_target(
ProtocolKind.LL,
LL_ALL_REDUCE_GB300_TP16_H8192,
),
),
),
),
)
)
BT_ONLY_CONFIG = MNNVLCuteDSLConfig(
profiles=(
StaticProfile(
tp_size=8,
hidden_size=8192,
top_k=10,
dtype=torch.bfloat16,
finalize_routes=MRangeDispatch(
upper_bounds=(48, 1024),
targets=(
_target(
ProtocolKind.BT,
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0,
),
_target(
ProtocolKind.BT,
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1,
),
),
),
all_reduce_routes=MRangeDispatch(
upper_bounds=(256, 1024),
targets=(
_target(
ProtocolKind.BT,
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0,
),
_target(
ProtocolKind.BT,
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1,
),
),
),
),
StaticProfile(
tp_size=16,
hidden_size=8192,
top_k=10,
dtype=torch.bfloat16,
finalize_routes=MRangeDispatch(
upper_bounds=(52, 1024),
targets=(
_target(
ProtocolKind.BT,
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0,
),
_target(
ProtocolKind.BT,
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1,
),
),
),
all_reduce_routes=MRangeDispatch(
upper_bounds=(512, 1024),
targets=(
_target(
ProtocolKind.BT,
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0,
),
_target(
ProtocolKind.BT,
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1,
),
),
),
),
)
)
HT_ONLY_CONFIG = MNNVLCuteDSLConfig(
profiles=(
StaticProfile(
tp_size=8,
hidden_size=8192,
top_k=10,
dtype=torch.bfloat16,
finalize_routes=MRangeDispatch(
upper_bounds=(None,),
targets=(
_target(
ProtocolKind.HT,
HT_FINALIZE_GB300_TP8_H8192_K10,
),
),
),
all_reduce_routes=MRangeDispatch(
upper_bounds=(None,),
targets=(
_target(
ProtocolKind.HT,
HT_ALL_REDUCE_GB300_TP8_H8192,
),
),
),
),
StaticProfile(
tp_size=16,
hidden_size=8192,
top_k=10,
dtype=torch.bfloat16,
finalize_routes=MRangeDispatch(
upper_bounds=(None,),
targets=(
_target(
ProtocolKind.HT,
HT_FINALIZE_GB300_TP16_H8192_K10,
),
),
),
all_reduce_routes=MRangeDispatch(
upper_bounds=(None,),
targets=(
_target(
ProtocolKind.HT,
HT_ALL_REDUCE_GB300_TP16_H8192,
),
),
),
),
)
)
DEFAULT_CONFIG = MNNVLCuteDSLConfig(
profiles=(
StaticProfile(
tp_size=8,
hidden_size=8192,
top_k=10,
dtype=torch.bfloat16,
finalize_routes=MRangeDispatch(
upper_bounds=(23, 48, 703, None),
targets=(
_target(
ProtocolKind.LL,
LL_FINALIZE_GB300_TP8_H8192_K10,
),
_target(
ProtocolKind.BT,
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_0,
),
_target(
ProtocolKind.BT,
BT_FINALIZE_GB300_TP8_H8192_K10_PRESET_1,
),
_target(
ProtocolKind.HT,
HT_FINALIZE_GB300_TP8_H8192_K10,
),
),
),
all_reduce_routes=MRangeDispatch(
upper_bounds=(15, 256, 1024, None),
targets=(
_target(
ProtocolKind.LL,
LL_ALL_REDUCE_GB300_TP8_H8192,
),
_target(
ProtocolKind.BT,
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_0,
),
_target(
ProtocolKind.BT,
BT_ALL_REDUCE_GB300_TP8_H8192_PRESET_1,
),
_target(
ProtocolKind.HT,
HT_ALL_REDUCE_GB300_TP8_H8192,
),
),
),
),
StaticProfile(
tp_size=16,
hidden_size=8192,
top_k=10,
dtype=torch.bfloat16,
finalize_routes=MRangeDispatch(
upper_bounds=(7, 52, 703, None),
targets=(
_target(
ProtocolKind.LL,
LL_FINALIZE_GB300_TP16_H8192_K10,
),
_target(
ProtocolKind.BT,
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_0,
),
_target(
ProtocolKind.BT,
BT_FINALIZE_GB300_TP16_H8192_K10_PRESET_1,
),
_target(
ProtocolKind.HT,
HT_FINALIZE_GB300_TP16_H8192_K10,
),
),
),
all_reduce_routes=MRangeDispatch(
upper_bounds=(5, 512, 959, None),
targets=(
_target(
ProtocolKind.LL,
LL_ALL_REDUCE_GB300_TP16_H8192,
),
_target(
ProtocolKind.BT,
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_0,
),
_target(
ProtocolKind.BT,
BT_ALL_REDUCE_GB300_TP16_H8192_PRESET_1,
),
_target(
ProtocolKind.HT,
HT_ALL_REDUCE_GB300_TP16_H8192,
),
),
),
),
)
)
python/sglang/kernels/ops/communication/mnnvl_cutedsl/runtime.py
+70
View File
@@ -0,0 +1,70 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Framework-side facilities shared by production Kernel wrappers."""
import cuda.bindings.driver as cuda
import cutlass.cute as cute
import torch
from cutlass.cute.runtime import from_dlpack, make_fake_compact_tensor
class _GraphSafeDLPack:
__slots__ = ("tensor",)
def __init__(self, tensor: torch.Tensor) -> None:
self.tensor = tensor
def __dlpack__(self, stream=None):
# stream=-1 skips producer sync; CuTe launches on the current captured stream.
return self.tensor.__dlpack__(stream=-1)
def __dlpack_device__(self):
return self.tensor.__dlpack_device__()
def to_cute(tensor: torch.Tensor, alignment: int) -> cute.Tensor:
return from_dlpack(
_GraphSafeDLPack(tensor.detach()),
assumed_align=alignment,
)
def to_cute_dynamic(
tensor: torch.Tensor,
alignment: int,
*,
divisibility: int,
) -> cute.Tensor:
return to_cute(tensor, alignment).mark_compact_shape_dynamic(
mode=0,
divisibility=divisibility,
)
def make_fake_dynamic_compact_tensor(
dtype,
*,
alignment: int,
divisibility: int,
) -> cute.Tensor:
return make_fake_compact_tensor(
dtype,
(cute.sym_int32(divisibility=divisibility),),
assumed_align=alignment,
)
def current_cu_stream() -> cuda.CUstream:
return cuda.CUstream(torch.cuda.current_stream().cuda_stream)
python/sglang/kernels/ops/communication/mnnvl_cutedsl/symmetric_buffer.py
+107
View File
@@ -0,0 +1,107 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Typed ownership for one rendezvoused symmetric Tensor."""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass, field
import torch
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem
# Upstream reads this helper through a package-relative import; every
# FlashInfer this tree supports ships it, so only the spelling differs here.
from flashinfer.comm.torch_symmetric_memory import _enable_symm_mem_for_group
@dataclass(frozen=True, slots=True)
class SymmetricBuffer:
"""A symmetric Tensor and the mapping resources derived at rendezvous."""
tensor: torch.Tensor
# Keep the rendezvous mapping alive without exposing the backend handle as
# part of a Kernel State's public surface.
_handle: object = field(repr=False)
multicast_address: int | None = field(default=None, repr=False)
peer_addresses: torch.Tensor | None = field(default=None, repr=False)
@classmethod
def allocate(
cls,
shape: Sequence[int],
dtype: torch.dtype,
device: torch.device,
group: dist.ProcessGroup,
*,
require_multicast: bool = False,
materialize_peer_addresses: bool = False,
) -> SymmetricBuffer:
"""Allocate with the current SymmMem backend and verify requested mappings."""
if symm_mem.get_backend(device) is None:
raise RuntimeError(
"PyTorch Symmetric Memory has no backend for the current device"
)
_enable_symm_mem_for_group(group.group_name)
return cls.rendezvous(
symm_mem.empty(shape, dtype=dtype, device=device),
group,
require_multicast=require_multicast,
materialize_peer_addresses=materialize_peer_addresses,
)
@classmethod
def rendezvous(
cls,
tensor: torch.Tensor,
group: dist.ProcessGroup,
*,
require_multicast: bool = False,
materialize_peer_addresses: bool = False,
) -> SymmetricBuffer:
_enable_symm_mem_for_group(group.group_name)
handle = symm_mem.rendezvous(tensor, group)
multicast_address = None
if require_multicast:
multicast_address = int(handle.multicast_ptr or 0)
if not multicast_address:
raise RuntimeError("NVLink multicast mapping is unavailable")
peer_addresses = None
if materialize_peer_addresses:
# Preserve the rendezvous offset for SymmMem Pool suballocations.
addresses = [
handle.get_remote_tensor(
peer,
tensor.shape,
tensor.dtype,
).data_ptr()
for peer in range(dist.get_world_size(group))
]
if any(not address for address in addresses):
raise RuntimeError("Symmetric peer mapping is unavailable")
peer_addresses = torch.tensor(
addresses,
dtype=torch.int64,
device=tensor.device,
)
return cls(
tensor=tensor,
_handle=handle,
multicast_address=multicast_address,
peer_addresses=peer_addresses,
)
python/sglang/kernels/ops/communication/mnnvl_cutedsl_ar.py
+584
View File
@@ -0,0 +1,584 @@
# Copyright (c) 2026 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""MNNVL AllReduce fusion backend implemented with CuTe DSL.
Ported from flashinfer-ai/flashinfer#4358 at main commit 906181e (with
sibling package mnnvl_cutedsl/), pending an installable FlashInfer release
that ships flashinfer.comm.mnnvl_cutedsl. The base communication
infrastructure (mnnvl probing, pattern enum, workspace ABC) still comes
from the installed FlashInfer; the pinned release provides it.
"""
from __future__ import annotations
import logging
from typing import Any, Optional, cast
import torch
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem
# Keep the copied backend and kernel package self-contained while reusing the
# stable communication infrastructure already supplied by the serving image.
from flashinfer.comm.mnnvl import is_multicast_supported
from flashinfer.comm.trtllm_ar import AllReduceFusionPattern
from flashinfer.comm.workspace_base import AllReduceFusionWorkspace
from torch.distributed import ProcessGroup
from .mnnvl_cutedsl import DEFAULT_CONFIG, MNNVLCuteDSLConfig, ProtocolKind
from .mnnvl_cutedsl.config import StaticProfile
from .mnnvl_cutedsl.kernel_bt import BTAllReduceTuning, BTFinalizeTuning
from .mnnvl_cutedsl.kernel_bt.protocol import BTProtocol
from .mnnvl_cutedsl.kernel_ht import HTAllReduceTuning, HTFinalizeTuning
from .mnnvl_cutedsl.kernel_ht.protocol import HTProtocol
from .mnnvl_cutedsl.kernel_ll import LLAllReduceTuning, LLFinalizeTuning
from .mnnvl_cutedsl.kernel_ll.protocol import LLProtocol
logger = logging.getLogger(__name__)
__all__ = [
"MNNVLCuteDSLAllReduceFusionWorkspace",
"mnnvl_cutedsl_allreduce_fusion",
]
def _check_tensor(
tensor: torch.Tensor,
name: str,
*,
shape: tuple[int | None, ...],
dtype: torch.dtype,
device: torch.device,
alignment: int,
) -> None:
if tensor.device != device:
raise ValueError(f"{name} must be on {device}")
if tensor.dtype != dtype:
raise ValueError(f"{name} must have dtype {dtype}")
if tensor.ndim != len(shape) or any(
expected is not None and actual != expected
for actual, expected in zip(tensor.shape, shape, strict=True)
):
raise ValueError(f"{name} has an unsupported shape")
if not tensor.is_contiguous():
raise ValueError(f"{name} must be contiguous")
if tensor.data_ptr() % alignment:
raise ValueError(f"{name} must be {alignment}-byte aligned")
def _warn_pdl_mismatch(
workspace: MNNVLCuteDSLAllReduceFusionWorkspace,
pattern: int,
m: int,
launch_with_pdl: bool,
) -> None:
preset_pdl = workspace._uses_pdl(pattern, m)
if launch_with_pdl != preset_pdl:
logger.warning(
"launch_with_pdl does not match the selected MNNVL CuTe DSL "
"preset; using enable_pdl=%s",
preset_pdl,
)
class MNNVLCuteDSLAllReduceFusionWorkspace(AllReduceFusionWorkspace):
"""Compiled LL, BT, and HT protocols for one static problem shape.
Workspace construction compiles the selected kernels and must finish before
the first invocation. Calls using the same workspace must not overlap.
Feature-disabled tensor slots use internal placeholders that are not read.
"""
_destroyed: bool
def __init__(
self,
tp_size: int,
tp_rank: int,
max_token_num: int,
hidden_dim: int,
dtype: torch.dtype,
*,
group: Optional[ProcessGroup] = None,
top_k: int = 10,
rms_eps: float = 1e-6,
routed_scaling_factor: float = 1.0,
weight_bias: float = 0.0,
include_shared_expert: bool = True,
add_residual: bool = True,
write_residual_output: bool = True,
config: MNNVLCuteDSLConfig = DEFAULT_CONFIG,
) -> None:
if tp_size not in (2, 4, 8, 16):
raise ValueError("tp_size must be 2, 4, 8, or 16")
if not 0 <= tp_rank < tp_size:
raise ValueError("tp_rank must be in [0, tp_size)")
if max_token_num <= 0:
raise ValueError("max_token_num must be positive")
if dtype != torch.bfloat16:
raise ValueError("MNNVL CuTe DSL kernels only support torch.bfloat16")
if not torch.cuda.is_available():
raise RuntimeError("MNNVL CuTe DSL kernels require CUDA")
device = torch.device("cuda", torch.cuda.current_device())
if torch.cuda.get_device_capability(device)[0] < 10:
raise RuntimeError("MNNVL CuTe DSL kernels require a Blackwell GPU")
if symm_mem.get_backend(device) is None:
raise RuntimeError("PyTorch Symmetric Memory is unavailable")
if not is_multicast_supported(device.index):
raise RuntimeError("NVLink multicast is unavailable")
if group is None:
if not dist.is_initialized():
raise ValueError("A ProcessGroup is required before initialization")
group = dist.group.WORLD
if dist.get_world_size(group) != tp_size:
raise ValueError("ProcessGroup size does not match tp_size")
if dist.get_rank(group) != tp_rank:
raise ValueError("ProcessGroup rank does not match tp_rank")
super().__init__(tp_size, tp_rank)
self._protocols: dict[ProtocolKind, LLProtocol | BTProtocol | HTProtocol] = {}
self.max_token_num = max_token_num
self.hidden_dim = hidden_dim
self.top_k = top_k
self.dtype = dtype
self.group = group
self.rms_eps = rms_eps
self.routed_scaling_factor = routed_scaling_factor
self.weight_bias = weight_bias
self.include_shared_expert = include_shared_expert
self.add_residual = add_residual
self.write_residual_output = write_residual_output
self.config = config
self.profile = config.resolve(
tp_size=tp_size,
hidden_size=hidden_dim,
top_k=top_k,
dtype=dtype,
capacity_m=max_token_num,
)
for protocol in (ProtocolKind.LL, ProtocolKind.BT, ProtocolKind.HT):
capacity = self.profile.protocol_capacity(
protocol, capacity_m=max_token_num
)
if capacity is None:
continue
finalize_tunings = self._tunings(
self.profile, protocol, finalize=True, capacity_m=capacity
)
all_reduce_tunings = self._tunings(
self.profile, protocol, finalize=False, capacity_m=capacity
)
common = dict(
hidden_size=hidden_dim,
top_k=top_k,
tp_size=tp_size,
rank=tp_rank,
capacity_m=capacity,
rms_epsilon=rms_eps,
routed_scaling_factor=routed_scaling_factor,
weight_bias=weight_bias,
include_shared_expert=include_shared_expert,
add_residual=add_residual,
write_residual_output=write_residual_output,
group=group,
)
instance: LLProtocol | BTProtocol | HTProtocol
if protocol is ProtocolKind.LL:
instance = LLProtocol(
**common,
finalize_tunings=finalize_tunings,
all_reduce_tunings=all_reduce_tunings,
)
elif protocol is ProtocolKind.BT:
instance = BTProtocol(
**common,
finalize_tunings=finalize_tunings,
all_reduce_tunings=all_reduce_tunings,
)
else:
instance = HTProtocol(
**common,
finalize_tunings=finalize_tunings,
all_reduce_tunings=all_reduce_tunings,
)
self._protocols[protocol] = instance
torch.cuda.synchronize(device)
dist.barrier(group=group)
@staticmethod
def _tunings(
profile: StaticProfile,
protocol: ProtocolKind,
*,
finalize: bool,
capacity_m: int,
) -> tuple:
routes = profile.finalize_routes if finalize else profile.all_reduce_routes
tunings = tuple(
dict.fromkeys(
target.preset
for target in routes.targets_for_capacity(capacity_m)
if target.protocol is protocol
)
)
expected_type = {
(ProtocolKind.LL, True): LLFinalizeTuning,
(ProtocolKind.LL, False): LLAllReduceTuning,
(ProtocolKind.BT, True): BTFinalizeTuning,
(ProtocolKind.BT, False): BTAllReduceTuning,
(ProtocolKind.HT, True): HTFinalizeTuning,
(ProtocolKind.HT, False): HTAllReduceTuning,
}[(protocol, finalize)]
if not all(isinstance(tuning, expected_type) for tuning in tunings):
path = "finalize" if finalize else "all-reduce"
raise TypeError(f"Invalid {protocol.value} {path} preset")
return tunings
def _uses_pdl(self, pattern: int, m: int) -> bool:
if pattern == AllReduceFusionPattern.kMoEFinalizeARResidualRMSNorm:
target = self.profile.finalize_routes.select(m)
elif pattern == AllReduceFusionPattern.kARResidualRMSNorm:
target = self.profile.all_reduce_routes.select(m)
else:
raise NotImplementedError("Unsupported MNNVL CuTe DSL fusion pattern")
preset = cast(Any, target.preset)
enabled = getattr(preset, "enable_pdl", None)
if enabled is None:
enabled = preset.collective.enable_pdl
return bool(enabled)
@property
def backend(self) -> str:
return "mnnvl-cutedsl"
def is_buffer_size_sufficient(
self,
tp_size: int,
num_tokens: int,
hidden_dim: int,
dtype: torch.dtype,
use_oneshot=None,
) -> bool:
del use_oneshot
return (
tp_size == self.world_size
and num_tokens <= self.max_token_num
and hidden_dim == self.hidden_dim
and dtype == self.dtype
)
def _finalize_all_reduce_rms_norm(
self,
routed_output: torch.Tensor,
expert_weights: torch.Tensor,
permuted_indices: torch.Tensor,
shared_output: torch.Tensor | None,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
m: int,
*,
norm_output: torch.Tensor | None,
residual_output: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
target = self.profile.finalize_routes.select(m)
protocol = cast(Any, self._protocols[target.protocol])
kernel = protocol.finalize_kernels[target.preset]
return kernel(
routed_output,
expert_weights,
permuted_indices,
shared_output,
residual_source,
gamma,
m,
state=protocol.state,
norm_output=norm_output,
residual_output=residual_output,
)
def _all_reduce_rms_norm(
self,
local_contribution: torch.Tensor,
residual_source: torch.Tensor | None,
gamma: torch.Tensor,
m: int,
*,
norm_output: torch.Tensor | None,
residual_output: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
target = self.profile.all_reduce_routes.select(m)
protocol = cast(Any, self._protocols[target.protocol])
kernel = protocol.all_reduce_kernels[target.preset]
return kernel(
local_contribution,
residual_source,
gamma,
m,
state=protocol.state,
norm_output=norm_output,
residual_output=residual_output,
)
def destroy(self) -> None:
if self._destroyed:
return
self._protocols.clear()
self._destroyed = True
def _mnnvl_cutedsl_allreduce_fusion(
input: torch.Tensor,
workspace: MNNVLCuteDSLAllReduceFusionWorkspace,
pattern: int,
*,
launch_with_pdl: bool,
output: Optional[torch.Tensor] = None,
residual_in: Optional[torch.Tensor] = None,
residual_out: Optional[torch.Tensor] = None,
norm_out: Optional[torch.Tensor] = None,
quant_out: Optional[torch.Tensor] = None,
scale_out: Optional[torch.Tensor] = None,
rms_gamma: Optional[torch.Tensor] = None,
rms_eps: float = 1e-6,
scale_factor: Optional[torch.Tensor | float] = None,
layout_code: Optional[int] = None,
use_oneshot: Optional[bool] = None,
fp32_acc: bool = False,
moe_reduction_device_num_experts: Optional[int] = None,
moe_reduction_scale_input: Optional[torch.Tensor] = None,
moe_reduction_active_experts_token_input: Optional[torch.Tensor] = None,
moe_reduction_token_input: Optional[torch.Tensor] = None,
weight_bias: float = 0.0,
expanded_idx_to_permuted_idx: Optional[torch.Tensor] = None,
expert_scale_factor: Optional[torch.Tensor] = None,
shared_expert_output: Optional[torch.Tensor] = None,
block_quant_group_size: Optional[int] = None,
) -> torch.Tensor:
if workspace._destroyed:
raise RuntimeError(
"The MNNVLCuteDSLAllReduceFusionWorkspace has been destroyed"
)
if pattern not in (
AllReduceFusionPattern.kARResidualRMSNorm,
AllReduceFusionPattern.kMoEFinalizeARResidualRMSNorm,
):
raise NotImplementedError("Unsupported MNNVL CuTe DSL fusion pattern")
unsupported = [
name
for name, value in (
("output", output),
("quant_out", quant_out),
("scale_out", scale_out),
("scale_factor", scale_factor),
("layout_code", layout_code),
("use_oneshot", use_oneshot),
("block_quant_group_size", block_quant_group_size),
("moe_reduction_scale_input", moe_reduction_scale_input),
(
"moe_reduction_active_experts_token_input",
moe_reduction_active_experts_token_input,
),
("moe_reduction_token_input", moe_reduction_token_input),
)
if value is not None
]
if fp32_acc:
unsupported.append("fp32_acc")
if moe_reduction_device_num_experts is not None:
unsupported.append("moe_reduction_device_num_experts")
if unsupported:
raise ValueError("MNNVL CuTe DSL does not support: " + ", ".join(unsupported))
if rms_eps != workspace.rms_eps:
raise ValueError("rms_eps does not match the compiled workspace")
if weight_bias != workspace.weight_bias:
raise ValueError("weight_bias does not match the compiled workspace")
if rms_gamma is None:
raise ValueError("rms_gamma is required")
if workspace.add_residual and residual_in is None:
raise ValueError("residual_in is required by the compiled workspace")
if not workspace.add_residual and residual_in is not None:
raise ValueError("residual_in must be None for this compiled workspace")
if not workspace.write_residual_output and residual_out is not None:
raise ValueError("residual_out must be None for this compiled workspace")
device = torch.device("cuda", torch.cuda.current_device())
hidden = workspace.hidden_dim
_check_tensor(
rms_gamma,
"rms_gamma",
shape=(hidden,),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
if pattern == AllReduceFusionPattern.kARResidualRMSNorm:
if any(
value is not None
for value in (
expanded_idx_to_permuted_idx,
expert_scale_factor,
shared_expert_output,
)
):
raise ValueError("MoE finalize operands require the finalize pattern")
_check_tensor(
input,
"input",
shape=(None, hidden),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
m = input.shape[0]
if not 1 <= m <= workspace.max_token_num:
raise ValueError("input token count exceeds workspace capacity")
_warn_pdl_mismatch(workspace, pattern, m, launch_with_pdl)
if residual_in is not None:
_check_tensor(
residual_in,
"residual_in",
shape=(m, hidden),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
if norm_out is not None:
_check_tensor(
norm_out,
"norm_out",
shape=(m, hidden),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
if residual_out is not None:
_check_tensor(
residual_out,
"residual_out",
shape=(m, hidden),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
norm_out, _ = workspace._all_reduce_rms_norm(
input,
residual_in,
rms_gamma,
input.shape[0],
norm_output=norm_out,
residual_output=residual_out,
)
return norm_out
if pattern == AllReduceFusionPattern.kMoEFinalizeARResidualRMSNorm:
if expanded_idx_to_permuted_idx is None:
raise ValueError("expanded_idx_to_permuted_idx is required")
if expert_scale_factor is None:
raise ValueError("expert_scale_factor is required")
if workspace.include_shared_expert and shared_expert_output is None:
raise ValueError(
"shared_expert_output is required by the compiled workspace"
)
if not workspace.include_shared_expert and shared_expert_output is not None:
raise ValueError(
"shared_expert_output must be None for this compiled workspace"
)
m = expanded_idx_to_permuted_idx.shape[0]
if not 1 <= m <= workspace.max_token_num:
raise ValueError("input token count exceeds workspace capacity")
_warn_pdl_mismatch(workspace, pattern, m, launch_with_pdl)
_check_tensor(
input,
"input",
shape=(None, hidden),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
_check_tensor(
expert_scale_factor,
"expert_scale_factor",
shape=(m, workspace.top_k),
dtype=torch.bfloat16,
device=device,
alignment=2,
)
_check_tensor(
expanded_idx_to_permuted_idx,
"expanded_idx_to_permuted_idx",
shape=(m, workspace.top_k),
dtype=torch.int32,
device=device,
alignment=4,
)
if shared_expert_output is not None:
_check_tensor(
shared_expert_output,
"shared_expert_output",
shape=(m, hidden),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
if residual_in is not None:
_check_tensor(
residual_in,
"residual_in",
shape=(m, hidden),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
if norm_out is not None:
_check_tensor(
norm_out,
"norm_out",
shape=(m, hidden),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
if residual_out is not None:
_check_tensor(
residual_out,
"residual_out",
shape=(m, hidden),
dtype=torch.bfloat16,
device=device,
alignment=16,
)
norm_out, _ = workspace._finalize_all_reduce_rms_norm(
input,
expert_scale_factor,
expanded_idx_to_permuted_idx,
shared_expert_output,
residual_in,
rms_gamma,
m,
norm_output=norm_out,
residual_output=residual_out,
)
return norm_out
raise AssertionError("unreachable")
# Keep the upstream private name intact while exposing SGLang's backend entry point;
# upstream exports it through flashinfer.comm.allreduce_fusion.
mnnvl_cutedsl_allreduce_fusion = _mnnvl_cutedsl_allreduce_fusion
python/sglang/kernels/ops/elementwise/elementwise.py
+51 -18
View File
@@ -465,9 +465,10 @@ def _fused_gate_sigmoid_mul_add_kernel(
hidden_states_ptr, # [num_tokens, hidden_dim]
gate_weight_ptr, # [hidden_dim]
shared_output_ptr, # [num_tokens, hidden_dim]
final_hidden_states_ptr, # [num_tokens, hidden_dim]
output_ptr, # [num_tokens, hidden_dim], optionally also the addend
hidden_dim: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
DO_ADD: tl.constexpr = True,
USE_PDL: tl.constexpr = False,
):
pid = tl.program_id(axis=0).to(tl.int64)
@@ -487,41 +488,39 @@ def _fused_gate_sigmoid_mul_add_kernel(
s = tl.load(shared_output_ptr + row_offset + offsets, mask=mask, other=0.0).to(
tl.float32
)
f = tl.load(
final_hidden_states_ptr + row_offset + offsets, mask=mask, other=0.0
).to(tl.float32)
if DO_ADD:
f = tl.load(output_ptr + row_offset + offsets, mask=mask, other=0.0).to(
tl.float32
)
if USE_PDL:
tl.extra.cuda.gdc_launch_dependents()
gate_val = tl.sigmoid(tl.sum(h * w, axis=0))
result = f + gate_val * s
result = gate_val * s
if DO_ADD:
result += f
tl.store(final_hidden_states_ptr + row_offset + offsets, result, mask=mask)
tl.store(output_ptr + row_offset + offsets, result, mask=mask)
def fused_gate_sigmoid_mul_add(
def _launch_fused_gate_sigmoid_mul(
hidden_states: torch.Tensor,
gate_weight: torch.Tensor,
shared_output: torch.Tensor,
final_hidden_states: torch.Tensor,
output: torch.Tensor,
*,
do_add: bool,
) -> None:
"""
Fused gate-sigmoid-mul-add for MoE shared expert gating.
Equivalent to:
gate = hidden_states @ gate_weight
final_hidden_states += sigmoid(gate).unsqueeze(1) * shared_output
"""
assert hidden_states.is_contiguous(), "hidden_states must be contiguous"
assert gate_weight.is_contiguous(), "gate_weight must be contiguous"
assert shared_output.is_contiguous(), "shared_output must be contiguous"
assert final_hidden_states.is_contiguous(), "final_hidden_states must be contiguous"
assert output.is_contiguous(), "output must be contiguous"
num_tokens, hidden_dim = hidden_states.shape
assert gate_weight.shape == (hidden_dim,)
assert shared_output.shape == (num_tokens, hidden_dim)
assert final_hidden_states.shape == (num_tokens, hidden_dim)
assert output.shape == (num_tokens, hidden_dim)
max_warps = 16 if _is_hip else 32
config = {
@@ -541,9 +540,43 @@ def fused_gate_sigmoid_mul_add(
hidden_states,
gate_weight,
shared_output,
final_hidden_states,
output,
hidden_dim=hidden_dim,
DO_ADD=do_add,
USE_PDL=use_pdl,
**config,
**pdl_kwargs,
)
def fused_gate_sigmoid_mul(
hidden_states: torch.Tensor,
gate_weight: torch.Tensor,
shared_output: torch.Tensor,
) -> torch.Tensor:
"""Materialize the gated shared-expert contribution without an add/copy."""
output = torch.empty_like(shared_output)
_launch_fused_gate_sigmoid_mul(
hidden_states,
gate_weight,
shared_output,
output,
do_add=False,
)
return output
def fused_gate_sigmoid_mul_add(
hidden_states: torch.Tensor,
gate_weight: torch.Tensor,
shared_output: torch.Tensor,
final_hidden_states: torch.Tensor,
) -> None:
"""Add the gated shared-expert contribution to routed-expert output."""
_launch_fused_gate_sigmoid_mul(
hidden_states,
gate_weight,
shared_output,
final_hidden_states,
do_add=True,
)
python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_direct.py
+367
View File
@@ -0,0 +1,367 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
# Vendored from flashinfer-ai/flashinfer@629147317d4149a12e53bcef27808bac380c283f.
"""Register-prefetch BF16 GEMM for low-M, long-K decode shapes.
The kernel keeps a complete output dot product inside one CTA and reuses each
prefetched B value across several public-M rows. It is intentionally a
separate autotuner runner from the Blackwell tensor-core split-K kernel: the
two algorithms have different useful shape regions and tactic spaces.
"""
from __future__ import annotations
import dataclasses
import functools
import cuda.bindings.driver as _cuda
import cutlass
import cutlass.cute as cute
import torch as _torch
from cutlass import const_expr
from cutlass.cute import experimental as cute_ext
from cutlass.cute.runtime import from_dlpack
_VECTOR_WIDTH = 8
_SUPPORTED_BLOCK_SIZES = (32, 64, 96, 128, 192, 256, 384)
_SUPPORTED_OUTPUTS_PER_BLOCK = (1, 2, 4)
_MAX_M = 32
_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64"
@dataclasses.dataclass(frozen=True)
class DirectTactic:
"""One direct-kernel specialization."""
block_size: int
outputs_per_block: int
rows_per_block: int
def _default_rows_per_block(m: int) -> int:
if m <= 8:
return m
return next(rows for rows in (8, 4, 2, 1) if m % rows == 0)
def validate_tactic(tactic: DirectTactic, m: int, n: int, k: int) -> None:
"""Reject a direct tactic that cannot serve ``(m, n, k)``."""
if tactic.block_size not in _SUPPORTED_BLOCK_SIZES:
raise ValueError(f"unsupported block_size={tactic.block_size}")
if tactic.outputs_per_block not in _SUPPORTED_OUTPUTS_PER_BLOCK:
raise ValueError(f"unsupported outputs_per_block={tactic.outputs_per_block}")
if not 1 <= m <= _MAX_M:
raise ValueError(f"direct GEMM requires 1 <= M <= {_MAX_M}, got {m}")
if not 1 <= tactic.rows_per_block <= m or m % tactic.rows_per_block:
raise ValueError(f"rows_per_block={tactic.rows_per_block} must divide M={m}")
if n <= 0 or n % tactic.outputs_per_block:
raise ValueError(
f"N={n} must be divisible by outputs_per_block={tactic.outputs_per_block}"
)
k_tile = tactic.block_size * _VECTOR_WIDTH
if k <= 0 or k % k_tile:
raise ValueError(f"K={k} must be divisible by {k_tile}")
def default_tactic(m: int, n: int, k: int) -> DirectTactic:
"""Choose the measured register-prefetch fallback tactic."""
block_size = next(
(
block
for block in (256, 192, 128, 96, 64, 32)
if k % (block * _VECTOR_WIDTH) == 0
),
None,
)
if block_size is None:
raise ValueError("direct GEMM requires a supported 16-byte K tiling")
outputs_per_block = next(outputs for outputs in (2, 1) if n % outputs == 0)
tactic = DirectTactic(
block_size,
outputs_per_block,
_default_rows_per_block(m),
)
validate_tactic(tactic, m, n, k)
return tactic
def autotune_tactics(m: int, n: int, k: int) -> list[DirectTactic]:
"""Enumerate the compact tactic space used by FlashInfer autotuning.
Block sizes cover every configuration exercised in the H100/B200 sweep;
output grouping spans the measured 1/2/4-column choices. Row tiling stays
at the occupancy-oriented default to keep JIT cost bounded.
"""
try:
default = default_tactic(m, n, k)
except ValueError:
return []
tactics = [default]
for block_size in _SUPPORTED_BLOCK_SIZES:
for outputs_per_block in _SUPPORTED_OUTPUTS_PER_BLOCK:
tactic = DirectTactic(
block_size,
outputs_per_block,
default.rows_per_block,
)
try:
validate_tactic(tactic, m, n, k)
except ValueError:
continue
tactics.append(tactic)
return list(dict.fromkeys(tactics))
def prefer_direct_bf16_gemm_sm100(m: int, n: int, k: int) -> bool:
"""Return the conservative B200 no-autotune crossover heuristic.
The three bands are a compact fit to a warm/cold sweep over M=1..16,24,32,
18 N values, and 11 K values. This is deliberately not a blanket rule for
K=8192: direct wins only where public M and N leave the tensor-core path
with too little independent output work.
"""
return k == 8192 and (
(m == 1 and n <= 4608) or (m <= 4 and n <= 512) or (m <= 8 and n <= 256)
)
class DirectDenseGemmKernel:
"""K-specialized direct GEMM with whole-mainloop vector prefetch."""
def __init__(
self,
*,
element_type,
num_rows: int,
k_extent: int,
tactic: DirectTactic,
use_pdl: bool,
) -> None:
validate_tactic(tactic, num_rows, tactic.outputs_per_block, k_extent)
self.element_type = element_type
self.num_rows = num_rows
self.rows_per_block = tactic.rows_per_block
self.k_extent = k_extent
self.block_size = tactic.block_size
self.outputs_per_block = tactic.outputs_per_block
self.vector_width = _VECTOR_WIDTH
self.use_pdl = use_pdl
self.num_warps = tactic.block_size // cute.arch.WARP_SIZE
self.num_k_tiles = k_extent // (tactic.block_size * _VECTOR_WIDTH)
@cute.jit
def __call__(
self,
gA: cute.Tensor,
gB: cute.Tensor,
gC: cute.Tensor,
stream: _cuda.CUstream,
) -> None:
n = cute.size(gB, mode=[0])
copy_a = cute.make_copy_atom(
cute.nvgpu.CopyG2ROp(),
self.element_type,
num_bits_per_copy=self.vector_width * self.element_type.width,
load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS,
)
copy_b = cute.make_copy_atom(
cute.nvgpu.CopyG2ROp(),
self.element_type,
num_bits_per_copy=self.vector_width * self.element_type.width,
load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING,
)
self.kernel(gA, gB, gC, copy_a, copy_b).launch(
grid=[
cute.ceil_div(n, self.outputs_per_block),
self.num_rows // self.rows_per_block,
1,
],
block=[self.block_size, 1, 1],
smem=self.rows_per_block * self.outputs_per_block * self.num_warps * 4,
stream=stream,
use_pdl=self.use_pdl,
min_blocks_per_mp=1,
)
@cute.kernel
def kernel(
self,
gA: cute.Tensor,
gB: cute.Tensor,
gC: cute.Tensor,
copy_a: cute.CopyAtom,
copy_b: cute.CopyAtom,
) -> None:
tidx, _, _ = cute.arch.thread_idx()
block_idx, block_m, _ = cute.arch.block_idx()
warp_idx = cute.arch.warp_idx()
num_rows: cutlass.Constexpr = self.rows_per_block
outputs_per_block: cutlass.Constexpr = self.outputs_per_block
vector_width: cutlass.Constexpr = self.vector_width
block_size: cutlass.Constexpr = self.block_size
num_warps: cutlass.Constexpr = self.num_warps
num_k_tiles: cutlass.Constexpr = self.num_k_tiles
acc = cute.make_rmem_tensor(
cute.make_layout(
(num_rows, outputs_per_block), stride=(outputs_per_block, 1)
),
cutlass.Float32,
)
acc.fill(0.0)
if const_expr(self.use_pdl):
cute.arch.griddepcontrol_wait()
n_base = block_idx * outputs_per_block
m_base = block_m * num_rows
gA_vec = cute.logical_divide(gA, (None, vector_width))
gB_vec = cute.logical_divide(gB, (None, vector_width))
tA_all = cute.logical_divide(gA_vec, (None, (None, block_size)))
tB_all = cute.logical_divide(gB_vec, (None, (None, block_size)))
tA = tA_all[None, (None, (tidx, None))]
b_regs = cute.make_rmem_tensor(
cute.make_layout(
(outputs_per_block, num_k_tiles, vector_width),
stride=(num_k_tiles * vector_width, vector_width, 1),
),
self.element_type,
)
for ni in cutlass.range_constexpr(outputs_per_block):
tB = tB_all[n_base + ni, (None, (tidx, None))]
for k_tile in cutlass.range_constexpr(num_k_tiles):
cute.copy(copy_b, tB[None, k_tile], b_regs[ni, k_tile, None])
a_regs = cute.make_rmem_tensor(
cute.make_layout((num_k_tiles, vector_width), stride=(vector_width, 1)),
self.element_type,
)
for mi in cutlass.range_constexpr(num_rows):
for k_tile in cutlass.range_constexpr(num_k_tiles):
cute.copy(
copy_a,
tA[m_base + mi, None, k_tile],
a_regs[k_tile, None],
)
for k_tile in cutlass.range_constexpr(num_k_tiles):
for vi in cutlass.range_constexpr(vector_width):
a_value = a_regs[k_tile, vi].to(cutlass.Float32)
for ni in cutlass.range_constexpr(outputs_per_block):
acc[mi, ni] = acc[mi, ni] + a_value * b_regs[ni, k_tile, vi].to(
cutlass.Float32
)
for mi in cutlass.range_constexpr(num_rows):
for ni in cutlass.range_constexpr(outputs_per_block):
acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])
smem_layout = cute.make_layout(
(num_rows, outputs_per_block, num_warps),
stride=(outputs_per_block * num_warps, num_warps, 1),
)
smem = cutlass.utils.SmemAllocator()
partials = smem.allocate_tensor(cutlass.Float32, smem_layout, byte_alignment=16)
with cute.arch.elect_one():
for mi in cutlass.range_constexpr(num_rows):
for ni in cutlass.range_constexpr(outputs_per_block):
partials[mi, ni, warp_idx] = acc[mi, ni]
cute.arch.sync_threads()
if tidx == 0:
for mi in cutlass.range_constexpr(num_rows):
for ni in cutlass.range_constexpr(outputs_per_block):
total = cutlass.Float32(0.0)
for warp in cutlass.range_constexpr(num_warps):
total = total + partials[mi, ni, warp]
gC[m_base + mi, n_base + ni] = total.to(self.element_type)
if const_expr(self.use_pdl):
cute.arch.griddepcontrol_launch_dependents()
def _from_dlpack_static(tensor: _torch.Tensor):
# K is specialized and the row stride must retain its 16-byte divisibility
# for the verifier to accept vectorized G2R copies.
return from_dlpack(tensor, assumed_align=32)
def _make_compile_repr_tensors(dtype, m: int, n: int, k: int):
return tuple(
_from_dlpack_static(tensor)
for tensor in (
_torch.empty((m, k), dtype=dtype, device="cuda"),
_torch.empty((n, k), dtype=dtype, device="cuda"),
_torch.empty((m, n), dtype=dtype, device="cuda"),
)
)
@functools.cache
def _get_compiled_direct_kernel(
dtype,
m: int,
n: int,
k: int,
tactic: DirectTactic,
use_pdl: bool,
):
if dtype != _torch.bfloat16:
raise ValueError(f"direct GEMM supports BF16; got {dtype}")
kernel = DirectDenseGemmKernel(
element_type=cutlass.BFloat16,
num_rows=m,
k_extent=k,
tactic=tactic,
use_pdl=use_pdl,
)
tensors = _make_compile_repr_tensors(dtype, m, n, k)
stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream)
return cute_ext.compile(kernel, *tensors, stream, options=_COMPILE_OPTIONS)
def _validate_runtime_tensors(a, b, out, tactic: DirectTactic):
if any(not isinstance(tensor, _torch.Tensor) for tensor in (a, b, out)):
raise ValueError("a, b, and out must be torch tensors")
if a.ndim != 2 or b.ndim != 2 or out.ndim != 2:
raise ValueError("direct GEMM accepts only 2D tensors")
if a.device.type != "cuda" or b.device != a.device or out.device != a.device:
raise ValueError("a, b, and out must be on the same CUDA device")
if a.dtype != _torch.bfloat16 or b.dtype != a.dtype or out.dtype != a.dtype:
raise ValueError("a, b, and out must share BF16 dtype")
if not a.is_contiguous() or not b.T.is_contiguous() or not out.is_contiguous():
raise ValueError("direct GEMM requires row-major A/out and column-major B")
if any(tensor.data_ptr() % 32 for tensor in (a, b, out)):
raise ValueError("a, b, and out must be 32-byte aligned")
m, k = a.shape
if b.shape[0] != k:
raise ValueError(
f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}"
)
n = b.shape[1]
if out.shape != (m, n):
raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}")
validate_tactic(tactic, m, n, k)
return m, n, k
def run_direct_dense(a, b, out, pdl: bool, tactic: DirectTactic):
"""Run direct ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts."""
m, n, k = _validate_runtime_tensors(a, b, out, tactic)
compiled = _get_compiled_direct_kernel(a.dtype, m, n, k, tactic, pdl)
tensors = tuple(_from_dlpack_static(tensor) for tensor in (a, b.T, out))
stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream)
compiled(*tensors, stream)
return out
__all__ = [
"DirectTactic",
"autotune_tactics",
"default_tactic",
"prefer_direct_bf16_gemm_sm100",
"run_direct_dense",
]
python/sglang/kernels/ops/gemm/flashinfer_pr4266_dense_bf16_gemm_sm100_splitk.py
+1041
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/ops/moe/ep_moe_kernels.py
+40 -21
View File
@@ -911,9 +911,7 @@ def post_reorder_deepgemm_triton_kernel(
BLOCK_SIZE: tl.constexpr,
NUM_STAGES: tl.constexpr,
):
"""`expert_id >= 0` includes the shared expert at num_experts (padding=-1); don't
switch to the cutlass `!= num_local_experts` gate. routed_scaling_factor is folded into the store.
"""
"""Accumulate valid permuted rows; routed_scaling_factor is folded into the store."""
OutDtype = output_ptr.dtype.element_ty
offset = BLOCK_SIZE * tl.program_id(1) + tl.arange(0, BLOCK_SIZE)
@@ -935,9 +933,8 @@ def post_reorder_deepgemm_triton_kernel(
sum_vec = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
for idx in range(topk):
expert_id = tl.load(token_topk_ids_ptr + idx)
if expert_id >= 0:
dst_idx = tl.load(token_src2dst_ptr + idx).to(tl.int64)
dst_idx = tl.load(token_src2dst_ptr + idx).to(tl.int64)
if dst_idx >= 0:
weight_scale = tl.load(token_topk_weights_ptr + idx).to(tl.float32)
load_ptr_offs = down_output_ptr_offs + dst_idx * hidden_size
in_data = tl.load(load_ptr_offs, mask=mask).to(tl.float32)
@@ -1074,6 +1071,8 @@ def _fwd_kernel_ep_scatter_2(
output_index,
output_index_stride0,
output_index_stride1,
expert_start,
num_experts,
topk_num: tl.constexpr,
HIDDEN_SIZE: tl.constexpr,
HIDDEN_SIZE_PAD: tl.constexpr,
@@ -1105,17 +1104,24 @@ def _fwd_kernel_ep_scatter_2(
for topk_idx_int32 in tl.range(0, topk_num, 1, num_stages=4):
topk_index = topk_idx_int32.to(tl.int64)
expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index)
if expert_id >= 0:
global_expert_id = tl.load(
recv_topk + token_id * recv_topk_stride0 + topk_index
)
expert_id = global_expert_id - expert_start
output_index_ptr = (
output_index + token_id * output_index_stride0 + topk_index
)
valid = (expert_id >= 0) & (expert_id < num_experts)
# The post-permute path uses this index as the validity sentinel.
# Initialize non-local/padding lanes in this same kernel.
tl.store(output_index_ptr, -1)
if valid:
dest_token_index_int32 = tl.atomic_add(
expert_start_loc + expert_id, 1, sem=ATOMIC_ADD_SEM
)
dest_token_index = dest_token_index_int32.to(tl.int64)
tl.store(
output_index + token_id * output_index_stride0 + topk_index,
dest_token_index_int32,
)
tl.store(output_index_ptr, dest_token_index_int32)
output_tensor_ptr = (
output_tensor + dest_token_index * output_tensor_stride0
)
@@ -1148,6 +1154,7 @@ def ep_scatter(
output_index: torch.Tensor,
scale_ue8m0: bool = False,
quant_block_size: int = 128,
expert_start: int = 0,
):
BLOCK_E = 128 # token num of per expert is aligned to 128
BLOCK_D = quant_block_size # block size of quantization
@@ -1208,6 +1215,8 @@ def ep_scatter(
output_index,
output_index.stride(0),
output_index.stride(1),
expert_start,
num_experts,
topk_num=recv_topk.shape[1],
num_warps=num_warps,
HIDDEN_SIZE=hidden_size,
@@ -1540,12 +1549,13 @@ def tma_align_input_scale(input_scale: torch.Tensor):
@triton.jit
def fused_moe_dispatch_index_triton_kernel(
topk_ids_ptr, # flat (num_toks,) int32; -1 = padding (drives the `expert >= 0` gate)
topk_ids_ptr, # flat (num_toks,) int32; global or already-local expert IDs
src2dst_ptr,
masked_m_ptr,
m_max,
num_toks,
num_experts,
expert_start,
BLOCK_SIZE: tl.constexpr,
ZERO_INIT: tl.constexpr,
):
@@ -1563,19 +1573,23 @@ def fused_moe_dispatch_index_triton_kernel(
tl.debug_barrier()
offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offs < num_toks
expert = tl.load(topk_ids_ptr + offs, mask=mask, other=-1)
valid = mask & (expert >= 0)
global_expert = tl.load(topk_ids_ptr + offs, mask=mask, other=-1)
expert = global_expert - expert_start
valid = mask & (expert >= 0) & (expert < num_experts)
# Clamp masked lanes to bin 0 so the masked atomic's pointer stays in-bounds.
expert_safe = tl.where(valid, expert, 0)
offset = tl.atomic_add(masked_m_ptr + expert_safe, 1, mask=valid)
dst = expert_safe * m_max + offset
tl.store(src2dst_ptr + offs, dst, mask=valid)
# post_reorder checks src2dst, so mark padding/non-local lanes -1 here to
# prevent uninitialized offsets from adding bogus expert contributions.
tl.store(src2dst_ptr + offs, tl.where(valid, dst, -1), mask=mask)
def fused_moe_dispatch_index(
topk_ids: torch.Tensor,
num_local_experts: int,
m_max: int,
expert_start: int = 0,
) -> Tuple[torch.Tensor, torch.Tensor]:
num_toks = topk_ids.numel()
src2dst = torch.empty(num_toks, device=topk_ids.device, dtype=torch.int32)
@@ -1600,6 +1614,7 @@ def fused_moe_dispatch_index(
m_max,
num_toks,
num_local_experts,
expert_start,
BLOCK_SIZE=BLOCK_SIZE,
ZERO_INIT=single_block,
)
@@ -1635,9 +1650,8 @@ def fill_gateup_input_triton_kernel(
vec = tl.arange(0, BLOCK_SIZE)
for idx in range(topk):
expert_id = tl.load(topk_ids_ptr + idx)
if expert_id >= 0:
dst_idx_int32 = tl.load(src2dst_ptr + idx)
dst_idx_int32 = tl.load(src2dst_ptr + idx)
if dst_idx_int32 >= 0:
dst_idx = dst_idx_int32.to(tl.int64)
dst_ptr = gateup_input_ptr + dst_idx * hidden_size
for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
@@ -1679,6 +1693,7 @@ def moe_ep_deepgemm_preprocess(
block_shape,
output_dtype: torch.dtype = torch.float8_e4m3fn,
use_mxfp8: bool = False,
expert_start: int = 0,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
# For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m}) https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/jit_kernels/m_grouped_gemm.py#L165
m_max = (hidden_states.size(0) // 256 + 1) * 256
@@ -1698,12 +1713,16 @@ def moe_ep_deepgemm_preprocess(
# correctness is unconditional (m_cap >= max(masked_m) by
# construction, and the final src2dst below is built with the same
# capped stride).
masked_m_probe, _ = fused_moe_dispatch_index(topk_ids, num_local_experts, m_max)
masked_m_probe, _ = fused_moe_dispatch_index(
topk_ids, num_local_experts, m_max, expert_start=expert_start
)
m_cap = (int(masked_m_probe.max().item()) + 255) // 256 * 256
m_max = min(m_max, max(m_cap, 256))
expected_m = (topk_ids.numel() - 1) // num_local_experts + 1
masked_m, src2dst = fused_moe_dispatch_index(topk_ids, num_local_experts, m_max)
masked_m, src2dst = fused_moe_dispatch_index(
topk_ids, num_local_experts, m_max, expert_start=expert_start
)
gateup_input = torch.empty(
(num_local_experts, m_max, hidden_states.size(1)),
python/sglang/kernels/ops/quantization/minimax_quant_ue8m0.py
+3 -2
View File
@@ -89,8 +89,9 @@ def per_token_quant_fp8_ue8m0_scatter(
then writes them to each of the token's ``topk`` destination rows:
``gateup_input`` fp8 ``[E, m_max, hidden]`` (row ``src2dst[token, i]``)
``gateup_input_scale`` int32 ``[E, hidden//group//4, m_max]`` (MN-major; byte-scattered)
Slots with ``topk_ids[token, i] < 0`` are skipped. Byte-identical to the
two-kernel path on every written row.
Slots with ``src2dst[token, i] < 0`` are skipped. Byte-identical to the
two-kernel path on every written row. ``topk_ids`` remains in the ABI for
compatibility with already-compiled JIT modules.
"""
assert x.is_cuda and x.dtype == torch.bfloat16 and x.dim() == 2
assert x.is_contiguous()
python/sglang/srt/arg_groups/moe_hook.py
+7 -1
View File
@@ -268,8 +268,14 @@ def handle_a2a_moe(server_args: Any):
assert resolved_view(server_args).moe_runner_backend in [
"flashinfer_cutlass",
"flashinfer_cutedsl",
"flashinfer_trtllm",
"flashinfer_trtllm_routed",
], "Flashinfer MoE A2A is only supported with flashinfer_cutlass, flashinfer_cutedsl or flashinfer_trtllm_routed moe runner backend"
"deep_gemm",
], (
"FlashInfer MoE A2A is supported with flashinfer_cutlass, "
"flashinfer_cutedsl, flashinfer_trtllm, "
"flashinfer_trtllm_routed, or deep_gemm."
)
if a2a_backend == "mori":
if cfg.deepep_mode == "auto":
python/sglang/srt/arg_groups/overrides.py
+15
View File
@@ -2272,6 +2272,20 @@ def _flashinfer_allreduce_fusion_auto_enable(view: Any) -> dict:
moe_a2a_backend (after the DeepSeek CP and a2a declarations), exactly
like the legacy tail block."""
model_arch = view.get_model_config().hf_config.architectures[0]
if envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION.get() and model_arch in {
"Qwen3_5MoeForCausalLM",
"Qwen3_5MoeForConditionalGeneration",
}:
# The Qwen backend owns one workspace for ordinary AR and MoE finalize;
# do not allocate or fall back to the legacy TRTLLM/MNNVL workspace.
if view.flashinfer_allreduce_fusion_backend is not None:
logger.warning(
"SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION owns both Qwen3.5 "
"AllReduce fusion patterns; suppressing the separately configured "
"--flashinfer-allreduce-fusion-backend=%s",
view.flashinfer_allreduce_fusion_backend,
)
return {"flashinfer_allreduce_fusion_backend": None}
if (
view.flashinfer_allreduce_fusion_backend is None
and model_arch in _FLASHINFER_ALLREDUCE_FUSION_ARCHS
@@ -2862,6 +2876,7 @@ _A2A_EP_SPANNING_BACKENDS = frozenset(
"flashinfer",
"mori",
"pplx",
"deepep_v2",
}
)
python/sglang/srt/disaggregation/common/staging_handler.py
+44
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import dataclasses
import logging
import struct
import threading
import time
from typing import TYPE_CHECKING, List, Optional, Tuple
@@ -536,6 +537,44 @@ class StagingTransferInfo:
self.ends[idx] = end
@dataclasses.dataclass
class StagingRegisterInfo:
"""Staging buffer registration info attached to a KVArgsRegisterInfo."""
base_ptr: int = 0
total_size: int = 0
# Staging slots stay [all K, all V] after draft buffers alter kv_data_ptrs order;
# older peers leave this empty and callers fall back to kv_layer_ids.
slot_layer_ids: List[int] = dataclasses.field(default_factory=list)
@classmethod
def from_zmq_fields(
cls, msg: list, msg_start_offset: int, slot_ids_index: Optional[int] = None
) -> Optional[StagingRegisterInfo]:
i = msg_start_offset
base_ptr = (
struct.unpack("Q", msg[i])[0] if len(msg) > i and len(msg[i]) == 8 else 0
)
total_size = (
int(msg[i + 1].decode("ascii"))
if len(msg) > i + 1 and len(msg[i + 1]) > 0
else 0
)
if base_ptr == 0 and total_size == 0:
return None
slot_layer_ids: List[int] = []
if (
slot_ids_index is not None
and len(msg) > slot_ids_index
and len(msg[slot_ids_index]) > 0
):
raw = msg[slot_ids_index]
slot_layer_ids = list(struct.unpack(f"{len(raw) // 8}Q", raw))
return cls(
base_ptr=base_ptr, total_size=total_size, slot_layer_ids=slot_layer_ids
)
class PrefillStagingStrategy:
"""Prefill-side staging transfer: readiness check + gather-RDMA execution.
@@ -627,6 +666,11 @@ class PrefillStagingStrategy:
target_info.dst_kv_item_len,
target_info.dst_kv_layer_ids,
staging_buffer=self.staging_buffer,
dst_slot_layer_ids=(
target_info.staging.slot_layer_ids
if target_info.staging is not None
else None
),
)
except Exception as e:
raise RuntimeError(
python/sglang/srt/disaggregation/decode.py
+21 -7
View File
@@ -51,6 +51,8 @@ from sglang.srt.disaggregation.utils import (
ReqToMetadataIdxAllocator,
TransferBackend,
_is_fake_transfer,
build_kv_layer_ids,
build_staging_slot_metadata,
get_dsv4_c128_state_indices,
get_kv_class,
is_dsv4_c128_online_enabled,
@@ -541,6 +543,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
kv_data_lens += device_kv_data_lens[c4_layer_num:]
kv_item_lens += device_kv_item_lens[c4_layer_num:]
kv_data_mem_kinds += ["VRAM"] * len(device_kv_data_ptrs[c4_layer_num:])
num_draft_entries = 0
if self.draft_token_to_kv_pool is not None:
# We should also transfer draft model kv cache. The indices are
# always shared with a target model.
@@ -551,15 +554,16 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
kv_data_lens += draft_kv_data_lens
kv_item_lens += draft_kv_item_lens
kv_data_mem_kinds += ["VRAM"] * len(draft_kv_data_ptrs)
num_draft_entries = len(draft_kv_data_ptrs)
kv_args.kv_data_ptrs = kv_data_ptrs
kv_args.kv_data_lens = kv_data_lens
kv_args.kv_item_lens = kv_item_lens
kv_args.kv_layer_ids = (
self.token_to_kv_pool.get_kv_layer_ids()
if self.draft_token_to_kv_pool is None
and hasattr(self.token_to_kv_pool, "get_kv_layer_ids")
else []
kv_args.kv_layer_ids = build_kv_layer_ids(
token_to_kv_pool=self.token_to_kv_pool,
draft_token_to_kv_pool=self.draft_token_to_kv_pool,
num_draft_entries=num_draft_entries,
num_hidden_layers=self.scheduler.model_config.num_hidden_layers,
)
if self.transfer_backend == TransferBackend.NIXL:
kv_args.kv_data_mem_kinds = kv_data_mem_kinds
@@ -599,9 +603,19 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
)
if hasattr(kv_manager, "set_kv_buffer_tensors"):
kv_pool = kv_pool_for_heads
if hasattr(kv_pool, "k_buffer") and hasattr(kv_pool, "v_buffer"):
staging_slots = build_staging_slot_metadata(
kv_layer_ids=kv_args.kv_layer_ids,
num_draft_entries=num_draft_entries,
kv_pool=kv_pool,
draft_kv_pool=self.draft_token_to_kv_pool,
)
if staging_slots is not None:
k_buffers, v_buffers, slot_layer_ids = staging_slots
kv_manager.set_kv_buffer_tensors(
kv_pool.k_buffer, kv_pool.v_buffer, kv_pool.page_size
k_buffers,
v_buffers,
kv_pool.page_size,
slot_layer_ids=slot_layer_ids,
)
return kv_manager
python/sglang/srt/disaggregation/mooncake/conn.py
+78 -21
View File
@@ -28,6 +28,7 @@ from sglang.srt.disaggregation.common.staging_handler import (
DecodeStagingContext,
PrefillStagingContext,
StagingManagerMixin,
StagingRegisterInfo,
StagingTransferInfo,
handle_staging_rsp,
handle_watermark_msg,
@@ -148,6 +149,7 @@ class KVArgsRegisterInfo:
dcp_token_item_lens: Optional[List[int]] = None
staging_base_ptr: int = 0
staging_total_size: int = 0
staging: Optional[StagingRegisterInfo] = None
@classmethod
def from_zmq(cls, msg: List[bytes]):
@@ -192,6 +194,8 @@ class KVArgsRegisterInfo:
dst_dcp_rank=(
int(msg[17].decode("ascii")) if len(msg) > 17 and msg[17] != b"" else 0
),
# Note: always put the staging field at the final
staging=StagingRegisterInfo.from_zmq_fields(msg, 14, slot_ids_index=18),
)
@@ -333,11 +337,20 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
self._staging_ctx.room_bootstrap[room] = bootstrap_infos
self._staging_ctx.room_receivers[room] = receiver
def set_kv_buffer_tensors(self, k_buffers: list, v_buffers: list, page_size: int):
def set_kv_buffer_tensors(
self,
k_buffers: list,
v_buffers: list,
page_size: int,
slot_layer_ids: Optional[List[int]] = None,
):
# slot_layer_ids follows the staging slot order (every k_buffer, then
# every v_buffer), which is not kv_args.kv_layer_ids once a draft exists.
self.kv_buffer_tensors = {
"k_buffers": k_buffers,
"v_buffers": v_buffers,
"page_size": page_size,
"slot_layer_ids": list(slot_layer_ids or []),
}
def _init_staging_buffers(self, count: int):
@@ -494,6 +507,7 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
dst_kv_item_len: int,
dst_layer_ids: List[int],
staging_buffer=None,
dst_slot_layer_ids: Optional[List[int]] = None,
) -> int:
"""Transfer KV cache via staging buffers (gather -> bulk RDMA -> scatter on decode)."""
from sglang.srt.disaggregation.common.staging_buffer import (
@@ -528,13 +542,20 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
local_bytes = per_layer_bytes * num_layers * 2
if self.pp_size > 1:
pairs = build_transfer_entry_pairs(
self.kv_args.kv_layer_ids,
dst_layer_ids,
num_layers * 2,
len(dst_layer_ids),
# Pair staging slots in [all K, all V] order; draft buffers make
# kv_data_ptrs diverge from this layout.
src_slot_ids = (
self.kv_buffer_tensors.get("slot_layer_ids")
or self.kv_args.kv_layer_ids
)
dst_num_layers = len(dst_layer_ids) // 2
dst_slot_ids = dst_slot_layer_ids or dst_layer_ids
pairs = build_transfer_entry_pairs(
src_slot_ids,
dst_slot_ids,
num_layers * 2,
len(dst_slot_ids),
)
dst_num_layers = len(dst_slot_ids) // 2
else:
pairs = None
dst_num_layers = num_layers
@@ -654,10 +675,18 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
layers_params = None
# Decode pp size should be equal to prefill pp size or 1
if self.is_mla_backend or self.is_hybrid_mla_backend or force_flat:
# Published layer IDs give exact pairing; plain-MHA peers publish none
# and keep positional slicing.
has_layer_ids = bool(src_layer_ids or dst_layer_ids)
if (
self.is_mla_backend
or self.is_hybrid_mla_backend
or force_flat
or has_layer_ids
):
# Layer IDs map PP-local buffers to global decode entries.
# Registrations without them retain the existing PP mapping.
if src_layer_ids or dst_layer_ids:
if has_layer_ids:
pairs = build_transfer_entry_pairs(
src_layer_ids,
dst_layer_ids,
@@ -970,6 +999,7 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
dst_attn_tp_size: int,
dst_kv_item_len: int,
executor: concurrent.futures.ThreadPoolExecutor,
dst_layer_ids: Optional[List[int]] = None,
):
"""
Sends KV cache slices from this Prefill rank to a target Decode rank,
@@ -1022,9 +1052,34 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
num_heads_to_send = dst_heads_per_rank
dst_head_start_offset = 0
src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
)
src_data_ptrs = self.kv_args.kv_data_ptrs
src_layer_ids = self.kv_args.kv_layer_ids
if src_layer_ids or dst_layer_ids:
# Draft buffers break the flat [K block, V block] layout, so pair by
# layer ID instead of the half-split used by get_mha_kv_ptrs_with_pp.
if any(l != src_kv_item_len for l in self.kv_args.kv_item_lens):
logger.error(
f"[{mooncake_session_id}] head-sliced transfer assumes one item "
f"length for every KV entry, got {set(self.kv_args.kv_item_lens)}"
)
return -1
layer_ptr_pairs = [
(src_data_ptrs[i], dst_kv_ptrs[j])
for i, j in build_transfer_entry_pairs(
src_layer_ids,
dst_layer_ids or [],
len(src_data_ptrs),
len(dst_kv_ptrs),
allow_positional_fallback=self.pp_size == 1,
)
]
else:
src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_kv_ptrs)
)
layer_ptr_pairs = [
(src_k_ptrs[i], dst_k_ptrs[i]) for i in range(layers_current_pp_stage)
] + [(src_v_ptrs[i], dst_v_ptrs[i]) for i in range(layers_current_pp_stage)]
# Calculate precise byte offset and length for the sub-slice within the token
src_head_slice_offset = src_head_start_offset * bytes_per_head_slice_to_send
@@ -1069,15 +1124,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
mooncake_session_id, src_addr_list, dst_addr_list, length_list
)
futures = []
for i in range(layers_current_pp_stage):
futures.append(
executor.submit(process_layer_tp_aware, src_k_ptrs[i], dst_k_ptrs[i])
)
for i in range(layers_current_pp_stage):
futures.append(
executor.submit(process_layer_tp_aware, src_v_ptrs[i], dst_v_ptrs[i])
)
futures = [
executor.submit(process_layer_tp_aware, src_layer_ptr, dst_layer_ptr)
for src_layer_ptr, dst_layer_ptr in layer_ptr_pairs
]
return self._await_transfer_futures(futures)
@@ -1786,6 +1836,7 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
target_rank_registration_info.dst_attn_tp_size,
target_rank_registration_info.dst_kv_item_len,
executor,
target_rank_registration_info.dst_kv_layer_ids,
)
if ret != 0:
with self.session_lock:
@@ -2418,6 +2469,11 @@ class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver):
else:
packed_staging_base_ptr = b""
staging_total_size_str = b""
staging_slots = getattr(self.kv_mgr, "kv_buffer_tensors", None) or {}
packed_staging_slot_layer_ids = b"".join(
struct.pack("Q", layer_id)
for layer_id in (staging_slots.get("slot_layer_ids") or [])
)
try:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
@@ -2442,6 +2498,7 @@ class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver):
staging_total_size_str,
dst_dcp_size,
dst_dcp_rank,
packed_staging_slot_layer_ids,
]
)
except zmq.ZMQError:
python/sglang/srt/disaggregation/prefill.py
+53 -17
View File
@@ -43,6 +43,8 @@ from sglang.srt.disaggregation.utils import (
MetadataBuffers,
ReqToMetadataIdxAllocator,
TransferBackend,
build_kv_layer_ids,
build_staging_slot_metadata,
get_dsv4_c128_state_indices,
get_kv_class,
is_aborted,
@@ -68,6 +70,7 @@ from sglang.srt.mem_cache.common import (
release_kv_cache,
)
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
from sglang.srt.observability.req_time_stats import set_schedule_time_batch
from sglang.srt.runtime_context import (
get_disagg,
@@ -99,6 +102,16 @@ def should_force_retry(req: Req) -> bool:
return int.from_bytes(digest[:8], "big") < retry_prob * 2**64
def _transfer_start_layer(*, pool, hf_text_config) -> int:
# Hybrid pools count all layers in start_layer, but peer KV lists contain only
# full-attention layers, so translate to a full-attention-relative offset.
if not isinstance(pool, HybridLinearKVPool):
return pool.start_layer
return sum(
1 for lid in hf_text_config.full_attention_layer_ids if lid < pool.start_layer
)
def maybe_release_metadata_buffer(
req: Req, allocator: ReqToMetadataIdxAllocator
) -> None:
@@ -212,7 +225,10 @@ class PrefillBootstrapQueue:
self.token_to_kv_pool.start_layer,
)
if layer_shard_enabled
else self.token_to_kv_pool.start_layer
else _transfer_start_layer(
pool=self.token_to_kv_pool,
hf_text_config=self.scheduler.model_config.hf_text_config,
)
)
kv_args.mla_compression_ratios = None
kv_data_ptrs, kv_data_lens, kv_item_lens = (
@@ -224,24 +240,27 @@ class PrefillBootstrapQueue:
else getattr(self.token_to_kv_pool, "end_layer", None)
)
if self.draft_token_to_kv_pool is not None and transfer_draft_cache:
draft_kv_pool = self.draft_token_to_kv_pool if transfer_draft_cache else None
num_draft_entries = 0
if draft_kv_pool is not None:
# We should also transfer draft model kv cache. The indices are
# always shared with a target model.
draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = (
self.draft_token_to_kv_pool.get_contiguous_buf_infos()
draft_kv_pool.get_contiguous_buf_infos()
)
kv_data_ptrs += draft_kv_data_ptrs
kv_data_lens += draft_kv_data_lens
kv_item_lens += draft_kv_item_lens
num_draft_entries = len(draft_kv_data_ptrs)
kv_args.kv_data_ptrs = kv_data_ptrs
kv_args.kv_data_lens = kv_data_lens
kv_args.kv_item_lens = kv_item_lens
kv_args.kv_layer_ids = (
self.token_to_kv_pool.get_kv_layer_ids()
if self.draft_token_to_kv_pool is None
and hasattr(self.token_to_kv_pool, "get_kv_layer_ids")
else []
kv_args.kv_layer_ids = build_kv_layer_ids(
token_to_kv_pool=self.token_to_kv_pool,
draft_token_to_kv_pool=draft_kv_pool,
num_draft_entries=num_draft_entries,
num_hidden_layers=self.scheduler.model_config.num_hidden_layers,
)
if not self.is_mla_backend:
kv_args.kv_head_num = self.token_to_kv_pool.head_num
@@ -288,11 +307,19 @@ class PrefillBootstrapQueue:
kv_pool = self.token_to_kv_pool
if hasattr(kv_pool, "full_kv_pool"):
kv_pool = kv_pool.full_kv_pool
if hasattr(kv_pool, "k_buffer") and hasattr(kv_pool, "v_buffer"):
staging_slots = build_staging_slot_metadata(
kv_layer_ids=kv_args.kv_layer_ids,
num_draft_entries=num_draft_entries,
kv_pool=kv_pool,
draft_kv_pool=draft_kv_pool,
)
if staging_slots is not None:
k_buffers, v_buffers, slot_layer_ids = staging_slots
kv_manager.set_kv_buffer_tensors(
kv_pool.k_buffer,
kv_pool.v_buffer,
k_buffers,
v_buffers,
kv_pool.page_size,
slot_layer_ids=slot_layer_ids,
)
return kv_manager
@@ -693,6 +720,16 @@ class SchedulerDisaggregationPrefillMixin:
logprob_pt = 0
assert batch.spec_info is result.next_draft_input
draft_input = result.next_draft_input
draft_hidden_states_cpu = None
draft_dsa_topk_indices_cpu = None
if self.spec_algorithm.is_eagle() and draft_input is not None:
draft_hidden_states_cpu = draft_input.hidden_states.to(
"cpu", non_blocking=False
)
if batch.spec_info.dsa_topk_indices is not None:
draft_dsa_topk_indices_cpu = batch.spec_info.dsa_topk_indices.to(
"cpu", non_blocking=False
)
# Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
next_token_ids = result.next_token_ids.tolist()
self.batch_result_processor.move_logprobs_to_cpu(
@@ -727,12 +764,11 @@ class SchedulerDisaggregationPrefillMixin:
if self.spec_algorithm.is_eagle() and draft_input is not None:
req.output_topk_p = draft_input.topk_p[i]
req.output_topk_index = draft_input.topk_index[i]
req.hidden_states_tensor = (
draft_input.hidden_states[i].cpu().clone()
)
dsa_topk_indices = batch.spec_info.dsa_topk_indices
if dsa_topk_indices is not None:
req.output_dsa_topk_indices = dsa_topk_indices[i].cpu().clone()
req.hidden_states_tensor = draft_hidden_states_cpu[i].clone()
if draft_dsa_topk_indices_cpu is not None:
req.output_dsa_topk_indices = draft_dsa_topk_indices_cpu[
i
].clone()
else:
req.output_dsa_topk_indices = None
else:
python/sglang/srt/disaggregation/utils.py
+104
View File
@@ -915,6 +915,64 @@ def build_transfer_entry_pairs(
return [(i, i) for i in range(n_src)]
def build_kv_layer_ids(
*,
token_to_kv_pool,
draft_token_to_kv_pool,
num_draft_entries: int,
num_hidden_layers: int,
) -> List[int]:
"""Global layer id for every entry in ``kv_args.kv_data_ptrs``.
Draft KV buffers are appended after the target's, so they need ids of their
own: build_transfer_entry_pairs requires the id list to cover every entry,
and a target-only list would be rejected. The draft numbers its layers from
zero, which would collide with the target's, so its entries are remapped
into a reserved band above the target's layer range. Both PD peers run this
against the same draft config and so agree on the band.
Returns [] for pools that cannot report ids, leaving the peers on positional
pairing.
"""
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
if not isinstance(token_to_kv_pool, HybridLinearKVPool):
return []
layer_ids = token_to_kv_pool.get_kv_layer_ids()
if draft_token_to_kv_pool is None:
return layer_ids
draft_ids = _draft_entry_layer_ids(
pool=draft_token_to_kv_pool, num_entries=num_draft_entries
)
# Rank the draft's own ids by first appearance, so the band stays dense and
# contiguous whatever the draft config numbers its layers.
band_index = {lid: i for i, lid in enumerate(dict.fromkeys(draft_ids))}
return layer_ids + [num_hidden_layers + band_index[lid] for lid in draft_ids]
def _draft_entry_layer_ids(*, pool, num_entries: int) -> List[int]:
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
if isinstance(pool, HybridLinearKVPool):
ids = pool.get_kv_layer_ids()
else:
# Pools register k0..k(L-1) then v0..v(L-1), so ids repeat once per
# group; derive the group count rather than assuming MHA vs MLA.
if pool.layer_num <= 0 or num_entries % pool.layer_num != 0:
raise RuntimeError(
"Draft KV buffers must register a whole number of per-layer "
f"groups: entries={num_entries}, layers={pool.layer_num}"
)
ids = list(range(pool.layer_num)) * (num_entries // pool.layer_num)
if len(ids) != num_entries:
raise RuntimeError(
"Draft KV layer ids must cover every registered entry: "
f"ids={len(ids)}, entries={num_entries}"
)
return ids
def resolve_dcp_dst_entry_indices(
src_layer_ids: List[int],
dst_layer_ids: List[int],
@@ -940,6 +998,52 @@ def resolve_dcp_dst_entry_indices(
]
def build_staging_slot_metadata(
*,
kv_layer_ids: List[int],
num_draft_entries: int,
kv_pool,
draft_kv_pool,
):
"""Buffers and per-slot layer ids for the staging gather.
The gather writes every k_buffer and then every v_buffer, while kv_layer_ids
follows kv_data_ptrs ([K target, V target, K draft, V draft]), so the two
orders diverge as soon as a draft pool is registered.
Returns (k_buffers, v_buffers, slot_layer_ids), or None for a pool that has
no contiguous K/V tensors to stage.
"""
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MHATokenToKVPool
# A hybrid pool keeps its contiguous K/V tensors on the inner full-attention
# pool, and the draft pool is wrapped the same way.
if isinstance(kv_pool, HybridLinearKVPool):
kv_pool = kv_pool.full_kv_pool
if isinstance(draft_kv_pool, HybridLinearKVPool):
draft_kv_pool = draft_kv_pool.full_kv_pool
if not isinstance(kv_pool, MHATokenToKVPool):
return None
ids = list(kv_layer_ids or [])
num_target = len(ids) - num_draft_entries
half = num_target // 2
k_buffers, k_ids = list(kv_pool.k_buffer), ids[:half]
v_buffers, v_ids = list(kv_pool.v_buffer), ids[half:num_target]
draft_half = num_draft_entries // 2
if draft_half:
if not isinstance(draft_kv_pool, MHATokenToKVPool):
# An empty id list puts the sender back on kv_data_ptrs order, which
# is what staging did before draft KV existed.
return k_buffers, v_buffers, []
k_buffers += list(draft_kv_pool.k_buffer)
v_buffers += list(draft_kv_pool.v_buffer)
k_ids += ids[num_target : num_target + draft_half]
v_ids += ids[num_target + draft_half :]
return k_buffers, v_buffers, k_ids + v_ids
def append_state_component(
kv_args: KVArgs,
state_type: StateType,
python/sglang/srt/distributed/bootstrap.py
+3 -1
View File
@@ -138,10 +138,12 @@ def init_torch_distributed(
host=server_args.host, port=server_args.gated_launch_port
)
# Draft workers reuse the target pool config and may exist on only one PP stage;
# including them in this WORLD reduction would deadlock on absent peers.
pre_model_load_memory = get_available_gpu_memory(
device,
ps.gpu_id,
distributed=get_world_group().world_size > 1,
distributed=get_world_group().world_size > 1 and not is_draft_worker,
cpu_group=get_world_group().cpu_group,
)
tp_group = get_tp_group()
python/sglang/srt/distributed/parallel_state.py
+40
View File
@@ -2052,6 +2052,12 @@ def get_moe_tp_group() -> GroupCoordinator:
get_tensor_model_parallel_group = get_tp_group
_PP: Optional[GroupCoordinator] = None
_SELF_PP: Optional[GroupCoordinator] = None
def get_self_pp_group() -> GroupCoordinator:
assert _SELF_PP is not None, "self pipeline group is not initialized"
return _SELF_PP
def get_pp_group() -> GroupCoordinator:
@@ -2716,6 +2722,18 @@ def initialize_model_parallel(
max_world_size=max_world_size,
)
# The one-layer draft uses a singleton PP group; every rank creates all groups
# because new_group is collective.
global _SELF_PP
if _SELF_PP is None:
_SELF_PP = init_model_parallel_group(
[[r] for r in range(world_size)],
get_world_group().local_rank,
backend,
use_custom_allreduce=False,
group_name="self_pp",
)
get_parallel().stamp_derived_widths(**derived_widths)
@@ -2814,6 +2832,28 @@ def model_parallel_is_initialized():
_TP_STATE_PATCHED = False
_PP_STATE_PATCHED = False
@contextmanager
def patch_pipeline_parallel_group(pp_group: GroupCoordinator):
"""Patch the pp group temporarily until this function ends.
This method is for draft workers of speculative decoding, whose model does not
span pipeline stages and must not read the target's pp topology.
"""
global _PP_STATE_PATCHED
assert not _PP_STATE_PATCHED, "Should not call when it's already patched"
_PP_STATE_PATCHED = True
old_pp_group = get_pp_group()
global _PP
_PP = pp_group
try:
yield
finally:
_PP_STATE_PATCHED = False
_PP = old_pp_group
@contextmanager
python/sglang/srt/distributed/utils.py
+4
View File
@@ -101,6 +101,10 @@ def get_pp_indices(
"""
# partition_list_str can be set to None in sglang
partition_list_str = os.getenv("SGLANG_PP_LAYER_PARTITION", None)
if pp_size == 1:
# A singleton draft PP group owns every layer and must ignore the target's
# process-global pipeline partition list.
partition_list_str = None
if partition_list_str is not None:
try:
partitions = [int(layer) for layer in partition_list_str.split(",")]
python/sglang/srt/entrypoints/http_server.py
+7
View File
@@ -674,6 +674,13 @@ async def health_generate(request: Request) -> Response:
if _global_state.tokenizer_manager.server_status == ServerStatus.Starting:
return Response(status_code=503)
# Let an external E2E driver establish a balanced DP batch before health traffic.
if envs.SGLANG_DIAG_BYPASS_HEALTH_GENERATE.get() and request.url.path in (
"/health",
"/health_generate",
):
return Response(status_code=200)
if (
not envs.SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION.get()
and request.url.path == "/health"
python/sglang/srt/environ.py
+29
View File
@@ -323,6 +323,7 @@ class Envs:
SGLANG_UVICORN_WORKER_HEALTHCHECK_TIMEOUT = EnvInt(10)
SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION = EnvBool(True)
SGLANG_EXPOSE_OWN_ENV_VARS = EnvBool(False)
SGLANG_DIAG_BYPASS_HEALTH_GENERATE = EnvBool(False)
# ===================================================================
# Logging
@@ -441,12 +442,21 @@ class Envs:
SGLANG_TRACE_ASYNC_FLUSH_THRESHOLD = EnvInt(100)
SGLANG_ENABLE_METRICS_DEVICE_TIMER = EnvBool(False)
SGLANG_ENABLE_METRICS_DP_ATTENTION = EnvBool(False)
SGLANG_TRACE_LOGITS_E2E = EnvBool(False)
SGLANG_TRACE_LOGITS_E2E_SYNC = EnvBool(False)
SGLANG_TRACE_SAMPLER_E2E = EnvBool(False)
SGLANG_TRACE_QWEN_MOE_DEEPEP_E2E = EnvBool(False)
SGLANG_DEEPEP_V2_TRACE_CONTIG = EnvBool(False)
SGLANG_DEEPEP_V2_TRACE_MASKED = EnvBool(False)
# ===================================================================
# Debugging and invariant checks
# ===================================================================
SGLANG_DETECT_SLOW_RANK = EnvBool(False)
SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
SGLANG_VALIDATE_MAMBA_REPLAY_STATE_INDICES = EnvBool(False)
SGLANG_GDN_DECODE_FUSION_LOG_LAYER_HITS = EnvBool(False)
SGLANG_GDN_DECODE_FUSION_VERIFY_REAL_TENSORS = EnvBool(False)
# NaN-fill the unified memory pool at boot (debug repro switch).
SGLANG_DEBUG_POISON_POOL = EnvBool(False)
SGLANG_DEBUG_REVERT_PR = EnvInt(0)
@@ -746,6 +756,7 @@ class Envs:
# ===================================================================
# MoRI transport and expert dispatch
# ===================================================================
SGLANG_DEEPEP_V2_FORCE_MAX_LEN = EnvBool(False)
# Send CPU-resident AUX data via RDMA instead of ZMQ TCP (default: TCP).
SGLANG_MORI_SEND_AUX_RDMA = EnvBool(False)
# Number of RDMA Queue Pairs (QPs) used per transfer operation. Higher
@@ -1009,6 +1020,9 @@ class Envs:
# DeepGEMM
# ===================================================================
SGLANG_ENABLE_JIT_DEEPGEMM = EnvBool(True)
# Enable the allowlisted low-M BF16 Split-K GEMM path on Blackwell. Shapes
# outside the measured allowlist continue to use CuTe DSL/cuBLAS.
SGLANG_ENABLE_BF16_SPLITK_GEMM = EnvBool(True)
SGLANG_DEEPGEMM_STANDARD_LAYOUT = EnvStr("auto")
SGLANG_DEEPGEMM_MASKED_MEMORY_BUDGET_FRACTION = EnvFloat(0.25)
# Cap the DeepGEMM masked grouped-GEMM per-expert padded capacity at
@@ -1017,6 +1031,9 @@ class Envs:
# load imbalance (they otherwise OOM saturated --moe-runner-backend
# deep_gemm serving). Costs one D2H sync per MoE layer.
SGLANG_OPT_DG_MASKED_M_CAP = EnvBool(False)
# Wide-DP eager prefill uses compact routing storage; masked storage scales
# with num_local_experts and can OOM on skewed batches.
SGLANG_OPT_DG_COMPACT_EAGER = EnvBool(False)
# Drop dp-attention MAX_LEN pad rows from MoE dispatch (StandardDispatcher
# post-translation topk_ids -> -1): pad rows otherwise run the router on
# stale hidden values and burn expert compute whose outputs are discarded;
@@ -1532,6 +1549,16 @@ class Envs:
SGLANG_SYMM_MEM_PREALLOC_GB_SIZE = EnvInt(-1)
SGLANG_DEBUG_SYMM_MEM = EnvBool(False)
# Qwen3.5 and GDN
SGLANG_ENABLE_GDN_DECODE_FUSED_PROJ_CONV = EnvBool(True)
SGLANG_TRACE_QWEN35_FINAL_NORM = EnvBool(False)
SGLANG_QWEN35_NATIVE_FINAL_NORM = EnvBool(False)
# One switch enables deferred MoE finalize and AR + residual + RMSNorm.
SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION = EnvBool(False)
# Distinct workspace configurations allowed in one process. Production
# uses one model/configuration per rank, so fail closed on accidental reuse.
SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION_MAX_INSTANCES = EnvInt(1)
# ===================================================================
# Plugin system
# ===================================================================
@@ -1685,6 +1712,8 @@ _DEPRECATED_ENVS: Dict[str, _DeprecatedEnv] = {
"SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN": _DeprecatedEnv(),
# sconv-family kernels always use the CUDA-JIT ports when supported; no toggle.
"SGLANG_OPT_USE_CUDA_SCONV": _DeprecatedEnv(),
# The direct dense BF16 GEMM source is vendored in-tree.
"SGLANG_FLASHINFER_PR4266_SOURCE": _DeprecatedEnv(),
# DSV4 compressor V2 is always used.
"SGLANG_OPT_USE_COMPRESSOR_V2": _DeprecatedEnv(),
# Replaced by CLI flags.
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+18
View File
@@ -14,6 +14,7 @@ from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
track_mamba_states_all_layers,
track_mamba_states_if_needed,
)
from sglang.srt.environ import envs
from sglang.srt.layers.attention.base_attn_backend import (
AttentionBackend,
SharedReadEnds,
@@ -23,6 +24,9 @@ from sglang.srt.layers.attention.mamba.mamba2_metadata import (
ForwardMetadata,
Mamba2Metadata,
)
from sglang.srt.layers.attention.mamba.replay_state_indices_validator import (
validate_replay_state_indices_cpu,
)
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -40,6 +44,9 @@ if TYPE_CHECKING:
from sglang.srt.layers.attention.verify_mask import VerifyMask
logger = logging.getLogger(__name__)
_validate_mamba_replay_state_indices = (
envs.SGLANG_VALIDATE_MAMBA_REPLAY_STATE_INDICES.get()
)
class MambaAttnBackendBase(AttentionBackend):
@@ -598,6 +605,17 @@ class MambaAttnBackendBase(AttentionBackend):
mamba_indices = self._translate_mamba_indices(mamba_indices)
mamba_indices[bs - num_padding :] = -1
self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices)
if _validate_mamba_replay_state_indices and not in_capture:
# This pre-replay diagnostic intentionally syncs to reject malformed
# live or padded indices before a captured state update uses them.
valid_bs = bs - int(num_padding)
validate_replay_state_indices_cpu(
mamba_indices.detach().cpu(),
valid_bs=valid_bs,
total_bs=bs,
num_state_slots=self.req_to_token_pool.mamba_pool.size + 1,
pad_slot_id=self.pad_slot_id,
)
# Refresh the static track-dest buffer in-place (translated); the captured
# track-save reads it, leaving the handed-in InputBuffer slot read-only.
# Hand out only the refreshed [:bs] prefix — Mamba2's track-save slices
python/sglang/srt/layers/attention/linear/gdn_backend.py
+177 -11
View File
@@ -8,6 +8,7 @@ from sglang.kernels.ops.mamba.causal_conv1d_triton import (
causal_conv1d_update,
)
from sglang.srt.configs.hybrid_arch import hybrid_gdn_config
from sglang.srt.environ import envs
from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
from sglang.srt.layers.attention.linear.utils import (
@@ -32,10 +33,21 @@ if not is_cpu():
if is_cuda() or is_hip() or is_xpu():
from sglang.kernels.ops.attention.triton_gdn_fused_proj import (
can_use_fused_qkvzba_causal_conv1d_update_contiguous,
fused_qkv_split_gdn_prefill,
fused_qkvzba_causal_conv1d_update_contiguous,
fused_qkvzba_split_reshape_cat_contiguous,
)
MAX_FUSED_QKV_SPLIT_DIM = 8192
_fused_decode_proj_conv_logged = False
_fused_decode_proj_conv_fallback_logged = False
_fused_decode_proj_conv_layers_logged: set[int] = set()
_fused_decode_real_tensor_verified_layers: set[int] = set()
_fused_decode_log_layer_hits = envs.SGLANG_GDN_DECODE_FUSION_LOG_LAYER_HITS.get()
_fused_decode_verify_real_tensors = (
envs.SGLANG_GDN_DECODE_FUSION_VERIFY_REAL_TENSORS.get()
)
if is_cuda():
from sglang.srt.layers.attention.mamba.causal_conv1d import (
@@ -433,6 +445,11 @@ class GDNAttnBackend(MambaAttnBackendBase):
b: torch.Tensor,
**kwargs,
):
global _fused_decode_proj_conv_fallback_logged
global _fused_decode_proj_conv_logged
global _fused_decode_proj_conv_layers_logged
global _fused_decode_real_tensor_verified_layers
if _is_hip and isinstance(mixed_qkv, torch.Tensor) and mixed_qkv.shape[0] == 0:
return mixed_qkv.new_zeros((1, 0, layer.num_v_heads, layer.head_v_dim))
@@ -453,15 +470,163 @@ class GDNAttnBackend(MambaAttnBackendBase):
replayssm_k = layer_cache.replayssm_k
replayssm_g = layer_cache.replayssm_g
assert isinstance(mixed_qkv, torch.Tensor)
mixed_qkv = causal_conv1d_update(
mixed_qkv,
conv_states,
layer.conv_weights,
layer.bias,
layer.activation,
conv_state_indices=cache_indices,
)
return_z = False
conv_already_applied = False
if isinstance(mixed_qkv, tuple):
if len(mixed_qkv) != 2:
raise ValueError(
"Fused GDN decode projection input must be "
"(projected_qkvz, projected_ba)"
)
projected_qkvz, projected_ba = mixed_qkv
eligible, eligibility_reason = (
can_use_fused_qkvzba_causal_conv1d_update_contiguous(
projected_qkvz,
projected_ba,
conv_states,
layer.conv_weights,
layer.bias,
cache_indices,
qkv_dim=layer.q_dim + layer.k_dim + layer.v_dim,
v_dim=layer.v_dim,
num_v_heads=layer.num_v_heads,
activation=layer.activation,
)
)
if eligible:
qkv_dim = layer.q_dim + layer.k_dim + layer.v_dim
fused_backend = "triton_direct_oracle_exact"
if not _fused_decode_proj_conv_logged:
rank0_log("Using fused GDN decode QKVZ/BA unpack + indexed Conv1D.")
_fused_decode_proj_conv_logged = True
if (
_fused_decode_log_layer_hits or _fused_decode_verify_real_tensors
) and layer.layer_id not in _fused_decode_proj_conv_layers_logged:
rank0_log(
"GDN_FUSED_DECODE_BACKEND "
f"layer_id={layer.layer_id} backend={fused_backend} "
f"batch={projected_qkvz.shape[0]} "
f"qkv_dim={qkv_dim} state_shape={tuple(conv_states.shape)} "
f"state_indices_dtype={cache_indices.dtype}"
)
_fused_decode_proj_conv_layers_logged.add(layer.layer_id)
# Compare real activations against the direct-Triton update on a
# compact state copy, leaving the live cache to the candidate.
verify_real_tensors = (
_fused_decode_verify_real_tensors
and layer.layer_id not in _fused_decode_real_tensor_verified_layers
)
if verify_real_tensors:
if bool(torch.any(cache_indices < 0).item()):
raise AssertionError(
"Real-tensor GDN fusion verification requires "
"non-padding cache indices"
)
ref_indices = torch.arange(
cache_indices.numel(),
device=cache_indices.device,
dtype=torch.int32,
)
ref_state = torch.index_select(
conv_states, 0, cache_indices.to(torch.int64)
)
ref_mixed_qkv, ref_z, ref_b, ref_a = (
fused_qkvzba_split_reshape_cat_contiguous(
projected_qkvz,
projected_ba,
layer.num_q_heads,
layer.num_v_heads,
layer.head_q_dim,
layer.head_v_dim,
)
)
ref_mixed_qkv = causal_conv1d_update(
ref_mixed_qkv,
ref_state,
layer.conv_weights,
layer.bias,
layer.activation,
conv_state_indices=ref_indices,
)
mixed_qkv, z, b, a = fused_qkvzba_causal_conv1d_update_contiguous(
projected_qkvz,
projected_ba,
conv_states,
layer.conv_weights,
layer.bias,
cache_indices,
qkv_dim=qkv_dim,
v_dim=layer.v_dim,
num_v_heads=layer.num_v_heads,
head_v_dim=layer.head_v_dim,
activation=layer.activation,
)
if verify_real_tensors:
candidate_state = torch.index_select(
conv_states, 0, cache_indices.to(torch.int64)
)
named_pairs = (
("qkv", mixed_qkv, ref_mixed_qkv),
("z", z, ref_z),
("b", b, ref_b),
("a", a, ref_a),
("state", candidate_state, ref_state),
)
report = []
mismatch = False
for tensor_name, candidate, reference in named_pairs:
diff = (candidate.float() - reference.float()).abs()
nonzero = int(torch.count_nonzero(diff).item())
mismatch |= nonzero != 0
report.append(
f"{tensor_name}_nonzero={nonzero}/"
f"{diff.numel()} {tensor_name}_max="
f"{diff.max().item()}"
)
rank0_log(
"GDN_FUSED_REAL_TENSOR_PARITY "
f"layer_id={layer.layer_id} backend={fused_backend} "
+ " ".join(report)
)
_fused_decode_real_tensor_verified_layers.add(layer.layer_id)
if mismatch:
raise AssertionError(
"GDN fused real-tensor parity failed at "
f"layer_id={layer.layer_id}; " + " ".join(report)
)
conv_already_applied = True
else:
# Explicit correctness fallback for an unexpected runtime
# tensor/state contract. This still returns Z to the model.
if not _fused_decode_proj_conv_fallback_logged:
rank0_log(
"Falling back from fused GDN decode projection/Conv1D: "
f"{eligibility_reason}"
)
_fused_decode_proj_conv_fallback_logged = True
mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous(
projected_qkvz,
projected_ba,
layer.num_q_heads,
layer.num_v_heads,
layer.head_q_dim,
layer.head_v_dim,
)
return_z = True
else:
assert isinstance(mixed_qkv, torch.Tensor)
if not conv_already_applied:
mixed_qkv = causal_conv1d_update(
mixed_qkv,
conv_states,
layer.conv_weights,
layer.bias,
layer.activation,
conv_state_indices=cache_indices,
)
# Skip split + reshape + separate gating kernel by consuming
# the packed mixed_qkv directly in a single fused Triton kernel.
@@ -486,7 +651,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
self._track_mamba_state_decode(
forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id
)
return core_attn_out
return (core_attn_out, z) if return_z else core_attn_out
query, key, value = torch.split(
mixed_qkv,
@@ -516,7 +681,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
forward_batch, conv_states, ssm_states, cache_indices, layer.layer_id
)
return core_attn_out
return (core_attn_out, z) if return_z else core_attn_out
def forward_extend(
self,
@@ -735,6 +900,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
state_checkpoint_every_n_tokens=(
forward_metadata.state_checkpoint_every_n_tokens
),
output=kwargs.get("linear_attn_output"),
)
if is_npu() and last_recurrent_state is not None:
python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py
+274 -17
View File
@@ -30,6 +30,29 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_FLASHINFER_GDN_ALIGNMENT = 32
def _empty_aligned_like(
tensor: torch.Tensor, alignment: int = _FLASHINFER_GDN_ALIGNMENT
) -> torch.Tensor:
element_size = tensor.dtype.itemsize
alignment_elements = max(1, alignment // element_size)
storage = torch.empty(
tensor.numel() + alignment_elements - 1,
dtype=tensor.dtype,
device=tensor.device,
)
start_bytes = (-storage.data_ptr()) % alignment
if start_bytes % element_size:
raise RuntimeError(
f"Cannot align {tensor.dtype} storage at {storage.data_ptr()} "
f"to {alignment} bytes"
)
start = start_bytes // element_size
return storage[start : start + tensor.numel()].view(tensor.shape)
# ---------------------------------------------------------------------------
# Lazy import for FlashInfer GDN kernels
# ---------------------------------------------------------------------------
@@ -167,6 +190,19 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
# states; SM100 accepts the state-pool dtype directly.
self._prefill_needs_fp32_state = sm_major >= 12
self.supports_target_verify = sm_major in (9, 10)
self._aligned_input_buffers: dict[tuple, torch.Tensor] = {}
self._aligned_parameter_cache: dict[
tuple, tuple[torch.Tensor, torch.Tensor]
] = {}
self._verify_intermediate_buffers: dict[tuple, torch.Tensor] = {}
self._alignment_fallback_warned = False
# FlashInfer writes through mutable state/workspace pointers, so misaligned
# inputs must fall back to Triton rather than use temporary copies.
from sglang.srt.layers.attention.linear.kernels.gdn_triton import (
TritonGDNKernel,
)
self._alignment_fallback_kernel = TritonGDNKernel()
if sm_major == 9 and self._prefill_fn is None:
raise RuntimeError("FlashInfer GDN prefill kernel is unavailable.")
@@ -209,6 +245,122 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
logger.info("Using FlashInfer GDN kernels")
def _prepare_dynamic_input(self, name: str, tensor: torch.Tensor) -> torch.Tensor:
# Reuse per-stream aligned scratch for uncommon read-only views; stable
# addresses avoid allocator churn and remain safe for CUDA graph capture.
if tensor.data_ptr() % _FLASHINFER_GDN_ALIGNMENT == 0:
return tensor
stream_key = (
torch.cuda.current_stream(tensor.device).cuda_stream
if tensor.device.type == "cuda"
else None
)
key = (
name,
tensor.device,
tensor.dtype,
tuple(tensor.shape),
stream_key,
)
aligned = self._aligned_input_buffers.get(key)
if aligned is None:
aligned = _empty_aligned_like(tensor)
self._aligned_input_buffers[key] = aligned
aligned.copy_(tensor)
return aligned
def _prepare_parameter(
self,
name: str,
tensor: torch.Tensor,
*,
dtype: Optional[torch.dtype] = None,
) -> torch.Tensor:
key = (
name,
id(tensor),
dtype,
)
cached_entry = self._aligned_parameter_cache.get(key)
if cached_entry is not None and cached_entry[0] is tensor:
return cached_entry[1]
prepared = tensor.detach().reshape(-1)
if dtype is not None:
prepared = prepared.to(dtype=dtype, copy=False)
if (
not prepared.is_contiguous()
or prepared.data_ptr() % _FLASHINFER_GDN_ALIGNMENT
):
aligned = _empty_aligned_like(prepared)
aligned.copy_(prepared)
prepared = aligned
# Retaining the source also prevents a recycled Python id from
# colliding with an older cache entry.
self._aligned_parameter_cache[key] = (tensor, prepared)
return prepared
def _prepare_gate_parameters(
self,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
*,
A_log_dtype: Optional[torch.dtype] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
return (
self._prepare_parameter("A_log", A_log, dtype=A_log_dtype),
self._prepare_parameter("dt_bias", dt_bias),
)
def _mutable_inputs_are_aligned(
self, *named_tensors: tuple[str, Optional[torch.Tensor]]
) -> bool:
for name, tensor in named_tensors:
if tensor is None or tensor.data_ptr() % _FLASHINFER_GDN_ALIGNMENT == 0:
continue
if not self._alignment_fallback_warned:
logger.warning(
"FlashInfer GDN mutable buffer %r has data_ptr %d "
"(mod 32 = %d); falling back to Triton for this call.",
name,
tensor.data_ptr(),
tensor.data_ptr() % _FLASHINFER_GDN_ALIGNMENT,
)
self._alignment_fallback_warned = True
return False
return True
def _prepare_verify_intermediate_buffer(
self,
intermediate_states_buffer: torch.Tensor,
batch_size: int,
) -> tuple[torch.Tensor, bool]:
# FlashInfer requires exact capture B, which may exceed the pool-scoped
# buffer; padded tiers use stable scratch and copy owned rows back.
direct = intermediate_states_buffer[:batch_size]
if direct.shape[0] == batch_size:
return direct, False
stream_key = (
torch.cuda.current_stream(intermediate_states_buffer.device).cuda_stream
if intermediate_states_buffer.device.type == "cuda"
else None
)
shape = (batch_size, *intermediate_states_buffer.shape[1:])
key = (
intermediate_states_buffer.device,
intermediate_states_buffer.dtype,
shape,
stream_key,
)
scratch = self._verify_intermediate_buffers.get(key)
if scratch is None:
template = intermediate_states_buffer.new_empty(shape)
scratch = _empty_aligned_like(template)
self._verify_intermediate_buffers[key] = scratch
return scratch, True
# ---- decode ----
def decode(
@@ -226,6 +378,21 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
if not self._mutable_inputs_are_aligned(("ssm_states", ssm_states)):
return self._alignment_fallback_kernel.decode(
q,
k,
v,
a,
b,
A_log=A_log,
dt_bias=dt_bias,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
**kwargs,
)
batch_size = cache_indices.shape[0]
num_heads = q.shape[2]
head_k_dim = q.shape[3]
@@ -237,20 +404,35 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
value_fi = v.view(batch_size, 1, num_v_heads, head_v_dim)
a_fi = a.view(batch_size, 1, num_v_heads)
b_fi = b.view(batch_size, 1, num_v_heads)
query_fi = self._prepare_dynamic_input("decode_q", query_fi)
key_fi = self._prepare_dynamic_input("decode_k", key_fi)
value_fi = self._prepare_dynamic_input("decode_v", value_fi)
a_fi = self._prepare_dynamic_input("decode_a", a_fi)
b_fi = self._prepare_dynamic_input("decode_b", b_fi)
A_log_fi, dt_bias_fi = self._prepare_gate_parameters(
A_log,
dt_bias,
# Preserve the original backend contract: the SM100 state-pool
# kernel consumes float32 A_log, while SM90 uses the source dtype.
A_log_dtype=torch.float32 if self.use_state_pool else None,
)
if self.use_state_pool:
cache_indices_fi = self._prepare_dynamic_input(
"decode_cache_indices", cache_indices
)
output_fi, _ = self._decode_fn(
q=query_fi,
k=key_fi,
v=value_fi,
state=None,
A_log=A_log.detach().float(),
A_log=A_log_fi,
a=a_fi,
dt_bias=dt_bias.detach(),
dt_bias=dt_bias_fi,
b=b_fi,
use_qk_l2norm=True,
initial_state=ssm_states,
initial_state_indices=cache_indices,
initial_state_indices=cache_indices_fi,
)
else:
# TODO: Once FlashInfer PR#2521 is merged for SM90, gather/scatter
@@ -261,9 +443,9 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
k=key_fi,
v=value_fi,
state=state_batch,
A_log=A_log.detach(),
A_log=A_log_fi,
a=a_fi,
dt_bias=dt_bias.detach(),
dt_bias=dt_bias_fi,
b=b_fi,
scale=None,
output=None,
@@ -301,6 +483,30 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
k_fi = l2norm_fwd(k[0].contiguous())
v_fi = v[0].contiguous()
output = kwargs.get("output")
output_fi = None
if output is not None:
expected_output_shape = (
1,
total_seq_len,
num_v_heads,
head_v_dim,
)
if tuple(output.shape) != expected_output_shape:
raise ValueError(
"FlashInfer GDN prefill output shape mismatch: "
f"expected {expected_output_shape}, got {tuple(output.shape)}"
)
if output.dtype != v.dtype or output.device != v.device:
raise ValueError(
"FlashInfer GDN prefill output must match v dtype/device, "
f"got output=({output.dtype}, {output.device}) and "
f"v=({v.dtype}, {v.device})"
)
output_fi = output[0]
if not output_fi.is_contiguous():
raise ValueError("FlashInfer GDN prefill output must be contiguous")
# g (alpha) and beta: [1, seq, HV] -> [seq, HV], float32 for FlashInfer
alpha_fi = torch.exp(g[0].to(torch.float32))
beta_fi = beta[0].to(torch.float32)
@@ -347,6 +553,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
output_final_state=True,
cu_seqlens=cu_seqlens,
use_qk_l2norm_in_kernel=False,
output=output_fi,
output_state=output_state_fi,
state_checkpoints=state_checkpoints,
checkpoint_cu_starts=state_checkpoint_cu_starts,
@@ -404,34 +611,79 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
num_v_heads = v.shape[2]
head_v_dim = v.shape[3]
query_mtp = q.view(batch_size, draft_token_num, num_heads, head_k_dim)
key_mtp = k.view(batch_size, draft_token_num, num_heads, head_k_dim)
value_mtp = v.view(batch_size, draft_token_num, num_v_heads, head_v_dim)
if a is None or b is None or A_log is None or dt_bias is None:
raise RuntimeError(
"FlashInfer GDN MTP kernel requires a, b, A_log, dt_bias."
)
a_mtp = a.view(batch_size, draft_token_num, num_v_heads)
b_mtp = b.view(batch_size, draft_token_num, num_v_heads)
intermediate_states_buffer_mtp = intermediate_states_buffer
copy_verify_intermediate_back = False
if self.use_state_pool and intermediate_states_buffer is not None:
# The SM100 bf16 MTP kernel indexes this scratch buffer by the
# per-call batch id, while SGLang's speculative state cache is
# pool-scoped and may include an extra dummy slot.
intermediate_states_buffer_mtp = intermediate_states_buffer[:batch_size]
# Graph padding can exceed the pool-scoped scratch; use exact-B storage
# and copy owned rows back before post-verify commit reads the pool.
(
intermediate_states_buffer_mtp,
copy_verify_intermediate_back,
) = self._prepare_verify_intermediate_buffer(
intermediate_states_buffer, batch_size
)
if not self._mutable_inputs_are_aligned(
("ssm_states", ssm_states),
("intermediate_states_buffer", intermediate_states_buffer_mtp),
):
return self._alignment_fallback_kernel.target_verify(
A_log=A_log,
dt_bias=dt_bias,
q=q,
k=k,
v=v,
a=a,
b=b,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
intermediate_states_buffer=intermediate_states_buffer,
intermediate_state_indices=intermediate_state_indices,
cache_steps=cache_steps,
retrieve_parent_token=retrieve_parent_token,
**kwargs,
)
query_mtp = self._prepare_dynamic_input(
"verify_q",
q.view(batch_size, draft_token_num, num_heads, head_k_dim),
)
key_mtp = self._prepare_dynamic_input(
"verify_k",
k.view(batch_size, draft_token_num, num_heads, head_k_dim),
)
value_mtp = self._prepare_dynamic_input(
"verify_v",
v.view(batch_size, draft_token_num, num_v_heads, head_v_dim),
)
a_mtp = self._prepare_dynamic_input(
"verify_a", a.view(batch_size, draft_token_num, num_v_heads)
)
b_mtp = self._prepare_dynamic_input(
"verify_b", b.view(batch_size, draft_token_num, num_v_heads)
)
A_log_fi, dt_bias_fi = self._prepare_gate_parameters(A_log, dt_bias)
cache_indices_fi = self._prepare_dynamic_input(
"verify_cache_indices", cache_indices
)
output_fi, _ = self._mtp_fn(
q=query_mtp,
k=key_mtp,
v=value_mtp,
initial_state=ssm_states,
initial_state_indices=cache_indices,
A_log=A_log.detach(),
initial_state_indices=cache_indices_fi,
A_log=A_log_fi,
a=a_mtp,
dt_bias=dt_bias.detach(),
dt_bias=dt_bias_fi,
b=b_mtp,
scale=None,
output=None,
@@ -440,4 +692,9 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
use_qk_l2norm=True,
)
if copy_verify_intermediate_back:
intermediate_states_buffer.copy_(
intermediate_states_buffer_mtp[: intermediate_states_buffer.shape[0]]
)
return output_fi.view(1, seq_len, num_v_heads, head_v_dim)
python/sglang/srt/layers/attention/mamba/replay_state_indices_validator.py
+82
View File
@@ -0,0 +1,82 @@
"""Debug validation for Mamba CUDA-graph replay state indices.
The validator is intentionally CPU-only. Replay metadata is prepared outside
the captured graph, so a diagnostic run may synchronize once here without
putting host reads or dynamic assertions into CUDA graph capture/replay.
"""
from __future__ import annotations
import torch
def validate_replay_state_indices_cpu(
state_indices: torch.Tensor,
*,
valid_bs: int,
total_bs: int,
num_state_slots: int,
pad_slot_id: int = -1,
) -> None:
"""Validate live and padded rows of a replay state-index buffer.
Live rows must own distinct in-range slots in ``[0, num_state_slots)``.
Slot zero is reserved for CUDA-graph dummy/idle traffic but is still a
valid storage row. All padded rows must carry exactly ``pad_slot_id`` so
indexed state kernels skip them.
"""
if state_indices.device.type != "cpu":
raise ValueError("state_indices must be copied to CPU before validation")
if state_indices.ndim != 1:
raise ValueError("state_indices must be rank-1")
if not 0 <= valid_bs <= total_bs <= state_indices.numel():
raise ValueError(
"expected 0 <= valid_bs <= total_bs <= state_indices.numel(), got "
f"valid_bs={valid_bs} total_bs={total_bs} "
f"numel={state_indices.numel()}"
)
if num_state_slots <= 1:
raise ValueError(
f"num_state_slots must include real slots, got {num_state_slots}"
)
indices = state_indices[:total_bs].to(dtype=torch.int64)
live = indices[:valid_bs]
padded = indices[valid_bs:]
errors: list[str] = []
live_in_range = (live >= 0) & (live < num_state_slots)
if not bool(torch.all(live_in_range)):
bad_rows = torch.nonzero(~live_in_range, as_tuple=False).flatten()
errors.append(
"live rows must contain in-range slots in "
f"[0, {num_state_slots}); bad_rows={bad_rows.tolist()} "
f"bad_values={live[bad_rows].tolist()}"
)
if live.numel() > 1:
unique_live, counts = torch.unique(live, sorted=True, return_counts=True)
duplicate_mask = counts > 1
if bool(torch.any(duplicate_mask)):
errors.append(
"live rows must own unique slots; "
f"duplicate_slots={unique_live[duplicate_mask].tolist()} "
f"counts={counts[duplicate_mask].tolist()}"
)
bad_padding = padded != pad_slot_id
if bool(torch.any(bad_padding)):
bad_offsets = torch.nonzero(bad_padding, as_tuple=False).flatten()
errors.append(
f"padded rows must equal pad_slot_id={pad_slot_id}; "
f"bad_rows={(bad_offsets + valid_bs).tolist()} "
f"bad_values={padded[bad_offsets].tolist()}"
)
if errors:
raise AssertionError(
"Invalid Mamba replay state indices: "
+ "; ".join(errors)
+ f"; valid_bs={valid_bs} total_bs={total_bs} "
f"indices={indices.tolist()}"
)
python/sglang/srt/layers/dp_attention.py
+6 -3
View File
@@ -28,6 +28,7 @@ from sglang.srt.distributed import (
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.runtime_context import (
derive_attention_widths,
get_device,
@@ -103,7 +104,11 @@ class DpPaddingMode(IntEnum):
# Force MAX_LEN so all ranks are padded to equal token counts.
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
if get_moe_a2a_backend().is_pplx():
moe_a2a_backend = get_moe_a2a_backend()
if moe_a2a_backend.is_pplx():
return DpPaddingMode.MAX_LEN
if moe_a2a_backend.is_deepep_v2() and envs.SGLANG_DEEPEP_V2_FORCE_MAX_LEN.get():
return DpPaddingMode.MAX_LEN
# When is_extend_in_batch and dp_size > 1, use SUM_LEN to avoid padding
@@ -588,8 +593,6 @@ _dp_gather_fp8_bufs: dict = {}
@functools.lru_cache(maxsize=1)
def _use_dp_gather_fp8() -> bool:
from sglang.srt.environ import envs
return envs.SGLANG_ENABLE_DP_GATHER_FP8.get()
python/sglang/srt/layers/flashinfer_mnnvl_cutedsl.py
+392
View File
@@ -0,0 +1,392 @@
"""Process-local access to FlashInfer's MNNVL CuTe DSL fusion workspace."""
from __future__ import annotations
import importlib.util
import logging
import threading
from dataclasses import dataclass, replace
from functools import lru_cache
from typing import TYPE_CHECKING
import torch
import torch.distributed as dist
if TYPE_CHECKING:
from torch.distributed import ProcessGroup
logger = logging.getLogger(__name__)
def _import_kernel_backend():
try:
from flashinfer.comm import AllReduceFusionPattern
except ImportError as error:
raise RuntimeError(
"MNNVL CuTe DSL fusion requires FlashInfer's communication "
"infrastructure (flashinfer >= 0.6.16)"
) from error
try:
from sglang.kernels.ops.communication.mnnvl_cutedsl import DEFAULT_CONFIG
from sglang.kernels.ops.communication.mnnvl_cutedsl_ar import (
MNNVLCuteDSLAllReduceFusionWorkspace,
mnnvl_cutedsl_allreduce_fusion,
)
except ImportError as error:
raise RuntimeError(
"SGLang's in-tree MNNVL CuTe DSL kernels failed to import; check "
"their dependencies, including nvidia-cutlass-dsl and cuda-python"
) from error
if importlib.util.find_spec("flashinfer.comm.mnnvl_cutedsl") is not None:
logger.warning(
"The installed FlashInfer now ships flashinfer.comm.mnnvl_cutedsl; "
"SGLang is still running its in-tree port "
"(sglang.kernels.ops.communication.mnnvl_cutedsl), which can now "
"be retired in favor of the upstream backend."
)
return (
MNNVLCuteDSLAllReduceFusionWorkspace,
mnnvl_cutedsl_allreduce_fusion,
AllReduceFusionPattern,
DEFAULT_CONFIG,
)
def _with_early_finalize_shared_load(config):
profiles = []
updated_presets = 0
for profile in config.profiles:
targets = []
for target in profile.finalize_routes.targets:
preset = target.preset
if hasattr(preset, "load_shared_expert_before_pdl"):
preset = replace(preset, load_shared_expert_before_pdl=True)
target = replace(target, preset=preset)
updated_presets += 1
targets.append(target)
profiles.append(
replace(
profile,
finalize_routes=replace(
profile.finalize_routes,
targets=tuple(targets),
),
)
)
if updated_presets == 0:
raise RuntimeError(
"FlashInfer MNNVL config does not expose the finalize shared-load "
"PDL ordering option"
)
return replace(config, profiles=tuple(profiles))
@dataclass(frozen=True, slots=True)
class _WorkspaceSignature:
hidden_size: int
top_k: int
rms_epsilon: float
weight_bias: float
max_m: int
device_index: int
process_group_identity: int
class FlashInferMNNVLCuteDSLARFusion:
"""One graph-stable workspace serving both supported fusion patterns."""
def __init__(
self,
*,
hidden_size: int,
top_k: int,
max_m: int,
rms_epsilon: float,
weight_bias: float,
process_group: ProcessGroup,
device: torch.device,
) -> None:
if hidden_size <= 0 or top_k <= 0 or max_m <= 0:
raise ValueError("hidden_size, top_k, and max_m must be positive")
if device.type != "cuda":
raise ValueError(f"MNNVL CuTe DSL fusion requires CUDA, got {device}")
self.hidden_size = int(hidden_size)
self.top_k = int(top_k)
self.max_m = int(max_m)
self.rms_epsilon = float(rms_epsilon)
self.weight_bias = float(weight_bias)
self.process_group = process_group
self.device = torch.device(device)
self._destroyed = False
with torch.cuda.device(self.device):
self.device = torch.device("cuda", torch.cuda.current_device())
# CuTe DSL obtains NVLS storage through PyTorch symmetric memory, whose
# process-local backend must be selected before workspace construction.
import torch.distributed._symmetric_memory as symm_mem
symmetric_memory_backend = symm_mem.get_backend(self.device)
if symmetric_memory_backend is None:
symm_mem.set_backend("NCCL")
symmetric_memory_backend = symm_mem.get_backend(self.device)
if symmetric_memory_backend is None:
raise RuntimeError(
"PyTorch symmetric memory has no backend for the current device"
)
logger.info(
"Using PyTorch symmetric-memory backend %s for %s",
symmetric_memory_backend,
self.device,
)
(
workspace_type,
self._allreduce_fusion,
self._patterns,
default_config,
) = _import_kernel_backend()
# Only fused finalize launches have a completed shared-expert handoff;
# standalone AllReduce kernels retain the safe load ordering.
from sglang.srt.runtime_context import get_spec
if get_spec().speculative_algorithm is None:
self.workspace_config = _with_early_finalize_shared_load(default_config)
else:
# Early shared load is safe only for a single looping decode graph;
# alternating draft/verify replays can read an unfinished buffer.
logger.info(
"Speculative decoding active: keeping the FlashInfer MNNVL "
"CuTe DSL finalize presets on the safe (non-early-load) "
"ordering."
)
self.workspace_config = default_config
self.workspace = workspace_type(
tp_size=dist.get_world_size(process_group),
tp_rank=dist.get_rank(process_group),
max_token_num=self.max_m,
hidden_dim=self.hidden_size,
dtype=torch.bfloat16,
group=process_group,
top_k=self.top_k,
rms_eps=self.rms_epsilon,
routed_scaling_factor=1.0,
weight_bias=self.weight_bias,
include_shared_expert=True,
add_residual=True,
write_residual_output=True,
config=self.workspace_config,
)
# Publish only after the mailbox barrier; older FlashInfer workspace
# classes may not provide it and would desynchronize Lamport stages.
torch.cuda.synchronize(self.device)
dist.barrier(group=process_group)
def supports(self, m: int) -> bool:
if self._destroyed or not 1 <= int(m) <= self.max_m:
return False
return self.workspace.is_buffer_size_sufficient(
tp_size=dist.get_world_size(self.process_group),
num_tokens=int(m),
hidden_dim=self.hidden_size,
dtype=torch.bfloat16,
)
def moe_finalize_all_reduce_rms_norm(
self,
*,
routed_output: torch.Tensor,
expert_weights: torch.Tensor,
permuted_indices: torch.Tensor,
gated_shared_output: torch.Tensor,
residual: torch.Tensor,
gamma: torch.Tensor,
norm_output: torch.Tensor | None = None,
residual_output: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
m = int(permuted_indices.shape[0])
if not self.supports(m):
raise ValueError(f"workspace does not support M={m}")
shape = (m, self.hidden_size)
if norm_output is None:
norm_output = torch.empty(shape, dtype=torch.bfloat16, device=self.device)
if residual_output is None:
residual_output = torch.empty(
shape, dtype=torch.bfloat16, device=self.device
)
pattern = self._patterns.kMoEFinalizeARResidualRMSNorm
self._allreduce_fusion(
input=routed_output,
workspace=self.workspace,
pattern=pattern,
# The public API carries the caller's PDL intent. The backend's
# routing profile owns the compiled choice and validates it.
launch_with_pdl=True,
residual_in=residual,
residual_out=residual_output,
norm_out=norm_output,
rms_gamma=gamma,
rms_eps=self.rms_epsilon,
weight_bias=self.weight_bias,
expanded_idx_to_permuted_idx=permuted_indices,
expert_scale_factor=expert_weights,
shared_expert_output=gated_shared_output,
)
return norm_output, residual_output
def all_reduce_residual_rms_norm(
self,
*,
local_contribution: torch.Tensor,
residual: torch.Tensor,
gamma: torch.Tensor,
norm_output: torch.Tensor | None = None,
residual_output: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
m = int(local_contribution.shape[0])
if not self.supports(m):
raise ValueError(f"workspace does not support M={m}")
if norm_output is None:
norm_output = torch.empty_like(local_contribution)
if residual_output is None:
residual_output = torch.empty_like(local_contribution)
pattern = self._patterns.kARResidualRMSNorm
self._allreduce_fusion(
input=local_contribution,
workspace=self.workspace,
pattern=pattern,
launch_with_pdl=True,
residual_in=residual,
residual_out=residual_output,
norm_out=norm_output,
rms_gamma=gamma,
rms_eps=self.rms_epsilon,
weight_bias=self.weight_bias,
)
return norm_output, residual_output
def destroy(self) -> None:
if self._destroyed:
return
self.workspace.destroy()
self._destroyed = True
_WORKSPACES: dict[_WorkspaceSignature, FlashInferMNNVLCuteDSLARFusion] = {}
_WORKSPACES_LOCK = threading.RLock()
@lru_cache(maxsize=1)
def _max_workspace_instances() -> int:
from sglang.srt.environ import envs
value = int(envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION_MAX_INSTANCES.get())
if value < 1:
raise ValueError(
"SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION_MAX_INSTANCES must be positive"
)
return value
def get_flashinfer_mnnvl_cutedsl_ar_fusion(
*,
hidden_size: int | None = None,
top_k: int | None = None,
max_m: int | None = None,
rms_epsilon: float | None = None,
weight_bias: float | None = None,
) -> FlashInferMNNVLCuteDSLARFusion:
"""Lookup, or before graph capture create, the process-local workspace."""
supplied = (hidden_size, top_k, max_m, rms_epsilon, weight_bias)
if all(value is None for value in supplied):
with _WORKSPACES_LOCK:
if len(_WORKSPACES) == 1:
return next(iter(_WORKSPACES.values()))
if not _WORKSPACES:
raise RuntimeError(
"MNNVL CuTe DSL fusion workspace was not initialized before use"
)
raise RuntimeError(
"multiple MNNVL CuTe DSL fusion workspaces exist; configuration "
"arguments are required"
)
if any(value is None for value in supplied):
raise TypeError(
"hidden_size, top_k, max_m, rms_epsilon, and weight_bias must be "
"supplied together"
)
if not torch.cuda.is_available():
raise RuntimeError("MNNVL CuTe DSL fusion requires CUDA")
assert hidden_size is not None
assert top_k is not None
assert max_m is not None
assert rms_epsilon is not None
assert weight_bias is not None
from sglang.srt.distributed.parallel_state import get_tp_group
device = torch.device("cuda", torch.cuda.current_device())
process_group = get_tp_group().device_group
domain = (
int(hidden_size),
int(top_k),
float(rms_epsilon),
float(weight_bias),
int(device.index),
id(process_group),
)
with _WORKSPACES_LOCK:
compatible = [
(signature.max_m, instance)
for signature, instance in _WORKSPACES.items()
if (
signature.hidden_size,
signature.top_k,
signature.rms_epsilon,
signature.weight_bias,
signature.device_index,
signature.process_group_identity,
)
== domain
and signature.max_m >= int(max_m)
]
if compatible:
return min(compatible, key=lambda item: item[0])[1]
if torch.cuda.is_current_stream_capturing():
raise RuntimeError(
"creating an MNNVL CuTe DSL fusion workspace during CUDA Graph "
"capture is forbidden"
)
if len(_WORKSPACES) >= _max_workspace_instances():
raise RuntimeError(
"MNNVL CuTe DSL fusion workspace instance limit exceeded; "
"increase SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION_MAX_INSTANCES "
"only when multiple model configurations intentionally coexist"
)
signature = _WorkspaceSignature(
hidden_size=int(hidden_size),
top_k=int(top_k),
rms_epsilon=float(rms_epsilon),
weight_bias=float(weight_bias),
max_m=int(max_m),
device_index=int(device.index),
process_group_identity=id(process_group),
)
instance = FlashInferMNNVLCuteDSLARFusion(
hidden_size=hidden_size,
top_k=top_k,
max_m=max_m,
rms_epsilon=rms_epsilon,
weight_bias=weight_bias,
process_group=process_group,
device=device,
)
_WORKSPACES[signature] = instance
return instance
python/sglang/srt/layers/logits_processor.py
+129
View File
@@ -27,6 +27,7 @@ from sglang.kernels.ops.activation.softcap import (
from sglang.srt.beam_search.logits_capture import BeamLogitsCapture
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
from sglang.srt.environ import envs
from sglang.srt.layers.aux_hidden_states import (
AuxHiddenStates,
pack_aux_hidden_states,
@@ -81,6 +82,83 @@ _UNQUANTIZED_LM_HEAD_METHODS = {
_autotune_run_lm_head: Optional[bool] = None
def _trace_e2e_logits(stage: str, **fields) -> None:
if not envs.SGLANG_TRACE_LOGITS_E2E.get():
return
try:
parallel = get_parallel()
rank = f"dp={parallel.attn_dp_rank} " f"tp={parallel.tp_rank}"
except Exception:
rank = "rank=unknown"
details = " ".join(f"{key}={value}" for key, value in fields.items())
print(f"SGLANG_TRACE_LOGITS_E2E {rank} stage={stage} {details}", flush=True)
def _has_lm_head_runtime_attrs(lm_head, attr_names: Tuple[str, ...]) -> bool:
return all(hasattr(lm_head, attr_name) for attr_name in attr_names)
def should_apply_lm_head_quant_method(lm_head, quant_method) -> bool:
if (
quant_method is None
or not hasattr(lm_head, "weight")
or not callable(getattr(quant_method, "apply", None))
):
return False
method_name = type(quant_method).__name__
if method_name in _UNQUANTIZED_LM_HEAD_METHODS:
return False
# A shared target lm_head can retain the draft's stale ModelOpt method; use it
# only when the runtime tensor layout matches that method.
if method_name == "ModelOptFp4LinearMethod":
if lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs(
lm_head,
(
"weight_scale",
"weight_global_scale",
"workspace",
"input_size_per_partition",
"output_size_per_partition",
),
):
return True
return lm_head.weight.dtype == torch.uint8 and _has_lm_head_runtime_attrs(
lm_head,
(
"weight_scale_interleaved",
"alpha",
"input_scale_inv",
"input_size_per_partition",
"output_size_per_partition",
),
)
if method_name == "ModelOptNvFp4A16LinearMethod":
return lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs(
lm_head,
(
"weight_scale",
"weight_global_scale",
"workspace",
"input_size_per_partition",
"output_size_per_partition",
),
)
if method_name == "ModelOptFp8LinearMethod":
return (
lm_head.weight.dtype == torch.float8_e4m3fn
and _has_lm_head_runtime_attrs(lm_head, ("weight_scale", "input_scale"))
)
return True
# FlashInfer autotune skips the unprofiled LM-head all-gather; its
# [batch * dp_size, vocab] output can OOM under tight DP-attention memory.
_in_autotune_dummy_run = False
def get_in_autotune_dummy_run() -> bool:
return _autotune_run_lm_head is not None
@@ -668,17 +746,42 @@ class LogitsProcessor(nn.Module):
last position (e.g., extend without input logprobs). The caller should
guarantee the given hidden_states follow this constraint.
"""
_trace_e2e_logits(
"get_logits_enter",
hidden_shape=tuple(hidden_states.shape),
dp_gather=self.do_tensor_parallel_all_gather_dp_attn,
tp_gather=self.do_tensor_parallel_all_gather,
)
hidden_states, local_hidden_states = self._gather_dp_attn_hidden_states(
hidden_states, logits_metadata
)
_trace_e2e_logits(
"dp_hidden_gather_returned",
global_shape=tuple(hidden_states.shape),
local_shape=tuple(local_hidden_states.shape),
)
if envs.SGLANG_TRACE_LOGITS_E2E_SYNC.get():
_trace_e2e_logits("pre_lm_head_sync_enter")
torch.cuda.synchronize()
_trace_e2e_logits("pre_lm_head_sync_returned")
_trace_e2e_logits("lm_head_enter", hidden_shape=tuple(hidden_states.shape))
logits = self._compute_lm_head(hidden_states, lm_head, embedding_bias)
_trace_e2e_logits("lm_head_returned", logits_shape=tuple(logits.shape))
if envs.SGLANG_TRACE_LOGITS_E2E_SYNC.get():
_trace_e2e_logits("post_lm_head_sync_enter")
torch.cuda.synchronize()
_trace_e2e_logits("post_lm_head_sync_returned")
if self.logit_scale is not None:
logits.mul_(self.logit_scale)
used_tp_lm_head_all_to_all = False
if self.do_tensor_parallel_all_gather:
_trace_e2e_logits(
"tp_logits_gather_enter", logits_shape=tuple(logits.shape)
)
if self.use_attn_tp_group:
logits = self._gather_attn_tp_logits(logits)
elif self._can_use_tp_lm_head_all_to_all(
@@ -688,11 +791,20 @@ class LogitsProcessor(nn.Module):
used_tp_lm_head_all_to_all = True
else:
logits = self._logits_gatherer(logits)
_trace_e2e_logits(
"tp_logits_gather_returned", logits_shape=tuple(logits.shape)
)
if not used_tp_lm_head_all_to_all:
_trace_e2e_logits(
"dp_logits_scatter_enter", logits_shape=tuple(logits.shape)
)
logits = self._scatter_dp_attn_logits(
logits, local_hidden_states, logits_metadata
)
_trace_e2e_logits(
"dp_logits_scatter_returned", logits_shape=tuple(logits.shape)
)
logits = self._copy_logits_to_buffer(
logits, logits_metadata, use_buffer=use_logits_buffer
@@ -776,10 +888,27 @@ class LogitsProcessor(nn.Module):
self, hidden_states: torch.Tensor, logits_metadata: LogitsMetadata
) -> Tuple[torch.Tensor, torch.Tensor]:
if self.do_tensor_parallel_all_gather_dp_attn:
_trace_e2e_logits(
"dp_metadata_enter",
local_shape=tuple(hidden_states.shape),
global_counts_cpu=logits_metadata.global_num_tokens_for_logprob_cpu,
)
logits_metadata.compute_dp_attention_metadata()
_trace_e2e_logits(
"dp_metadata_returned",
buffer_shape=tuple(logits_metadata.gathered_buffer.shape),
local_start=logits_metadata.dp_local_start_pos,
local_tokens=logits_metadata.dp_local_num_tokens,
)
local_hidden_states = hidden_states
hidden_states = logits_metadata.gathered_buffer
_trace_e2e_logits(
"dp_hidden_gather_enter",
global_shape=tuple(hidden_states.shape),
local_shape=tuple(local_hidden_states.shape),
)
dp_gather_replicate(hidden_states, local_hidden_states, logits_metadata)
_trace_e2e_logits("dp_hidden_gather_collective_returned")
return hidden_states, local_hidden_states
return hidden_states, hidden_states
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+14 -3
View File
@@ -206,6 +206,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
num_experts=moe_runner_config.num_experts,
num_local_experts=moe_runner_config.num_local_experts,
hidden_size=moe_runner_config.hidden_size,
moe_runner_config=moe_runner_config,
)
else:
raise NotImplementedError(f"Unsupported a2a backend: {a2a_backend}")
@@ -283,6 +284,8 @@ class FusedMoE(torch.nn.Module):
reduce_results: Whether to apply all_reduce on the output of the layer
quant_config: Quantization configuration.
inplace: suggestion to compute inplace (modify input activation).
enable_qwen35_fp8_deferred_finalize: Whether this concrete Qwen3.5
layer may expose FlashInfer's block-FP8 deferred MoE output.
"""
# True on shared-expert FusedMoE subclasses (e.g. Inkling's sink); lets
@@ -321,6 +324,7 @@ class FusedMoE(torch.nn.Module):
routing_method_type: Optional[RoutingMethodType] = None,
is_gated: bool = True,
gate_up_interleaved: bool = True,
enable_qwen35_fp8_deferred_finalize: bool = False,
):
super().__init__()
if params_dtype is None:
@@ -446,10 +450,17 @@ class FusedMoE(torch.nn.Module):
)
_validate_hpc_ops_quant_method(self.quant_method)
_validate_deepep_v2_quant_method(self.quant_method)
nvfp4_deferred = envs.SGLANG_ENABLE_MOE_DEFERRED_FINALIZE.get() and isinstance(
self.quant_method, ModelOptNvFp4FusedMoEMethod
)
qwen35_fp8_deferred = (
enable_qwen35_fp8_deferred_finalize
and isinstance(self.quant_method, Fp8MoEMethod)
and self.quant_method.block_quant
)
self.supports_deferred_finalize = (
envs.SGLANG_ENABLE_MOE_DEFERRED_FINALIZE.get()
and get_moe_runner_backend().is_flashinfer_trtllm()
and isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod)
get_moe_runner_backend().is_flashinfer_trtllm()
and (nvfp4_deferred or qwen35_fp8_deferred)
)
global _deferred_finalize_info_logged
if not _deferred_finalize_info_logged:
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
+134 -3
View File
@@ -30,7 +30,7 @@ from sglang.srt.layers.moe.moe_runner.base import (
register_pre_permute,
)
from sglang.srt.layers.moe.utils import MoeRunnerBackend, get_moe_a2a_backend
from sglang.srt.runtime_context import get_exec
from sglang.srt.runtime_context import get_exec, get_flags
from sglang.srt.utils import (
ceil_div,
dispose_tensor,
@@ -53,6 +53,9 @@ if TYPE_CHECKING:
DeepEPv2CombineInput,
DeepEPv2DispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.flashinfer import (
FlashinferDispatchOutput,
)
from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardCombineInput,
StandardDispatchOutput,
@@ -161,6 +164,14 @@ def _should_use_masked_standard_layout(
quant_info: DeepGemmMoeQuantInfo,
hidden_states: torch.Tensor,
) -> bool:
# Preserve the Oakhaven WideEP escape hatch while adopting upstream's
# memory-budget-based auto policy. CUDA graph capture remains masked.
if (
envs.SGLANG_OPT_DG_COMPACT_EAGER.get()
and not get_flags().capture.disable_dispose_tensor
):
return False
mode = envs.SGLANG_DEEPGEMM_STANDARD_LAYOUT.get().lower()
if mode not in ("auto", "masked", "compact"):
raise ValueError(
@@ -306,11 +317,24 @@ class DeepGemmRunnerCore(MoeRunnerCore):
hidden_states_dtype = running_state["hidden_states_dtype"]
hidden_states_shape = running_state["hidden_states_shape"]
m_indices = runner_input.m_indices
trace_deepep_v2_contig = (
envs.SGLANG_DEEPEP_V2_TRACE_CONTIG.get()
and running_state.get("deepep_v2_expanded", False)
)
N = quant_info.w13_weight.size(1)
K = hidden_states_shape[1]
scale_block_size = 128
if all_tokens == 0:
if trace_deepep_v2_contig:
logger.warning("DeepEP v2 expanded contig runner empty return")
dispose_tensor(hidden_states)
dispose_tensor(hidden_states_scale)
return torch.empty(
(0, K), device=hidden_states_device, dtype=torch.bfloat16
)
recipe_a, recipe_b = (
((1, 128), (1, 32)) if quant_info.is_fp4_experts else (None, None)
)
@@ -340,6 +364,9 @@ class DeepGemmRunnerCore(MoeRunnerCore):
recipe_a=recipe_a,
recipe_b=recipe_b,
)
if trace_deepep_v2_contig:
torch.cuda.synchronize()
logger.warning("DeepEP v2 expanded contig gateup GEMM returned")
dispose_tensor(hidden_states)
dispose_tensor(hidden_states_scale)
@@ -400,6 +427,9 @@ class DeepGemmRunnerCore(MoeRunnerCore):
del down_input
elif self.use_swizzle:
swiglu_limit_arg: Optional[float] = self.swiglu_limit
use_contig_swizzle = self.use_swizzle and not running_state.get(
"deepep_v2_disable_contig_swizzle", False
)
down_input_fp8 = torch.empty(
(all_tokens, N // 2),
@@ -422,7 +452,7 @@ class DeepGemmRunnerCore(MoeRunnerCore):
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
transposed=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
swiglu_limit=swiglu_limit_arg,
swizzle=self.use_swizzle,
swizzle=use_contig_swizzle,
)
del gateup_output
else:
@@ -454,6 +484,9 @@ class DeepGemmRunnerCore(MoeRunnerCore):
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
)
del down_input
if trace_deepep_v2_contig:
torch.cuda.synchronize()
logger.warning("DeepEP v2 expanded contig activation returned")
# Allocate the MoE output in the NCCL symmetric memory pool when symmetric
# allocation is required, so the downstream all-reduce takes the low-latency
@@ -478,6 +511,9 @@ class DeepGemmRunnerCore(MoeRunnerCore):
recipe_a=recipe_a,
recipe_b=recipe_b,
)
if trace_deepep_v2_contig:
torch.cuda.synchronize()
logger.warning("DeepEP v2 expanded contig down GEMM returned")
return down_output
@@ -568,6 +604,22 @@ class DeepGemmRunnerCore(MoeRunnerCore):
w2_scale = quant_info.w2_scale
hidden_states_device = running_state["hidden_states_device"]
trace_deepep_v2_masked = envs.SGLANG_DEEPEP_V2_TRACE_MASKED.get()
if trace_deepep_v2_masked:
logger.warning(
"DeepEP v2 masked runner enter: hidden=%s hidden_stride=%s "
"scale=%s scale_stride=%s masked_m=%s expected_m=%s",
tuple(hidden_states.shape),
hidden_states.stride(),
(
None
if hidden_states_scale is None
else tuple(hidden_states_scale.shape)
),
None if hidden_states_scale is None else hidden_states_scale.stride(),
masked_m.detach().cpu().tolist(),
expected_m,
)
use_mxfp8 = quant_info.use_mxfp8
scale_block_size = quant_info.block_shape[1] if quant_info.block_shape else 128
@@ -632,6 +684,9 @@ class DeepGemmRunnerCore(MoeRunnerCore):
recipe_a=recipe_a,
recipe_b=recipe_b,
)
if trace_deepep_v2_masked:
torch.cuda.synchronize()
logger.warning("DeepEP v2 masked runner gateup GEMM returned")
dispose_tensor(hidden_states)
dispose_tensor(hidden_states_scale)
@@ -672,6 +727,9 @@ class DeepGemmRunnerCore(MoeRunnerCore):
gemm1_clamp_limit=self.config.gemm1_clamp_limit,
num_real_tokens=num_real_tokens,
)
if trace_deepep_v2_masked:
torch.cuda.synchronize()
logger.warning("DeepEP v2 masked runner activation returned")
del gateup_output
# Down activation is quantised locally at scale_block_size (never DeepEP-LL),
@@ -730,6 +788,9 @@ class DeepGemmRunnerCore(MoeRunnerCore):
recipe_b=recipe_b,
**gemm_overlap_args_dict,
)
if trace_deepep_v2_masked:
torch.cuda.synchronize()
logger.warning("DeepEP v2 masked runner down GEMM returned")
meta_overlap_args = running_state.get("meta_overlap_args", None)
# Returns (block_m, threshold) only with down-gemm overlap, else None;
# meta_overlap_args may be set without overlap, so guard the unpack.
@@ -818,6 +879,7 @@ def pre_permute_standard_to_deep_gemm(
quant_info: DeepGemmMoeQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
expert_start: int = 0,
) -> DeepGemmRunnerInput:
from sglang.kernels.ops.moe.ep_moe_kernels import (
ep_scatter,
@@ -853,6 +915,7 @@ def pre_permute_standard_to_deep_gemm(
quant_info.block_shape,
output_dtype=output_dtype,
use_mxfp8=quant_info.use_mxfp8,
expert_start=expert_start,
)
)
# Use the global expert count because expected_m is a tuning hint, not
@@ -894,7 +957,7 @@ def pre_permute_standard_to_deep_gemm(
all_tokens = _get_compact_all_tokens(num_assignments, num_experts, block_e)
tokens_per_expert, unused_masked_dst = fused_moe_dispatch_index(
topk_ids, num_experts, 1
topk_ids, num_experts, 1, expert_start=expert_start
)
dispose_tensor(unused_masked_dst)
valid_tokens_per_expert = tokens_per_expert
@@ -974,6 +1037,7 @@ def pre_permute_standard_to_deep_gemm(
src2dst,
scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
quant_block_size=(quant_info.block_shape[1] if quant_info.block_shape else 128),
expert_start=expert_start,
)
if packed_input_source is not hidden_states:
dispose_tensor(packed_input_source)
@@ -1003,6 +1067,49 @@ def pre_permute_standard_to_deep_gemm(
)
@register_pre_permute("flashinfer", "deep_gemm")
def pre_permute_flashinfer_to_deep_gemm(
dispatch_output: FlashinferDispatchOutput,
quant_info: DeepGemmMoeQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> DeepGemmRunnerInput:
"""Feed one-sided A2A output into DeepGEMM with fused expert remapping."""
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
from sglang.srt.runtime_context import get_parallel
if dispatch_output.hidden_states.dtype != torch.bfloat16:
raise TypeError(
"FlashInfer A2A + DeepGEMM requires a BF16 dispatch payload, got "
f"{dispatch_output.hidden_states.dtype}."
)
if dispatch_output.hidden_states_scale is not None:
raise ValueError(
"FlashInfer A2A + DeepGEMM expects unquantized BF16 dispatch; "
"hidden_states_scale must be None."
)
if dispatch_output.topk_output.topk_ids.dtype != torch.int32:
raise TypeError(
"FlashInfer A2A expert IDs must be int32 before DeepGEMM, got "
f"{dispatch_output.topk_output.topk_ids.dtype}."
)
standard_output = StandardDispatchOutput(
hidden_states=dispatch_output.hidden_states,
hidden_states_scale=None,
topk_output=dispatch_output.topk_output,
)
expert_start = get_parallel().moe_ep_rank * runner_config.num_local_experts
return pre_permute_standard_to_deep_gemm(
standard_output,
quant_info,
runner_config,
running_state,
expert_start=expert_start,
)
@register_post_permute("deep_gemm", "standard")
def post_permute_deep_gemm_to_standard(
runner_output: DeepGemmRunnerOutput,
@@ -1047,6 +1154,30 @@ def post_permute_deep_gemm_to_standard(
)
@register_post_permute("deep_gemm", "flashinfer")
def post_permute_deep_gemm_to_flashinfer(
runner_output: DeepGemmRunnerOutput,
quant_info: DeepGemmMoeQuantInfo,
runner_config: MoeRunnerConfig,
running_state: dict,
):
"""Reuse DeepGEMM's weighted post-permute and hand BF16 to A2A combine."""
from sglang.srt.layers.moe.token_dispatcher.flashinfer import (
FlashinferCombineInput,
)
standard_input = post_permute_deep_gemm_to_standard(
runner_output, quant_info, runner_config, running_state
)
if standard_input.hidden_states.dtype != torch.bfloat16:
raise TypeError(
"FlashInfer A2A + DeepGEMM combine payload must be BF16, got "
f"{standard_input.hidden_states.dtype}."
)
return FlashinferCombineInput(hidden_states=standard_input.hidden_states)
@register_pre_permute("deepep_ll", "deep_gemm")
def pre_permute_deepep_ll_to_deep_gemm(
dispatch_output: DeepEPLLDispatchOutput,
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py
+9 -2
View File
@@ -496,10 +496,10 @@ def fused_experts_none_to_flashinfer_cutedsl_fp4(
@register_fused_func("flashinfer", "flashinfer_cutedsl")
def fused_experts_flashinfer_to_flashinfer_cutedsl_fp4(
dispatch_output: FlashinferDispatchOutput,
dispatch_output: FlashinferDispatchOutput | StandardDispatchOutput,
quant_info: CuteDslFp4MoeQuantInfo,
runner_config: MoeRunnerConfig,
) -> FlashinferCombineInput:
) -> FlashinferCombineInput | StandardCombineInput:
"""CuteDSL fused func for flashinfer alltoall dispatcher.
Two cases depending on whether the dispatcher did FP4 quantization:
@@ -509,6 +509,10 @@ def fused_experts_flashinfer_to_flashinfer_cutedsl_fp4(
from sglang.srt.layers.moe.token_dispatcher.flashinfer import (
FlashinferCombineInput,
)
from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardCombineInput,
StandardDispatchOutput,
)
from sglang.srt.layers.moe.topk import TopKOutputChecker
from sglang.srt.layers.quantization.fp4_utils import fp4_quantize
@@ -580,6 +584,9 @@ def fused_experts_flashinfer_to_flashinfer_cutedsl_fp4(
# Note: output contains routed expert results; shared_expert is handled separately
if isinstance(dispatch_output, StandardDispatchOutput):
return StandardCombineInput(hidden_states=output)
# Write into pre-allocated workspace buffer if available
if dispatch_output.moe_output is not None:
dispatch_output.moe_output.copy_(output)
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+120 -26
View File
@@ -92,6 +92,41 @@ def finalize_flashinfer_trtllm_deferred_output(
)
def _make_deferred_finalize_output(
result,
*,
top_k: int,
) -> FlashInferTrtllmDeferredFinalizeOutput:
"""Validate and adapt FlashInfer's ``do_finalize=False`` output ABI."""
gemm2_out, expert_weights, expanded_idx_to_permuted_idx = result[:3]
# Some FlashInfer versions size this buffer from routing_logits dtype while
# writing BF16 weights into it. Reinterpret only the live BF16 prefix.
if expert_weights.dtype == torch.float32:
n, k = expert_weights.shape
expert_weights = expert_weights.view(torch.bfloat16).view(-1, k)[:n]
if expert_weights.dtype != torch.bfloat16:
raise RuntimeError(
"FlashInfer deferred finalize must return BF16 expert weights, got "
f"{expert_weights.dtype}"
)
if gemm2_out.dtype != torch.bfloat16:
raise RuntimeError(
"FlashInfer deferred finalize must return BF16 GEMM2 output, got "
f"{gemm2_out.dtype}"
)
if expanded_idx_to_permuted_idx.dtype != torch.int32:
raise RuntimeError(
"FlashInfer deferred finalize must return Int32 permuted indices, got "
f"{expanded_idx_to_permuted_idx.dtype}"
)
return FlashInferTrtllmDeferredFinalizeOutput(
gemm2_out=gemm2_out,
expert_weights=expert_weights,
expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx,
top_k=top_k,
)
def round_up_to_multiple(x: int, m: int) -> int:
"""Round up *x* to the nearest multiple of *m*."""
return (x + m - 1) // m * m
@@ -714,6 +749,16 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
else Fp8QuantizationType.DeepSeekFp8
)
use_shuffled_weight = quant_info.use_mxfp8
defer_finalize = _deferred_finalize_enabled.get()
if defer_finalize and (
not quant_info.block_quant
or use_routed_topk
or not TopKOutputChecker.format_is_bypassed(topk_output)
):
raise RuntimeError(
"FP8 deferred finalize requires block quantization, the logits-based "
"FlashInfer TRTLLM backend, and bypassed TopK"
)
if quant_info.block_quant:
assert quant_info.weight_block_k is not None
@@ -736,16 +781,19 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
)
a_sf_t = a_sf.t()
# Allocate output inside symmetric memory context
with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
):
symm_output = torch.empty(
hidden_states.shape[0],
hidden_states.shape[1],
dtype=hidden_states.dtype,
device=hidden_states.device,
)
symm_output = None
if not defer_finalize:
# The deferred path returns FlashInfer's permuted/padded GEMM2
# materialization and must not allocate the ordinary final output.
with use_symmetric_memory(
get_tp_group(), disabled=not is_allocation_symmetric()
):
symm_output = torch.empty(
hidden_states.shape[0],
hidden_states.shape[1],
dtype=hidden_states.dtype,
device=hidden_states.device,
)
# Move kernel call outside context manager to avoid graph breaks
# during torch.compile for piecewise cuda graph.
@@ -791,10 +839,10 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
fp8_quantization_type=int(fp8_quantization_type),
activation_type=quant_info.activation_type,
)
output = cast(torch.Tensor, symm_output)
else:
assert TopKOutputChecker.format_is_bypassed(topk_output)
trtllm_fp8_block_scale_moe_out_wrapper(
common_kwargs = dict(
routing_logits=router_logits,
routing_bias=correction_bias,
hidden_states=a_q,
@@ -806,7 +854,6 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
gemm1_clamp_limit=quant_info.gemm1_clamp_limit,
gemm2_weights=quant_info.w2_weight,
gemm2_weights_scale=quant_info.w2_weight_scale_inv,
output=symm_output,
num_experts=quant_info.global_num_experts,
top_k=topk_config.top_k,
n_group=topk_config.num_expert_group,
@@ -822,10 +869,34 @@ def fused_experts_none_to_flashinfer_trtllm_fp8(
routing_method_type=routing_method_type,
use_shuffled_weight=use_shuffled_weight,
tune_max_num_tokens=next_power_of_2(a_q.shape[0]),
fp8_quantization_type=int(fp8_quantization_type),
activation_type=quant_info.activation_type,
)
output = symm_output
if defer_finalize:
from flashinfer.fused_moe import trtllm_fp8_block_scale_moe
deferred_kwargs = dict(
**common_kwargs,
do_finalize=False,
fp8_quantization_type=fp8_quantization_type,
enable_pdl=a_q.shape[0] <= _TRTLLM_MOE_PDL_MAX_TOKENS,
)
if quant_info.activation_type is not None:
from flashinfer.fused_moe.core import ActivationType
deferred_kwargs["activation_type"] = ActivationType(
quant_info.activation_type
)
output = _make_deferred_finalize_output(
trtllm_fp8_block_scale_moe(**deferred_kwargs),
top_k=topk_config.top_k,
)
else:
trtllm_fp8_block_scale_moe_out_wrapper(
**common_kwargs,
output=cast(torch.Tensor, symm_output),
fp8_quantization_type=int(fp8_quantization_type),
activation_type=quant_info.activation_type,
)
output = cast(torch.Tensor, symm_output)
else:
assert TopKOutputChecker.format_is_bypassed(topk_output)
assert quant_info.w13_input_scale is not None
@@ -1344,23 +1415,26 @@ def fused_experts_none_to_flashinfer_trtllm_routed(
)
@register_fused_func("flashinfer", "flashinfer_trtllm")
@register_fused_func("flashinfer", "flashinfer_trtllm_routed")
def fused_experts_flashinfer_to_flashinfer_trtllm_routed(
dispatch_output: FlashinferDispatchOutput,
def fused_experts_flashinfer_to_flashinfer_trtllm(
dispatch_output: FlashinferDispatchOutput | StandardDispatchOutput,
quant_info: MoeQuantInfo,
runner_config: MoeRunnerConfig,
) -> FlashinferCombineInput:
"""Fused function for flashinfer A2A + flashinfer_trtllm_routed runner.
) -> FlashinferCombineInput | StandardCombineInput:
"""Fused function for FlashInfer A2A + TRT-LLM Gen MoE.
FlashinferDispatchOutput and StandardDispatchOutput share the same field
layout (hidden_states, hidden_states_scale, topk_output), so the existing
FP8/FP4/BF16 implementations work unchanged. We wrap the returned
StandardCombineInput into a FlashinferCombineInput for the FlashinferDispatcher
combine path.
Both one-sided decode and AG+RS prefill materialize routing IDs and weights,
so the regular and explicitly-routed backend names enter TRT-LLM's routed
kernel. The dispatch formats share the fields consumed by the implementation;
only the combine wrapper differs.
"""
from sglang.srt.layers.moe.token_dispatcher.flashinfer import (
FlashinferCombineInput,
)
from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardDispatchOutput,
)
if isinstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo):
result = fused_experts_none_to_flashinfer_trtllm_fp4(
@@ -1370,6 +1444,16 @@ def fused_experts_flashinfer_to_flashinfer_trtllm_routed(
use_routed_topk=True,
)
elif isinstance(quant_info, FlashInferTrtllmFp8MoeQuantInfo):
if dispatch_output.hidden_states.dtype != torch.bfloat16:
raise TypeError(
"FlashInfer A2A + TRT-LLM Gen FP8 MoE requires a BF16 "
f"dispatch payload, got {dispatch_output.hidden_states.dtype}."
)
if dispatch_output.hidden_states_scale is not None:
raise ValueError(
"FlashInfer A2A + TRT-LLM Gen FP8 MoE quantizes locally; "
"the BF16 dispatch payload must not carry activation scales."
)
result = fused_experts_none_to_flashinfer_trtllm_fp8(
dispatch_output,
quant_info,
@@ -1385,8 +1469,18 @@ def fused_experts_flashinfer_to_flashinfer_trtllm_routed(
)
else:
raise TypeError(
f"Unexpected quant_info type for flashinfer a2a + flashinfer_trtllm_routed: {type(quant_info)}"
f"Unexpected quant_info type for flashinfer a2a + flashinfer_trtllm: {type(quant_info)}"
)
if (
isinstance(quant_info, FlashInferTrtllmFp8MoeQuantInfo)
and result.hidden_states.dtype != torch.bfloat16
):
raise TypeError(
"FlashInfer A2A + TRT-LLM Gen FP8 MoE must return a BF16 combine "
f"payload, got {result.hidden_states.dtype}."
)
if isinstance(dispatch_output, StandardDispatchOutput):
return result
return FlashinferCombineInput(hidden_states=result.hidden_states)
python/sglang/srt/layers/moe/qwen35_flashinfer_fusion.py
+347
View File
@@ -0,0 +1,347 @@
"""Qwen3.5 integration for FlashInfer MNNVL CuTe DSL AllReduce fusion."""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import Optional
import torch
from sglang.srt.arg_groups.overrides import resolving_view
from sglang.srt.layers.communicator import (
CommunicateWithAllReduceAndLayerNormFn,
LayerCommunicator,
ScatterMode,
get_attn_tp_context,
)
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.moe import get_moe_a2a_backend
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.runtime_context import get_exec, get_parallel
logger = logging.getLogger(__name__)
def is_supported_forward_mode(forward_mode: ForwardMode) -> bool:
return forward_mode in (
ForwardMode.DECODE,
ForwardMode.EXTEND,
ForwardMode.TARGET_VERIFY,
)
def resolve_max_m(model_runner) -> int:
"""Use framework token bounds as the workspace-capacity source of truth."""
server_args = resolving_view(model_runner.server_args)
decode_config = server_args.cuda_graph_config.decode
prefill_config = server_args.cuda_graph_config.prefill
candidates = [
server_args.cutedsl_moe_max_num_tokens(),
model_runner.max_running_requests,
decode_config.max_bs,
prefill_config.max_bs,
*(decode_config.bs or []),
*(prefill_config.bs or []),
]
positive = [
int(value) for value in candidates if value is not None and int(value) > 0
]
if not positive:
raise RuntimeError("framework reported no positive fusion workspace M bound")
return max(positive)
@dataclass(frozen=True)
class Qwen35MoeFinalizeHandoff:
"""Unfinalized routed output plus the separately gated shared contribution."""
routed_output: torch.Tensor
expert_weights: torch.Tensor
permuted_indices: torch.Tensor
gated_shared_output: torch.Tensor
m: int
@classmethod
def from_flashinfer(
cls,
deferred_output,
*,
gated_shared_output: torch.Tensor,
m: int,
) -> Qwen35MoeFinalizeHandoff:
top_k = int(deferred_output.top_k)
return cls(
routed_output=deferred_output.gemm2_out.view(
-1, deferred_output.gemm2_out.shape[-1]
),
expert_weights=deferred_output.expert_weights.view(-1, top_k)[:m],
permuted_indices=deferred_output.expanded_idx_to_permuted_idx.view(
-1, top_k
)[:m],
gated_shared_output=gated_shared_output,
m=int(m),
)
class Qwen35FlashInferFusionService:
"""A lightweight model handle for the process-local FlashInfer workspace."""
def __init__(
self,
*,
hidden_size: int,
top_k: int,
rms_epsilon: float,
) -> None:
self.hidden_size = int(hidden_size)
self.top_k = int(top_k)
self.rms_epsilon = float(rms_epsilon)
self.max_m: int | None = None
self._workspace = None
@property
def is_prepared(self) -> bool:
return self._workspace is not None
def prepare(self, *, max_m: int) -> None:
if self._workspace is not None:
assert self.max_m is not None
if int(max_m) > self.max_m:
raise RuntimeError(
f"fusion workspace is already prepared for M_max={self.max_m}; "
f"refusing M_max={max_m}"
)
return
from sglang.srt.layers.flashinfer_mnnvl_cutedsl import (
get_flashinfer_mnnvl_cutedsl_ar_fusion,
)
workspace = get_flashinfer_mnnvl_cutedsl_ar_fusion(
hidden_size=self.hidden_size,
top_k=self.top_k,
max_m=int(max_m),
rms_epsilon=self.rms_epsilon,
# GemmaRMSNorm.gemma_weight is already checkpoint weight + 1.
weight_bias=0.0,
)
self._workspace = workspace
self.max_m = workspace.max_m
def supports(self, m: int) -> bool:
if self._workspace is None or self.max_m is None:
return False
return 1 <= int(m) <= self.max_m and self._workspace.supports(m)
def finalize(
self,
handoff: Qwen35MoeFinalizeHandoff,
residual: torch.Tensor,
gamma: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
self._validate_finalize(handoff, residual, gamma)
assert self._workspace is not None
return self._workspace.moe_finalize_all_reduce_rms_norm(
routed_output=handoff.routed_output,
expert_weights=handoff.expert_weights,
permuted_indices=handoff.permuted_indices,
gated_shared_output=handoff.gated_shared_output,
residual=residual,
gamma=gamma,
)
def all_reduce_residual_rms_norm(
self,
local_contribution: torch.Tensor,
residual: torch.Tensor,
gamma: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
self._validate_matrix(local_contribution, "local_contribution")
self._validate_matrix(residual, "residual", m=local_contribution.shape[0])
self._validate_gamma(gamma)
if not self.supports(local_contribution.shape[0]):
raise ValueError(f"unsupported M={local_contribution.shape[0]}")
assert self._workspace is not None
return self._workspace.all_reduce_residual_rms_norm(
local_contribution=local_contribution,
residual=residual,
gamma=gamma,
)
def _validate_finalize(
self,
handoff: Qwen35MoeFinalizeHandoff,
residual: torch.Tensor,
gamma: torch.Tensor,
) -> None:
if not self.supports(handoff.m):
raise ValueError(f"unsupported M={handoff.m}")
self._validate_matrix(handoff.routed_output, "routed_output", exact_m=False)
expected_metadata = (handoff.m, self.top_k)
if tuple(handoff.expert_weights.shape) != expected_metadata:
raise ValueError("expert_weights must have shape [M, top_k]")
if tuple(handoff.permuted_indices.shape) != expected_metadata:
raise ValueError("permuted_indices must have shape [M, top_k]")
if handoff.expert_weights.dtype != torch.bfloat16:
raise ValueError("expert_weights must be BF16")
if handoff.permuted_indices.dtype != torch.int32:
raise ValueError("permuted_indices must be Int32")
self._validate_matrix(
handoff.gated_shared_output, "gated_shared_output", m=handoff.m
)
self._validate_matrix(residual, "residual", m=handoff.m)
self._validate_gamma(gamma)
def _validate_matrix(
self,
tensor: torch.Tensor,
name: str,
*,
m: int | None = None,
exact_m: bool = True,
) -> None:
if tensor.ndim != 2 or tensor.shape[1] != self.hidden_size:
raise ValueError(f"{name} must have shape [M, hidden_size]")
if m is not None and exact_m and tensor.shape[0] != int(m):
raise ValueError(f"{name} has the wrong M dimension")
if tensor.dtype != torch.bfloat16 or not tensor.is_contiguous():
raise ValueError(f"{name} must be contiguous BF16")
def _validate_gamma(self, gamma: torch.Tensor) -> None:
if (
tuple(gamma.shape) != (self.hidden_size,)
or gamma.dtype != torch.bfloat16
or not gamma.is_contiguous()
):
raise ValueError("gamma must be contiguous BF16 [hidden_size]")
class Qwen35FlashInferLayerCommunicator(LayerCommunicator):
"""Qwen-only hooks; generic LayerCommunicator remains backend agnostic."""
fusion_service: Qwen35FlashInferFusionService | None = None
def prepare_attn(
self,
hidden_states,
residual,
forward_batch,
quant_format: str = "",
post_residual_addition=None,
):
if isinstance(hidden_states, Qwen35MoeFinalizeHandoff):
if not self.should_use_finalize(forward_batch, hidden_states.m):
raise RuntimeError("received deferred MoE output on an ineligible path")
if residual is None:
raise RuntimeError("deferred MoE finalize requires residual input")
if not hasattr(self.input_layernorm, "gemma_weight"):
raise RuntimeError("deferred Qwen finalize requires GemmaRMSNorm")
if post_residual_addition is not None:
residual = residual + post_residual_addition
assert self.fusion_service is not None
return self.fusion_service.finalize(
hidden_states, residual, self.input_layernorm.gemma_weight
)
return super().prepare_attn(
hidden_states,
residual,
forward_batch,
quant_format=quant_format,
post_residual_addition=post_residual_addition,
)
def prepare_mlp(
self,
hidden_states: torch.Tensor,
residual: torch.Tensor,
forward_batch: ForwardBatch,
cache=None,
):
if cache is not None:
self._context.cache = cache
if self.should_use_all_reduce_rms_norm(
forward_batch, int(hidden_states.shape[0]), residual
):
assert self.fusion_service is not None and residual is not None
return self.fusion_service.all_reduce_residual_rms_norm(
hidden_states,
residual,
self.post_attention_layernorm.gemma_weight,
)
return super().prepare_mlp(hidden_states, residual, forward_batch, cache=cache)
def should_use_all_reduce_rms_norm(
self,
forward_batch: ForwardBatch,
m: int,
residual: Optional[torch.Tensor],
) -> bool:
communicate_fn = self._communicate_with_all_reduce_and_layer_norm_fn
norm_fn = getattr(communicate_fn, "func", communicate_fn)
residual_input_mode = getattr(communicate_fn, "keywords", {}).get(
"residual_input_mode"
)
parallel = get_parallel()
return (
self._common_eligible(forward_batch, m)
and residual is not None
and hasattr(self.post_attention_layernorm, "gemma_weight")
and norm_fn
is CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual
and residual_input_mode is ScatterMode.TP_ATTN_FULL
and self._context.attn_dp_size == 1
and parallel.attn_tp_size == parallel.tp_size
and not get_exec().comm.enable_quant_communications
)
def should_use_finalize(self, forward_batch: ForwardBatch, m: int) -> bool:
parallel = get_parallel()
return (
self._common_eligible(forward_batch, m)
and self.layer_scatter_modes.mlp_mode is not ScatterMode.SCATTERED
and parallel.moe_ep_size == 1
)
def _common_eligible(self, forward_batch: ForwardBatch, m: int) -> bool:
parallel = get_parallel()
return bool(
self.fusion_service is not None
and self.fusion_service.is_prepared
and is_supported_forward_mode(forward_batch.forward_mode)
and self.fusion_service.supports(m)
and not is_dp_attention_enabled()
and parallel.attn_cp_size == 1
and not get_attn_tp_context().input_scattered
and get_moe_a2a_backend().is_none()
and self._context.tp_size > 1
)
def should_fuse_mlp_allreduce_with_next_layer(
self, forward_batch: ForwardBatch
) -> bool:
m = (
int(forward_batch.input_ids.shape[0])
if getattr(forward_batch, "input_ids", None) is not None
else 0
)
if self.should_use_finalize(forward_batch, m):
# The Qwen model consumes the final layer's handoff with its final
# GemmaRMSNorm, so this is intentionally also true for that layer.
return True
return super().should_fuse_mlp_allreduce_with_next_layer(forward_batch)
def prepare_qwen35_flashinfer_fusion(model, model_runner) -> None:
service = getattr(model, "flashinfer_mnnvl_cutedsl_fusion", None)
if service is None:
return
if model_runner.server_args.enable_pdmux:
raise RuntimeError(
"FlashInfer MNNVL CuTe DSL fusion does not support concurrent PDMux "
"streams sharing one mutable workspace"
)
service.prepare(max_m=resolve_max_m(model_runner))
logger.info(
"Prepared Qwen3.5 FlashInfer MNNVL CuTe DSL fusion workspace for M_max=%d",
service.max_m,
)
python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py
+202 -26
View File
@@ -9,6 +9,7 @@ from sglang.kernels.kernel_api_logging import debug_kernel_api
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
get_dp_global_num_tokens,
get_is_extend_in_batch,
is_dp_attention_enabled,
)
from sglang.srt.layers.moe.token_dispatcher import (
@@ -21,13 +22,18 @@ from sglang.srt.layers.moe.token_dispatcher import (
from sglang.srt.layers.moe.token_dispatcher.flashinfer_utils import (
TorchDistributedCommBackend,
)
from sglang.srt.layers.moe.token_dispatcher.standard import (
StandardCombineInput,
StandardDispatcher,
StandardDispatchOutput,
)
from sglang.srt.layers.moe.topk import (
StandardTopKOutput,
TopKOutput,
TopKOutputChecker,
)
from sglang.srt.layers.moe.utils import get_moe_runner_backend
from sglang.srt.runtime_context import get_schedule, get_spec
from sglang.srt.runtime_context import get_flags, get_parallel, get_schedule, get_spec
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
try:
@@ -46,6 +52,36 @@ logger = logging.getLogger(__name__)
MOE_NVFP4_DISPATCH = envs.SGLANG_MOE_NVFP4_DISPATCH.get()
# FlashInfer keys MNNVL allocations by workspace size; aligned tail padding gives
# concurrently live paths distinct persistent workspaces without extra token work.
_WORKSPACE_NAMESPACE_ALIGNMENT = 128
def _max_tokens_per_scattered_source(
dp_global_num_tokens: list[int], attn_tp_size: int
) -> int:
assert attn_tp_size > 0
max_dp_tokens = max(dp_global_num_tokens)
return (max_dp_tokens + attn_tp_size - 1) // attn_tp_size
def _scattered_source_token_counts(
dp_global_num_tokens: list[int], attn_tp_size: int
) -> list[int]:
assert attn_tp_size > 0
counts = []
for num_tokens in dp_global_num_tokens:
base, remainder = divmod(num_tokens, attn_tp_size)
counts.extend(
base + int(attn_tp_rank < remainder) for attn_tp_rank in range(attn_tp_size)
)
return counts
def _workspace_size_for_namespace(workspace_size: int, *, speculative: bool) -> int:
slot = int(speculative)
return workspace_size + slot * _WORKSPACE_NAMESPACE_ALIGNMENT
class FlashinferDispatchOutput(NamedTuple):
"""Flashinfer EP dispatch output."""
@@ -88,6 +124,7 @@ class FlashinferDispatcher(BaseDispatcher):
num_local_experts: int = None, # Unused
hidden_size: int = None,
params_dtype: torch.dtype = None, # Unused
moe_runner_config=None,
):
super().__init__()
if not use_flashinfer:
@@ -102,13 +139,28 @@ class FlashinferDispatcher(BaseDispatcher):
self.hidden_size = hidden_size
self.num_experts = num_experts
self.num_local_experts = num_local_experts
runner_backend = get_moe_runner_backend()
self.invalid_token_expert_id = (
-1
if get_moe_runner_backend().is_flashinfer_trtllm_routed()
if (
runner_backend.is_deep_gemm()
or runner_backend.is_flashinfer_trtllm()
or runner_backend.is_flashinfer_trtllm_routed()
)
else self.num_experts
)
# TODO: Can other moe runners use payload_in_workspace too?
self.payload_in_workspace = get_moe_runner_backend().is_flashinfer_cutlass()
if moe_runner_config is None:
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
moe_runner_config = MoeRunnerConfig(
num_experts=num_experts,
num_local_experts=num_local_experts,
hidden_size=hidden_size,
top_k=router_topk,
)
self.prefill_dispatcher = StandardDispatcher(moe_runner_config)
# FlashInfer sizes the workspace from the maximum dispatched tokens per
# EP rank. See FlashInfer's moe_a2a_get_workspace_size_per_rank(),
@@ -167,19 +219,106 @@ class FlashinferDispatcher(BaseDispatcher):
pp_size=1,
cp_size=1,
)
self.moe_a2a = MoeAlltoAll(
mapping=self.mapping,
max_num_tokens=self.max_num_tokens,
top_k=self.router_topk,
num_experts=self.num_experts,
workspace_size_per_rank=self.workspace_size,
mnnvl_config=MnnvlConfig(comm_backend=TorchDistributedCommBackend(group)),
mnnvl_config = MnnvlConfig(comm_backend=TorchDistributedCommBackend(group))
is_speculative_model = get_flags().moe.speculative_context
def make_moe_a2a() -> MoeAlltoAll:
# Target and draft decode graphs can coexist; prefill/mixed extend use
# AG+RS and do not lease MNNVL A2A workspaces.
workspace_size = _workspace_size_for_namespace(
self.workspace_size,
speculative=is_speculative_model,
)
return MoeAlltoAll(
mapping=self.mapping,
max_num_tokens=self.max_num_tokens,
top_k=self.router_topk,
num_experts=self.num_experts,
workspace_size_per_rank=workspace_size,
mnnvl_config=mnnvl_config,
)
self.moe_a2a = make_moe_a2a()
def set_quant_config(self, quant_config: dict) -> None:
super().set_quant_config(quant_config)
self.prefill_dispatcher.set_quant_config(quant_config)
def _dispatch_prefill_allgather(
self, hidden_states: torch.Tensor, topk_output: TopKOutput
) -> StandardDispatchOutput:
# Eager extend can overlap another stream, so use BF16 all-gatherv instead
# of reusing pure-decode A2A signal state across streams.
if hidden_states.dtype != torch.bfloat16:
raise TypeError(
"FlashInfer WideEP prefill AG requires BF16 hidden states, got "
f"{hidden_states.dtype}."
)
if TopKOutputChecker.format_is_bypassed(topk_output):
topk_output = topk_output.to_standard()
if not TopKOutputChecker.format_is_standard(topk_output):
raise TypeError(
"FlashInfer WideEP prefill AG requires materialized top-k "
f"routing, got {type(topk_output).__name__}."
)
dp_global = get_dp_global_num_tokens()
if dp_global is None:
source_sizes = [hidden_states.shape[0]] * self.ep_size
else:
source_sizes = _scattered_source_token_counts(
dp_global, get_parallel().attn_tp_size
)
if len(source_sizes) != self.ep_size:
raise RuntimeError(
"FlashInfer WideEP prefill AG source geometry does not match "
f"EP: len(source_sizes)={len(source_sizes)}, ep_size={self.ep_size}."
)
if source_sizes[self.ep_rank] != hidden_states.shape[0]:
raise RuntimeError(
"FlashInfer WideEP prefill AG local source geometry mismatch: "
f"source_sizes[{self.ep_rank}]={source_sizes[self.ep_rank]} != "
f"hidden_states.shape[0]={hidden_states.shape[0]}."
)
topk_ids = topk_output.topk_ids.to(torch.int32)
hidden_states, topk_ids, topk_weights = get_parallel().tp_group.all_gatherv(
[hidden_states, topk_ids, topk_output.topk_weights],
sizes=source_sizes,
)
self.prefill_source_sizes = source_sizes
return self.prefill_dispatcher.dispatch(
hidden_states,
StandardTopKOutput(topk_weights, topk_ids, topk_output.router_logits),
)
@debug_kernel_api
def dispatch(
self, hidden_states: torch.Tensor, topk_output: TopKOutput
) -> FlashinferDispatchOutput:
) -> FlashinferDispatchOutput | StandardDispatchOutput:
if get_is_extend_in_batch():
return self._dispatch_prefill_allgather(hidden_states, topk_output)
self.active_moe_a2a = self.moe_a2a
# Block-wise FP8 runners quantize before GEMM, so keep dispatch/combine BF16;
# FP4 retains its packed wire path keyed by input_global_scale.
runner_backend = get_moe_runner_backend()
weight_dtype = self.quant_config.get("weight_dtype")
uses_bf16_fp8_payload = weight_dtype in (
torch.float8_e4m3fn,
torch.float8_e5m2,
) and (
runner_backend.is_deep_gemm()
or runner_backend.is_flashinfer_trtllm()
or runner_backend.is_flashinfer_trtllm_routed()
)
if uses_bf16_fp8_payload and hidden_states.dtype != torch.bfloat16:
raise TypeError(
"FlashInfer A2A with an FP8 DeepGEMM/TRT-LLM Gen MoE runner "
"requires BF16 dispatch and combine payloads, but received "
f"{hidden_states.dtype}."
)
output_dtype = hidden_states.dtype
x = hidden_states
x_sf = None
@@ -218,16 +357,8 @@ class FlashinferDispatcher(BaseDispatcher):
# CUDA-graph *capture*; on *replay* dispatch() is not re-executed and the
# value baked at capture is reused. Two cases, both rank-invariant:
#
# Case 1 — max(dp_global): DP attention feeding EP. The scheduler
# all-gathers per-DP-rank token counts into dp_global (length dp_size,
# identical on every rank), which differ across ranks, so we must take
# the max. FlashInfer A2A forces require_mlp_tp_gather=True (see
# require_mlp_tp_gather()), so: eager reads the live list; capture sees
# [num_tokens] * dp_size (uniform capture bs) and bakes max() == the
# bucket; replay reuses that baked value and every rank replays the same
# bucket because the decode graph runner sizes it from the cross-rank
# max. Without this, per-rank buckets could diverge -> geometry mismatch
# -> illegal memory access (issue #30242).
# Each EP source owns ceil(max(dp_global) / attn_tp_size) tokens; the
# shared maximum keeps graph geometry rank-uniform (issue #30242).
#
# Case 2 — x.shape[0]: no per-rank DP list (dp_global absent or scalar).
# This is SP attention feeding EP (tokens are sequence-parallel scattered
@@ -237,7 +368,10 @@ class FlashinferDispatcher(BaseDispatcher):
dp_global = get_dp_global_num_tokens()
if dp_global is not None and len(dp_global) > 1:
# Case 1
self.runtime_max_tokens_per_rank = max(dp_global)
attn_tp_size = get_parallel().attn_tp_size
self.runtime_max_tokens_per_rank = _max_tokens_per_scattered_source(
dp_global, attn_tp_size
)
else:
# Case 2. Guard against the #30242 failure mode: DP attention must
# never land here with ep_size > 1, because there x.shape[0] differs
@@ -252,9 +386,20 @@ class FlashinferDispatcher(BaseDispatcher):
)
self.runtime_max_tokens_per_rank = x.shape[0]
# MoeAlltoAll does not resize its max_num_tokens workspace; reject larger
# runtime geometry here before it becomes an illegal memory access.
assert self.runtime_max_tokens_per_rank <= self.max_num_tokens, (
"FlashInfer A2A runtime token geometry exceeds its fixed workspace: "
f"runtime_max_tokens_per_rank={self.runtime_max_tokens_per_rank} > "
f"max_num_tokens={self.max_num_tokens}. Increase "
"SGLANG_FLASHINFER_NUM_MAX_DISPATCH_TOKENS_PER_RANK to cover the "
"largest mixed prefill and speculative-verify batch."
)
# The recv buffer reserves runtime_max_tokens_per_rank slots for THIS
# rank, so it must cover this rank's own tokens. This holds in both cases
# (Case 1: max(dp_global) >= the local count; Case 2: exactly x.shape[0]),
# (Case 1: ceil(max(dp_global) / attn_tp_size) covers every token-scatter
# shard; Case 2: exactly x.shape[0]),
# so a violation signals a sizing/plumbing bug (e.g. an un-adjusted spec
# count) rather than a benign case.
assert self.runtime_max_tokens_per_rank >= x.shape[0], (
@@ -268,7 +413,7 @@ class FlashinferDispatcher(BaseDispatcher):
# padding slots whose expert_id would otherwise route to a real expert
# and waste downstream MoE compute. Sanitizing the padding to a
# sentinel id is structural, not optional.
recv_tensors = self.moe_a2a.dispatch(
recv_tensors = self.active_moe_a2a.dispatch(
topk_ids,
payloads,
self.runtime_max_tokens_per_rank,
@@ -290,7 +435,7 @@ class FlashinferDispatcher(BaseDispatcher):
# Provide an output tensor to fused_moe so it writes directly to our buffer
moe_output = None
if self.payload_in_workspace:
moe_output = self.moe_a2a.get_combine_payload_tensor_in_workspace(
moe_output = self.active_moe_a2a.get_combine_payload_tensor_in_workspace(
self.runtime_max_tokens_per_rank, self.hidden_size, output_dtype
).view(-1, self.hidden_size)
return FlashinferDispatchOutput(
@@ -301,10 +446,40 @@ class FlashinferDispatcher(BaseDispatcher):
)
@debug_kernel_api
def combine(self, combine_input: FlashinferCombineInput) -> torch.Tensor:
def combine(
self, combine_input: FlashinferCombineInput | StandardCombineInput
) -> torch.Tensor:
hidden_states = combine_input.hidden_states
if combine_input.format == CombineInputFormat.STANDARD:
if hidden_states.dtype != torch.bfloat16:
raise TypeError(
"FlashInfer WideEP prefill RS requires BF16 expert output, "
f"got {hidden_states.dtype}."
)
source_sizes = self.prefill_source_sizes
hidden_states = get_parallel().tp_group.reduce_scatterv(
hidden_states, sizes=source_sizes
)
del self.prefill_source_sizes
return hidden_states
weight_dtype = self.quant_config.get("weight_dtype")
runner_backend = get_moe_runner_backend()
if (
weight_dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
and (
runner_backend.is_deep_gemm()
or runner_backend.is_flashinfer_trtllm()
or runner_backend.is_flashinfer_trtllm_routed()
)
and hidden_states.dtype != torch.bfloat16
):
raise TypeError(
"FlashInfer A2A FP8 MoE combine payload must be BF16, but "
f"received {hidden_states.dtype}."
)
output_hidden_size = hidden_states.shape[-1]
hidden_states = self.moe_a2a.combine(
hidden_states = self.active_moe_a2a.combine(
hidden_states.view(
self.ep_size, self.runtime_max_tokens_per_rank, output_hidden_size
),
@@ -313,4 +488,5 @@ class FlashinferDispatcher(BaseDispatcher):
)
del self.runtime_max_tokens_per_rank
del self.active_moe_a2a
return hidden_states
python/sglang/srt/layers/moe/utils.py
+3
View File
@@ -704,14 +704,17 @@ def speculative_moe_a2a_backend_context():
moe = get_flags().moe
original_backend = moe.a2a_backend
original_disable_fp4_allgather = moe.disable_fp4_allgather
original_speculative_context = moe.speculative_context
try:
moe.a2a_backend = get_speculative_moe_a2a_backend()
# Disable FP4 allgather for spec decode since MTP layers are unquantized
moe.disable_fp4_allgather = True
moe.speculative_context = True
yield
finally:
moe.a2a_backend = original_backend
moe.disable_fp4_allgather = original_disable_fp4_allgather
moe.speculative_context = original_speculative_context
# The type of method in top-K routing, for use in torch custom op
python/sglang/srt/layers/quantization/unquant.py
+128 -21
View File
@@ -75,6 +75,7 @@ if _use_aiter:
class Bf16GemmBackend(Enum):
AUTO = "auto"
CUTEDSL = "cutedsl"
FLASHINFER_PR4266 = "flashinfer_pr4266"
GEMV = "gemv"
TORCH = "torch"
@@ -87,16 +88,77 @@ class Bf16GemmBackend(Enum):
def is_gemv(self) -> bool:
return self == Bf16GemmBackend.GEMV
def is_flashinfer_pr4266(self) -> bool:
return self == Bf16GemmBackend.FLASHINFER_PR4266
def is_optimized(self) -> bool:
return self.is_cutedsl() or self.is_flashinfer_pr4266()
_BF16_GEMM_BACKEND: Optional[Bf16GemmBackend] = None
_cutedsl_bf16_gemm = None
_use_cutedsl_bf16_gemm = None
_hopper_bf16_gemv = None
_use_hopper_bf16_gemv = None
_flashinfer_pr4266_splitk_tactic = None
_flashinfer_pr4266_run_splitk_dense = None
_flashinfer_pr4266_direct_default_tactic = None
_flashinfer_pr4266_prefer_direct = None
_flashinfer_pr4266_run_direct_dense = None
_enable_bf16_splitk_gemm = False
# GB300 TP16 tactics measured under CUDA graph replay with PDL and cold weights.
# Unlisted shapes, including M=64, retain the existing TGV/cuBLAS path.
_FLASHINFER_PR4266_TUNED_TACTICS = {
(1, 256, 8192): (64, 8, 4, 11),
(2, 256, 8192): (64, 8, 4, 11),
(4, 256, 8192): (64, 8, 4, 11),
(8, 256, 8192): (64, 8, 4, 11),
(16, 256, 8192): (64, 8, 4, 10),
(24, 256, 8192): (64, 8, 4, 11),
(32, 256, 8192): (64, 8, 4, 12),
(1, 512, 8192): (64, 8, 4, 11),
(2, 512, 8192): (64, 8, 4, 12),
(4, 512, 8192): (64, 8, 4, 10),
(8, 512, 8192): (64, 8, 4, 12),
(16, 512, 8192): (64, 8, 4, 12),
(24, 512, 8192): (64, 8, 4, 12),
(32, 512, 8192): (64, 16, 4, 9),
(1, 2304, 8192): (128, 8, 4, 6),
(2, 2304, 8192): (64, 8, 2, 12),
(4, 2304, 8192): (128, 8, 4, 6),
(8, 2304, 8192): (64, 8, 4, 10),
(16, 2304, 8192): (64, 16, 4, 9),
(24, 2304, 8192): (64, 32, 2, 9),
(32, 2304, 8192): (64, 32, 2, 9),
(1, 2560, 8192): (64, 8, 2, 10),
(2, 2560, 8192): (64, 8, 2, 10),
(4, 2560, 8192): (64, 8, 2, 10),
(8, 2560, 8192): (64, 8, 2, 10),
(16, 2560, 8192): (64, 16, 2, 11),
(24, 2560, 8192): (64, 32, 2, 9),
(32, 2560, 8192): (64, 32, 2, 9),
}
def use_flashinfer_pr4266_bf16_gemm(m: int, n: int, k: int) -> bool:
return (m, n, k) in _FLASHINFER_PR4266_TUNED_TACTICS
def should_enable_bf16_splitk_gemm(backend: Bf16GemmBackend) -> bool:
"""Return whether the optional Split-K path should be initialized."""
return backend.is_optimized() and envs.SGLANG_ENABLE_BF16_SPLITK_GEMM.get()
def initialize_bf16_gemm_config(server_args: ServerArgs) -> None:
global _BF16_GEMM_BACKEND, _cutedsl_bf16_gemm, _use_cutedsl_bf16_gemm
global _BF16_GEMM_BACKEND
global _cutedsl_bf16_gemm, _use_cutedsl_bf16_gemm
global _flashinfer_pr4266_splitk_tactic
global _flashinfer_pr4266_run_splitk_dense
global _flashinfer_pr4266_direct_default_tactic
global _flashinfer_pr4266_prefer_direct
global _flashinfer_pr4266_run_direct_dense
global _enable_bf16_splitk_gemm
from sglang.srt.utils import is_sm100_supported
@@ -122,14 +184,17 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None:
_hopper_bf16_gemv = hopper_bf16_gemv
_use_hopper_bf16_gemv = use_hopper_bf16_gemv
elif backend.is_cutedsl():
elif backend.is_optimized():
if get_exec().deterministic.enable_deterministic_inference:
raise ValueError(
"--bf16-gemm-backend cutedsl is batch-size dependent and cannot "
"be combined with --enable-deterministic-inference"
)
if not is_sm100_supported():
raise ValueError("--bf16-gemm-backend cutedsl requires an SM10x GPU")
raise ValueError(
f"--bf16-gemm-backend {backend.value} requires "
"SM100/SM103 (Blackwell)"
)
from sglang.kernels.ops.gemm.cutedsl_bf16_gemm import (
cutedsl_bf16_gemm,
@@ -139,6 +204,25 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None:
_cutedsl_bf16_gemm = cutedsl_bf16_gemm
_use_cutedsl_bf16_gemm = use_cutedsl_bf16_gemm
_enable_bf16_splitk_gemm = False
if should_enable_bf16_splitk_gemm(backend):
from sglang.kernels.ops.gemm.flashinfer_pr4266_dense_bf16_gemm_sm100_direct import (
default_tactic,
prefer_direct_bf16_gemm_sm100,
run_direct_dense,
)
from sglang.kernels.ops.gemm.flashinfer_pr4266_dense_bf16_gemm_sm100_splitk import (
SplitKTactic,
run_splitk_dense,
)
_flashinfer_pr4266_splitk_tactic = SplitKTactic
_flashinfer_pr4266_run_splitk_dense = run_splitk_dense
_flashinfer_pr4266_direct_default_tactic = default_tactic
_flashinfer_pr4266_prefer_direct = prefer_direct_bf16_gemm_sm100
_flashinfer_pr4266_run_direct_dense = run_direct_dense
_enable_bf16_splitk_gemm = True
_BF16_GEMM_BACKEND = backend
@@ -148,22 +232,48 @@ def _bf16_gemm_dispatch_fake(
return x.new_empty((*x.shape[:-1], weight.shape[0]))
@register_custom_op(fake_impl=_bf16_gemm_dispatch_fake)
def bf16_gemm_dispatch(
def _flashinfer_pr4266_bf16_gemm(
x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]
) -> torch.Tensor:
x_2d = x.view(-1, x.shape[-1])
out = torch.empty((x_2d.shape[0], weight.shape[0]), dtype=x.dtype, device=x.device)
m, n, k = x_2d.shape[0], weight.shape[0], weight.shape[1]
if bias is None and _flashinfer_pr4266_prefer_direct(m, n, k):
tactic = _flashinfer_pr4266_direct_default_tactic(m, n, k)
_flashinfer_pr4266_run_direct_dense(x_2d, weight.T, out, True, tactic)
else:
tactic = _flashinfer_pr4266_splitk_tactic(
*_FLASHINFER_PR4266_TUNED_TACTICS[(m, n, k)]
)
_flashinfer_pr4266_run_splitk_dense(
x_2d,
weight.T,
bias,
out,
True,
tactic,
)
return out.view(*x.shape[:-1], weight.shape[0])
def _bf16_gemm_dispatch_impl(
x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]
) -> torch.Tensor:
m = x.numel() // x.shape[-1]
if _enable_bf16_splitk_gemm and use_flashinfer_pr4266_bf16_gemm(
m, weight.shape[0], weight.shape[1]
):
return _flashinfer_pr4266_bf16_gemm(x, weight, bias)
if (
_use_hopper_bf16_gemv is not None
and bias is None
and _use_hopper_bf16_gemv(
x.numel() // x.shape[-1], weight.shape[0], weight.shape[1]
)
and _use_hopper_bf16_gemv(m, weight.shape[0], weight.shape[1])
):
return _hopper_bf16_gemv(x.view(-1, x.shape[-1]), weight).view(
*x.shape[:-1], -1
)
if _use_cutedsl_bf16_gemm is not None and _use_cutedsl_bf16_gemm(
x.numel() // x.shape[-1], weight.shape[0], weight.shape[1]
m, weight.shape[0], weight.shape[1]
):
return _cutedsl_bf16_gemm(x.view(-1, x.shape[-1]), weight, bias).view(
*x.shape[:-1], -1
@@ -171,6 +281,13 @@ def bf16_gemm_dispatch(
return F.linear(x, weight, bias)
@register_custom_op(fake_impl=_bf16_gemm_dispatch_fake)
def bf16_gemm_dispatch(
x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]
) -> torch.Tensor:
return _bf16_gemm_dispatch_impl(x, weight, bias)
def get_bf16_gemm_backend() -> Bf16GemmBackend:
global _BF16_GEMM_BACKEND
if _BF16_GEMM_BACKEND is None:
@@ -269,7 +386,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
return tgemm.mm(x, layer.weight, bias, otype=x.dtype)
elif (
get_bf16_gemm_backend().is_cutedsl()
get_bf16_gemm_backend().is_optimized()
and x.is_cuda
and x.dtype == torch.bfloat16
and layer.weight.dtype == torch.bfloat16
@@ -283,17 +400,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
# opaque op resolves it at runtime with concrete shapes,
# keeping the per-shape kernel choice.
return bf16_gemm_dispatch(x, layer.weight, bias)
if _use_cutedsl_bf16_gemm(
x.numel() // x.shape[-1],
layer.weight.shape[0],
layer.weight.shape[1],
):
x_shapes = x.shape
output = _cutedsl_bf16_gemm(
x.view(-1, x_shapes[-1]), layer.weight, bias
)
return output.view(*x_shapes[:-1], -1)
return F.linear(x, layer.weight, bias)
return _bf16_gemm_dispatch_impl(x, layer.weight, bias)
return F.linear(x, layer.weight, bias)
python/sglang/srt/layers/radix_linear_attention.py
+101 -32
View File
@@ -84,10 +84,8 @@ class RadixLinearAttention(nn.Module):
a: torch.Tensor,
b: torch.Tensor,
) -> torch.Tensor:
if (
forward_batch.forward_mode.is_extend()
and get_tc_piecewise_forward_context() is not None
):
is_extend = forward_batch.forward_mode.is_extend()
if is_extend and get_tc_piecewise_forward_context() is not None:
# Output shape from linear attention: (1, seq_len, num_v_heads, head_v_dim)
seq_len = mixed_qkv.shape[0]
output = torch.empty(
@@ -112,14 +110,102 @@ class RadixLinearAttention(nn.Module):
self.layer_id,
)
return output
else:
return get_attn_backend().forward(
layer=self,
forward_batch=forward_batch,
# Target verify rebuilds query_start_loc from the physical padded input,
# unlike ordinary extend where it retains the logical sequence ends.
should_trim_padded_extend = (
is_extend and not forward_batch.forward_mode.is_target_verify()
)
real_num_tokens = (
getattr(forward_batch, "num_token_non_padded_cpu", None)
if should_trim_padded_extend
else None
)
if real_num_tokens is not None and real_num_tokens < mixed_qkv.shape[0]:
# DP synchronization may append physical rows beyond the logical varlen
# layout; compute its real prefix, then restore the downstream shape.
output = torch.empty(
(1, mixed_qkv.shape[0], self.num_v_heads, self.head_v_dim),
dtype=mixed_qkv.dtype,
device=mixed_qkv.device,
)
_linear_attention_with_output_impl(
mixed_qkv=mixed_qkv,
a=a,
b=b,
output=output,
attention_layer=self,
forward_batch=forward_batch,
)
return output
return get_attn_backend().forward(
layer=self,
forward_batch=forward_batch,
mixed_qkv=mixed_qkv,
a=a,
b=b,
)
def _linear_attention_with_output_impl(
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
output: torch.Tensor,
attention_layer: RadixLinearAttention,
forward_batch: ForwardBatch,
) -> None:
"""Run linear attention on the real prefix and initialize physical padding."""
real_num_tokens = min(forward_batch.num_token_non_padded_cpu, mixed_qkv.shape[0])
original_out_cache_loc = forward_batch.out_cache_loc
# Keep the original ForwardBatch object and only narrow cache locations for
# this backend call so model/backend state is still written to the same batch.
forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens]
logical_output = output[:, :real_num_tokens]
try:
ret = get_attn_backend().forward(
layer=attention_layer,
forward_batch=forward_batch,
mixed_qkv=mixed_qkv[:real_num_tokens],
a=a[:real_num_tokens],
b=b[:real_num_tokens],
linear_attn_output=logical_output,
)
finally:
forward_batch.out_cache_loc = original_out_cache_loc
# FlashInfer GDN can write directly into the physical output's logical
# prefix. Other backends return their own tensor and keep the copy fallback.
if ret.data_ptr() != logical_output.data_ptr():
logical_output.copy_(ret)
# Physical padding participates in following residual, router, expert/MoE,
# and collective operations. Keep those inputs finite and deterministic.
output[:, real_num_tokens:].zero_()
def _unified_linear_attention_with_output_impl(
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
output: torch.Tensor,
layer_id: int,
) -> None:
"""Eager implementation kept separate for backend-independent tests."""
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
_linear_attention_with_output_impl(
mixed_qkv=mixed_qkv,
a=a,
b=b,
output=output,
attention_layer=attention_layer,
forward_batch=forward_batch,
)
return
@register_custom_op(mutates_args=["output"])
@@ -131,31 +217,14 @@ def unified_linear_attention_with_output(
output: torch.Tensor,
layer_id: int,
) -> None:
"""
Custom op wrapper for linear attention computation only.
"""
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
real_num_tokens = forward_batch.num_token_non_padded_cpu
original_out_cache_loc = forward_batch.out_cache_loc
# Keep the original ForwardBatch object and only narrow cache locations for
# this backend call so model/backend state is still written to the same batch.
forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens]
ret = get_attn_backend().forward(
layer=attention_layer,
forward_batch=forward_batch,
mixed_qkv=mixed_qkv[:real_num_tokens],
a=a[:real_num_tokens],
b=b[:real_num_tokens],
"""Custom op wrapper for linear attention computation only."""
_unified_linear_attention_with_output_impl(
mixed_qkv=mixed_qkv,
a=a,
b=b,
output=output,
layer_id=layer_id,
)
forward_batch.out_cache_loc = original_out_cache_loc
output[:, :real_num_tokens].copy_(ret)
return
bcg_unified_linear_attention_with_output = eager_on_graph(True)(
python/sglang/srt/layers/sampler.py
+27
View File
@@ -8,6 +8,7 @@ from torch import nn
from sglang.kernels.ops.sampling.murmur_hash import murmur_hash32
from sglang.srt.distributed import get_tp_group
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled,
)
@@ -68,6 +69,18 @@ _CUSTOM_SAMPLER_FACTORIES: Dict[str, Callable[[], "Sampler"]] = {}
_BUILT_IN_SAMPLING_BACKENDS = {"flashinfer", "pytorch", "ascend"}
def _trace_e2e_sampler(stage: str, **fields) -> None:
if not envs.SGLANG_TRACE_SAMPLER_E2E.get():
return
try:
parallel = get_parallel()
rank = f"dp={parallel.attn_dp_rank} tp={parallel.tp_rank}"
except Exception:
rank = "rank=unknown"
details = " ".join(f"{key}={value}" for key, value in fields.items())
print(f"SGLANG_TRACE_SAMPLER_E2E {rank} stage={stage} {details}", flush=True)
class Sampler(nn.Module):
def __init__(self):
super().__init__()
@@ -119,15 +132,23 @@ class Sampler(nn.Module):
to get the unique seed for each position.
"""
logits = logits_output.next_token_logits
_trace_e2e_sampler(
"forward_enter",
logits_shape=tuple(logits.shape),
all_greedy=sampling_info.is_all_greedy,
)
if _is_hip and logits.shape[0] == 0:
return torch.empty((0,), dtype=torch.int64, device=logits.device)
# Preprocess logits (custom processors and NaN handling)
_trace_e2e_sampler("preprocess_enter")
logits = self._preprocess_logits(logits, sampling_info)
_trace_e2e_sampler("preprocess_returned")
return_sampling_mask = any(sampling_info.return_sampling_masks or [])
if sampling_info.is_all_greedy:
_trace_e2e_sampler("greedy_enter")
if _use_aiter and not _disable_aiter_greedy_sample:
batch_next_token_ids = torch.empty(
logits.shape[0], device=logits.device, dtype=torch.int32
@@ -135,6 +156,9 @@ class Sampler(nn.Module):
_aiter_greedy_sample(batch_next_token_ids, logits)
else:
batch_next_token_ids = torch.argmax(logits, -1)
_trace_e2e_sampler(
"greedy_returned", output_shape=tuple(batch_next_token_ids.shape)
)
if return_sampling_mask:
self._attach_greedy_sampling_mask_to_output(
logits_output, sampling_info, batch_next_token_ids
@@ -243,8 +267,11 @@ class Sampler(nn.Module):
)
logprob_result.write_output_to(logits_output)
_trace_e2e_sampler("token_sync_enter")
self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info)
_trace_e2e_sampler("token_sync_returned")
_trace_e2e_sampler("forward_returned")
return batch_next_token_ids
def _sample_from_probs(
python/sglang/srt/managers/overlap_utils.py
+23 -19
View File
@@ -295,8 +295,12 @@ class FutureMap:
else:
self.new_seq_lens_cpu_pinned = None
self.fwd_prepare_d2h_stream = None
# Lazy-inited on the first non-empty stash (peeks tensor shapes); non-spec's is a no-op.
self._forward_buf_initialized = False
self.need_topk = False
self.need_hidden_states = False
self.topk_p_buf = None
self.topk_index_buf = None
self.hidden_states_buf = None
self.draft_probs_buf = None
self.dsa_topk_indices_buf = None
# ngram-only relay bufs
@@ -314,22 +318,18 @@ class FutureMap:
pool=req_to_token_pool,
)
def _lazy_init_forward_buf(self, payload: RelayPayload):
def _maybe_init_forward_bufs(self, payload: RelayPayload) -> None:
# Local import (see decide_needs_cpu_seq_lens): keep module-level deps leaf.
from sglang.srt.speculative.spec_utils import spec_need_hidden_states
self._forward_buf_initialized = True
# Spec extras are gated by spec_algo, not by the payload's shape, so a
# non-spec stash allocates no extra bufs (only output_tokens_buf).
self.need_topk = self.spec_algo.is_some() and self.spec_algo.need_topk()
self.need_hidden_states = (
# Prefill can omit spec extras; initialize each buffer when decode first
# carries it instead of fixing the layout from the first payload.
if not self.need_topk and (
self.spec_algo.is_some()
and spec_need_hidden_states()
and payload.hidden_states is not None
)
if self.need_topk:
and self.spec_algo.need_topk()
and payload.topk_p is not None
):
self.need_topk = True
topk_p0 = payload.topk_p[0]
topk_index0 = payload.topk_index[0]
self.topk_p_buf = torch.empty(
@@ -342,7 +342,13 @@ class FutureMap:
dtype=topk_index0.dtype,
device=self.device,
)
if self.need_hidden_states:
if not self.need_hidden_states and (
self.spec_algo.is_some()
and spec_need_hidden_states()
and payload.hidden_states is not None
):
self.need_hidden_states = True
hidden_states0 = payload.hidden_states[0]
self.hidden_states_buf = torch.empty(
(self.req_pool_size, *hidden_states0.shape),
@@ -350,8 +356,7 @@ class FutureMap:
device=self.device,
)
self.draft_probs_buf = None
if payload.draft_probs is not None:
if self.draft_probs_buf is None and payload.draft_probs is not None:
draft_probs0 = payload.draft_probs[0]
self.draft_probs_buf = torch.empty(
(self.req_pool_size, *draft_probs0.shape),
@@ -545,8 +550,7 @@ class FutureMap:
self.accept_tokens_buf[indices] = payload.accept_tokens
self.accept_lens_buf[indices] = payload.accept_lens
return
if not self._forward_buf_initialized:
self._lazy_init_forward_buf(payload)
self._maybe_init_forward_bufs(payload)
self._maybe_init_dsa_topk_indices_buf(payload)
self.output_tokens_buf[indices] = payload.bonus_tokens.to(
self.output_tokens_buf.dtype
python/sglang/srt/managers/scheduler.py
+25 -12
View File
@@ -292,14 +292,16 @@ from sglang.srt.observability.trace import process_tracing_init, trace_set_threa
from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.platforms import current_platform
from sglang.srt.plugins import load_plugins
from sglang.srt.runtime_context import get_context, publish
from sglang.srt.runtime_context import get_context, get_spec, publish
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.sampling.sampling_params import TOP_K_ALL
from sglang.srt.server_args import PortArgs, ServerArgs, compute_world_size
from sglang.srt.session.session_controller import SessionController
from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
from sglang.srt.speculative.dflash_utils import validate_dflash_request
from sglang.srt.speculative.eagle_utils import get_draft_recurrent_hidden_state_spec
from sglang.srt.speculative.eagle_utils import (
get_draft_recurrent_hidden_state_spec_from_config,
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.utils import (
DynamicGradMode,
@@ -1380,11 +1382,18 @@ class Scheduler(
)
if self.spec_algorithm.carries_draft_hidden_states():
# `draft_runner` aliases `draft_runner_list[0]` in the multi-layer
# worker, so a single accessor covers both shapes.
draft_runner = self.draft_worker.draft_worker.draft_runner
# Derive the rank-uniform PD wire schema from config because only the
# last prefill PP stage owns a draft runner.
draft_model_config = ModelConfig.from_server_args(
self.server_args,
model_path=get_spec().speculative_draft_model_path,
model_revision=get_spec().speculative_draft_model_revision,
is_draft_model=True,
)
disagg_hidden_size, disagg_hidden_states_dtype = (
get_draft_recurrent_hidden_state_spec(draft_runner)
get_draft_recurrent_hidden_state_spec_from_config(
draft_model_config, self.spec_algorithm
)
)
else:
disagg_hidden_size = 16 # minimal padding size for RDMA
@@ -3989,7 +3998,9 @@ class Scheduler(
# future_map relay / on_publish).
resolve_forward_inputs(batch, self.future_map)
with self._forward_isolation(batch, overlap=False):
batch_result = self.model_worker.forward_batch_generation(batch)
batch_result = self.model_worker.forward_batch_generation(
batch, pp_proxy_tensors=pp_proxy_tensors
)
# The isolation restore reverted the worker's in-forward SB edits;
# re-apply what must carry to the next iter.
batch.spec_info = batch_result.next_draft_input
@@ -4000,12 +4011,14 @@ class Scheduler(
batch.seq_lens_sum = int(batch.seq_lens_cpu.sum())
batch.input_ids = None # rebuilt next iter from draft_token
self.update_cache_from_scheduler(batch, batch_result)
# Sync D2H so the result processor can read CPU tensors.
# Only the last PP rank owns real results requiring D2H; other ranks
# consume device tensors rebuilt from the output ring.
batch_result.copy_done = self.device_module.Event()
batch_result.copy_to_cpu(
return_logprob=batch.return_logprob,
return_hidden_states=batch.return_hidden_states,
)
if batch_result.has_sampled_token_ids and self.ps.pp_size == 1:
batch_result.copy_to_cpu(
return_logprob=batch.return_logprob,
return_hidden_states=batch.return_hidden_states,
)
else:
kwargs = (
{"pp_proxy_tensors": pp_proxy_tensors}
python/sglang/srt/managers/scheduler_pp_mixin.py
+37 -1
View File
@@ -1025,6 +1025,14 @@ class SchedulerPPMixin:
"next_token_ids": result.next_token_ids,
}
# Draft extend runs only on the last stage, but every rank needs its relayed
# output to fill PD auxiliary buffers.
draft_input = result.next_draft_input
if draft_input is not None and draft_input.topk_p is not None:
tensor_dict["draft_topk_p"] = draft_input.topk_p.contiguous()
tensor_dict["draft_topk_index"] = draft_input.topk_index.contiguous()
tensor_dict["draft_hidden_states"] = draft_input.hidden_states.contiguous()
if batch.return_logprob:
logprob_dict = get_logprob_dict_from_result(result)
tensor_dict = {
@@ -1172,17 +1180,45 @@ class SchedulerPPMixin:
logits_output = LogitsProcessorOutput(next_token_logits=None)
logits_output.auxiliary_device_output = auxiliary_output
next_token_ids = pp_outputs["next_token_ids"].to(torch.int64)
# Rebind the last stage's ring proposal as batch.spec_info so the PD result
# processor sees the same object on every rank.
next_draft_input = None
if "draft_topk_p" in pp_outputs.tensors:
from sglang.srt.speculative.eagle_info import EagleDraftInput
next_draft_input = EagleDraftInput(
topk_p=pp_outputs["draft_topk_p"],
topk_index=pp_outputs["draft_topk_index"],
hidden_states=pp_outputs["draft_hidden_states"],
bonus_tokens=next_token_ids,
num_tokens_per_req=1,
num_tokens_for_logprob_per_req=1,
)
batch.spec_info = next_draft_input
# PP rank 0 also relays into output_tokens_buf so the next iter's
# resolve_forward_inputs finds these tokens for the decode portion
# of mixed-chunk batches (which gather via mix_running_indices).
self.future_map.stash(
batch.req_pool_indices, RelayPayload(bonus_tokens=next_token_ids)
batch.req_pool_indices,
RelayPayload(
bonus_tokens=next_token_ids,
topk_p=None if next_draft_input is None else next_draft_input.topk_p,
topk_index=(
None if next_draft_input is None else next_draft_input.topk_index
),
hidden_states=(
None if next_draft_input is None else next_draft_input.hidden_states
),
),
)
batch.input_ids = None
output_result = GenerationBatchResult(
logits_output=logits_output,
pp_hidden_states_proxy_tensors=None,
next_token_ids=pp_outputs["next_token_ids"],
next_draft_input=next_draft_input,
extend_input_len_per_req=extend_input_len_per_req,
extend_logprob_start_len_per_req=extend_logprob_start_len_per_req,
can_run_cuda_graph=mb_metadata.can_run_cuda_graph,
python/sglang/srt/managers/tp_worker.py
+6 -2
View File
@@ -325,6 +325,7 @@ class TpModelWorker(BaseTpWorker):
is_multi_layer_eagle: bool = False,
context_length: Optional[int] = None,
draft_attention_backend: Optional[str] = None,
random_seed: Optional[int] = None,
):
# Parse args
self.server_args = server_args
@@ -384,8 +385,11 @@ class TpModelWorker(BaseTpWorker):
self.world_group = get_world_group()
# Sync random seed across TP workers.
# Elastic joiners cannot enter the launch-time WORLD broadcast.
if server_args.is_ep_joiner:
# Elastic joiners and last-stage-only draft workers cannot enter the WORLD
# broadcast, so they reuse the target's already-broadcast seed.
if random_seed is not None:
self.random_seed = random_seed
elif server_args.is_ep_joiner:
self.random_seed = get_device().random_seed
else:
self.random_seed = broadcast_pyobj(
python/sglang/srt/mem_cache/kv_cache_builder.py
+24
View File
@@ -23,6 +23,7 @@ class KVCacheBuildResult:
from typing import TYPE_CHECKING
from sglang.srt.arg_groups.overrides import resolving_view
from sglang.srt.configs.hybrid_arch import (
hybrid_gdn_config,
hybrid_lightning_config,
@@ -62,6 +63,29 @@ if TYPE_CHECKING:
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
def get_draft_kv_pool(
*,
draft_worker: BaseTpWorker,
spec_algorithm: SpeculativeAlgorithm,
server_args: ServerArgs,
):
"""Return the draft token-to-KV pool for the current draft worker,
or None when no draft KV pool is available."""
if draft_worker is None or spec_algorithm.is_ngram():
return None
# V2 draft workers exist only on their hosting PP stage; other ranks own no
# nested draft worker or draft KV pool.
if draft_worker.draft_worker is None:
return None
if resolving_view(server_args).enable_multi_layer_eagle:
draft_runner = draft_worker.draft_worker.draft_runner_list[0]
else:
draft_runner = draft_worker.draft_worker.draft_runner
return draft_runner.token_to_kv_pool
def maybe_register_hicache_draft(
*,
tree_cache,
python/sglang/srt/mem_cache/kv_cache_configurator.py
+22
View File
@@ -161,6 +161,28 @@ MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP_LAZY = 1
MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1
MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_BUFFER = 1
def _pp_local_per_request_bytes(
total_bytes: int,
layer_ids: list[int],
start_layer: int,
end_layer: int,
) -> int:
# BaseLinearStateParams reports global bytes, but PP pools allocate only local
# layers; charge this stage its proportional per-request share.
if not layer_ids:
return 0
if total_bytes % len(layer_ids) != 0:
raise ValueError(
"Linear-state bytes must be uniform per layer: "
f"total_bytes={total_bytes}, num_layers={len(layer_ids)}"
)
local_layer_count = sum(
start_layer <= layer_id < end_layer for layer_id in layer_ids
)
return total_bytes // len(layer_ids) * local_layer_count
if TYPE_CHECKING:
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.mem_cache.unified_memory_pool import (
python/sglang/srt/model_executor/cuda_graph_buffer_registry.py
+6 -1
View File
@@ -939,6 +939,8 @@ def build_prefill_registry(
)
reg.register_slot(slot, bind=bind)
# PP stage inputs live outside ForwardBatch; adopt runner-owned buffers for
# stable addresses and clear padding because prefill executes every bucket row.
if source is not None:
pp = getattr(source, "pp_proxy_tensors", None)
if pp is not None:
@@ -954,7 +956,10 @@ def build_prefill_registry(
reg.register_slot(
GraphSlot(
name=f"pp_proxy_tensors.{_key}",
shape_fn=lambda _bs, _mt, _s=tuple(_backing.shape): _s,
shape_fn=lambda _bs, mt, _tail=tuple(_backing.shape[1:]): (
mt,
*_tail,
),
dtype=_backing.dtype,
axis="tokens",
padding_policy=PaddingPolicy.ZERO,
python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py
+55 -9
View File
@@ -3,7 +3,7 @@ from __future__ import annotations
import logging
import time
from collections import defaultdict
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING, Any, Optional
import msgspec
@@ -80,6 +80,45 @@ def _align_pipeline_layers(layers: list, layer_model) -> list:
)
def has_standard_gqa_for_all_local_layers(
*, attention_layer_count: int, start_layer: int, end_layer: int
) -> bool:
"""Check the layers materialized on this pipeline rank, not the full model."""
return attention_layer_count >= end_layer - start_layer
def index_attention_layers_by_global_id(
attention_layers: list[Any],
mha_companion_layers: list[Any],
layer_model=None,
) -> tuple[list[Any], list[Any]]:
"""Pad PP-local attention metadata so global layer_id remains a valid index."""
if len(attention_layers) != len(mha_companion_layers):
raise ValueError("attention and MHA companion metadata must be parallel")
populated = [layer for layer in attention_layers if layer is not None]
if not populated or any(not hasattr(layer, "layer_id") for layer in populated):
if layer_model is not None:
return (
_align_pipeline_layers(attention_layers, layer_model),
_align_pipeline_layers(mha_companion_layers, layer_model),
)
return attention_layers, mha_companion_layers
max_layer_id = max(int(layer.layer_id) for layer in populated)
indexed_attention = [None] * (max_layer_id + 1)
indexed_companions = [None] * (max_layer_id + 1)
for attention, companion in zip(attention_layers, mha_companion_layers):
if attention is None:
if companion is not None:
raise ValueError("MHA companion has no primary attention layer")
continue
layer_id = int(attention.layer_id)
if layer_id < 0 or indexed_attention[layer_id] is not None:
raise ValueError(f"invalid or duplicate attention layer_id: {layer_id}")
indexed_attention[layer_id] = attention
indexed_companions[layer_id] = companion
return indexed_attention, indexed_companions
class GraphCapture(msgspec.Struct, frozen=True, kw_only=True):
runner: Optional[BaseRunner]
memory_phase: str
@@ -392,11 +431,22 @@ def capture_prefill_graph(
model_runner.dsa_indexers,
model_runner.mha_companion_layers,
) = compute_attention_and_moe_layers(layer_model)
model_runner.attention_layers = _align_pipeline_layers(
model_runner.attention_layers, layer_model
(
model_runner.attention_layers,
model_runner.mha_companion_layers,
) = index_attention_layers_by_global_id(
model_runner.attention_layers,
model_runner.mha_companion_layers,
layer_model,
)
if len(model_runner.attention_layers) < model_runner.model_config.num_hidden_layers:
if not has_standard_gqa_for_all_local_layers(
attention_layer_count=sum(
layer is not None for layer in model_runner.attention_layers
),
start_layer=model_runner.layer_info.start_layer,
end_layer=model_runner.layer_info.end_layer,
):
# TODO(yuwei): support Non-Standard GQA
log_info_on_rank0(
logger,
@@ -404,10 +454,6 @@ def capture_prefill_graph(
)
return result(None)
model_runner.mha_companion_layers = _align_pipeline_layers(
model_runner.mha_companion_layers, layer_model
)
tic = time.perf_counter()
before_mem = get_available_gpu_memory(model_runner.device, model_runner.gpu_id)
role = "draft" if model_runner.is_draft_worker else "target"
python/sglang/srt/model_executor/runner/base_runner.py
+9 -1
View File
@@ -134,7 +134,7 @@ def _allocate_decode_buffers(
residual_shape = (
(max_num_token, pp_proxy_residual_num_blocks, hidden_size)
if pp_proxy_residual_num_blocks is not None
else (max_bs, hidden_size)
else (max_num_token, hidden_size)
)
pp_proxy_tensors["residual"] = torch.zeros(residual_shape, dtype=dtype)
if pp_proxy_topk_size is not None:
@@ -244,6 +244,14 @@ class BaseRunner(ABC):
self._pre_initialize_flashinfer_allreduce_workspace()
self._pre_initialize_fi_a2a_workspace()
# Model-owned communication resources may depend on the resolved
# request pool and must be compiled/allocated before graph capture.
prepare_model_resources = getattr(
mr.model, "prepare_before_cuda_graph_capture", None
)
if prepare_model_resources is not None:
prepare_model_resources(mr)
if should_run_flashinfer_autotune(self.model_runner):
buffers, batch_size = self._autotune_buffers()
assert (
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+38 -22
View File
@@ -38,6 +38,7 @@ Backend selection comes from cuda_graph_config.prefill:
from __future__ import annotations
import copy
import dataclasses
import inspect
import logging
from collections.abc import Sequence
@@ -328,10 +329,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
dtype=self.model_runner.dtype,
enable_mamba_track=self.mamba_track_enabled,
pp_size=self.pp_size,
hc_hidden_size=model_runner.model_config.hc_hidden_size,
pp_proxy_topk_size=model_runner.get_pp_proxy_topk_size(),
hc_hidden_size=getattr(
self.model_runner.model_config, "hc_hidden_size", None
),
pp_proxy_topk_size=self.model_runner.get_pp_proxy_topk_size(),
pp_proxy_residual_num_blocks=(
model_runner.get_pp_proxy_residual_num_blocks()
self.model_runner.get_pp_proxy_residual_num_blocks()
),
)
self.buffers.share_buffers()
@@ -385,11 +388,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self._capture_lora = False
self.enable_cp_v2_bcg_capture = False
self.prefill_cp_bcg_input: Optional[PrefillCPBCGInput] = None
self._static_pp_proxy_tensors = (
PPProxyTensors(self.buffers.pp_proxy_tensors)
if self.buffers.pp_proxy_tensors is not None
else None
)
# TcPiecewise does its compile pass during backend construction.
# Wrap only that path with the prefill CUDA graph failure hint.
try:
@@ -646,6 +644,14 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
return forward_batch.positions
def _static_pp_proxy_tensors(self, num_tokens: int) -> Optional[PPProxyTensors]:
buffers = self.buffers.pp_proxy_tensors
if buffers is None:
return None
return PPProxyTensors(
{key: value[:num_tokens] for key, value in buffers.items()}
)
@contextmanager
def _prefill_forward_context(
self,
@@ -702,31 +708,25 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
set_is_extend_in_batch(False)
with self._prefill_forward_context(forward_batch):
pp_kwargs = self.model_runner._pp_kwargs(
self._static_pp_proxy_tensors(num_tokens)
)
if self._uses_eager_prefill_tail():
# BCG / Full: capture the transformer body only.
positions = self._get_layer_model_positions(forward_batch)
input_ids = forward_batch.input_ids
input_embeds = forward_batch.input_embeds
layer_kwargs = {}
if self._static_pp_proxy_tensors is not None:
layer_kwargs["pp_proxy_tensors"] = self._static_pp_proxy_tensors[
:num_tokens
]
if not self.model_runner.pp_group.is_first_rank:
input_ids = None
input_embeds = None
return self.layer_model.forward(
input_ids,
forward_batch.input_ids,
positions,
forward_batch,
input_embeds,
**layer_kwargs,
forward_batch.input_embeds,
**pp_kwargs,
)
# tc_piecewise: compile/capture the outer model.forward path.
return self.model_runner.model.forward(
forward_batch.input_ids,
forward_batch.positions,
forward_batch,
**pp_kwargs,
)
def _run_dummy_forward(self, num_tokens: int) -> None:
@@ -1568,6 +1568,19 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
else forward_batch.global_forward_mode
)
# The draft tail concatenates hidden states with padded input embeddings;
# expose the bucket-sized static view and refresh its live prefix below.
padded_spec_info = forward_batch.spec_info
if (
self.static_draft_hidden_states is not None
and padded_spec_info is not None
and getattr(padded_spec_info, "hidden_states", None) is not None
):
padded_spec_info = dataclasses.replace(
padded_spec_info,
hidden_states=self.static_draft_hidden_states[:static_num_tokens],
)
static_forward_batch = ForwardBatch(
forward_mode=pcg_forward_mode,
batch_size=bs,
@@ -1611,7 +1624,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
global_dp_buffer_len=forward_batch.global_dp_buffer_len,
mrope_positions=mrope_positions,
spec_algorithm=forward_batch.spec_algorithm,
spec_info=forward_batch.spec_info,
spec_info=padded_spec_info,
capture_hidden_mode=forward_batch.capture_hidden_mode,
num_token_non_padded=num_token_non_padded,
num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu,
@@ -1619,6 +1632,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
lora_ids=forward_batch.lora_ids,
sampling_info=forward_batch.sampling_info,
mm_inputs=forward_batch.mm_inputs,
# Multimodal preparation consumes mm_inputs but retains embeddings for
# later MTP draft extend, so copy that side channel into the replay view.
mm_input_embeds=forward_batch.mm_input_embeds,
temperature=forward_batch.temperature,
top_p=forward_batch.top_p,
dimensions=forward_batch.dimensions,
python/sglang/srt/model_executor/runner_utils/buffers.py
+3 -3
View File
@@ -133,7 +133,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
is_mhc = hc_hidden_size is not None
hs = hc_hidden_size if is_mhc else hidden_size
pp_proxy_tensors = {
"hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
"hidden_states": torch.zeros((max_num_token, hs), dtype=dtype),
}
if not is_mhc:
# Only Kimi K3 supplies num_blocks: its PP bank is token-major
@@ -141,7 +141,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
residual_shape = (
(max_num_token, pp_proxy_residual_num_blocks, hidden_size)
if pp_proxy_residual_num_blocks is not None
else (max_bs, hidden_size)
else (max_num_token, hidden_size)
)
pp_proxy_tensors["residual"] = torch.zeros(
residual_shape, dtype=dtype
@@ -356,7 +356,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
hidden_size: int,
dtype: torch.dtype,
enable_mamba_track: bool,
pp_size: int,
pp_size: int = 1,
hc_hidden_size: Optional[int] = None,
pp_proxy_topk_size: Optional[int] = None,
pp_proxy_residual_num_blocks: Optional[int] = None,
python/sglang/srt/models/deepseek_v2.py
+6 -1
View File
@@ -2740,7 +2740,12 @@ class DeepseekV2Model(nn.Module):
for i in range(len(self.layers)):
if isinstance(self.layers[i].mlp, DeepseekV2MoE):
# tp_size = get_parallel().tp_size
is_a2a_moe = is_deepep_class_backend()
# Keep the original deepep-class scope here and only add DeepEP v2,
# so unrelated backends' allocator sizing is unchanged.
is_a2a_moe = (
is_deepep_class_backend()
or get_moe_a2a_backend().is_deepep_v2()
)
tp_size = 1 if is_a2a_moe else get_parallel().tp_size
intermediate_size = (
config.moe_intermediate_size * config.n_shared_experts
python/sglang/srt/models/qwen2_moe.py
+125 -7
View File
@@ -27,7 +27,10 @@ import torch.nn.functional as F
from torch import nn
from transformers import PretrainedConfig
from sglang.kernels.ops.elementwise.elementwise import fused_gate_sigmoid_mul_add
from sglang.kernels.ops.elementwise.elementwise import (
fused_gate_sigmoid_mul,
fused_gate_sigmoid_mul_add,
)
from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
from sglang.srt.distributed import (
get_pp_group,
@@ -154,6 +157,7 @@ def can_fuse_shared_expert(
or getattr(config, "shared_expert_intermediate_size", 0) <= 0
or config.shared_expert_intermediate_size != config.moe_intermediate_size
or get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_deepep_v2()
or get_moe_a2a_backend().is_mori()
):
return False
@@ -341,6 +345,10 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
routing_method_type=RoutingMethodType.RenormalizeNaive,
num_fused_shared_experts=self.num_fused_shared_experts,
inplace=not _needs_hidden_after_experts,
enable_qwen35_fp8_deferred_finalize=(
config.model_type == "qwen3_5_moe_text"
and envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION.get()
),
)
self.gate = ReplicatedLinear(
@@ -369,6 +377,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
if (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mori()
or get_moe_a2a_backend().is_deepep_v2()
or get_moe_a2a_backend().is_flashinfer()
)
else {}
@@ -387,7 +396,11 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
else:
self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori():
if (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_deepep_v2()
or get_moe_a2a_backend().is_mori()
):
# TODO: we will support tp < ep in the future
self.ep_size = get_parallel().moe_ep_size
self.num_experts = (
@@ -532,6 +545,23 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
return shared_output
def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch):
trace_e2e = envs.SGLANG_TRACE_QWEN_MOE_DEEPEP_E2E.get()
def trace_sync(stage: str):
if not trace_e2e:
return
print(
"SGLANG_TRACE_QWEN_MOE_DEEPEP_E2E "
f"stage={stage}_sync_enter tokens={hidden_states.shape[0]}",
flush=True,
)
torch.cuda.synchronize()
print(
"SGLANG_TRACE_QWEN_MOE_DEEPEP_E2E "
f"stage={stage}_sync_returned tokens={hidden_states.shape[0]}",
flush=True,
)
enable_dual_stream = (
is_npu()
and envs.SGLANG_NPU_USE_MULTI_STREAM.get()
@@ -564,6 +594,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
shared_event = self.alt_stream.record_event()
else:
shared_output = self._forward_shared_experts(hidden_states)
trace_sync("shared_expert")
topk_output = self.topk(
hidden_states,
router_logits,
@@ -578,36 +609,95 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
)
else:
topk_output = self.topk.empty_topk_output(hidden_states.device)
trace_sync("pre_experts")
final_hidden_states = self.experts(
hidden_states=hidden_states,
topk_output=topk_output,
)
trace_sync("post_experts")
if enable_dual_stream:
wait_share_stream()
elif enable_cuda_shared_overlap:
torch.cuda.current_stream().wait_event(shared_event)
if shared_output is not None:
trace_sync("pre_shared_add")
final_hidden_states.add_(shared_output)
trace_sync("post_shared_add")
return final_hidden_states
def _forward_router_experts(self, hidden_states: torch.Tensor):
@property
def supports_deferred_finalize(self) -> bool:
return bool(
self.experts.supports_deferred_finalize and self.shared_expert is not None
)
def _forward_router_experts(
self,
hidden_states: torch.Tensor,
*,
defer_finalize: bool = False,
):
# router_logits: (num_tokens, n_experts)
router_logits, _ = self.gate(hidden_states)
topk_output = self.topk(hidden_states, router_logits)
if defer_finalize:
if not self.supports_deferred_finalize:
raise RuntimeError(
"Qwen deferred finalize requires a compatible FlashInfer "
"TRTLLM MoE producer and a separate shared expert"
)
if not TopKOutputChecker.format_is_bypassed(topk_output):
raise RuntimeError(
"Qwen deferred finalize requires logits-based bypassed TopK"
)
return self.experts.forward_deferred_finalize(hidden_states, topk_output)
if self.enable_shared_expert_fusion and TopKOutputChecker.format_is_standard(
topk_output
):
topk_output = self._append_shared_to_topk_output(topk_output, hidden_states)
return self.experts(hidden_states, topk_output)
def _gate_shared_output_out_of_place(
self,
hidden_states: torch.Tensor,
shared_output: torch.Tensor,
) -> torch.Tensor:
if self.shared_expert_gate is None:
return shared_output
return fused_gate_sigmoid_mul(
hidden_states,
self.shared_expert_gate.weight.squeeze(0),
shared_output,
)
def forward_normal_dual_stream(
self,
hidden_states: torch.Tensor,
use_fused_gate: bool = False,
defer_finalize: bool = False,
) -> torch.Tensor:
current_stream = torch.cuda.current_stream()
if defer_finalize:
# Keep routed FC2 on the current stream as finalize's PDL dependency;
# the shared branch reads the same input and writes a separate output.
self.alt_stream.wait_stream(current_stream)
router_output = self._forward_router_experts(
hidden_states, defer_finalize=True
)
with torch.cuda.stream(self.alt_stream):
shared_output = self._forward_shared_experts(
hidden_states, apply_gate=False
)
if shared_output is not None:
shared_output = self._gate_shared_output_out_of_place(
hidden_states, shared_output
)
current_stream.wait_stream(self.alt_stream)
return router_output, shared_output
self.alt_stream.wait_stream(current_stream)
shared_output = (
self._forward_shared_experts(
@@ -648,11 +738,18 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
self,
hidden_states: torch.Tensor,
forward_batch: Optional[ForwardBatch] = None,
defer_finalize: bool = False,
) -> torch.Tensor:
num_tokens, hidden_dim = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_dim)
if defer_finalize and num_tokens == 0:
raise RuntimeError("Qwen deferred finalize does not support M=0")
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori():
if (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_deepep_v2()
or get_moe_a2a_backend().is_mori()
):
return self._forward_deepep(hidden_states, forward_batch)
use_fused_gate = (
@@ -674,13 +771,34 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
and not torch.compiler.is_compiling()
):
final_hidden_states, shared_output = self.forward_normal_dual_stream(
hidden_states, use_fused_gate=use_fused_gate
hidden_states,
use_fused_gate=use_fused_gate,
defer_finalize=defer_finalize,
)
else:
shared_output = self._forward_shared_experts(
hidden_states, apply_gate=not use_fused_gate
hidden_states, apply_gate=not use_fused_gate and not defer_finalize
)
if defer_finalize and shared_output is not None:
shared_output = self._gate_shared_output_out_of_place(
hidden_states, shared_output
)
final_hidden_states = self._forward_router_experts(
hidden_states, defer_finalize=defer_finalize
)
if defer_finalize:
if shared_output is None:
raise RuntimeError("Qwen deferred finalize requires shared output")
from sglang.srt.layers.moe.qwen35_flashinfer_fusion import (
Qwen35MoeFinalizeHandoff,
)
return Qwen35MoeFinalizeHandoff.from_flashinfer(
final_hidden_states,
gated_shared_output=shared_output,
m=num_tokens,
)
final_hidden_states = self._forward_router_experts(hidden_states)
if shared_output is not None:
if use_fused_gate:
python/sglang/srt/models/qwen3_5.py
+238 -13
View File
@@ -39,6 +39,7 @@ from sglang.srt.configs.qwen3_5 import (
# Distributed
from sglang.srt.distributed import get_pp_group
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader
@@ -135,6 +136,9 @@ _gdn_use_alt_stream = _is_cuda or (
_qknorm_use_alt_stream = _is_cuda or (
get_bool_env_var("SGLANG_QK_NORM_ALT_STREAM", "False") and _hip_use_alt_stream
)
_gdn_decode_fused_proj_conv = (
_is_cuda and envs.SGLANG_ENABLE_GDN_DECODE_FUSED_PROJ_CONV.get()
)
_is_amx_available = cpu_has_amx_support()
_is_xpu = is_xpu()
@@ -156,7 +160,12 @@ def _disable_shared_experts_fusion() -> bool:
# Resolved lazily: the flag is written by the owning model's gate before
# its layers build (per runner); models without a gate see the config
# intent through the accessor's fallback.
return is_shared_experts_fusion_disabled()
# The deferred-finalize ABI needs the shared expert as a separate, gated
# local contribution; it cannot consume a shared slot fused into routed MoE.
return bool(
envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION.get()
or is_shared_experts_fusion_disabled()
)
def _maybe_enable_silu_fp4_quant_fusion(mlp: nn.Module) -> None:
@@ -185,6 +194,24 @@ def _maybe_enable_silu_fp4_quant_fusion(mlp: nn.Module) -> None:
logger.info("Enabled fused SiLU+mul+FP4-quant for dense MLP down_proj input.")
def _use_mnnvl_cutedsl_fusion(config: Qwen3_5TextConfig, is_nextn: bool) -> bool:
return bool(
not is_nextn
and config.model_type == "qwen3_5_moe_text"
and envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION.get()
)
def _layer_communicator_class(config: Qwen3_5TextConfig, is_nextn: bool):
if _use_mnnvl_cutedsl_fusion(config, is_nextn):
from sglang.srt.layers.moe.qwen35_flashinfer_fusion import (
Qwen35FlashInferLayerCommunicator,
)
return Qwen35FlashInferLayerCommunicator
return LayerCommunicator
if _is_cuda:
from sglang.kernels.ops.attention.fused_qk_rmsnorm_rope_gate import (
fused_qk_gemma_rmsnorm_rope_gate,
@@ -257,6 +284,27 @@ def _select_fused_ar_input_for_linear(hidden_states, linear: nn.Module):
)
def _finish_mlp_output(hidden_states, *, expect_deferred: bool):
if not expect_deferred:
if not isinstance(hidden_states, torch.Tensor):
from sglang.srt.layers.moe.qwen35_flashinfer_fusion import (
Qwen35MoeFinalizeHandoff,
)
if isinstance(hidden_states, Qwen35MoeFinalizeHandoff):
raise RuntimeError("unexpected deferred-finalize handoff")
hidden_states._sglang_needs_allreduce_fusion = True
return hidden_states
from sglang.srt.layers.moe.qwen35_flashinfer_fusion import (
Qwen35MoeFinalizeHandoff,
)
if not isinstance(hidden_states, Qwen35MoeFinalizeHandoff):
raise RuntimeError("Qwen3.5 expected a FlashInfer deferred-finalize handoff")
return hidden_states
if _is_npu:
from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import (
split_qkvgate_gemma_rmsnorm_rope,
@@ -730,10 +778,23 @@ class Qwen3_5GatedDeltaNet(nn.Module):
backend, projected_states_qkvz, projected_states_ba, forward_batch
)
if (
use_fused_decode_proj_conv = (
_gdn_decode_fused_proj_conv
and forward_batch.forward_mode.is_decode()
and isinstance(projected_states_qkvz, torch.Tensor)
and isinstance(projected_states_ba, torch.Tensor)
)
use_fused_contiguous_unpack = (
self.num_v_heads // self.num_k_heads in _GDN_FUSED_QKVZBA_RATIOS
and not _is_npu
):
)
if use_fused_decode_proj_conv:
# GDN owns indexed Conv1D state and the safe unpack/Conv boundary;
# it replaces these temporary B/A placeholders before recurrence.
mixed_qkv = (projected_states_qkvz, projected_states_ba)
z = None
b = projected_states_ba
a = projected_states_ba
elif use_fused_contiguous_unpack and not _is_npu:
if _is_cpu:
num_k_heads_tp = self.num_k_heads // self.attn_tp_size
num_v_heads_tp = self.num_v_heads // self.attn_tp_size
@@ -760,12 +821,22 @@ class Qwen3_5GatedDeltaNet(nn.Module):
)
mixed_qkv = torch.cat((query, key, value), dim=-1)
core_attn_out = self.attn(
attn_result = self.attn(
forward_batch,
mixed_qkv=mixed_qkv,
a=a,
b=b,
)
if use_fused_decode_proj_conv:
if not isinstance(attn_result, tuple) or len(attn_result) != 2:
raise RuntimeError(
"Fused GDN decode projection/Conv1D backend must return "
"(core_attn_out, z)"
)
core_attn_out, z = attn_result
else:
core_attn_out = attn_result
assert z is not None
z_shape_og = z.shape
# reshape input data into 2D tensor
@@ -862,7 +933,7 @@ class Qwen3_5LinearDecoderLayer(nn.Module):
_enable_qwen35_fused_ar_quant()
and _linear_accepts_fp8_tuple(self.linear_attn.in_proj_qkvz)
)
self.layer_communicator = LayerCommunicator(
self.layer_communicator = _layer_communicator_class(config, is_nextn)(
layer_scatter_modes=self.layer_scatter_modes,
input_layernorm=self.input_layernorm,
post_attention_layernorm=self.post_attention_layernorm,
@@ -913,6 +984,23 @@ class Qwen3_5LinearDecoderLayer(nn.Module):
forward_batch
)
)
defer_moe_finalize = (
fuse_mlp_allreduce
and isinstance(hidden_states, torch.Tensor)
and isinstance(self.mlp, Qwen2MoeSparseMoeBlock)
and hasattr(self.layer_communicator, "should_use_finalize")
and self.layer_communicator.should_use_finalize(
forward_batch, int(hidden_states.shape[0])
)
)
if (
fuse_mlp_allreduce
and self.layer_communicator.is_last_layer
and not defer_moe_finalize
):
# The last layer has no prepare_attn consumer for deferred AllReduce;
# fall back before MLP so postprocess_layer performs the collective.
fuse_mlp_allreduce = False
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_reduce_scatter,
@@ -921,11 +1009,14 @@ class Qwen3_5LinearDecoderLayer(nn.Module):
hidden_states = self.mlp(
hidden_states,
forward_batch,
defer_finalize=defer_moe_finalize,
)
else:
hidden_states = self.mlp(hidden_states)
if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True
hidden_states = _finish_mlp_output(
hidden_states, expect_deferred=defer_moe_finalize
)
else:
hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch
@@ -1084,7 +1175,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
enable_fused_ar_quant = (
_enable_qwen35_fused_ar_quant() and _linear_accepts_fp8_tuple(self.qkv_proj)
)
self.layer_communicator = LayerCommunicator(
self.layer_communicator = _layer_communicator_class(config, is_nextn)(
layer_scatter_modes=self.layer_scatter_modes,
input_layernorm=self.input_layernorm,
post_attention_layernorm=self.post_attention_layernorm,
@@ -1319,6 +1410,21 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
forward_batch
)
)
defer_moe_finalize = (
fuse_mlp_allreduce
and isinstance(hidden_states, torch.Tensor)
and isinstance(self.mlp, Qwen2MoeSparseMoeBlock)
and hasattr(self.layer_communicator, "should_use_finalize")
and self.layer_communicator.should_use_finalize(
forward_batch, int(hidden_states.shape[0])
)
)
if (
fuse_mlp_allreduce
and self.layer_communicator.is_last_layer
and not defer_moe_finalize
):
fuse_mlp_allreduce = False
with get_forward().scoped(
fuse_mlp_allreduce=fuse_mlp_allreduce,
mlp_reduce_scatter=mlp_reduce_scatter,
@@ -1327,11 +1433,14 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
hidden_states = self.mlp(
hidden_states,
forward_batch,
defer_finalize=defer_moe_finalize,
)
else:
hidden_states = self.mlp(hidden_states)
if fuse_mlp_allreduce:
hidden_states._sglang_needs_allreduce_fusion = True
hidden_states = _finish_mlp_output(
hidden_states, expect_deferred=defer_moe_finalize
)
else:
hidden_states, residual = self.layer_communicator.postprocess_layer(
hidden_states, residual, forward_batch
@@ -1479,6 +1588,48 @@ class Qwen3_5ForCausalLM(nn.Module):
prefix=f"{prefix}.layers",
)
self.flashinfer_mnnvl_cutedsl_fusion = None
if _use_mnnvl_cutedsl_fusion(config, is_nextn):
if self.pp_group.world_size != 1:
raise RuntimeError(
"Qwen3.5 FlashInfer MNNVL CuTe DSL fusion currently requires PP=1"
)
unsupported_layers = [
layer.layer_id
for layer in self.layers
if not isinstance(layer.mlp, Qwen2MoeSparseMoeBlock)
or not layer.mlp.supports_deferred_finalize
]
if unsupported_layers:
raise RuntimeError(
"Qwen3.5 FlashInfer MNNVL CuTe DSL fusion currently "
"requires block-FP8 MoE weights with FlashInfer TRTLLM "
"deferred-finalize support on every layer; unsupported "
"layers: "
f"{unsupported_layers}"
)
from sglang.srt.layers.moe.qwen35_flashinfer_fusion import (
Qwen35FlashInferFusionService,
Qwen35FlashInferLayerCommunicator,
)
self.flashinfer_mnnvl_cutedsl_fusion = Qwen35FlashInferFusionService(
hidden_size=config.hidden_size,
top_k=config.num_experts_per_tok,
rms_epsilon=config.rms_norm_eps,
)
for layer in self.layers:
communicator = layer.layer_communicator
if not isinstance(communicator, Qwen35FlashInferLayerCommunicator):
raise RuntimeError(
"Qwen3.5 fusion-enabled layer has the wrong communicator"
)
communicator.fusion_service = self.flashinfer_mnnvl_cutedsl_fusion
logger.info(
"Installed one Qwen3.5 FlashInfer fusion handle for %d layers",
len(self.layers),
)
# Final normalization
if self.pp_group.is_last_rank:
self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -1490,6 +1641,15 @@ class Qwen3_5ForCausalLM(nn.Module):
def get_input_embeddings(self):
return self.embed_tokens
def prepare_before_cuda_graph_capture(self, model_runner) -> None:
if self.flashinfer_mnnvl_cutedsl_fusion is None:
return
from sglang.srt.layers.moe.qwen35_flashinfer_fusion import (
prepare_qwen35_flashinfer_fusion,
)
prepare_qwen35_flashinfer_fusion(self, model_runner)
def set_dflash_layers_to_capture(self, layers_to_capture: list[int]):
self.layers_to_capture = layers_to_capture
for layer_id in self.layers_to_capture:
@@ -1513,6 +1673,16 @@ class Qwen3_5ForCausalLM(nn.Module):
pp_proxy_tensors: Optional[PPProxyTensors] = None,
input_deepstack_embeds: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, PPProxyTensors]:
if (
self.flashinfer_mnnvl_cutedsl_fusion is not None
and input_deepstack_embeds is not None
and input_deepstack_embeds.numel() > 0
):
raise RuntimeError(
"Qwen3.5 FlashInfer MNNVL CuTe DSL fusion currently supports "
"the text-only path, not deepstack visual inputs"
)
# Initialize hidden states
if self.pp_group.is_first_rank:
if input_embeds is None:
@@ -1564,12 +1734,62 @@ class Qwen3_5ForCausalLM(nn.Module):
}
)
# Apply final normalization
if hidden_states.shape[0] != 0:
# The final layer has no successor to consume its deferred MoE tail.
trace_final_norm = envs.SGLANG_TRACE_QWEN35_FINAL_NORM.get()
use_native_final_norm = envs.SGLANG_QWEN35_NATIVE_FINAL_NORM.get()
is_deferred_finalize = False
if self.flashinfer_mnnvl_cutedsl_fusion is not None:
from sglang.srt.layers.moe.qwen35_flashinfer_fusion import (
Qwen35MoeFinalizeHandoff,
)
is_deferred_finalize = isinstance(hidden_states, Qwen35MoeFinalizeHandoff)
if is_deferred_finalize:
if residual is None or self.flashinfer_mnnvl_cutedsl_fusion is None:
raise RuntimeError("invalid final deferred MoE handoff")
hidden_states, _ = self.flashinfer_mnnvl_cutedsl_fusion.finalize(
hidden_states, residual, self.norm.gemma_weight
)
elif hidden_states.shape[0] != 0:
if trace_final_norm:
print(
"SGLANG_TRACE_QWEN35_FINAL_NORM "
f"stage=pre_sync_enter hidden={tuple(hidden_states.shape)} "
f"hidden_stride={hidden_states.stride()} "
f"hidden_dtype={hidden_states.dtype} "
f"hidden_contiguous={hidden_states.is_contiguous()} "
f"residual={None if residual is None else tuple(residual.shape)} "
f"native={use_native_final_norm}",
flush=True,
)
torch.cuda.synchronize()
print(
"SGLANG_TRACE_QWEN35_FINAL_NORM stage=pre_sync_returned",
flush=True,
)
if residual is None:
hidden_states = self.norm(hidden_states)
hidden_states = (
self.norm.forward_native(hidden_states)
if use_native_final_norm
else self.norm(hidden_states)
)
else:
hidden_states, _ = self.norm(hidden_states, residual)
hidden_states, _ = (
self.norm.forward_native(hidden_states, residual)
if use_native_final_norm
else self.norm(hidden_states, residual)
)
if trace_final_norm:
print(
"SGLANG_TRACE_QWEN35_FINAL_NORM stage=post_sync_enter",
flush=True,
)
torch.cuda.synchronize()
print(
"SGLANG_TRACE_QWEN35_FINAL_NORM stage=post_sync_returned",
flush=True,
)
if len(aux_hidden_states) == 0:
return hidden_states
@@ -2064,6 +2284,11 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
def get_hidden_dim(self, module_name: str, layer_idx: int):
return self.model.get_hidden_dim(module_name, layer_idx)
def prepare_before_cuda_graph_capture(self, model_runner) -> None:
prepare = getattr(self.model, "prepare_before_cuda_graph_capture", None)
if prepare is not None:
prepare(model_runner)
def should_apply_lora(self, module_name: str) -> bool:
# Accept all language model layer modules (attention, linear_attn, mlp).
return module_name.startswith("model.layers.")
python/sglang/srt/models/qwen3_5_mtp.py
+17 -5
View File
@@ -154,12 +154,14 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
return self.model.embed_tokens.weight, self.lm_head.weight
def set_embed_and_head(self, embed, head):
del self.model.embed_tokens.weight
if not self.config.tie_word_embeddings:
# A last-stage draft can share only the target lm_head under PP; retain its
# own embedding for the first-stage half it cannot receive.
if embed is not None:
del self.model.embed_tokens.weight
self.model.embed_tokens.weight = embed
if head is not None and not self.config.tie_word_embeddings:
del self.lm_head.weight
self.model.embed_tokens.weight = embed
self.lm_head.weight = head
self.lm_head.weight = head
torch.cuda.empty_cache()
torch.cuda.synchronize()
@@ -214,6 +216,16 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
if not forward_batch.forward_mode.is_idle():
input_embeds = self.pre_fc_norm_embedding(input_embeds)
hidden_states = self.pre_fc_norm_hidden(hidden_states)
# Captured prefill gives padded embeddings but real-height target states;
# place the real rows in an equal-height slot whose padding stays unread.
if hidden_states.shape[0] != input_embeds.shape[0]:
rows = min(hidden_states.shape[0], input_embeds.shape[0])
slot = hidden_states.new_zeros(
(input_embeds.shape[0], hidden_states.shape[1])
)
slot[:rows] = hidden_states[:rows]
hidden_states = slot
hidden_states = torch.cat([input_embeds, hidden_states], dim=-1)
hidden_states = self.fc(hidden_states)
python/sglang/srt/models/qwen3_5_text.py
+47 -15
View File
@@ -107,8 +107,36 @@ class Qwen3_5ForCausalLM(nn.Module):
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens
def prepare_before_cuda_graph_capture(self, model_runner) -> None:
"""Forward model-owned warmup to the text backbone."""
self.model.prepare_before_cuda_graph_capture(model_runner)
def set_dflash_layers_to_capture(self, layer_ids: list[int]) -> None:
if self.pp_group.world_size > 1:
raise NotImplementedError("DFLASH/DSPARK aux hidden capture requires PP=1.")
num_layers = len(self.model.layers)
if sorted(set(layer_ids)) != list(layer_ids) or not all(
0 <= layer_id < num_layers - 1 for layer_id in layer_ids
):
raise ValueError(
"target_layer_ids must be unique, strictly increasing, and in "
f"[0, {num_layers - 1}); got {layer_ids}"
)
self.capture_aux_hidden_states = True
self.model.set_dflash_layers_to_capture(
[layer_id + 1 for layer_id in layer_ids]
)
def get_embed_and_head(self):
return self.model.embed_tokens.weight, self.lm_head.weight
# PP splits embedding and lm_head across first/last stages; the draft keeps
# its own copy of whichever half its stage cannot receive.
embed = (
None
if isinstance(self.model.embed_tokens, PPMissingLayer)
else self.model.embed_tokens.weight
)
head = None if isinstance(self.lm_head, PPMissingLayer) else self.lm_head.weight
return embed, head
def set_embed_and_head(self, embed, head):
del self.model.embed_tokens.weight
@@ -164,21 +192,25 @@ class Qwen3_5ForCausalLM(nn.Module):
params_dict = dict(self.named_parameters())
loaded_params: Set[str] = set()
body_weights = []
for name, loaded_weight in weights:
if name.startswith(_MODEL_PREFIX):
body_weights.append((name[len(_MODEL_PREFIX) :], loaded_weight))
elif name == "lm_head.weight":
if self.config.tie_word_embeddings:
continue
if "lm_head.weight" not in params_dict:
continue
param = params_dict["lm_head.weight"]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add("lm_head.weight")
# Keep prefix stripping lazy: materializing mmap-backed checkpoint tensors
# across model processes can exceed host memory before GPU copies finish.
def body_weights():
for name, loaded_weight in weights:
if name.startswith(_MODEL_PREFIX):
yield name[len(_MODEL_PREFIX) :], loaded_weight
elif name == "lm_head.weight":
if self.config.tie_word_embeddings:
continue
if "lm_head.weight" not in params_dict:
continue
param = params_dict["lm_head.weight"]
weight_loader = getattr(
param, "weight_loader", default_weight_loader
)
weight_loader(param, loaded_weight)
loaded_params.add("lm_head.weight")
body_loaded = self.model.load_weights(body_weights)
body_loaded = self.model.load_weights(body_weights())
loaded_params.update(f"{_MODEL_PREFIX}{n}" for n in body_loaded)
if self.config.tie_word_embeddings and self.pp_group.is_last_rank:
python/sglang/srt/runtime_context.py
+3
View File
@@ -497,6 +497,9 @@ class MoeFlags(_FlagGroupBase):
# speculative_moe_backend_context is active, so a draft gate's write also
# lands on the speculative leaf.
in_speculative_scope: bool = False
# Draft construction/execution uses a separate one-sided A2A workspace from
# the target model's concurrently live CUDA graphs.
speculative_context: bool = False
@dataclasses.dataclass
python/sglang/srt/server_args.py
+2 -2
View File
@@ -1743,9 +1743,9 @@ class ServerArgs:
bf16_gemm_backend: A[
str,
Arg(
help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM10x GPUs, except deterministic inference selects 'torch'; otherwise uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10x; dispatches between the CuTe DSL kernel and cuBLAS), 'torch' (always uses cuBLAS via torch.nn.functional.linear).",
help="Choose the backend for unquantized BF16 GEMM operations. Options: 'auto' (default; selects 'cutedsl' on SM10x GPUs, except deterministic inference selects 'torch'; otherwise uses cuBLAS via torch.nn.functional.linear), 'cutedsl' (SGLang JIT CuTe DSL TGV BF16 GEMM on SM10x; dispatches between the allowlisted low-M Split-K kernel, the CuTe DSL kernel, and cuBLAS; set SGLANG_ENABLE_BF16_SPLITK_GEMM=0 to disable Split-K), 'flashinfer_pr4266' (legacy compatibility alias for the optimized CuTe DSL path), 'gemv', 'torch' (always uses cuBLAS via torch.nn.functional.linear).",
cli_name="--bf16-gemm-backend",
choices=["auto", "cutedsl", "gemv", "torch"],
choices=["auto", "cutedsl", "flashinfer_pr4266", "gemv", "torch"],
),
NS("exec.kernel"),
] = "auto"
python/sglang/srt/speculative/dflash_worker_v2.py
+5 -2
View File
@@ -306,7 +306,7 @@ class DFlashWorkerV2(BaseSpecWorker):
bundle = build_draft_tp_worker(
server_args=server_args,
gpu_id=gpu_id,
ps=replace(ps, pp_rank=0),
ps=replace(ps, pp_rank=0, pp_size=1),
nccl_port=nccl_port,
target_model_config=target_worker.model_runner.model_config,
algo_label="DFLASH",
@@ -1662,13 +1662,16 @@ class DFlashWorkerV2(BaseSpecWorker):
batch: ScheduleBatch,
on_publish=None,
grammar_barrier=None,
pp_proxy_tensors=None,
) -> GenerationBatchResult:
self._validate_phase1_sampling_support(batch)
if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
# Target prefill: capture DFlash aux hidden states for prompt tokens.
batch_output = self.target_worker.forward_batch_generation(
batch, capture_hidden_mode=CaptureHiddenMode.FULL
batch,
pp_proxy_tensors=pp_proxy_tensors,
capture_hidden_mode=CaptureHiddenMode.FULL,
)
logits_output, next_token_ids = (
python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py
+1 -1
View File
@@ -131,7 +131,7 @@ class DSparkWorkerV2(BaseSpecWorker):
bundle = build_draft_tp_worker(
server_args=server_args,
gpu_id=gpu_id,
ps=replace(ps, pp_rank=0),
ps=replace(ps, pp_rank=0, pp_size=1),
nccl_port=nccl_port,
target_model_config=target_worker.model_runner.model_config,
algo_label="DSPARK",
python/sglang/srt/speculative/eagle_utils.py
+17 -3
View File
@@ -485,13 +485,27 @@ def get_draft_input_from_target_hidden_dim(model_runner: ModelRunner) -> int:
return target_hidden * num_aux
def get_draft_recurrent_hidden_state_spec_from_config(
model_config, spec_algorithm
) -> tuple[Optional[int], Optional[torch.dtype]]:
"""Return hidden_states width/dtype carried between draft decode steps.
Config-only so callers without a draft runner can reach it: prefill-side PP
builds the draft on the last stage alone, but the PD metadata wire schema it
feeds has to come out identical on every rank.
"""
if spec_algorithm.is_standalone():
return None, None
return model_config.spec_hidden_size, model_config.dtype
def get_draft_recurrent_hidden_state_spec(
model_runner: ModelRunner,
) -> tuple[Optional[int], Optional[torch.dtype]]:
"""Return hidden_states width/dtype carried between draft decode steps."""
if model_runner.spec_algorithm.is_standalone():
return None, None
return model_runner.model_config.spec_hidden_size, model_runner.model_config.dtype
return get_draft_recurrent_hidden_state_spec_from_config(
model_runner.model_config, model_runner.spec_algorithm
)
def eagle_prepare_for_verify(
python/sglang/srt/speculative/eagle_worker_v2.py
+33 -12
View File
@@ -7,6 +7,7 @@ from typing import List, Optional
import torch
from sglang.kernels.ops.speculative.topk1 import draft_topk1_postprocess
from sglang.srt.distributed import get_pp_group
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_runner import (
@@ -87,6 +88,7 @@ from sglang.srt.speculative.eagle_worker_common import (
)
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import (
draft_pp_context,
draft_tp_context,
fast_sample,
get_plan_stream,
@@ -168,16 +170,17 @@ class EagleDraftWorker(EagleDraftWorkerBase):
ctx = empty_context()
with (
ctx
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), draft_model_build_scope():
), draft_pp_context(), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), draft_model_build_scope():
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
# spec workers don't support pipeline parallelism
ps=replace(ps, pp_rank=0),
ps=replace(ps, pp_rank=0, pp_size=1),
nccl_port=nccl_port,
is_draft_worker=True,
# The draft runs at absolute target positions.
context_length=target_worker.model_runner.model_config.context_len,
random_seed=target_worker.random_seed,
)
# Alias for better readability
@@ -314,7 +317,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
)
else:
if self.hot_token_id is not None:
if self.hot_token_id is not None and head is not None:
head = head.clone()
self.hot_token_id = self.hot_token_id.to(head.device)
head.data = head.data[self.hot_token_id]
@@ -1078,17 +1081,24 @@ class EAGLEWorkerV2(BaseSpecWorker):
get_spec().speculative_algorithm
)
self._draft_worker = EagleDraftWorker(
server_args,
gpu_id,
ps,
nccl_port,
target_worker,
# Only the last PP stage runs the draft; other EAGLEWorkerV2 instances
# return proxies so scheduler dispatch remains rank-uniform.
self._hosts_draft = get_pp_group().is_last_rank
self._draft_worker = (
EagleDraftWorker(
server_args,
gpu_id,
ps,
nccl_port,
target_worker,
)
if self._hosts_draft
else None
)
# Adaptive speculative
self.adaptive_controller: Optional[AdaptiveController] = None
if get_spec().speculative_adaptive:
if get_spec().speculative_adaptive and self._hosts_draft:
self.adaptive_controller = AdaptiveController(
self,
config_path=get_spec().speculative_adaptive_config,
@@ -1151,7 +1161,11 @@ class EAGLEWorkerV2(BaseSpecWorker):
)
def forward_batch_generation(
self, batch: ScheduleBatch, on_publish=None, grammar_barrier=None
self,
batch: ScheduleBatch,
on_publish=None,
grammar_barrier=None,
pp_proxy_tensors=None,
):
if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
# Target prefill
@@ -1161,7 +1175,9 @@ class EAGLEWorkerV2(BaseSpecWorker):
else CaptureHiddenMode.FULL
)
batch_output = self.target_worker.forward_batch_generation(
batch, capture_hidden_mode=target_capture_mode
batch,
pp_proxy_tensors=pp_proxy_tensors,
capture_hidden_mode=target_capture_mode,
)
# Spec_v2 convention: batch.seq_lens = length BEFORE this iter's tokens.
@@ -1171,6 +1187,11 @@ class EAGLEWorkerV2(BaseSpecWorker):
if on_publish is not None:
on_publish(batch_output.new_seq_lens)
# A rank that does not host the draft (prefill-side PP builds it only on
# the last stage) forwards the target's proxy tensors and stops here.
if self._draft_worker is None:
return batch_output
# Draft prefill
with (
self.draft_worker.draft_tp_context(
python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py
+1 -1
View File
@@ -143,7 +143,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
server_args=server_args,
gpu_id=gpu_id,
# spec workers don't support pipeline parallelism
ps=replace(ps, pp_rank=0),
ps=replace(ps, pp_rank=0, pp_size=1),
nccl_port=nccl_port,
is_draft_worker=True,
# The draft runs at absolute target positions.
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+1 -1
View File
@@ -164,7 +164,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
server_args=server_args,
gpu_id=gpu_id,
# spec workers don't support pipeline parallelism
ps=replace(ps, pp_rank=0),
ps=replace(ps, pp_rank=0, pp_size=1),
nccl_port=nccl_port,
is_draft_worker=True,
is_multi_layer_eagle=True,
python/sglang/srt/speculative/ngram_worker.py
+5 -3
View File
@@ -427,7 +427,7 @@ class NGRAMWorker(BaseSpecWorker):
self.ngram_corpus.batch_put(batch_tokens)
def forward_batch_generation(
self, batch: ScheduleBatch, on_publish=None
self, batch: ScheduleBatch, on_publish=None, pp_proxy_tensors=None
) -> GenerationBatchResult:
fwd_stream = torch.get_device_module(self.device).current_stream()
record_stream_for_v2_verify(batch, None, fwd_stream)
@@ -442,7 +442,7 @@ class NGRAMWorker(BaseSpecWorker):
if batch.forward_mode.is_target_verify():
batch_result = self.target_worker.forward_batch_generation(
batch, is_verify=True
batch, pp_proxy_tensors=pp_proxy_tensors, is_verify=True
)
logits_output, can_run_cuda_graph = (
@@ -528,7 +528,9 @@ class NGRAMWorker(BaseSpecWorker):
batch.forward_mode = ForwardMode.DECODE
else:
batch_result = self.target_worker.forward_batch_generation(batch)
batch_result = self.target_worker.forward_batch_generation(
batch, pp_proxy_tensors=pp_proxy_tensors
)
logits_output, predict, can_run_cuda_graph = (
batch_result.logits_output,
batch_result.next_token_ids,
python/sglang/srt/speculative/spec_utils.py
+10
View File
@@ -35,6 +35,8 @@ from sglang.srt.configs.hybrid_arch import mambaish_config
from sglang.srt.constrained.base_grammar_backend import GrammarMask
from sglang.srt.distributed.parallel_state import (
GroupCoordinator,
get_self_pp_group,
patch_pipeline_parallel_group,
patch_tensor_parallel_group,
)
from sglang.srt.environ import envs
@@ -676,6 +678,14 @@ def load_token_map(token_map_path: str) -> List[int]:
return torch.tensor(hot_token_id, dtype=torch.int64)
@contextmanager
def draft_pp_context():
# The draft model is one layer and never spans pipeline stages; give it a
# single-member pp group so it initializes as if pp were off.
with patch_pipeline_parallel_group(get_self_pp_group()):
yield
@contextmanager
def draft_tp_context(tp_group: GroupCoordinator):
# Draft model doesn't use dp and has its own tp group.
python/sglang/srt/speculative/standalone_worker_v2.py
+1 -1
View File
@@ -84,7 +84,7 @@ class StandaloneDraftWorker(EagleDraftWorker):
server_args=server_args,
gpu_id=gpu_id,
# spec workers don't support pipeline parallelism
ps=replace(ps, pp_rank=0),
ps=replace(ps, pp_rank=0, pp_size=1),
nccl_port=nccl_port,
is_draft_worker=True,
# The draft runs at absolute target positions.