From 19c30dff56204d3c259410ae8819c8309e3985f5 Mon Sep 17 00:00:00 2001 From: eeecho <82921722+AliceChenyy@users.noreply.github.com> Date: Thu, 3 Sep 2026 06:49:59 +0800 Subject: [PATCH] [SM120] DeepSeek-V4: DeepGEMM paged-MQA indexer +FP4 MoE+ page-split (#29927) Co-authored-by: Baizhou Zhang --- .../kernels/ops/attention/flash_mla_sm120.py | 117 +++++++++-- python/sglang/kernels/ops/layernorm/mhc.py | 31 ++- python/sglang/srt/arg_groups/model_hook.py | 17 +- .../layers/attention/deepseek_v4_backend.py | 12 ++ .../srt/layers/attention/dsv4/indexer.py | 186 +++++++++------- .../srt/layers/attention/dsv4/metadata.py | 31 ++- .../layers/deep_gemm_wrapper/configurer.py | 12 +- .../srt/layers/moe/moe_runner/deep_gemm.py | 31 ++- .../layers/moe/moe_runner/deep_gemm_sm120.py | 198 ++++++++++++++++++ python/sglang/srt/model_loader/utils.py | 5 + python/sglang/srt/models/deepseek_v4.py | 19 +- 11 files changed, 537 insertions(+), 122 deletions(-) create mode 100644 python/sglang/srt/layers/moe/moe_runner/deep_gemm_sm120.py diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index 6f39fa05b..bf8dc234d 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/python/sglang/kernels/ops/attention/flash_mla_sm120.py @@ -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, diff --git a/python/sglang/kernels/ops/layernorm/mhc.py b/python/sglang/kernels/ops/layernorm/mhc.py index a6a56318e..d34cca1b4 100644 --- a/python/sglang/kernels/ops/layernorm/mhc.py +++ b/python/sglang/kernels/ops/layernorm/mhc.py @@ -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, diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index e513232dd..56bfd3b87 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -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) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index fcf42aaf0..67dc74e18 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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, diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index f9fccebe8..ea4608eb8 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -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[ diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index 56c762cf1..4c6c3d663 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -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 diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py index 17f04b2a6..a21073f5c 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py @@ -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) diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 421b25ab5..8eff8224f 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -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"] diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm_sm120.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm_sm120.py new file mode 100644 index 000000000..78d341b22 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm_sm120.py @@ -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, + ) diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index 199603321..c82e5d60f 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/python/sglang/srt/model_loader/utils.py @@ -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 ( diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index fba2702e7..e6ce967ca 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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, )