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

[SM120] DeepSeek-V4: DeepGEMM paged-MQA indexer +FP4 MoE+ page-split (#29927)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
eeecho
2026-09-02 15:49:59 -07:00
committed by GitHub
co-authored by Baizhou Zhang
parent 982aa8acfc
commit 19c30dff56
11 changed files with 537 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/flash_mla_sm120.py
+98 -19
View File
@@ -210,6 +210,64 @@ def _sm120_sparse_decode_fwd(
_sm120_default_backend = envs.SGLANG_SM120_FLASHMLA_BACKEND.get()
SM120_DECODE_MAX_TOKENS = 64
def _flash_mla_sm120_prefill(
q,
k_cache,
indices,
topk_length,
attn_sink,
head_dim_v,
softmax_scale,
extra_k_cache,
extra_indices,
extra_topk_length,
):
from flashinfer.mla._sparse_mla_sm120 import _sparse_mla_sm120_paged_attention
q2 = q.squeeze(1) if q.ndim == 4 else q
num_tokens, num_heads, _ = q2.shape
dev = q2.device
kv_u8 = k_cache.view(torch.uint8) if k_cache.dtype != torch.uint8 else k_cache
src_pbs = k_cache.shape[1] if k_cache.ndim >= 3 else _PBS_SRC
idx = indices.squeeze(1) if indices.dim() == 3 else indices
kv_64 = (
_split_kv_pages_to_64(kv_u8, src_pbs, touched_indices=idx)
if src_pbs != _PBS_DST
else kv_u8
)
extra_kv_u8 = (
extra_k_cache.view(torch.uint8)
if extra_k_cache is not None and extra_k_cache.dtype != torch.uint8
else extra_k_cache
)
extra_idx = (
extra_indices.squeeze(1)
if extra_indices is not None and extra_indices.dim() == 3
else extra_indices
)
output = q2.new_empty((num_tokens, num_heads, head_dim_v), dtype=torch.bfloat16)
out_lse = torch.empty((num_tokens, num_heads), dtype=torch.float32, device=dev)
_sparse_mla_sm120_paged_attention(
q2,
kv_64,
idx,
output,
out_lse,
softmax_scale,
topk_length=topk_length,
attn_sink=attn_sink,
extra_kv_cache=extra_kv_u8,
extra_indices=extra_idx,
extra_topk_length=extra_topk_length,
mid_out=None,
mid_lse=None,
)
return (output.unsqueeze(1), None)
def flash_mla_with_kvcache_sm120(**kwargs):
"""SM120 FlashMLA sparse decode entry point.
@@ -229,6 +287,19 @@ def flash_mla_with_kvcache_sm120(**kwargs):
extra_topk_length = kwargs.get("extra_topk_length")
if _sm120_default_backend == "flashinfer":
if q.shape[0] > SM120_DECODE_MAX_TOKENS:
return _flash_mla_sm120_prefill(
q,
k_cache,
indices,
topk_length,
attn_sink,
head_dim_v,
softmax_scale,
extra_k_cache,
extra_indices,
extra_topk_length,
)
return _flash_mla_flashinfer(
q,
k_cache,
@@ -326,24 +397,31 @@ def _page_split_kernel(
if tl.load(mask_ptr + page_idx) == 0:
return
src_base = src_ptr + page_idx * src_stride0
dst_base = dst_ptr + (page_idx * RATIO + sub) * dst_stride0
# All layout strides/offsets are 8-byte aligned (asserted at the call
# site), so copy in u64 lanes instead of single bytes.
src_u64 = src_ptr.to(tl.pointer_type(tl.uint64))
dst_u64 = dst_ptr.to(tl.pointer_type(tl.uint64))
src_base = src_u64 + page_idx * (src_stride0 // 8)
dst_base = dst_u64 + (page_idx * RATIO + sub) * (dst_stride0 // 8)
# Copy data region: DATA_PER_SUB bytes from src offset sub*DATA_PER_SUB
data_src_off = sub * DATA_PER_SUB
for start in tl.range(0, DATA_PER_SUB, BLOCK_SIZE):
DATA_U64: tl.constexpr = DATA_PER_SUB // 8
SCALE_U64: tl.constexpr = SCALE_PER_SUB // 8
# Copy data region: DATA_U64 u64 lanes from src offset sub*DATA_U64
data_src_off = sub * DATA_U64
for start in tl.range(0, DATA_U64, BLOCK_SIZE):
offs = start + tl.arange(0, BLOCK_SIZE)
mask = offs < DATA_PER_SUB
mask = offs < DATA_U64
vals = tl.load(src_base + data_src_off + offs, mask=mask)
tl.store(dst_base + offs, vals, mask=mask)
# Copy scale region: SCALE_PER_SUB bytes
scale_src_off = SRC_SCALE_OFF + sub * SCALE_PER_SUB
for start in tl.range(0, SCALE_PER_SUB, BLOCK_SIZE):
# Copy scale region: SCALE_U64 u64 lanes
scale_src_off = SRC_SCALE_OFF // 8 + sub * SCALE_U64
for start in tl.range(0, SCALE_U64, BLOCK_SIZE):
offs = start + tl.arange(0, BLOCK_SIZE)
mask = offs < SCALE_PER_SUB
mask = offs < SCALE_U64
vals = tl.load(src_base + scale_src_off + offs, mask=mask)
tl.store(dst_base + DST_SCALE_OFF + offs, vals, mask=mask)
tl.store(dst_base + DST_SCALE_OFF // 8 + offs, vals, mask=mask)
@triton.jit
@@ -361,13 +439,12 @@ def _page_mark_kernel(
the same value 1 are safe (no atomic needed).
"""
pid = tl.program_id(0)
if pid >= N_idx:
return
idx = tl.load(indices_ptr + pid)
if idx < 0:
return
offs = pid * BLOCK + tl.arange(0, BLOCK)
valid = offs < N_idx
idx = tl.load(indices_ptr + offs, mask=valid, other=-1)
keep = valid & (idx >= 0)
page = idx // SRC_PBS
tl.store(mask_ptr + page, 1)
tl.store(mask_ptr + page, tl.full((BLOCK,), 1, tl.int8), mask=keep)
def _split_kv_pages_to_64(
@@ -441,15 +518,17 @@ def _split_kv_pages_to_64(
idx_flat = touched_indices.reshape(-1).contiguous()
if idx_flat.dtype != torch.int32:
idx_flat = idx_flat.to(torch.int32)
_page_mark_kernel[(idx_flat.numel(),)](
_MARK_BLOCK = 1024
_page_mark_kernel[(triton.cdiv(idx_flat.numel(), _MARK_BLOCK),)](
idx_flat,
mask,
idx_flat.numel(),
src_pbs, # SRC_PBS
1024, # BLOCK (unused, kept for JIT signature)
_MARK_BLOCK,
)
mask_ptr = mask
assert src_stride0 % 8 == 0 and _BYTES_PER_DST_PAGE_PADDED % 8 == 0
grid = (N * ratio,)
_page_split_kernel[grid](
src_2d,
python/sglang/kernels/ops/layernorm/mhc.py
+28 -3
View File
@@ -583,6 +583,20 @@ def mhc_pre_gemm_sqrsum_tilelang(
T.pdl_trigger()
@functools.cache
def _mhc_pre_gemm_sqrsum_dispatch():
"""SM120's TileLang pipeline cannot warp-specialize this kernel (the role
marker fails on tirx.Bind), so re-wrap it there with warp specialization
disabled. Other archs keep the original compiled form."""
from sglang.srt.utils import is_sm120_supported
if not is_sm120_supported():
return mhc_pre_gemm_sqrsum_tilelang
_tl = _load_tilelang()
cfg = {_tl.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True}
return _tl.jit(pass_configs=cfg)(mhc_pre_gemm_sqrsum_tilelang.__wrapped__)
@functools.cache
def mhc_pre_gemm_sqrsum_splitk_kernel(
hc_mult3: int,
@@ -603,7 +617,18 @@ def mhc_pre_gemm_sqrsum_splitk_kernel(
ENABLE_PDL = is_arch_support_pdl()
@tilelang.jit
from sglang.srt.utils import is_sm120_supported
_tl = _load_tilelang()
# See _mhc_pre_gemm_sqrsum_dispatch: SM120 cannot compile the
# warp-specialized form of these kernels.
_cfg = (
{_tl.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True}
if is_sm120_supported()
else None
)
@tilelang.jit(pass_configs=_cfg)
def mhc_pre_gemm_sqrsum_splitk_stage_0(
x: T.Tensor[(num_tokens, hc_hidden_size), T.bfloat16],
fn: T.Tensor[(hc_mult3, hc_hidden_size), T.float32],
@@ -1081,7 +1106,7 @@ def mhc_pre(
assert (
n_splits == 1
), "The simple TileLang version gemm_sqrsum doesn't support split-k"
mhc_pre_gemm_sqrsum_tilelang(
_mhc_pre_gemm_sqrsum_dispatch()(
residual_flat.view(num_tokens, hc_mult * hidden_size),
fn_flat,
gemm_out_mul.squeeze(0),
@@ -1637,7 +1662,7 @@ def mhc_fused_post_pre(
gemm_out_sqrsum_1d = torch.empty(
num_tokens, dtype=torch.float32, device=residual.device
)
mhc_pre_gemm_sqrsum_tilelang(
_mhc_pre_gemm_sqrsum_dispatch()(
residual_cur.view(num_tokens, hc_hidden_size),
fn,
gemm_out_mul_2d,
python/sglang/srt/arg_groups/model_hook.py
+12 -5
View File
@@ -345,13 +345,20 @@ def handle_model_specific_adjustments(server_args: Any):
# DeepGEMM or require >99KB SMEM (topk_v2).
envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
envs.SGLANG_OPT_USE_TOPK_V2.set(False)
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.set(False)
if not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.is_set():
envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.set(False)
if not envs.SGLANG_OPT_FUSE_MHC_POST_PRE.is_set():
envs.SGLANG_OPT_FUSE_MHC_POST_PRE.set(True)
envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False)
envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.set(True)
# Prefer TileLang over the Torch fallback.
envs.SGLANG_OPT_USE_TILELANG_INDEXER.set(True)
if not envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.is_set():
envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False)
# Out of the box the indexer runs the TileLang kernel (works on
# stock DeepGEMM); both knobs stay env-overridable so a DeepGEMM
# build with SM120 attention support can opt into
# fp8_paged_mqa_logits by setting them to 0.
if not envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.is_set():
envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.set(True)
if not envs.SGLANG_OPT_USE_TILELANG_INDEXER.is_set():
envs.SGLANG_OPT_USE_TILELANG_INDEXER.set(True)
elif get_platform().is_hip:
envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.set(False)
envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+12
View File
@@ -1789,9 +1789,21 @@ class DeepseekV4AttnBackend(
if get_platform().is_sm120:
from sglang.kernels.ops.attention.flash_mla_sm120 import (
SM120_DECODE_MAX_TOKENS,
flash_mla_with_kvcache_sm120,
)
# The pad to 64 heads only serves the decode kernel's h_q
# specialization; the prefill kernel takes arbitrary h_q, so
# drop it instead of attending on garbage heads (4x the work
# at attn-TP 4).
real_heads = layer.tp_q_head_num
if q.shape[0] > SM120_DECODE_MAX_TOKENS:
if q.shape[-2] > real_heads:
q = q[..., :real_heads, :].contiguous()
if attn_sink is not None and attn_sink.shape[0] > real_heads:
attn_sink = attn_sink[:real_heads]
o = flash_mla_with_kvcache_sm120(
q=q,
k_cache=swa_k_cache,
python/sglang/srt/layers/attention/dsv4/indexer.py
+105 -81
View File
@@ -33,6 +33,7 @@ from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend
from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.attention.dsv4.compressor import Compressor
from sglang.srt.layers.attention.dsv4.metadata import (
_SM120_INDEXER_M_CHUNK,
NonPagedIndexerPlan,
PagedIndexerMetadata,
)
@@ -822,55 +823,6 @@ class C4IndexerBackendMixin:
c4_seq_lens=c4_seq_lens,
query_rows=query_rows,
)
if use_aiter_fp4:
q_fp4, q_scale = q
logits = aiter_fp4_paged_mqa_logits(
q_fp4=q_fp4,
q_scale=q_scale,
k_payload=token_to_kv_pool.get_index_k_fp4_payload_buffer(
c4_indexer.layer_id
),
k_scale=token_to_kv_pool.get_index_k_fp4_scale_buffer(
c4_indexer.layer_id
),
weights=weights,
page_table=page_table,
c4_seq_lens=c4_seq_lens,
weight_scale=c4_indexer.weight_scale,
is_decode=forward_batch.forward_mode.is_decode(),
decode_workspace=metadata.fp4_decode_workspace,
prefill_workspace=metadata.fp4_prefill_workspace,
)
elif nonpaged_plan is not None:
assert isinstance(q_indexer, torch.Tensor)
logits = self._forward_nonpaged_indexer(
q_indexer=q_indexer,
weights=weights,
c4_indexer=c4_indexer,
token_to_kv_pool=token_to_kv_pool,
plan=nonpaged_plan,
)
else:
c4_indexer_kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(
layer_id=c4_indexer.layer_id,
)
assert c4_indexer_kv_cache.dim() == 2
head_dim_with_sf = 68 if use_fp4_indexer else 132
c4_indexer_kv_cache = c4_indexer_kv_cache.view(
c4_indexer_kv_cache.shape[0], 64, 1, head_dim_with_sf
)
logits = fn(
q,
c4_indexer_kv_cache,
weights,
_c4sl,
page_table,
indexer_metadata.deep_gemm_metadata,
indexer_metadata.max_c4_seq_len,
False,
)
assert indexer_metadata.page_table is core_metadata.page_table
if self.debug_use_external_c4_sparse_indices:
return
@@ -892,42 +844,114 @@ class C4IndexerBackendMixin:
elif core_metadata.c4_sparse_raw_indices is not None:
raw_indices = core_metadata.c4_sparse_raw_indices
if self.dsa_topk_backend.is_torch():
topk_transform_pytorch_vectorized(
logits,
c4_seq_lens,
page_table,
c4_sparse_page_indices,
indexer_metadata.c4_page_size,
raw_indices,
all_rows = slice(0, _c4sl.shape[0])
def run_topk_transform(rows: slice, logits: torch.Tensor) -> None:
row_raw_indices = raw_indices[rows] if raw_indices is not None else None
if self.dsa_topk_backend.is_torch():
topk_transform_pytorch_vectorized(
logits,
c4_seq_lens[rows],
page_table[rows],
c4_sparse_page_indices[rows],
indexer_metadata.c4_page_size,
row_raw_indices,
)
elif self.dsa_topk_backend.is_flashinfer():
self.flashinfer_topk_transform(
logits,
c4_seq_lens[rows],
page_table[rows],
c4_sparse_page_indices[rows],
indexer_metadata.c4_page_size,
row_raw_indices,
)
elif self.dsa_topk_backend.should_use_topk_v2() and raw_indices is None:
topk_transform_paged_v2(
logits,
c4_seq_lens[rows],
page_table[rows],
c4_sparse_page_indices[rows],
indexer_metadata.c4_page_size,
indexer_metadata.topk_metadata,
)
else:
topk_transform_paged(
logits,
c4_seq_lens[rows],
page_table[rows],
c4_sparse_page_indices[rows],
indexer_metadata.c4_page_size,
row_raw_indices,
)
if nonpaged_plan is not None:
assert isinstance(q_indexer, torch.Tensor)
logits = self._forward_nonpaged_indexer(
q_indexer=q_indexer,
weights=weights,
c4_indexer=c4_indexer,
token_to_kv_pool=token_to_kv_pool,
plan=nonpaged_plan,
)
elif self.dsa_topk_backend.is_flashinfer():
self.flashinfer_topk_transform(
logits,
c4_seq_lens,
page_table,
c4_sparse_page_indices,
indexer_metadata.c4_page_size,
raw_indices,
)
elif self.dsa_topk_backend.should_use_topk_v2() and raw_indices is None:
topk_transform_paged_v2(
logits,
c4_seq_lens,
page_table,
c4_sparse_page_indices,
indexer_metadata.c4_page_size,
indexer_metadata.topk_metadata,
run_topk_transform(all_rows, logits)
elif use_aiter_fp4:
q_fp4, q_scale = q
logits = aiter_fp4_paged_mqa_logits(
q_fp4=q_fp4,
q_scale=q_scale,
k_payload=token_to_kv_pool.get_index_k_fp4_payload_buffer(
c4_indexer.layer_id
),
k_scale=token_to_kv_pool.get_index_k_fp4_scale_buffer(
c4_indexer.layer_id
),
weights=weights,
page_table=page_table,
c4_seq_lens=c4_seq_lens,
weight_scale=c4_indexer.weight_scale,
is_decode=forward_batch.forward_mode.is_decode(),
decode_workspace=metadata.fp4_decode_workspace,
prefill_workspace=metadata.fp4_prefill_workspace,
)
run_topk_transform(all_rows, logits)
else:
topk_transform_paged(
logits,
c4_seq_lens,
page_table,
c4_sparse_page_indices,
indexer_metadata.c4_page_size,
raw_indices,
c4_indexer_kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(
layer_id=c4_indexer.layer_id,
)
assert c4_indexer_kv_cache.dim() == 2
head_dim_with_sf = 68 if use_fp4_indexer else 132
c4_indexer_kv_cache = c4_indexer_kv_cache.view(
c4_indexer_kv_cache.shape[0], 64, 1, head_dim_with_sf
)
def run_paged_indexer(rows: slice, metadata: torch.Tensor) -> None:
row_q = (q[0][rows], q[1][rows]) if isinstance(q, tuple) else q[rows]
logits = fn(
row_q,
c4_indexer_kv_cache,
weights[rows],
_c4sl[rows],
page_table[rows],
metadata,
indexer_metadata.max_c4_seq_len,
False,
)
run_topk_transform(rows, logits)
deep_gemm_metadata = indexer_metadata.deep_gemm_metadata
if isinstance(deep_gemm_metadata, list):
# SM120 only: DeepGEMM's metadata kernel caps the row count, so
# PagedIndexerMetadata split it; run indexer + topk per chunk.
num_rows = _c4sl.shape[0]
for chunk_idx, start in enumerate(
range(0, num_rows, _SM120_INDEXER_M_CHUNK)
):
rows = slice(start, min(start + _SM120_INDEXER_M_CHUNK, num_rows))
run_paged_indexer(rows, deep_gemm_metadata[chunk_idx])
else:
run_paged_indexer(all_rows, deep_gemm_metadata)
if hisparse_coordinator is not None:
if hisparse_decode:
compress_layer_id = token_to_kv_pool.layer_mapping[
python/sglang/srt/layers/attention/dsv4/metadata.py
+24 -7
View File
@@ -7,7 +7,9 @@ from typing import TYPE_CHECKING, Any, List, Optional
import torch
from sglang.srt.environ import envs
from sglang.srt.utils import is_hip, is_xpu
from sglang.srt.utils import is_hip, is_sm120_supported, is_xpu
_IS_SM120 = is_sm120_supported()
if TYPE_CHECKING:
pass
@@ -50,6 +52,8 @@ Some other notes:
"""
_LARGE_INDEXER_QUERY_THRESHOLD = 11673
_SM120_INDEXER_M_CHUNK = 4096
def copy_metadata(
*,
@@ -143,13 +147,26 @@ class PagedIndexerMetadata:
_c4 = self.c4_seq_lens.to(torch.int32)
if _c4.dim() == 1:
_c4 = _c4.unsqueeze(-1)
self.deep_gemm_metadata = get_paged_mqa_logits_metadata(
_c4,
self.c4_page_size,
deep_gemm.get_num_sms(),
)
if _IS_SM120 and _c4.shape[0] > _SM120_INDEXER_M_CHUNK:
# Chunk metadata is identical for every layer in the forward
# pass; compute the per-chunk list once here instead of per
# layer in the indexer.
self.deep_gemm_metadata = [
get_paged_mqa_logits_metadata(
_c4[_s : _s + _SM120_INDEXER_M_CHUNK],
self.c4_page_size,
deep_gemm.get_num_sms(),
)
for _s in range(0, _c4.shape[0], _SM120_INDEXER_M_CHUNK)
]
else:
self.deep_gemm_metadata = get_paged_mqa_logits_metadata(
_c4,
self.c4_page_size,
deep_gemm.get_num_sms(),
)
assert isinstance(self.deep_gemm_metadata, torch.Tensor)
assert isinstance(self.deep_gemm_metadata, (torch.Tensor, list))
if self.use_topk_v2:
from sglang.kernels.ops.attention.dsv4 import plan_topk_v2
python/sglang/srt/layers/deep_gemm_wrapper/configurer.py
+9 -3
View File
@@ -18,9 +18,13 @@ def _compute_enable_deep_gemm():
sm_version = get_device_sm()
if (_is_cuda and sm_version < 90) or (_is_musa and sm_version < 31):
return False
# DeepGEMM requires TMEM/tcgen05 (SM100+datacenter), not available on SM120
# SM120 support (mma.sync block-scale, no TMEM) landed in DeepGEMM#324;
# probe the entry point since installed builds may predate it.
if sm_version == 120:
return False
try:
from deep_gemm import m_grouped_fp8_fp4_gemm_nt_contiguous # noqa: F401
except (ImportError, AttributeError):
return False
if not (_is_cuda or _is_musa):
return False
@@ -35,5 +39,7 @@ def _compute_enable_deep_gemm():
ENABLE_JIT_DEEPGEMM = _compute_enable_deep_gemm()
DEEPGEMM_BLACKWELL = ENABLE_JIT_DEEPGEMM and get_platform().is_sm100
DEEPGEMM_SCALE_UE8M0 = DEEPGEMM_BLACKWELL
DEEPGEMM_SCALE_UE8M0 = ENABLE_JIT_DEEPGEMM and (
get_platform().is_sm100 or get_device_sm() == 120
)
DEEPGEMM_NEED_TMA_ALIGNED_SCALES = not (DEEPGEMM_SCALE_UE8M0 or _is_musa)
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
+29 -2
View File
@@ -20,6 +20,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
from sglang.srt.environ import envs
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.moe.moe_runner import deep_gemm_sm120
from sglang.srt.layers.moe.moe_runner.base import (
MoeQuantInfo,
MoeRunnerConfig,
@@ -272,7 +273,11 @@ class DeepGemmRunnerCore(MoeRunnerCore):
assert self.config.activation in ("silu", "situ")
assert self.config.is_gated
self.swiglu_limit = self.config.swiglu_limit
self.use_swizzle = get_moe_a2a_backend().is_megamoe()
# SM120's contiguous GEMM only consumes standard-layout activations, so
# it opts out of swizzle regardless of the a2a backend.
self.use_swizzle = (
get_moe_a2a_backend().is_megamoe() and deep_gemm_sm120.use_swizzle()
)
def run(
self,
@@ -896,6 +901,19 @@ def pre_permute_standard_to_deep_gemm(
dispatch_output.topk_output,
)
topk_weights, topk_ids, _ = topk_output
# SM120's DeepGEMM grouped GEMM consumes standard-layout activations only.
# Feeding it the shared masked/swizzled layout does not raise -- it silently
# returns wrong results (GSM8K 0.96 -> 0.06, measured on 4x RTX 6000D), so
# refuse the combination rather than corrupt output.
assert deep_gemm_sm120.is_supported(), (
"--moe-runner-backend deep_gemm on consumer Blackwell (SM120) requires "
"the standard-layout MoE path, which is unavailable in this build."
)
sm120_input = deep_gemm_sm120.maybe_pre_permute(
hidden_states, topk_ids, topk_weights, quant_info, runner_config, running_state
)
if sm120_input is not None:
return sm120_input
hidden_states_shape = hidden_states.shape
hidden_states_dtype = hidden_states.dtype
@@ -904,7 +922,10 @@ def pre_permute_standard_to_deep_gemm(
topk_weights, topk_ids = topk_weights, topk_ids
if _should_use_masked_standard_layout(runner_config, quant_info, hidden_states):
if (
deep_gemm_sm120.allows_masked_standard_layout()
and _should_use_masked_standard_layout(runner_config, quant_info, hidden_states)
):
output_dtype = (
torch.bfloat16
if quant_info.w13_weight.dtype == torch.bfloat16
@@ -1123,6 +1144,12 @@ def post_permute_deep_gemm_to_standard(
from sglang.kernels.ops.moe.ep_moe_kernels import post_reorder_deepgemm
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
sm120_output = deep_gemm_sm120.maybe_post_permute(
runner_output, runner_config, running_state
)
if sm120_output is not None:
return sm120_output
hidden_states_shape = running_state["hidden_states_shape"]
hidden_states_dtype = running_state["hidden_states_dtype"]
hidden_states_device = running_state["hidden_states_device"]
python/sglang/srt/layers/moe/moe_runner/deep_gemm_sm120.py
+198
View File
@@ -0,0 +1,198 @@
"""SM120-specific DeepGEMM MoE path.
Consumer Blackwell's DeepGEMM contiguous grouped GEMM takes standard-layout
activations, so SM120 uses a contiguous scatter/gather permute instead of the
masked layout the other architectures use. Kept out of ``deep_gemm.py`` so the
shared runner keeps a single code path.
Entry points, all no-ops off SM120:
* ``use_swizzle`` — layout choice for the shared runner
* ``maybe_pre_permute`` — contiguous scatter, or ``None`` to fall through
* ``maybe_post_permute`` — matching gather, or ``None`` to fall through
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
from sglang.srt.utils import ceil_div, dispose_tensor, is_sm120_supported
if TYPE_CHECKING:
from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmRunnerInput
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
_is_sm120 = is_sm120_supported()
# Above this the standard path uses the contiguous grouped GEMM; the masked
# layout's [num_local_experts, capacity, k] transients OOM at prefill sizes.
_STANDARD_CONTIG_MIN_TOKENS = 1024
def is_supported() -> bool:
"""True when this build can serve the SM120 standard-layout MoE path.
Off SM120 the shared masked path is correct, so this is vacuously true;
on SM120 it reports whether the contiguous implementation is present.
"""
return True
def use_swizzle() -> bool:
"""SM120's contiguous GEMM consumes standard-layout activations only."""
return not _is_sm120
def allows_masked_standard_layout() -> bool:
"""SM120 cannot serve the masked-standard layout for DSV4 shapes: its varlen
activation kernel requires ``D // 8 >= num_experts`` (512 // 8 = 64 < 256),
so keep upstream's memory-budget heuristic off this path.
"""
return not _is_sm120
def _eligible(hidden_states, quant_info, runner_config) -> bool:
return (
_is_sm120
and hidden_states.shape[0] >= _STANDARD_CONTIG_MIN_TOKENS
and quant_info.w13_weight.dtype != torch.bfloat16
# Under EP the standard dispatcher maps non-local experts to -1,
# which the contiguous path does not handle; keep the masked path.
and runner_config.num_local_experts == runner_config.num_experts
)
def maybe_pre_permute(
hidden_states: torch.Tensor,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
quant_info,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> Optional[DeepGemmRunnerInput]:
"""Contiguous scatter for SM120; ``None`` means use the shared masked path."""
if not _eligible(hidden_states, quant_info, runner_config):
return None
return _pre_permute_standard_contig(
hidden_states, topk_ids, topk_weights, runner_config, running_state
)
def maybe_post_permute(
runner_output,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> Optional[StandardCombineInput]:
"""Gather matching ``maybe_pre_permute``; ``None`` means fall through."""
if not running_state.get("contig_mode"):
return None
from sglang.kernels.ops.moe.ep_moe_kernels import ep_gather
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
gather_out = torch.empty(
running_state["hidden_states_shape"],
device=running_state["hidden_states_device"],
dtype=running_state["hidden_states_dtype"],
)
ep_gather(
runner_output.hidden_states,
running_state["topk_ids"],
running_state["topk_weights"],
running_state["output_index"],
gather_out,
)
dispose_tensor(runner_output.hidden_states)
if runner_config.routed_scaling_factor is not None:
gather_out *= runner_config.routed_scaling_factor
return StandardCombineInput(hidden_states=gather_out)
def _pre_permute_standard_contig(
hidden_states: torch.Tensor,
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
runner_config: MoeRunnerConfig,
running_state: dict,
) -> DeepGemmRunnerInput:
from sglang.kernels.ops.moe.ep_moe_kernels import ep_scatter
from sglang.kernels.ops.quantization.fp8_kernel import (
sglang_per_token_group_quant_fp8,
)
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmRunnerInput
num_tokens, K = hidden_states.shape
num_experts = runner_config.num_local_experts
device = hidden_states.device
running_state["topk_ids"] = topk_ids
running_state["topk_weights"] = topk_weights
running_state["hidden_states_shape"] = hidden_states.shape
running_state["hidden_states_device"] = device
running_state["hidden_states_dtype"] = hidden_states.dtype
running_state["contig_mode"] = True
ue8m0 = deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
q, q_scale = sglang_per_token_group_quant_fp8(
hidden_states,
128,
column_major_scales=ue8m0,
scale_tma_aligned=ue8m0,
scale_ue8m0=ue8m0,
)
dispose_tensor(hidden_states)
# ep_scatter fills expert slots in blocks of 128; size from a static upper
# bound and accumulate counts on-device to avoid a GPU->CPU sync.
flat_ids = topk_ids.flatten().to(torch.int64)
counts = torch.zeros(num_experts, device=device, dtype=torch.int32)
counts.index_add_(0, flat_ids.clamp_min(0), (flat_ids >= 0).to(torch.int32))
counts_aligned = (counts + 127) // 128 * 128
all_tokens = ceil_div(num_tokens * runner_config.top_k, 128) * 128 + (
num_experts * 128
)
running_state["all_tokens"] = all_tokens
# Pad slots (m_indices == -1) are skipped by the grouped GEMM, so the
# buffer needs no zero fill.
input_tensor = torch.empty((all_tokens, K), device=device, dtype=q.dtype)
if ue8m0:
input_tensor_scale = torch.zeros(
(ceil_div(K // 128, 4), all_tokens), device=device, dtype=torch.int
).transpose(0, 1)
else:
input_tensor_scale = torch.empty(
(all_tokens, K // 128), device=device, dtype=torch.float32
)
m_indices = torch.full((all_tokens,), -1, device=device, dtype=torch.int32)
output_index = torch.empty_like(topk_ids)
expert_start_loc = torch.empty_like(counts_aligned)
ep_scatter(
q,
q_scale,
topk_ids,
counts_aligned,
counts,
expert_start_loc,
input_tensor,
input_tensor_scale,
m_indices,
output_index,
scale_ue8m0=ue8m0,
)
dispose_tensor(q)
if q_scale is not None:
dispose_tensor(q_scale)
running_state["output_index"] = output_index
return DeepGemmRunnerInput(
hidden_states=input_tensor,
hidden_states_scale=input_tensor_scale,
use_masked_gemm=False,
m_indices=m_indices,
)
python/sglang/srt/model_loader/utils.py
+5
View File
@@ -16,6 +16,7 @@ from transformers.dynamic_module_utils import get_class_from_dynamic_module
from sglang.srt.configs.model_config import ModelConfig, ModelImpl
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.utils import get_device_sm
logger = logging.getLogger(__name__)
@@ -277,6 +278,10 @@ def should_deepgemm_weight_requant_ue8m0(
and weight_block_size is not None
):
return False
# SM120 routes dense block-FP8 GEMMs to CUTLASS/Triton (fp32 scales);
# only the grouped MoE GEMM consumes DeepGEMM layouts there.
if get_device_sm() == 120:
return False
if output_dtype is not None and output_dtype != torch.bfloat16:
return False
if weight_shape is not None and (
python/sglang/srt/models/deepseek_v4.py
+17 -2
View File
@@ -29,6 +29,7 @@ from sglang.kernels.ops.attention.dsv4 import (
fused_rope_inplace,
sglang_per_token_group_quant_fp8_dsv4_wo_a,
)
from sglang.kernels.ops.attention.flash_mla_sm120 import SM120_DECODE_MAX_TOKENS
from sglang.kernels.ops.quantization.fp8_kernel import (
sglang_per_token_group_quant_fp8,
)
@@ -176,6 +177,7 @@ from sglang.srt.utils import (
is_gfx95_supported,
is_gfx942_supported,
is_gfx1250_supported,
is_sm120_supported,
log_info_on_rank0,
make_layers,
)
@@ -222,6 +224,7 @@ logger = logging.getLogger(__name__)
_FP8_WO_A_GEMM = envs.SGLANG_OPT_FP8_WO_A_GEMM.get()
_MHC_POST_MULT_VALUE = 2.0
_HC_PRENORM_DEEPGEMM_MIN_TOKENS = 1024
DEEPSEEK_V4_STACKED_PARAMS_MAPPING: List[Tuple[str, str, int]] = [
("gate_up_proj", "gate_proj", 0),
@@ -1581,12 +1584,19 @@ class MQALayer(MqaAttentionBase):
)
tp_slice, q_padded, q_out = slice(None), None, None
# Above this the SM120 route is the prefill kernel, which takes
# arbitrary h_q, so the decode pad below would just be sliced back off.
skip_decode_pad = is_sm120_supported() and x.shape[0] > SM120_DECODE_MAX_TOKENS
if self.attn_tp_size > 1:
# FlashMLA's fp8 sparse decode kernel only specializes h_q for {64, 128}.
# Pad the per-rank heads to 64 (not the full n_heads) when they fit, to
# dispatch the cheaper decode::head64 variant; attn_sink is sliced to
# this rank and padded to match.
padded_num_heads = 64 if self.n_local_heads <= 64 else self.n_heads
padded_num_heads = (
self.n_local_heads
if skip_decode_pad
else (64 if self.n_local_heads <= 64 else self.n_heads)
)
# Only [0:n_local_heads] is written below. Uninitialized padded TP
# heads inject NaN into attention on gfx942 (fnuz), so zero-init
# there; other archs tolerate new_empty and skip the per-forward
@@ -1970,7 +1980,12 @@ class DeepseekV4DecoderLayer(nn.Module):
)
return y, post.squeeze(-1), comb, False
if envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get():
# The deepgemm tf32 gemm wins at large M (prefill) but its fixed
# dispatch cost dominates at small M (decode): dispatch by token count.
if (
envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get()
and x.shape[0] >= _HC_PRENORM_DEEPGEMM_MIN_TOKENS
):
from sglang.srt.layers.deep_gemm_wrapper.entrypoint import (
tf32_hc_prenorm_gemm,
)