chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

feat(deepgemm): migrate to sgl-deep-gemm 0.1.2 wheel API

Browse Source 
Port the DeepGEMM "deprecate-from-sgl-kernel, separate sgl-deep-gemm wheel"
migration (upstream PR #24268 and follow-ons) so our code runs against the
torch-2.11 dev-cu13 image, which ships DeepGEMM as the separate sgl-deep-gemm
0.1.2 wheel (import name still `deep_gemm`). Verified against upstream/main HEAD,
NOT the introducing PR — the wheel API drifted between 0.0.1 and 0.1.2.

Ports (all verified against HEAD = wheel 0.1.2):
- nsa_indexer/nsa_backend: paged-MQA context_lens must be (N_total, 1). The
  0.0.1 form (batch_size, next_n) DEADLOCKS fp8_paged_mqa_logits on next_n>=2
  (our EAGLE deploy uses next_n=4 on SM90/H200, which does not take the SM100-
  only DG-native broadcast path). Matches HEAD's _to_2d_context_lens.
- compile_utils warmup: hasattr-guard the dropped get/set_compile_mode API;
  pass m_indices positionally to m_grouped_fp8_gemm_nt_contiguous.
- fp8_utils.transform_scale_ue8m0: restore TMA-aligned stride when the DLPack
  round-trip collapses a size-1 trailing dim.
- moe_runner/deep_gemm: guard the SBO masked-gemm return unpack when overlap is
  inactive (#26839) — reachable via our --enable-single-batch-overlap.
- entrypoint: enable DeepGEMM PDL by default, hasattr-guarded (#23979).
- engine: require sglang-kernel >= 0.4.3 (first sgl-deep-gemm-era kernel); the
  migrated (N_total,1) paths would misbehave on the old bundled DeepGEMM.

Verified-unchanged symbols left as-is: fp8_mqa_logits, fp8_gemm_nt, bf16_gemm_*,
get_mk_alignment_for_contiguous_layout, transform_sf_into_required_layout,
get/set_num_sms, the masked-gemm signature. Runtime validation pending the
torch-2.11 image (tai-kernel rebuild + harness).

Refs: WI-2026-06-07-001

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 32818784d332b5c48faf8027247cd4cd9a0f48cd)
This commit is contained in:
孙远航
2026-06-07 17:33:32 +00:00
committed by laoyao0822
parent 430ed85083
commit db5ab48b19
9 changed files with 121 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/entrypoints/engine.py
View File

@@ -1141,9 +1141,12 @@ def _set_envs_and_config(server_args: ServerArgs):
"at https://docs.flashinfer.ai/installation.html.",
)
if _is_cuda:
# 0.4.3 is the first sgl-kernel that ships DeepGEMM as the separate
# sgl-deep-gemm wheel; our migrated deep_gemm call paths target that wheel's
# API and would misbehave against the older bundled DeepGEMM. [[WI-2026-06-07-001]]
assert_pkg_version(
"sglang-kernel",
"0.4.0",
"0.4.3",
"Please reinstall the latest version with `pip install sglang-kernel --force-reinstall`",
)

12
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -1012,10 +1012,18 @@ class Indexer(MultiPlatformOp):
# Reuse pre-computed schedule metadata if available (from init_forward_metadata),
# otherwise fall back to computing it here.
schedule_metadata = getattr(metadata, "paged_mqa_schedule_metadata", None)
# sgl-deep-gemm (DeepGEMM release-0426+) requires context_lens of shape
# [batch_size, next_n] to match q.shape = [batch_size, next_n, heads, head_dim].
# The indexer uses next_n=1 with batch_size=N_total via q_fp8.unsqueeze(1) below,
# so mirror that layout here. See [[WI-2026-06-07-001]] (upstream PR #24268).
if seqlens_32.dim() == 2:
seqlens_32_2d = seqlens_32
else:
seqlens_32_2d = seqlens_32.unsqueeze(-1)
if _is_cuda:
if schedule_metadata is None:
schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
seqlens_32, blocksize, self.sm_count
seqlens_32_2d, blocksize, self.sm_count
)
assert len(q_fp8.shape) == 3
@@ -1067,7 +1075,7 @@ class Indexer(MultiPlatformOp):
q_fp8[:q_offset],
kv_cache_fp8,
weights[:q_offset],
seqlens_32,
seqlens_32_2d,
block_tables,
schedule_metadata,
max_seq_len,

