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

[Comm] Drop the in-tree MNNVL CuTe DSL port in favor of FlashInfer 0.6.18 (#37206)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2026-09-03 21:49:48 -07:00
committed by GitHub
parent e3305b3b87
commit 3ad3f23ed5
17 changed files with 9 additions and 7121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/communication/mnnvl_cutedsl/__init__.py
-52
View File
@@ -1,52 +0,0 @@
# 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
@@ -1,231 +0,0 @@
# 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
@@ -1,876 +0,0 @@
# 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
@@ -1,43 +0,0 @@
# 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
@@ -1,511 +0,0 @@
# 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
@@ -1,37 +0,0 @@
# 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
@@ -1,978 +0,0 @@
# 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
@@ -1,469 +0,0 @@
# 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
@@ -1,47 +0,0 @@
# 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
@@ -1,461 +0,0 @@
# 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
@@ -1,338 +0,0 @@
# 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
@@ -1,70 +0,0 @@
# 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
@@ -1,107 +0,0 @@
# 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
@@ -1,584 +0,0 @@
# 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/srt/layers/flashinfer_mnnvl_cutedsl.py
+9 -29
View File
@@ -2,7 +2,6 @@
from __future__ import annotations from __future__ import annotations
import importlib.util
import logging import logging
import threading import threading
from dataclasses import dataclass, replace from dataclasses import dataclass, replace
@@ -21,34 +20,15 @@ logger = logging.getLogger(__name__)
def _import_kernel_backend(): def _import_kernel_backend():
try: # Imported here rather than at module scope: the CuTe DSL backend drags in
from flashinfer.comm import AllReduceFusionPattern # CUDA-only dependencies that CPU-side importers of this module never need.
except ImportError as error: from flashinfer.comm import AllReduceFusionPattern, allreduce_fusion
raise RuntimeError( from flashinfer.comm.mnnvl_cutedsl import DEFAULT_CONFIG
"MNNVL CuTe DSL fusion requires FlashInfer's communication " from flashinfer.comm.mnnvl_cutedsl_ar import MNNVLCuteDSLAllReduceFusionWorkspace
"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 ( return (
MNNVLCuteDSLAllReduceFusionWorkspace, MNNVLCuteDSLAllReduceFusionWorkspace,
mnnvl_cutedsl_allreduce_fusion, allreduce_fusion,
AllReduceFusionPattern, AllReduceFusionPattern,
DEFAULT_CONFIG, DEFAULT_CONFIG,
) )
@@ -180,8 +160,8 @@ class FlashInferMNNVLCuteDSLARFusion:
config=self.workspace_config, config=self.workspace_config,
) )
# Publish only after the mailbox barrier; older FlashInfer workspace # Publish only after the mailbox barrier; without it the ranks
# classes may not provide it and would desynchronize Lamport stages. # would desynchronize their Lamport stages.
torch.cuda.synchronize(self.device) torch.cuda.synchronize(self.device)
dist.barrier(group=process_group) dist.barrier(group=process_group)