diff --git a/docker/Dockerfile b/docker/Dockerfile index e6884d3b5..fe7c2680a 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -11,7 +11,7 @@ ARG GRACE_BLACKWELL_DEEPEP_BRANCH=gb200_blog_part_2 ARG HOPPER_SBO_DEEPEP_COMMIT=9f2fc4b3182a51044ae7ecb6610f7c9c3258c4d6 ARG DEEPEP_COMMIT=9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee ARG BUILD_AND_DOWNLOAD_PARALLEL=8 -ARG SGL_KERNEL_VERSION=0.4.2 +ARG SGL_KERNEL_VERSION=0.4.2.post1 ARG SGL_VERSION ARG USE_LATEST_SGLANG=0 ARG GDRCOPY_VERSION=2.5.1 diff --git a/python/pyproject.toml b/python/pyproject.toml index 9da94e25f..19eab8a8b 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -59,7 +59,8 @@ dependencies = [ "sentencepiece", "setproctitle", "flash-attn-4>=4.0.0b9", - "sglang-kernel==0.4.2", + "sgl-deep-gemm==0.0.1", + "sglang-kernel==0.4.2.post1", "soundfile==0.13.1", "tiktoken", "timm==1.0.16", diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 0b7849a6b..1f7e79276 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1170,7 +1170,7 @@ def _set_envs_and_config(server_args: ServerArgs): if _is_cuda: assert_pkg_version( "sglang-kernel", - "0.4.2", + "0.4.2.post1", "Please reinstall the latest version with `pip install sglang-kernel --force-reinstall`", ) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 87bfecc75..78dcf0aec 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -456,10 +456,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) + # 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. + 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 @@ -508,7 +516,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, diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index f583509a7..f515228d4 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -67,6 +67,17 @@ else: ) +def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tensor: + if seqlens_32.dim() == 2: + return seqlens_32 + n = seqlens_32.numel() + assert ( + n % batch_size == 0 + ), f"seqlens_32 size {n} is not a multiple of batch_size {batch_size}" + next_n = n // batch_size + return seqlens_32.view(batch_size, next_n) + + # Reuse this workspace buffer across all NSA backend instances global_workspace_buffer = None @@ -647,8 +658,11 @@ class NativeSparseAttnBackend( ) else cache_seqlens_int32 ) + seqlens_32_2d = _to_2d_context_lens( + seqlens_32, forward_batch.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 @@ -932,8 +946,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 @@ -1082,8 +1097,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 diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index a30bf5e14..02c31cf15 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -196,12 +196,18 @@ 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) + 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() @@ -300,7 +306,7 @@ 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], + self.m_indices[:m], ) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py index 6ae5ee431..de433fe2d 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/configurer.py @@ -18,6 +18,8 @@ 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 + if not (_is_cuda or _is_musa): + return False try: import deep_gemm # noqa: F401 diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index b39da5267..3bb83c111 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1286,6 +1286,19 @@ 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. + 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 diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index 66a37310d..5d8e0cd69 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -294,6 +294,12 @@ install_sglang_kernel() { else echo "CUSTOM_BUILD_SGL_KERNEL=true: keeping freshly built sgl-kernel wheel." fi + SGL_DEEP_GEMM_VERSION=$(grep -Po -m1 '(?<=sgl-deep-gemm==)[0-9A-Za-z\.\-]+' python/pyproject.toml) + if [ "$CU_MAJOR" = "13" ]; then + $PIP_CMD install "sgl-deep-gemm==${SGL_DEEP_GEMM_VERSION}" --force-reinstall $PIP_INSTALL_SUFFIX + else + $PIP_CMD install "https://github.com/sgl-project/whl/releases/download/v${SGL_DEEP_GEMM_VERSION}/sgl_deep_gemm-${SGL_DEEP_GEMM_VERSION}+cu129-py3-none-manylinux2014_$(uname -m).whl" --force-reinstall $PIP_INSTALL_SUFFIX + fi mark_step_done "${FUNCNAME[0]}" } diff --git a/scripts/ci/cuda/warmup_deep_gemm.py b/scripts/ci/cuda/warmup_deep_gemm.py index b1844bc59..0e8a0442d 100644 --- a/scripts/ci/cuda/warmup_deep_gemm.py +++ b/scripts/ci/cuda/warmup_deep_gemm.py @@ -235,8 +235,11 @@ def compile_one_shape(kernel_type, n, k, num_groups, m_list): ) m_list = [m for m in m_list if m <= max_m] - old_mode = deep_gemm.get_compile_mode() - deep_gemm.set_compile_mode(1) + get_compile_mode = getattr(deep_gemm, "get_compile_mode", None) + set_compile_mode = getattr(deep_gemm, "set_compile_mode", None) + old_mode = get_compile_mode() if get_compile_mode is not None else None + if set_compile_mode is not None: + set_compile_mode(1) try: if kernel_type == "NORMAL": lhs_q, lhs_s = _empty_token_fp8((max_m, k)) @@ -255,7 +258,7 @@ def compile_one_shape(kernel_type, n, k, num_groups, m_list): (lhs_q[:m], lhs_s[:m]), (rhs_q, rhs_s), out[:m], - m_indices=m_indices[:m], + m_indices[:m], ) elif kernel_type == "MASKED": @@ -274,7 +277,8 @@ def compile_one_shape(kernel_type, n, k, num_groups, m_list): expected_m=m, ) finally: - deep_gemm.set_compile_mode(old_mode) + if set_compile_mode is not None and old_mode is not None: + set_compile_mode(old_mode) torch.cuda.current_stream().synchronize() torch.cuda.empty_cache() diff --git a/test/registered/8-gpu-models/test_dsa_models_mtp.py b/test/registered/8-gpu-models/test_dsa_models_mtp.py index 05cd47632..9bb05f36c 100644 --- a/test/registered/8-gpu-models/test_dsa_models_mtp.py +++ b/test/registered/8-gpu-models/test_dsa_models_mtp.py @@ -16,7 +16,11 @@ from sglang.test.test_utils import ( write_github_step_summary, ) -register_cuda_ci(est_time=1048, suite="stage-c-test-8-gpu-h200") +register_cuda_ci( + est_time=1048, + suite="stage-c-test-8-gpu-h200", + disabled="Disabled due to #24268. Should be fixed soon.", +) FULL_DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2" GLM5_MODEL_PATH = "zai-org/GLM-5-FP8" diff --git a/test/registered/cp/test_deepseek_v32_cp_single_node.py b/test/registered/cp/test_deepseek_v32_cp_single_node.py index 55fa21907..dcfcb4d78 100644 --- a/test/registered/cp/test_deepseek_v32_cp_single_node.py +++ b/test/registered/cp/test_deepseek_v32_cp_single_node.py @@ -13,7 +13,11 @@ from sglang.test.test_utils import ( write_github_step_summary, ) -register_cuda_ci(est_time=616, suite="stage-c-test-deepep-8-gpu-h200") +register_cuda_ci( + est_time=616, + suite="stage-c-test-deepep-8-gpu-h200", + disabled="Disabled due to #24268. Should be fixed soon.", +) DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2" diff --git a/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py b/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py index 1cf38cdf7..14793c17a 100644 --- a/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py +++ b/test/registered/quant/test_deepseek_v32_fp4_mtp_4gpu.py @@ -15,7 +15,11 @@ from sglang.test.test_utils import ( write_github_step_summary, ) -register_cuda_ci(est_time=1060, suite="stage-c-test-4-gpu-b200") +register_cuda_ci( + est_time=1060, + suite="stage-c-test-4-gpu-b200", + disabled="Disabled due to #24268. Should be fixed soon.", +) FULL_DEEPSEEK_V3_FP4_MODEL_PATH = "nvidia/DeepSeek-V3.2-NVFP4" SERVER_LAUNCH_TIMEOUT = 1200