32
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -87,6 +87,29 @@ _EAGLE_ACCEPT_DRAFT_MLA_PATH_DEBUG_COUNTS: Dict[Tuple[int, int], int] = {}
_CP_SHARED_KV_MLA_PREFETCH_FALLBACK_LOG_COUNTS: Dict[str, int] = {}
def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tensor:
"""Normalize context_lens to (N_total, 1) for sgl-deep-gemm.
sgl-deep-gemm 0.1.2's ``get_paged_mqa_logits_metadata`` /
``fp8_paged_mqa_logits`` require an (N_total, 1) context_lens layout. PR #24268
originally introduced a (batch_size, next_n) 2-D requirement for the 0.0.1 wheel,
but the wheel later changed the required layout: passing a (batch_size, next_n>=2)
view DEADLOCKS ``fp8_paged_mqa_logits`` (matters here — EAGLE runs next_n=4 on
SM90/H200, which does not take the SM100-only DG-native broadcast path).
Matches upstream HEAD's ``_to_2d_context_lens``. See [[WI-2026-06-07-001]].
``batch_size`` is retained for signature parity with upstream but is unused.
"""
# Always normalize to (N_total, 1) layout, to avoid deadlock at
# deep_gemm.fp8_paged_mqa_logits.
if seqlens_32.dim() == 2:
if seqlens_32.size(1) == 1:
return seqlens_32
# Re-flatten a (bs, next_n) view — we want (N_total, 1) regardless.
seqlens_32 = seqlens_32.reshape(-1)
return seqlens_32.contiguous().view(-1, 1)
def _log_cp_shared_kv_mla_prefetch_fallback(
reason: str,
message: str,
@@ -928,8 +951,9 @@ class NativeSparseAttnBackend(
)
else cache_seqlens_int32
)
seqlens_32_2d = _to_2d_context_lens(seqlens_32, batch_size)
paged_mqa_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
seqlens_32, 64, deep_gemm.get_num_sms()
seqlens_32_2d, 64, deep_gemm.get_num_sms()
)
except (ImportError, ModuleNotFoundError):
paged_mqa_schedule_metadata = None
@@ -1308,8 +1332,9 @@ class NativeSparseAttnBackend(
)
else cache_seqlens_int32
)
seqlens_32_2d = _to_2d_context_lens(seqlens_32, bs)
paged_mqa_schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata(
seqlens_32, 64, deep_gemm.get_num_sms()
seqlens_32_2d, 64, deep_gemm.get_num_sms()
)
except (ImportError, ModuleNotFoundError):
paged_mqa_schedule_metadata = None
@@ -1457,8 +1482,9 @@ class NativeSparseAttnBackend(
)
else metadata.cache_seqlens_int32
)
seqlens_32_2d = _to_2d_context_lens(seqlens_32, bs)
new_schedule = deep_gemm.get_paged_mqa_logits_metadata(
seqlens_32, 64, deep_gemm.get_num_sms()
seqlens_32_2d, 64, deep_gemm.get_num_sms()
)
if metadata.paged_mqa_schedule_metadata is None:
metadata.paged_mqa_schedule_metadata = new_schedule

16
python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py
View File

@@ -198,12 +198,19 @@ def _compile_deep_gemm_one_type_all(
kernel_type, max_m=max_m, n=n, k=k, num_groups=num_groups
)
old_compile_mode = deep_gemm.get_compile_mode()
deep_gemm.set_compile_mode(1)
# sgl-deep-gemm (DeepGEMM release-0426+) dropped the compile-mode API; guard it.
# See [[WI-2026-06-07-001]] (upstream PR #24268).
has_compile_mode_api = hasattr(deep_gemm, "get_compile_mode") and hasattr(
deep_gemm, "set_compile_mode"
)
if has_compile_mode_api:
old_compile_mode = deep_gemm.get_compile_mode()
deep_gemm.set_compile_mode(1)
# TODO can use multi thread
for m in tqdm(m_list, desc=f"DeepGEMM warmup"):
executor.execute(m=m)
deep_gemm.set_compile_mode(old_compile_mode)
if has_compile_mode_api:
deep_gemm.set_compile_mode(old_compile_mode)
# clean up input buffers
torch.cuda.current_stream().synchronize()
@@ -302,7 +309,8 @@ class _GroupedContWarmupExecutor(_BaseWarmupExecutor):
(self.lhs_q[:m], self.lhs_s[:m]),
(self.rhs_q, self.rhs_s),
self.out[:m],
m_indices=self.m_indices[:m],
# sgl-deep-gemm takes m_indices positionally. [[WI-2026-06-07-001]] (PR #24268)
self.m_indices[:m],
)

7
python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py
View File

@@ -19,6 +19,13 @@ if ENABLE_JIT_DEEPGEMM:
import deep_gemm
from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor # noqa: F401
# Enable DeepGEMM PDL (Programmatic Dependent Launch) by default to overlap
# kernel launches, matching upstream's new-wheel default. Guarded: no-op if the
# sgl-deep-gemm wheel does not expose set_pdl. Upstream PR #23979.
# [[WI-2026-06-07-001]]
if get_bool_env_var("SGLANG_DEEPGEMM_PDL", "true") and hasattr(deep_gemm, "set_pdl"):
deep_gemm.set_pdl(True)
_SANITY_CHECK = get_bool_env_var("SGLANG_DEEPGEMM_SANITY_CHECK")

5
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
View File

@@ -385,7 +385,10 @@ class DeepGemmRunnerCore(MoeRunnerCore):
**gemm_overlap_args_dict,
)
meta_overlap_args = running_state.get("meta_overlap_args", None)
if meta_overlap_args is not None:
# sgl-deep-gemm's masked gemm returns (block_m, threshold) only with down-gemm
# overlap, else None; meta_overlap_args may be set without overlap when SBO is
# enabled, so guard the unpack. Upstream PR #26839. [[WI-2026-06-07-001]]
if meta_overlap_args is not None and deep_gemm_return_value is not None:
block_m, threshold = deep_gemm_return_value
meta_overlap_args["block_m"] = block_m
meta_overlap_args["threshold"] = threshold

14
python/sglang/srt/layers/quantization/fp8_utils.py
View File

@@ -1198,6 +1198,20 @@ def transform_scale_ue8m0(sf, mn, use_torch_impl: bool = False):
sf = sf.index_select(-2, torch.arange(mn, device=sf.device) // 128)
sf = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf)
# In sgl-deep-gemm, the C++ deepgemm path returns through DLPack which collapses the
# stride of size-1 trailing dims to 1 (happens when packed_sf_k == 1, i.e.
# K <= block_k * 4). Restore the TMA-aligned stride so the deepgemm assertion
# sf.stride(-1) == get_tma_aligned_size(mn, element_size) holds.
# See [[WI-2026-06-07-001]] (upstream PR #24268).
if not use_torch_impl and sf.shape[-1] == 1:
from deep_gemm.utils import get_tma_aligned_size
aligned_mn = get_tma_aligned_size(sf.shape[-2], sf.element_size())
if sf.stride(-1) != aligned_mn:
new_stride = list(sf.stride())
new_stride[-1] = aligned_mn
sf = sf.as_strided(sf.shape, tuple(new_stride))
return sf