From 7ffc265de386c7ea0589a6b1d01ef158eb3e90b4 Mon Sep 17 00:00:00 2001 From: Zilin Zhu Date: Wed, 12 Aug 2026 05:24:41 +0000 Subject: [PATCH] [docker] Upgrade to sglang v0.5.17 --- docker/Dockerfile | 17 +- docker/README.md | 25 +- docker/justfile | 22 +- docker/patch/latest/megatron.patch | 23 +- .../patch/latest/sglang-deterministic.patch | 1536 ++++++++--------- docker/patch/latest/sglang-pull_weights.patch | 40 +- .../patch/latest/sglang-release_hicache.patch | 324 ++-- docker/patch/latest/sglang-top_p.patch | 870 +++++----- docker/patch/latest/sglang.patch | 711 +++----- .../alignment/deepgemm_moe_forward.py | 27 +- slime/backends/sglang_utils/compat.py | 15 + slime/backends/sglang_utils/sglang_engine.py | 7 +- slime/utils/routing_replay.py | 84 + slime_plugins/models/glm5/glm5.py | 17 +- slime_plugins/models/glm5/ops/indexer.py | 46 +- tests/test_deepgemm_moe_forward.py | 150 +- tests/test_glm52_6layer_deterministic_e2e.py | 10 +- tests/utils/test_sglang_arguments.py | 13 +- 18 files changed, 2020 insertions(+), 1917 deletions(-) create mode 100644 slime/backends/sglang_utils/compat.py diff --git a/docker/Dockerfile b/docker/Dockerfile index 142d9df866..e34b367cda 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,12 +1,12 @@ -ARG SGLANG_IMAGE_TAG=v0.5.15.post1-cu129 +ARG SGLANG_IMAGE_TAG=v0.5.17-cu129 FROM slimerl/sglang:${SGLANG_IMAGE_TAG} AS sglang # ======================================== Arguments ============================================= ARG PATCH_VERSION=latest ARG MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087 -ARG DEEPGEMM_COMMIT=b38a77cd193cf38f670caae192310521d24343be -ARG DEEPEP_COMMIT=6845ffd9d59126ec0030c13e0e155935a61e5b5a +ARG DEEPGEMM_COMMIT=7ad54cbd80ebe24ad36cb5cd3729f37beed420bb +ARG DEEPEP_COMMIT=c5c2b0cfa767b1afeb4d237c4ed7e5ff062a550d ARG DEEPEP_CUDA_ARCH_LIST= ARG FLASH_QLA_COMMIT=821fd9d37ede18fdc2a4e707fefe3770bfc32e58 ARG TRANSFORMER_ENGINE_COMMIT=c9877beb87ad7e711e1869dd0b5062167ede447a @@ -39,7 +39,7 @@ RUN if [ "$(python -c 'import torch; print(torch.version.cuda.split(".")[0])')" fi RUN pip install flash-linear-attention==0.4.2 -# FlashQLA currently requires TileLang 0.1.9, while SGLang 0.5.15.post1 on +# FlashQLA currently requires TileLang 0.1.9, while SGLang 0.5.17 on # CUDA 13 requires 0.1.11. Keep the existing CUDA 12 installation and use # the default FLA backend on CUDA 13 instead of downgrading SGLang's runtime. RUN if [ "$(python -c 'import torch; print(torch.version.cuda.split(".")[0])')" = "12" ]; then \ @@ -69,9 +69,9 @@ RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \ cd Megatron-LM && git checkout ${MEGATRON_COMMIT} && \ pip install -e . -# zhuzilin fork builds, grouped together right after Megatron-LM: -# torch_memory_saver, plus the GLM-5 train/rollout alignment kernels -# (DeepGEMM batch-invariant selection and the DeepEP low-latency fork). +# Alignment builds, grouped together right after Megatron-LM: +# torch_memory_saver, plus the GLM-5 train/rollout alignment kernels rebased +# onto the DeepGEMM and DeepEP revisions shipped with SGLang 0.5.17. RUN TMS_CUDA_MAJOR="${TMS_CUDA_MAJOR:-$(python -c 'import torch; print(torch.version.cuda.split(".")[0])')}" && \ export TMS_CUDA_MAJOR && \ pip install git+https://github.com/zhuzilin/torch_memory_saver.git@${TMS_COMMIT} --no-cache-dir --force-reinstall @@ -84,7 +84,8 @@ RUN git clone https://github.com/zhuzilin/DeepGEMM.git --recursive && \ pip install --force-reinstall --no-deps dist/sgl_deep_gemm-*.whl && \ cd /root/ && rm -rf DeepGEMM -# DeepEP low-latency alignment fork. CUDA 13 relocated the CCCL/libcu++ headers +# DeepEP fork with low-latency FP8 alignment for SGLang 0.5.15.post1 and 0.5.17. +# CUDA 13 relocated the CCCL/libcu++ headers # () under include/cccl, which nvshmem_tensor.h includes, so add that # directory to the compile search path (nvcc + host compiler) when it exists. RUN git clone https://github.com/zhuzilin/DeepEP.git /root/DeepEP && \ diff --git a/docker/README.md b/docker/README.md index efb49fbefd..08f9c9db81 100644 --- a/docker/README.md +++ b/docker/README.md @@ -5,9 +5,10 @@ We will publish 2 kinds of docker images: 2. latest version, which aligns to `lmsysorg/sglang:latest`. current stable version is: -- sglang v0.5.15.post1 (0b3bb0cbe31873994c9f989fddfe2f87ca839fdd), megatron dev 1dcf0dafa884ad52ffb243625717a3471643e087 +- sglang v0.5.17 (29481685462732237d80d86076d6563e1f658102), megatron dev 1dcf0dafa884ad52ffb243625717a3471643e087 history versions: +- sglang v0.5.15.post1 (0b3bb0cbe31873994c9f989fddfe2f87ca839fdd), megatron dev 1dcf0dafa884ad52ffb243625717a3471643e087 - sglang v0.5.13 (28b095c01005d4a3a2a5b637b7d028b07fba31b2), megatron dev 1dcf0dafa884ad52ffb243625717a3471643e087 - sglang v0.5.12.post1 (5a15cde858ea09b77116212a39356f2fc51b8584), megatron dev 1dcf0dafa884ad52ffb243625717a3471643e087 - sglang v0.5.10.post1 (7c35342c10e201899e22fe2972d40e60da19ff3e), megatron dev 1dcf0dafa884ad52ffb243625717a3471643e087 @@ -20,13 +21,14 @@ history versions: The commands to build and publish: ```bash -just release-primary # CUDA 12 (cu129 base): publishes `latest`, `latest-cu129`, `-cu129` -just release-cu13 # CUDA 13 (cu130 base, Blackwell): publishes `latest-cu130`, `-cu130` +just release-primary # publishes both CUDA 12 and CUDA 13 variants below +just release-cu12 # cu129 base: `latest`, `latest-cu129`, `-cu129` +just release-cu13 # cu130 base: `latest-cu130`, `-cu130` ``` `slimerl/slime:latest` tracks the CUDA 12 build. The tag suffixes (`-cu129` / `-cu130`) match the SGLang base image. `docker/Dockerfile` branches on the base -image's CUDA version; it defaults to the cu130 SGLang base, and the cu129 base +image's CUDA version; it defaults to the cu129 SGLang base, and the cu130 base is selected via build args (see `docker/justfile`). To build a single image directly without publishing: @@ -34,13 +36,13 @@ To build a single image directly without publishing: ```bash # CUDA 12 docker build -f docker/Dockerfile . \ - --build-arg SGLANG_IMAGE_TAG=v0.5.15.post1-cu129 \ + --build-arg SGLANG_IMAGE_TAG=v0.5.17-cu129 \ -t slimerl/slime:latest-cu129 # CUDA 13 (Blackwell) docker build -f docker/Dockerfile . \ --build-arg DEEPEP_CUDA_ARCH_LIST='10.0 10.3' \ - --build-arg SGLANG_IMAGE_TAG=v0.5.15.post1-cu130 \ + --build-arg SGLANG_IMAGE_TAG=v0.5.17-cu130 \ -t slimerl/slime:latest-cu130 ``` @@ -48,11 +50,12 @@ The following components are pinned and rebuilt in the image: - Megatron-LM `1dcf0dafa884ad52ffb243625717a3471643e087`, plus `docker/patch//megatron.patch`. -- DeepGEMM `b38a77cd193cf38f670caae192310521d24343be` from the - `zhuzilin/DeepGEMM` batch-invariant branch, rebuilt as an SGLang-compatible wheel. -- DeepEP `6845ffd9d59126ec0030c13e0e155935a61e5b5a` from the - `zhuzilin/DeepEP` `align_fp8_quantization` branch (GLM-5 low-latency alignment). +- DeepGEMM `7ad54cbd80ebe24ad36cb5cd3729f37beed420bb` from the + `zhuzilin/DeepGEMM` fork. It adds batch-invariant FP8 kernels directly to + the 0.5.17 `v0.1.5` source baseline. +- DeepEP `c5c2b0cfa767b1afeb4d237c4ed7e5ff062a550d` from the + `zhuzilin/DeepEP` fork. It merges the 0.5.17 baseline and provides aligned + low-latency FP8 modes for both SGLang 0.5.15.post1 and 0.5.17. For a non-default GPU architecture list, pass `--build-arg DEEPEP_CUDA_ARCH_LIST=''`. - diff --git a/docker/justfile b/docker/justfile index 508758922d..8481528af9 100644 --- a/docker/justfile +++ b/docker/justfile @@ -1,17 +1,13 @@ -# Primary release stays on CUDA 12 (the cu129 SGLang base). Publishes the -# unqualified `latest`, plus the rolling `latest-cu129` and the versioned -# `-cu129` tags. The tag suffix matches the SGLang base image. -release-primary: - ARG_TAG_POSTFIX="-cu129" ARG_TAG_LATEST=1 ARG_BUILD_EXTRA_ARGS='--build-arg SGLANG_IMAGE_TAG=v0.5.15.post1-cu129' just _release-raw - -# B200/B300 (Blackwell, x86) CUDA 13 build on the cu130 SGLang base. Publishes -# `latest-cu130` and `-cu130` (tag suffix matches the SGLang base -# image). DeepEP must be compiled for the Blackwell arch list; the Dockerfile -# default (9.0) targets Hopper only. The semicolon separator is normalized to -# spaces by PyTorch's TORCH_CUDA_ARCH_LIST parser and survives the unquoted -# ARG_BUILD_EXTRA_ARGS expansion as one token. +# Publish both supported CUDA variants; CUDA 12 owns the unqualified `latest` tag. +release-primary: release-cu12 release-cu13 + +# Publish CUDA 12 tags: `latest`, `latest-cu129`, and `-cu129`. +release-cu12: + ARG_TAG_POSTFIX="-cu129" ARG_TAG_LATEST=1 ARG_BUILD_EXTRA_ARGS='--build-arg SGLANG_IMAGE_TAG=v0.5.17-cu129' just _release-raw + +# Publish B200/B300 CUDA 13 tags; compile DeepEP for Blackwell via the build arg. release-cu13: - ARG_TAG_POSTFIX="-cu130" ARG_BUILD_EXTRA_ARGS="--build-arg DEEPEP_CUDA_ARCH_LIST=10.0;10.3 --build-arg SGLANG_IMAGE_TAG=v0.5.15.post1-cu130" just _release-raw + ARG_TAG_POSTFIX="-cu130" ARG_BUILD_EXTRA_ARGS="--build-arg DEEPEP_CUDA_ARCH_LIST=10.0;10.3 --build-arg SGLANG_IMAGE_TAG=v0.5.17-cu130" just _release-raw _release-raw: #!/bin/bash diff --git a/docker/patch/latest/megatron.patch b/docker/patch/latest/megatron.patch index 4b4c74d28e..5045d96146 100644 --- a/docker/patch/latest/megatron.patch +++ b/docker/patch/latest/megatron.patch @@ -735,14 +735,31 @@ diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transfor index 75825cd37..445b3fb84 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py -@@ -711,6 +711,9 @@ def topk_routing_with_score_function( +@@ -711,6 +711,25 @@ def topk_routing_with_score_function( scores, topk, num_groups, group_topk, _compute_topk ) -+ from slime.utils.routing_replay import get_routing_replay_compute_topk ++ from slime.utils.routing_replay import ( ++ get_routing_replay_compute_topk, ++ maybe_sglang_fused_sigmoid_topk, ++ ) + compute_topk = get_routing_replay_compute_topk(compute_topk) + - if score_function == "softmax": ++ aligned_topk = maybe_sglang_fused_sigmoid_topk( ++ logits, ++ topk, ++ use_pre_softmax=use_pre_softmax, ++ num_groups=num_groups, ++ group_topk=group_topk, ++ scaling_factor=scaling_factor, ++ score_function=score_function, ++ expert_bias=expert_bias, ++ replay_active=router_replay is not None, ++ ) ++ if aligned_topk is not None: ++ probs, top_indices = aligned_topk +- if score_function == "softmax": ++ elif score_function == "softmax": if use_pre_softmax: scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py diff --git a/docker/patch/latest/sglang-deterministic.patch b/docker/patch/latest/sglang-deterministic.patch index ad31f7a454..a6bc4e0340 100644 --- a/docker/patch/latest/sglang-deterministic.patch +++ b/docker/patch/latest/sglang-deterministic.patch @@ -1,7 +1,204 @@ -diff --git a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh b/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh +diff --git a/python/sglang/kernels/aot/csrc/gemm/per_token_group_quant_8bit_v2.cu b/python/sglang/kernels/aot/csrc/gemm/per_token_group_quant_8bit_v2.cu +index 4fbaf08efc..5bad403d0f 100644 +--- a/python/sglang/kernels/aot/csrc/gemm/per_token_group_quant_8bit_v2.cu ++++ b/python/sglang/kernels/aot/csrc/gemm/per_token_group_quant_8bit_v2.cu +@@ -35,7 +35,9 @@ __device__ __forceinline__ float silu(const float& val) { + float t = __tanhf(half); + return half * (1.0f + t); + #else +- return val / (1.0f + __expf(-val)); ++ // Match sgl_kernel.silu_and_mul on SM90. That kernel uses expf and ++ // exposes only the final SiLU*up BF16 rounding boundary. ++ return val / (1.0f + expf(-val)); + #endif + } + +@@ -266,7 +268,8 @@ __global__ void per_token_group_quant_8bit_kernel( + // TODO can this be removed? + const int scale_expert_stride, + const int scale_hidden_stride, +- const int num_tokens_per_expert) { ++ const int num_tokens_per_expert, ++ const float max_8bit) { + using dst_dtype_info = DtypeInfo; + using scale_element_t = std::conditional_t; + static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0); +@@ -347,8 +350,11 @@ __global__ void per_token_group_quant_8bit_kernel( + for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) { + float val; + if constexpr (FUSE_SILU_AND_MUL) { +- // TODO maybe vectorize +- T val_lowprec = static_cast(silu(static_cast(input_primary_vec[j]))) * input_secondary_vec[j]; ++ // Keep the same single visible low-precision boundary as the ++ // contiguous activation path. Rounding SiLU before multiplying ++ // changes the second MoE GEMM's FP8 inputs. ++ T val_lowprec = static_cast( ++ silu(static_cast(input_primary_vec[j])) * static_cast(input_secondary_vec[j])); + val = static_cast(val_lowprec); + input_primary_vec[j] = val_lowprec; + } else { +@@ -362,8 +368,15 @@ __global__ void per_token_group_quant_8bit_kernel( + local_absmax = GroupReduceMax(local_absmax, lane_id); + + float y_scale, y_scale_inv; +- calculate_fp8_scales(local_absmax, y_scale, y_scale_inv); +- float2 y_scale_repeated = {y_scale, y_scale}; ++ if constexpr (SCALE_UE8M0) { ++ calculate_fp8_scales(local_absmax, y_scale, y_scale_inv); ++ } else { ++ // Preserve the scale and division sequence of the legacy unmasked ++ // quantizer. Reciprocal reuse under -use_fast_math moves a few ++ // values across an FP8 rounding boundary. ++ y_scale_inv = __fdiv_rn(local_absmax, max_8bit); ++ y_scale = 1.0f / y_scale_inv; ++ } + + if (lane_id == 0) { + *scale_output = extract_required_scale_format(y_scale_inv); +@@ -373,19 +386,33 @@ __global__ void per_token_group_quant_8bit_kernel( + static_assert(sizeof(output_buf) == INPUT_PRIMARY_VEC_SIZE * sizeof(DST_DTYPE)); + + if constexpr (std::is_same_v) { +- const auto output_buf_ptr = reinterpret_cast<__nv_fp8x2_storage_t*>(&output_buf); +- static_assert(sizeof(output_buf) == INPUT_PRIMARY_VEC_SIZE / 2 * sizeof(__nv_fp8x2_storage_t)); +- static_assert(INPUT_PRIMARY_VEC_SIZE % 2 == 0); ++ if constexpr (SCALE_UE8M0) { ++ const auto output_buf_ptr = reinterpret_cast<__nv_fp8x2_storage_t*>(&output_buf); ++ static_assert(sizeof(output_buf) == INPUT_PRIMARY_VEC_SIZE / 2 * sizeof(__nv_fp8x2_storage_t)); ++ static_assert(INPUT_PRIMARY_VEC_SIZE % 2 == 0); + + #pragma unroll +- for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; j += 2) { +- float2 inputx2 = {static_cast(input_primary_vec[j]), static_cast(input_primary_vec[j + 1])}; +- float2 outputx2 = fmul2_rn(inputx2, y_scale_repeated); ++ for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; j += 2) { ++ float2 inputx2 = { ++ static_cast(input_primary_vec[j]), static_cast(input_primary_vec[j + 1])}; ++ float2 outputx2 = fmul2_rn(inputx2, {y_scale, y_scale}); + +- outputx2.x = fminf(fmaxf(outputx2.x, dst_dtype_info::MIN), dst_dtype_info::MAX); +- outputx2.y = fminf(fmaxf(outputx2.y, dst_dtype_info::MIN), dst_dtype_info::MAX); ++ outputx2.x = fminf(fmaxf(outputx2.x, dst_dtype_info::MIN), dst_dtype_info::MAX); ++ outputx2.y = fminf(fmaxf(outputx2.y, dst_dtype_info::MIN), dst_dtype_info::MAX); + +- output_buf_ptr[j / 2] = __nv_cvt_float2_to_fp8x2(outputx2, __NV_SATFINITE, __NV_E4M3); ++ output_buf_ptr[j / 2] = __nv_cvt_float2_to_fp8x2(outputx2, __NV_SATFINITE, __NV_E4M3); ++ } ++ } else { ++ // Match the legacy SM90 quantizer's scalar conversion exactly; ++ // packed conversion makes a different tie choice for rare values. ++ const auto output_buf_ptr = reinterpret_cast(&output_buf); ++#pragma unroll ++ for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) { ++ float val = static_cast(input_primary_vec[j]); ++ float q_val = ++ fminf(fmaxf(__fdiv_rn(val, y_scale_inv), dst_dtype_info::MIN), dst_dtype_info::MAX); ++ output_buf_ptr[j] = DST_DTYPE(q_val); ++ } + } + } else { + const auto output_buf_ptr = reinterpret_cast(&output_buf); +@@ -474,7 +501,8 @@ void sgl_per_token_group_quant_8bit_v2( + hidden_dim_num_groups, \ + scale_expert_stride, \ + scale_hidden_stride, \ +- num_tokens_per_expert); \ ++ num_tokens_per_expert, \ ++ static_cast(max_8bit)); \ + } while (0) + + #define LAUNCH_KERNEL(GROUP_SIZE, T, DST_DTYPE) \ +@@ -497,13 +525,51 @@ void sgl_per_token_group_quant_8bit_v2( + NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, true); \ + } \ + } else { \ +- LAUNCH_KERNEL_INNER(NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true); \ ++ if (masked_layout) { \ ++ LAUNCH_KERNEL_INNER( \ ++ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, false); \ ++ } else { \ ++ LAUNCH_KERNEL_INNER( \ ++ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, false); \ ++ } \ + } \ + } else { \ +- LAUNCH_KERNEL_INNER(NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true); \ ++ if (fuse_silu_and_mul) { \ ++ if (masked_layout) { \ ++ LAUNCH_KERNEL_INNER( \ ++ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, true); \ ++ } else { \ ++ LAUNCH_KERNEL_INNER( \ ++ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, true); \ ++ } \ ++ } else { \ ++ if (masked_layout) { \ ++ LAUNCH_KERNEL_INNER( \ ++ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, false); \ ++ } else { \ ++ LAUNCH_KERNEL_INNER( \ ++ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, false); \ ++ } \ ++ } \ + } \ + } else { \ +- LAUNCH_KERNEL_INNER(NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false); \ ++ if (fuse_silu_and_mul) { \ ++ if (masked_layout) { \ ++ LAUNCH_KERNEL_INNER( \ ++ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, true); \ ++ } else { \ ++ LAUNCH_KERNEL_INNER( \ ++ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, true); \ ++ } \ ++ } else { \ ++ if (masked_layout) { \ ++ LAUNCH_KERNEL_INNER( \ ++ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, false); \ ++ } else { \ ++ LAUNCH_KERNEL_INNER( \ ++ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, false); \ ++ } \ ++ } \ + } \ + } while (0) + +diff --git a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh +index 4b4ac822fe..2c4ddcb5d2 100644 +--- a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh ++++ b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh +@@ -26,7 +26,9 @@ SGL_DEVICE float silu(const float val) { + const float half = 0.5f * val; + return half * (1.0f + __tanhf(half)); + #else +- return val * __frcp_rn(1.0f + __expf(-val)); ++ // Match sgl_kernel.silu_and_mul on SM90. That kernel uses expf and exposes ++ // only the final SiLU*up BF16 rounding boundary. ++ return val / (1.0f + expf(-val)); + #endif + } + +@@ -268,8 +270,14 @@ struct QuantTrait { + #pragma unroll + for (uint32_t i = 0; i < kVecSize / 2; ++i) { + const auto gate = cast(in[i]); +- const auto act = cast(float2{details::silu(gate.x), details::silu(gate.y)}); +- in[i] = __hmul2(act, up[i]); ++ const auto up_fp32 = cast(up[i]); ++ // Keep the same single visible low-precision boundary as the ++ // contiguous activation path. Rounding SiLU before multiplying changes ++ // the second MoE GEMM's FP8 inputs. ++ in[i] = cast(float2{ ++ details::silu(gate.x) * up_fp32.x, ++ details::silu(gate.y) * up_fp32.y, ++ }); + } + } + +diff --git a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant_8bit_v2.cuh b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant_8bit_v2.cuh index 022567398a..8442ba0710 100644 ---- a/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh -+++ b/python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit_v2.cuh +--- a/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant_8bit_v2.cuh ++++ b/python/sglang/kernels/jit/csrc/gemm/per_token_group_quant_8bit_v2.cuh @@ -45,7 +45,9 @@ SGL_DEVICE float silu(const float& val) { float t = __tanhf(half); return half * (1.0f + t); @@ -12,7 +209,7 @@ index 022567398a..8442ba0710 100644 + return val / (1.0f + expf(-val)); #endif } - + @@ -297,7 +299,11 @@ __global__ void per_token_group_quant_8bit_v2_kernel( for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) { float val; @@ -28,7 +225,7 @@ index 022567398a..8442ba0710 100644 } else { @@ -309,12 +315,16 @@ __global__ void per_token_group_quant_8bit_v2_kernel( local_absmax = GroupReduceMax(local_absmax); - + float y_scale, y_scale_inv; - // When SCALE_UE8M0, always quantize with the rounded (power-of-2) scale - // — not with the exact scale followed by post-hoc rounding. @@ -50,7 +247,7 @@ index 022567398a..8442ba0710 100644 *scale_output = extract_required_scale_format < SCALE_UE8M0 && IS_COLUMN_MAJOR > (y_scale_inv); } @@ -322,14 +332,28 @@ __global__ void per_token_group_quant_8bit_v2_kernel( - + int4 output_buf; if constexpr (std::is_same_v) { - const auto output_buf_ptr = reinterpret_cast<__nv_fp8x2_storage_t*>(&output_buf); @@ -146,11 +343,11 @@ index 022567398a..8442ba0710 100644 } } } -diff --git a/python/sglang/jit_kernel/glm5_router_gemm.py b/python/sglang/jit_kernel/glm5_router_gemm.py +diff --git a/python/sglang/kernels/jit/glm5_router_gemm.py b/python/sglang/kernels/jit/glm5_router_gemm.py new file mode 100644 -index 0000000000..29076ade64 +index 0000000000..3c8307484c --- /dev/null -+++ b/python/sglang/jit_kernel/glm5_router_gemm.py ++++ b/python/sglang/kernels/jit/glm5_router_gemm.py @@ -0,0 +1,36 @@ +from __future__ import annotations + @@ -158,11 +355,11 @@ index 0000000000..29076ade64 + +import torch + -+from sglang.jit_kernel.utils import cache_once, load_jit ++from sglang.kernels.jit.utils import cache_once, load_jit +from sglang.kernel_api_logging import debug_kernel_api + +if TYPE_CHECKING: -+ from tvm_ffi.module import Module ++ from tvm_ffi import Module + + +@cache_once @@ -188,50 +385,13 @@ index 0000000000..29076ade64 + output, hidden_states, router_weight + ) + return output -diff --git a/python/sglang/jit_kernel/tests/test_glm5_router_gemm.py b/python/sglang/jit_kernel/tests/test_glm5_router_gemm.py -new file mode 100644 -index 0000000000..2455659739 ---- /dev/null -+++ b/python/sglang/jit_kernel/tests/test_glm5_router_gemm.py -@@ -0,0 +1,31 @@ -+import pytest -+import torch -+ -+from sglang.jit_kernel.glm5_router_gemm import glm5_router_gemm -+ -+ -+@pytest.mark.parametrize("num_tokens", [1, 2, 4, 8, 16]) -+def test_glm5_router_gemm_accuracy(num_tokens: int): -+ torch.manual_seed(42) -+ hidden_states = torch.randn((num_tokens, 6144), dtype=torch.bfloat16, device="cuda") -+ router_weight = torch.randn((256, 6144), dtype=torch.bfloat16, device="cuda") -+ actual = glm5_router_gemm(hidden_states, router_weight) -+ expected = torch.mm(hidden_states, router_weight.t(), out_dtype=torch.float32) -+ torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-3) -+ -+ -+def test_glm5_router_gemm_batch_invariance(): -+ torch.manual_seed(42) -+ token = torch.randn((1, 6144), dtype=torch.bfloat16, device="cuda") -+ router_weight = torch.randn((256, 6144), dtype=torch.bfloat16, device="cuda") -+ expected = glm5_router_gemm(token, router_weight) -+ -+ for num_tokens, token_index in [(2, 1), (8, 5), (16, 15)]: -+ hidden_states = torch.randn( -+ (num_tokens, 6144), dtype=torch.bfloat16, device="cuda" -+ ) -+ hidden_states[token_index].copy_(token[0]) -+ actual = glm5_router_gemm(hidden_states, router_weight)[ -+ token_index : token_index + 1 -+ ] -+ torch.testing.assert_close(actual, expected, rtol=0, atol=0) -diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py -index bdff4bb67a..ba11284fc9 100644 ---- a/python/sglang/jit_kernel/utils.py -+++ b/python/sglang/jit_kernel/utils.py -@@ -131,6 +131,32 @@ def _resolve_kernel_path() -> pathlib.Path: - - +diff --git a/python/sglang/kernels/jit/utils/compile.py b/python/sglang/kernels/jit/utils/compile.py +index 284ae33184..49ac430ec4 100644 +--- a/python/sglang/kernels/jit/utils/compile.py ++++ b/python/sglang/kernels/jit/utils/compile.py +@@ -97,6 +97,32 @@ def _resolve_kernel_path() -> pathlib.Path: + + KERNEL_PATH = _resolve_kernel_path() + + @@ -262,37 +422,47 @@ index bdff4bb67a..ba11284fc9 100644 DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] DEFAULT_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_LDFLAGS = [] -@@ -261,8 +287,8 @@ def load_jit( - extra_ldflags = extra_ldflags or [] - extra_include_paths = extra_include_paths or [] - -- cpp_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cpp_files] -- cuda_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cuda_files] -+ cpp_files = [_resolve_source_file(f) for f in cpp_files] -+ cuda_files = [_resolve_source_file(f) for f in cuda_files] - - for dep in set(extra_dependencies or []): - if dep not in _REGISTERED_DEPENDENCIES: +@@ -250,10 +276,10 @@ def load_jit( + if flag not in ("--use_fast_math", "-use_fast_math") + ] + +- cpp_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cpp_files] + [ ++ cpp_files = [_resolve_source_file(f) for f in cpp_files] + [ + str(pathlib.Path(f).resolve()) for f in external_cpp_files + ] +- cuda_files = [str((KERNEL_PATH / "csrc" / f).resolve()) for f in cuda_files] + [ ++ cuda_files = [_resolve_source_file(f) for f in cuda_files] + [ + str(pathlib.Path(f).resolve()) for f in external_cuda_files + ] + diff --git a/python/sglang/srt/batch_invariant_ops/__init__.py b/python/sglang/srt/batch_invariant_ops/__init__.py -index 3caa60f0b4..4a387c7884 100644 +index 3caa60f0b4..da16a8d3f4 100644 --- a/python/sglang/srt/batch_invariant_ops/__init__.py +++ b/python/sglang/srt/batch_invariant_ops/__init__.py -@@ -10,6 +10,7 @@ from .batch_invariant_ops import ( +@@ -4,12 +4,14 @@ from .batch_invariant_ops import ( + AttentionBlockSize, + disable_batch_invariant_mode, + enable_batch_invariant_mode, ++ fused_residual_rms_norm_batch_invariant, + get_batch_invariant_attention_block_size, + is_batch_invariant_mode_enabled, + log_softmax, matmul_persistent, mean_dim, rms_norm_batch_invariant, + router_gemm_batch_invariant, set_batch_invariant_mode, ) - -@@ -26,4 +27,5 @@ __all__ = [ + +@@ -26,4 +28,6 @@ __all__ = [ "get_batch_invariant_attention_block_size", "AttentionBlockSize", "rms_norm_batch_invariant", ++ "fused_residual_rms_norm_batch_invariant", + "router_gemm_batch_invariant", ] diff --git a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py -index 51a911eb18..bdee93be6f 100644 +index 51a911eb18..7f780c2037 100644 --- a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py +++ b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py @@ -214,7 +214,7 @@ def _matmul_persistent_triton( @@ -304,10 +474,99 @@ index 51a911eb18..bdee93be6f 100644 "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8, -@@ -315,6 +315,68 @@ def matmul_persistent( +@@ -315,6 +315,157 @@ def matmul_persistent( return _matmul_persistent_triton(a=a, b=b, bias=bias) - - + + ++@triton.jit ++def _fused_residual_rms_norm_prepare_kernel( ++ x_ptr, ++ residual_ptr, ++ summed_ptr, ++ square_ptr, ++ residual_out_ptr, ++ n_elements: tl.constexpr, ++ BLOCK_SIZE: tl.constexpr, ++): ++ offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) ++ mask = offsets < n_elements ++ x = tl.load(x_ptr + offsets, mask=mask).to(tl.float32) ++ residual = tl.load(residual_ptr + offsets, mask=mask).to(tl.float32) ++ summed = x + residual ++ tl.store(summed_ptr + offsets, summed, mask=mask) ++ tl.store(square_ptr + offsets, summed * summed, mask=mask) ++ tl.store(residual_out_ptr + offsets, summed, mask=mask) ++ ++ ++@triton.jit ++def _fused_residual_rms_norm_finish_kernel( ++ summed_ptr, ++ variance_ptr, ++ weight_ptr, ++ output_ptr, ++ hidden_size: tl.constexpr, ++ n_elements: tl.constexpr, ++ eps: tl.constexpr, ++ BLOCK_SIZE: tl.constexpr, ++): ++ offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) ++ mask = offsets < n_elements ++ summed = tl.load(summed_ptr + offsets, mask=mask) ++ variance = tl.load(variance_ptr + offsets // hidden_size, mask=mask) ++ weight = tl.load(weight_ptr + offsets % hidden_size, mask=mask).to(tl.float32) ++ inverse_rms = tl.rsqrt(variance + eps) ++ output = summed * inverse_rms * weight ++ tl.store(output_ptr + offsets, output, mask=mask) ++ ++ ++def fused_residual_rms_norm_batch_invariant( ++ x: torch.Tensor, ++ residual: torch.Tensor, ++ weight: torch.Tensor, ++ eps: float, ++) -> tuple[torch.Tensor, torch.Tensor]: ++ """Match the BF16 residual/RMSNorm boundary used by aligned training.""" ++ if x.shape != residual.shape: ++ raise ValueError( ++ "fused residual RMSNorm expects matching inputs, got " ++ f"{x.shape=} and {residual.shape=}" ++ ) ++ if x.shape[-1] != weight.numel(): ++ raise ValueError( ++ "fused residual RMSNorm weight does not match hidden size: " ++ f"{x.shape[-1]=}, {weight.numel()=}" ++ ) ++ ++ summed = torch.empty_like(x, dtype=torch.float32) ++ square = torch.empty_like(x, dtype=torch.float32) ++ residual_out = torch.empty_like(x) ++ n_elements = x.numel() ++ block_size = 256 ++ grid = (triton.cdiv(n_elements, block_size),) ++ _fused_residual_rms_norm_prepare_kernel[grid]( ++ x, ++ residual, ++ summed, ++ square, ++ residual_out, ++ n_elements=n_elements, ++ BLOCK_SIZE=block_size, ++ ) ++ variance = mean_dim(square, -1, keepdim=True) ++ output = torch.empty_like(x) ++ _fused_residual_rms_norm_finish_kernel[grid]( ++ summed, ++ variance, ++ weight, ++ output, ++ hidden_size=x.shape[-1], ++ n_elements=n_elements, ++ eps=eps, ++ BLOCK_SIZE=block_size, ++ ) ++ return output, residual_out ++ ++ +def router_gemm_batch_invariant( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, @@ -346,11 +605,11 @@ index 51a911eb18..bdee93be6f 100644 + ) + + if (router_weight.shape[0], router_weight.shape[1]) == (256, 6144): -+ from sglang.jit_kernel.glm5_router_gemm import ( ++ from sglang.kernels.jit.glm5_router_gemm import ( + glm5_router_gemm as fast_router_gemm, + ) + else: -+ from sgl_kernel import dsv3_router_gemm ++ from sglang.kernels.ops.gemm.dsv3_router_gemm import dsv3_router_gemm + + def fast_router_gemm(input_chunk, weight): + return dsv3_router_gemm(input_chunk, weight, out_dtype=torch.float32) @@ -373,6 +632,44 @@ index 51a911eb18..bdee93be6f 100644 @triton.jit def _log_softmax_kernel( input_ptr, +diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py +index 67661147f3..894c15fa43 100644 +--- a/python/sglang/srt/layers/layernorm.py ++++ b/python/sglang/srt/layers/layernorm.py +@@ -22,6 +22,7 @@ import torch.nn as nn + import torch.nn.functional as F + + from sglang.srt.batch_invariant_ops import ( ++ fused_residual_rms_norm_batch_invariant, + is_batch_invariant_mode_enabled, + rms_norm_batch_invariant, + ) +@@ -649,6 +650,25 @@ class RMSNorm(MultiPlatformOp): + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + if not x.is_contiguous(): + x = x.contiguous() ++ if ( ++ is_batch_invariant_mode_enabled() ++ and residual is not None ++ and post_residual_addition is None ++ and not self.fp32_residual ++ and self.variance_size_override is None ++ and not self.cast_x_before_out_mul ++ and self.override_orig_dtype is None ++ and x.is_cuda ++ and x.dtype == torch.bfloat16 ++ and residual.dtype == torch.bfloat16 ++ and self.weight.dtype == torch.bfloat16 ++ ): ++ return fused_residual_rms_norm_batch_invariant( ++ x, ++ residual.contiguous(), ++ self.weight, ++ self.variance_epsilon, ++ ) + orig_dtype = self.override_orig_dtype or x.dtype + x = x.to(torch.float32) + if residual is not None: diff --git a/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py b/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py index 66a2660755..d38f312c09 100644 --- a/python/sglang/srt/debug_utils/tensor_dump_forward_hook.py @@ -453,7 +750,7 @@ index 66a2660755..d38f312c09 100644 - torch.save(self._current_tensors, str(tensor_file_for_pass)) self._current_tensors = {} self._forward_pass_id += 1 - + @@ -100,21 +132,29 @@ class TensorDumper: top_level_model = True else: @@ -491,7 +788,7 @@ index 66a2660755..d38f312c09 100644 # self_attn.qkv_proj, self_attn.attn & self_attn.o_proj. # Therefore, we do not need to add output hooks for self_attn, @@ -126,14 +166,14 @@ class TensorDumper: - + def _dump_hook(self, tensor_name, do_dump): def inner_dump_hook(module, input, output): + if output is not None and not (do_dump and self._layer_outputs_only): @@ -504,38 +801,38 @@ index 66a2660755..d38f312c09 100644 self.dump_current_tensors() - if output is not None: - self.add_tensor(tensor_name, output) - + return inner_dump_hook - + diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py -index 790acce902..294cedc06c 100644 +index 81cbfc18dd..68383f23d3 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py -@@ -601,6 +601,7 @@ class Envs: +@@ -742,6 +742,7 @@ class Envs: SGLANG_DG_USE_NVRTC = EnvBool(False) SGLANG_USE_DEEPGEMM_BMM = EnvBool(False) SGLANG_DEEPGEMM_SANITY_CHECK = EnvBool(False) + SGLANG_DEEPGEMM_BATCH_INVARIANT = EnvBool(False) SGLANG_DEEPGEMM_PDL = EnvBool(True) SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP = EnvBool(False) - -@@ -611,6 +612,7 @@ class Envs: + +@@ -752,6 +753,7 @@ class Envs: # DeepEP SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False) # This argument is deprecated SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128) + SGLANG_DEEPEP_LL_PREFILL_STAGING = EnvBool(False) SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32) SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False) - # Force dynamic DeepEP Waterfill with runtime EP all-reduce instead of the + # Force dynamic Waterfill with runtime EP all-reduce instead of the default diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py -index 5e5f7aa2a5..f074a044bc 100644 +index a74aba3631..1b271c6e23 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py -@@ -505,6 +505,19 @@ class Indexer(MultiPlatformOp): +@@ -352,6 +352,19 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): if _use_aiter and _is_gfx95_supported and isinstance(x, tuple) and len(x) == 3: x = x[2] if _is_cuda: -+ if get_global_server_args().enable_deterministic_inference: ++ if get_exec().deterministic.enable_deterministic_inference: + flat_input = x.reshape(-1, x.shape[-1]).contiguous() + output = torch.empty( + (flat_input.shape[0], self.weights_proj.weight.shape[0]), @@ -549,17 +846,196 @@ index 5e5f7aa2a5..f074a044bc 100644 + ) + return output.view(*x.shape[:-1], self.weights_proj.weight.shape[0]) return torch.mm(x, self.weights_proj.weight.t(), out_dtype=torch.float32) - + weights, _ = self.weights_proj(x) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py -index f9db4b88e6..e9daf23d65 100644 +index c80315fd82..fa056efd7b 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_topk_backend.py -@@ -229,6 +229,15 @@ def _topk_unfused( +@@ -4,8 +4,11 @@ from enum import Enum, IntEnum, auto + from typing import Callable, Dict, List, Optional, Tuple + + import torch ++import triton ++import triton.language as tl + + from sglang.srt.environ import envs ++from sglang.srt.runtime_context import get_exec + + _FLASHINFER_TIE_BREAK_VALUES = { + "small": 1, +@@ -13,6 +16,151 @@ _FLASHINFER_TIE_BREAK_VALUES = { + } + + ++@triton.jit ++def _deterministic_topk_mask_logits_kernel( ++ score_ptr, ++ lengths_ptr, ++ row_starts_ptr, ++ masked_ptr, ++ score_row_stride, ++ score_column_stride, ++ num_columns: tl.constexpr, ++ HAS_ROW_STARTS: tl.constexpr, ++ BLOCK_SIZE: tl.constexpr, ++): ++ row = tl.program_id(0) ++ columns = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) ++ in_bounds = columns < num_columns ++ length = tl.load(lengths_ptr + row) ++ row_start = tl.load(row_starts_ptr + row) if HAS_ROW_STARTS else 0 ++ values = tl.load( ++ score_ptr ++ + row * score_row_stride ++ + columns * score_column_stride, ++ mask=in_bounds, ++ other=-float("inf"), ++ ) ++ values = tl.where(values == 0.0, 0.0, values) ++ values = tl.where( ++ (columns >= row_start) & (columns < row_start + length), ++ values, ++ -float("inf"), ++ ) ++ tl.store(masked_ptr + row * num_columns + columns, values, mask=in_bounds) ++ ++ ++@triton.jit ++def _deterministic_topk_canonicalize_kernel( ++ selected_ptr, ++ values_ptr, ++ lengths_ptr, ++ row_starts_ptr, ++ output_ptr, ++ topk: tl.constexpr, ++ HAS_ROW_STARTS: tl.constexpr, ++ BLOCK_SIZE: tl.constexpr, ++): ++ row = tl.program_id(0) ++ columns = tl.arange(0, BLOCK_SIZE) ++ selected_mask = columns < topk ++ indices = tl.load( ++ selected_ptr + row * topk + columns, ++ mask=selected_mask, ++ other=-1, ++ ) ++ values = tl.load( ++ values_ptr + row * topk + columns, ++ mask=selected_mask, ++ other=-float("inf"), ++ ) ++ length = tl.load(lengths_ptr + row) ++ row_start = tl.load(row_starts_ptr + row) if HAS_ROW_STARTS else 0 ++ indices -= row_start ++ invalid_key = 0x7FFFFFFF ++ valid = ( ++ selected_mask ++ & (indices >= 0) ++ & (indices < length) ++ & (values != -float("inf")) ++ ) ++ indices = tl.where(valid, indices, invalid_key) ++ indices = tl.sort(indices, dim=0) ++ indices = tl.where(indices == invalid_key, -1, indices) ++ tl.store(output_ptr + row * topk + columns, indices, mask=selected_mask) ++ ++ ++def deterministic_flashinfer_topk( ++ score: torch.Tensor, ++ lengths: torch.Tensor, ++ topk: int, ++ row_starts: Optional[torch.Tensor] = None, ++) -> torch.Tensor: ++ """Fast deterministic DSA top-k with canonical ascending indices.""" ++ from flashinfer.topk import get_topk_module ++ from flashinfer.utils import _get_cache_buf ++ ++ batch_size, max_score_len = score.shape ++ lengths_i32 = lengths.to(dtype=torch.int32, device=score.device) ++ has_row_starts = row_starts is not None ++ row_starts_i32 = ( ++ row_starts.to(dtype=torch.int32, device=score.device) ++ if has_row_starts ++ else lengths_i32 ++ ) ++ masked_logits = torch.empty( ++ (batch_size, max_score_len), dtype=score.dtype, device=score.device ++ ) ++ mask_block_size = 256 ++ _deterministic_topk_mask_logits_kernel[ ++ (batch_size, triton.cdiv(max_score_len, mask_block_size)) ++ ]( ++ score, ++ lengths_i32, ++ row_starts_i32, ++ masked_logits, ++ score_row_stride=score.stride(0), ++ score_column_stride=score.stride(1), ++ num_columns=max_score_len, ++ HAS_ROW_STARTS=has_row_starts, ++ BLOCK_SIZE=mask_block_size, ++ ) ++ ++ output_values = torch.empty( ++ (batch_size, topk), dtype=score.dtype, device=score.device ++ ) ++ row_states = _get_cache_buf( ++ f"radix_topk_row_states_{score.device}", ++ 1024 * 1024, ++ score.device, ++ zero_init=True, ++ ) ++ selected = get_topk_module().radix_topk( ++ masked_logits, ++ topk, ++ False, ++ True, ++ _FLASHINFER_TIE_BREAK_VALUES["small"], ++ row_states, ++ output_values, ++ True, ++ ) ++ output = torch.empty( ++ (batch_size, topk), dtype=torch.int32, device=score.device ++ ) ++ _deterministic_topk_canonicalize_kernel[(batch_size,)]( ++ selected, ++ output_values, ++ lengths_i32, ++ row_starts_i32, ++ output, ++ topk=topk, ++ HAS_ROW_STARTS=has_row_starts, ++ BLOCK_SIZE=triton.next_power_of_2(topk), ++ num_warps=8, ++ ) ++ return output ++ ++ + class TopkTransformMethod(IntEnum): + # Transform topk indices to indices to the page table (page_size = 1) + PAGED = auto() +@@ -44,6 +192,14 @@ class DSATopKBackend(Enum): + topk: int, + row_starts: Optional[torch.Tensor] = None, + ) -> torch.Tensor: ++ if ( ++ get_exec().deterministic.enable_deterministic_inference ++ and row_starts is None ++ and score.is_cuda ++ and score.shape[1] >= topk ++ and 0 < topk <= 2048 ++ ): ++ return deterministic_flashinfer_topk(score, lengths, topk) + if self.is_sgl_kernel(): + from sgl_kernel import fast_topk_v2 + +@@ -253,6 +409,15 @@ def _topk_unfused( topk_local_indices = topk_local_indices.masked_fill( topk_scores == float("-inf"), -1 ) -+ if envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get(): ++ if get_exec().deterministic.enable_deterministic_inference: + invalid_sort_key = torch.iinfo(torch.int32).max + topk_local_indices = torch.sort( + topk_local_indices.masked_fill(topk_local_indices < 0, invalid_sort_key), @@ -569,28 +1045,24 @@ index f9db4b88e6..e9daf23d65 100644 + topk_local_indices == invalid_sort_key, -1 + ) topk_indices[:, :valid_topk] = topk_local_indices - + return topk_indices diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py -index 75aa41d784..65482ed603 100644 +index 37fcbe4042..0c16290252 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py -@@ -20,7 +20,10 @@ from sglang.srt.runtime_context import get_parallel +@@ -20,6 +20,7 @@ from sglang.srt.runtime_context import get_parallel, get_spec logger = logging.getLogger(__name__) - from sglang.srt.environ import envs - from sglang.srt.layers.attention.base_attn_backend import AttentionBackend --from sglang.srt.layers.attention.dsa.dequant_k_cache import dequantize_k_cache_paged -+from sglang.srt.layers.attention.dsa.dequant_k_cache import ( + from sglang.kernels.ops.attention.dsa.dequant_k_cache import ( + concat_cast_kv_fp8_pad, + dequantize_k_cache, -+ dequantize_k_cache_paged, -+) - from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import ( - DeepseekSparseAttnBackendMTPPrecomputeMixin, - PrecomputedMetadata, -@@ -73,6 +76,45 @@ from sglang.srt.utils import ( + dequantize_k_cache_paged, + gather_dequant_requant_fp8_paged, + ) +@@ -86,6 +87,45 @@ from sglang.srt.utils import ( _DSA_TRITON_PREFILL = get_bool_env_var("SGLANG_DSA_TRITON_PREFILL") _IS_GFX95 = is_gfx95_supported() - + + +def _quantize_dequantize_fp8_sparse_kv(kv: torch.Tensor) -> torch.Tensor: + """Round fresh MLA KV through the packed FP8 cache representation.""" @@ -632,22 +1104,22 @@ index 75aa41d784..65482ed603 100644 + if is_cuda(): import deep_gemm - -@@ -1997,7 +2039,13 @@ class DeepseekSparseAttnBackend( + +@@ -2112,6 +2152,13 @@ class DeepseekSparseAttnBackend( ) else: kv_cache = _cat([k, k_rope], dim=-1) + if self.dsa_kv_cache_store_fp8: + kv_cache = _quantize_dequantize_fp8_sparse_kv(kv_cache) - page_table_1 = topk_indices ++ page_table_1 = topk_indices + elif self.dsa_kv_cache_store_fp8: + kv_cache, page_table_1 = _dequantize_fp8_sparse_paged_kv( + kv_cache, page_table_1 + ) - - return self._forward_flashmla_sparse( - q_all=q_all, -@@ -2142,6 +2190,10 @@ class DeepseekSparseAttnBackend( + + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) +@@ -2272,6 +2319,10 @@ class DeepseekSparseAttnBackend( if self.dsa_decode_impl == "flashmla_sparse": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) @@ -659,20 +1131,20 @@ index 75aa41d784..65482ed603 100644 q_all=q_all, kv_cache=kv_cache, diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py -index 003a716a78..c3448d7edd 100644 +index 6151b3cd7e..451c6a4907 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py -@@ -203,9 +203,31 @@ def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): +@@ -251,9 +251,31 @@ def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): if envs.SGLANG_DEEPGEMM_PDL.get() and hasattr(deep_gemm, "set_pdl"): deep_gemm.set_pdl(True) - + + configure_deep_gemm_batch_invariant( + server_args.enable_deterministic_inference + or envs.SGLANG_DEEPGEMM_BATCH_INVARIANT.get() + ) compile_utils.update_deep_gemm_config(gpu_id, server_args) - - + + +def configure_deep_gemm_batch_invariant(enabled: bool): + if not ENABLE_JIT_DEEPGEMM: + return @@ -695,10 +1167,10 @@ index 003a716a78..c3448d7edd 100644 def configure_deep_gemm_num_sms(num_sms): if num_sms is None or not ENABLE_JIT_DEEPGEMM: diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py -index 3d4cbfc445..af536ee105 100644 +index b3f1291f48..8f9df94594 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py -@@ -21,7 +21,10 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( +@@ -22,7 +22,10 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( ) from sglang.srt.environ import envs from sglang.srt.eplb.expert_location import get_global_expert_location_metadata @@ -710,10 +1182,10 @@ index 3d4cbfc445..af536ee105 100644 from sglang.srt.layers.moe import ( MoeRunnerConfig, get_deepep_mode, -@@ -84,6 +87,31 @@ _is_cpu = is_cpu() - _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip - - +@@ -134,6 +137,31 @@ def _get_deepep_comm_group(a2a_backend): + return group + + +def compute_deepep_ll_prefill_staging_slices( + *, + local_num_tokens: int, @@ -741,11 +1213,11 @@ index 3d4cbfc445..af536ee105 100644 + def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher: a2a_backend = get_moe_a2a_backend() - if ( -@@ -1228,7 +1256,76 @@ class FusedMoE(torch.nn.Module): - else: - return self.forward_impl(hidden_states, topk_output) - + if a2a_backend.is_none() and is_npu(): +@@ -1442,11 +1470,96 @@ class FusedMoE(torch.nn.Module): + hidden_states, topk_output, pre_quant_input=pre_quant_input + ) + + def _get_deepep_ll_prefill_staging_slices( + self, + hidden_states: torch.Tensor, @@ -757,7 +1229,7 @@ index 3d4cbfc445..af536ee105 100644 + return None + if not get_deepep_mode().resolve(True).is_low_latency(): + return None -+ if not TopKOutputChecker.format_is_standard(topk_output): ++ if not isinstance(topk_output, StandardTopKOutput): + raise TypeError( + "DeepEP LL prefill staging requires standard top-k output, " + f"got {topk_output.format}" @@ -798,53 +1270,76 @@ index 3d4cbfc445..af536ee105 100644 + ), + ) + - def forward_impl(self, hidden_states: torch.Tensor, topk_output: TopKOutput): + def forward_impl( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + pre_quant_input: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ++ ): + staging_slices = self._get_deepep_ll_prefill_staging_slices( + hidden_states, topk_output + ) + if staging_slices is None: -+ return self._forward_impl_once(hidden_states, topk_output) ++ return self._forward_impl_once( ++ hidden_states, ++ topk_output, ++ pre_quant_input=pre_quant_input, ++ ) + + final_hidden_states = hidden_states.new_empty(hidden_states.shape) + for token_slice in staging_slices: ++ staged_pre_quant_input = None ++ if pre_quant_input is not None: ++ staged_pre_quant_input = ( ++ pre_quant_input[0][token_slice], ++ pre_quant_input[1][token_slice], ++ ) + staged_hidden_states = self._forward_impl_once( + hidden_states[token_slice], + self._slice_standard_topk_output(topk_output, token_slice), ++ pre_quant_input=staged_pre_quant_input, + ) + if token_slice.stop > token_slice.start: + final_hidden_states[token_slice].copy_(staged_hidden_states) + return final_hidden_states + -+ def _forward_impl_once(self, hidden_states: torch.Tensor, topk_output: TopKOutput): ++ def _forward_impl_once( ++ self, ++ hidden_states: torch.Tensor, ++ topk_output: TopKOutput, ++ pre_quant_input: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ): origin_hidden_states_dim = hidden_states.shape[-1] assert self.quant_method is not None - 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 4c3fda886d..32552c5e02 100644 +index b769cce033..328621f10c 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py -@@ -59,6 +59,17 @@ else: - - _MASKED_GEMM_FAST_ACT = get_bool_env_var("SGLANG_MASKED_GEMM_FAST_ACT") +@@ -73,6 +73,20 @@ else: + + _DEEPGEMM_ON_H20 = get_bool_env_var("SGLANG_DEEPGEMM_ON_H20") +_DEEPGEMM_PAD_EXPERT_M = get_bool_env_var("SGLANG_DEEPGEMM_PAD_EXPERT_M") + + +def _should_pad_contiguous_expert_m() -> bool: -+ return _DEEPGEMM_PAD_EXPERT_M or envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get() ++ return ( ++ _DEEPGEMM_PAD_EXPERT_M ++ or get_exec().deterministic.enable_deterministic_inference ++ ) + + +def _pad_count_to_deepgemm_m_tile(count: int) -> int: + if count <= 0: + return 0 + return ceil_div(count, 128) * 128 - - + + # TODO(kaixih@nvidia): ideally we should merge this logic into -@@ -722,8 +733,17 @@ def pre_permute_deepep_normal_to_deep_gemm( +@@ -1046,8 +1060,17 @@ def pre_permute_deepep_normal_to_deep_gemm( ) = dispatch_output - assert runner_config.activation == "silu" - + assert runner_config.activation in ("silu", "situ") + - all_tokens = sum(num_recv_tokens_per_expert) + actual_all_tokens = sum(num_recv_tokens_per_expert) + if _should_pad_contiguous_expert_m(): @@ -857,56 +1352,25 @@ index 4c3fda886d..32552c5e02 100644 + all_tokens = sum(deepgemm_num_recv_tokens_per_expert) running_state["all_tokens"] = all_tokens + running_state["actual_all_tokens"] = actual_all_tokens - + K = hidden_states.shape[1] - -@@ -737,10 +757,18 @@ def pre_permute_deepep_normal_to_deep_gemm( - running_state["topk_ids"] = topk_ids - running_state["topk_weights"] = topk_weights - -- input_tensor = torch.empty( -- (all_tokens, K), -- device=hidden_states.device, -- dtype=hidden_states.dtype, -+ input_tensor = ( -+ torch.zeros( -+ (all_tokens, K), -+ device=hidden_states.device, -+ dtype=hidden_states.dtype, -+ ) -+ if all_tokens != actual_all_tokens -+ else torch.empty( -+ (all_tokens, K), -+ device=hidden_states.device, -+ dtype=hidden_states.dtype, -+ ) - ) - if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0: - # TODO check whether need `zeros` -@@ -750,21 +778,29 @@ def pre_permute_deepep_normal_to_deep_gemm( - dtype=torch.int, - ).transpose(0, 1) - else: -- input_tensor_scale = torch.empty( -- (all_tokens, K // 128), -- device=hidden_states.device, -- dtype=torch.float32, -+ input_tensor_scale = ( -+ torch.zeros( -+ (all_tokens, K // 128), -+ device=hidden_states.device, -+ dtype=torch.float32, -+ ) -+ if all_tokens != actual_all_tokens -+ else torch.empty( -+ (all_tokens, K // 128), -+ device=hidden_states.device, -+ dtype=torch.float32, -+ ) - ) - m_indices = torch.empty(all_tokens, device=hidden_states.device, dtype=torch.int32) - output_index = torch.empty_like(topk_ids) - + +@@ -1066,8 +1089,11 @@ def pre_permute_deepep_normal_to_deep_gemm( + # input_tensor would leak batch-dependent values into the grouped GEMM. + # The scale buffer only matters for FP8 activations sharing this + # pre-permute (ep_scatter skips scales entirely for BF16 dispatch). +- deterministic = get_exec().deterministic.enable_deterministic_inference +- buffer_init = torch.zeros if deterministic else torch.empty ++ needs_zero_fill = ( ++ get_exec().deterministic.enable_deterministic_inference ++ or all_tokens != actual_all_tokens ++ ) ++ buffer_init = torch.zeros if needs_zero_fill else torch.empty + + input_tensor = buffer_init( + (all_tokens, K), +@@ -1092,11 +1118,11 @@ def pre_permute_deepep_normal_to_deep_gemm( + if get_offloader().forbid_copy_engine_usage: num_recv_tokens_per_expert_gpu = copy_list_to_gpu_no_ce( - num_recv_tokens_per_expert @@ -920,12 +1384,12 @@ index 4c3fda886d..32552c5e02 100644 pin_memory=True, device="cpu", diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py -index 33a51a89da..ca9ca021aa 100644 +index f33c74f77e..214ce6be02 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -1,5 +1,6 @@ from __future__ import annotations - + +import inspect import logging from contextlib import nullcontext @@ -934,41 +1398,45 @@ index 33a51a89da..ca9ca021aa 100644 get_deepep_output_dtype, is_tbo_enabled, ) -+from sglang.srt.server_args import get_global_server_args ++from sglang.srt.runtime_context import get_exec from sglang.srt.utils import ( get_bool_env_var, get_cuda_version, -@@ -644,6 +646,19 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): +@@ -664,6 +666,21 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): self.return_recv_hook = return_recv_hook self.device_module = torch.get_device_module() self.quant_config = {} -+ self.align_fp8_quantization = self._supports_aligned_fp8_quantization( -+ Buffer.low_latency_dispatch -+ ) and bool(get_global_server_args().enable_deterministic_inference) -+ -+ @staticmethod -+ def _supports_aligned_fp8_quantization(dispatch_method) -> bool: -+ try: -+ return ( -+ "align_fp8_quantization" -+ in inspect.signature(dispatch_method).parameters -+ ) -+ except (TypeError, ValueError): -+ return False ++ self.fp8_quantization_mode = None ++ if get_exec().deterministic.enable_deterministic_inference: ++ try: ++ supports_fp8_quantization_mode = ( ++ "fp8_quantization_mode" ++ in inspect.signature(Buffer.low_latency_dispatch).parameters ++ ) ++ except (TypeError, ValueError): ++ supports_fp8_quantization_mode = False ++ if not supports_fp8_quantization_mode: ++ raise RuntimeError( ++ "Deterministic inference with SGLang 0.5.17 requires a " ++ "DeepEP build that supports fp8_quantization_mode" ++ ) ++ self.fp8_quantization_mode = "fast_pairwise" def dispatch_a( self, -@@ -725,6 +740,9 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): +@@ -747,6 +764,11 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): buffer = self._get_buffer() _deepep_precompile_tp_barrier() + alignment_kwargs = ( -+ {"align_fp8_quantization": True} if self.align_fp8_quantization else {} ++ {"fp8_quantization_mode": self.fp8_quantization_mode} ++ if self.fp8_quantization_mode is not None ++ else {} + ) packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = ( buffer.low_latency_dispatch( hidden_states, -@@ -741,6 +759,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): +@@ -768,6 +790,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): async_finish=not self.return_recv_hook, return_recv_hook=self.return_recv_hook, **fp8_deepgemm_scale_opts, @@ -976,62 +1444,32 @@ index 33a51a89da..ca9ca021aa 100644 ) ) return packed_recv_hidden, self.packed_recv_count, event, hook -diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py -index 27933521b1..a0dd2ddcdb 100644 ---- a/python/sglang/srt/layers/moe/topk.py -+++ b/python/sglang/srt/layers/moe/topk.py -@@ -1393,6 +1393,20 @@ def biased_grouped_topk_gpu( - routed_scaling_factor: Optional[float] = None, - apply_routed_scaling_factor_on_output: Optional[bool] = False, - ) -> Tuple[torch.Tensor, torch.Tensor]: -+ if envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get(): -+ return biased_grouped_topk_impl.__wrapped__( -+ hidden_states, -+ gating_output, -+ correction_bias, -+ topk, -+ renormalize, -+ num_expert_group, -+ topk_group, -+ num_fused_shared_experts=num_fused_shared_experts, -+ routed_scaling_factor=routed_scaling_factor, -+ apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, -+ ) -+ - num_tokens = gating_output.shape[0] - num_experts = gating_output.shape[1] - experts_per_group = ( -diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 159ef8fc5a..dd0193fa5c 100644 ---- a/python/sglang/srt/model_executor/model_runner.py -+++ b/python/sglang/srt/model_executor/model_runner.py -@@ -2220,6 +2220,14 @@ class ModelRunner(ModelRunnerKVCacheMixin): +diff --git a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py +index ecbfa7a713..0ccd60390c 100644 +--- a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py ++++ b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py +@@ -382,6 +382,12 @@ class WeightUpdater: # Load the reconstructed tensors using the standard method - self.model.load_weights(reconstructed_tensors) - + self.get_model().load_weights(reconstructed_tensors) + + # The flattened tensor can be backed by CUDA IPC memory owned by the -+ # training process. Weight loaders enqueue device-to-device copies, -+ # while the RPC response releases the producer-side IPC buffer. Make -+ # the copies complete before returning so a following bucket cannot -+ # reuse that storage and corrupt the just-loaded weights. ++ # training process. Complete the enqueued copies before the RPC returns ++ # and lets the producer reuse that storage. + if self.device == "cuda": + torch.cuda.synchronize() + return True, "Success" - - def get_weights_by_name( + + def update_weights_from_ipc(self: WeightUpdater, recv_req): diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py -index 587a132189..a173c92ce5 100644 +index 61a4be7356..c028fc1ab0 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py -@@ -492,7 +492,15 @@ class MoEGate(nn.Module): - True, # is_vnni +@@ -506,6 +506,13 @@ class MoEGate(nn.Module): ) - -- if get_global_server_args().enable_deterministic_inference: -+ server_args = get_global_server_args() -+ if server_args.enable_deterministic_inference: -+ if server_args.enable_fp32_moe_router: + + if get_exec().deterministic.enable_deterministic_inference: ++ if get_server_args().enable_fp32_moe_router: + from sglang.srt.batch_invariant_ops import router_gemm_batch_invariant + + return router_gemm_batch_invariant( @@ -1039,33 +1477,37 @@ index 587a132189..a173c92ce5 100644 + self.weight.contiguous(), + ) return F.linear(hidden_states, self.weight, None) - + if ( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index bd708d202f..de398a854e 100644 +index 9363e8c75f..0f47fc3e1b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py -@@ -195,9 +195,15 @@ ATTENTION_BACKEND_CHOICES = [ - "intel_xpu", +@@ -219,13 +219,20 @@ CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS = [ + + DETERMINISTIC_ATTENTION_BACKEND_CHOICES = [ + "ascend", ++ "dsa", + "fa3", + "fa4", + "flashinfer", + "triton", ] - --DETERMINISTIC_ATTENTION_BACKEND_CHOICES = ["flashinfer", "fa3", "triton", "ascend"] -+DETERMINISTIC_ATTENTION_BACKEND_CHOICES = [ -+ "flashinfer", + +-RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["ascend", "fa3", "fa4", "triton"] ++RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = [ ++ "ascend", ++ "dsa", + "fa3", ++ "fa4", + "triton", -+ "dsa", -+ "ascend", +] - --RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["fa3", "triton", "ascend"] -+RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["fa3", "triton", "dsa", "ascend"] - + DISAGG_TRANSFER_BACKEND_CHOICES = [ "mooncake", -@@ -2517,6 +2523,10 @@ class ServerArgs: - bool, +@@ -3257,6 +3264,10 @@ class ServerArgs: "Enable deterministic inference mode with batch invariant ops.", + NS("exec.deterministic"), ] = False + enable_fp32_moe_router: A[ + bool, @@ -1074,547 +1516,3 @@ index bd708d202f..de398a854e 100644 rl_on_policy_target: A[ Optional[str], Arg( -@@ -5860,7 +5870,7 @@ class ServerArgs: - - attention_backend = resolved_view(self).attention_backend - if is_deepseek_model: -- if attention_backend not in ["fa3", "triton"]: -+ if attention_backend not in ["fa3", "triton", "dsa"]: - raise ValueError( - f"Currently only {RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with DeepSeek models. But you're using {attention_backend}." - ) -diff --git a/sgl-kernel/csrc/gemm/per_token_group_quant_8bit_v2.cu b/sgl-kernel/csrc/gemm/per_token_group_quant_8bit_v2.cu -index 4fbaf08efc..5bad403d0f 100644 ---- a/sgl-kernel/csrc/gemm/per_token_group_quant_8bit_v2.cu -+++ b/sgl-kernel/csrc/gemm/per_token_group_quant_8bit_v2.cu -@@ -35,7 +35,9 @@ __device__ __forceinline__ float silu(const float& val) { - float t = __tanhf(half); - return half * (1.0f + t); - #else -- return val / (1.0f + __expf(-val)); -+ // Match sgl_kernel.silu_and_mul on SM90. That kernel uses expf and -+ // exposes only the final SiLU*up BF16 rounding boundary. -+ return val / (1.0f + expf(-val)); - #endif - } - -@@ -266,7 +268,8 @@ __global__ void per_token_group_quant_8bit_kernel( - // TODO can this be removed? - const int scale_expert_stride, - const int scale_hidden_stride, -- const int num_tokens_per_expert) { -+ const int num_tokens_per_expert, -+ const float max_8bit) { - using dst_dtype_info = DtypeInfo; - using scale_element_t = std::conditional_t; - static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0); -@@ -347,8 +350,11 @@ __global__ void per_token_group_quant_8bit_kernel( - for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) { - float val; - if constexpr (FUSE_SILU_AND_MUL) { -- // TODO maybe vectorize -- T val_lowprec = static_cast(silu(static_cast(input_primary_vec[j]))) * input_secondary_vec[j]; -+ // Keep the same single visible low-precision boundary as the -+ // contiguous activation path. Rounding SiLU before multiplying -+ // changes the second MoE GEMM's FP8 inputs. -+ T val_lowprec = static_cast( -+ silu(static_cast(input_primary_vec[j])) * static_cast(input_secondary_vec[j])); - val = static_cast(val_lowprec); - input_primary_vec[j] = val_lowprec; - } else { -@@ -362,8 +368,15 @@ __global__ void per_token_group_quant_8bit_kernel( - local_absmax = GroupReduceMax(local_absmax, lane_id); - - float y_scale, y_scale_inv; -- calculate_fp8_scales(local_absmax, y_scale, y_scale_inv); -- float2 y_scale_repeated = {y_scale, y_scale}; -+ if constexpr (SCALE_UE8M0) { -+ calculate_fp8_scales(local_absmax, y_scale, y_scale_inv); -+ } else { -+ // Preserve the scale and division sequence of the legacy unmasked -+ // quantizer. Reciprocal reuse under -use_fast_math moves a few -+ // values across an FP8 rounding boundary. -+ y_scale_inv = __fdiv_rn(local_absmax, max_8bit); -+ y_scale = 1.0f / y_scale_inv; -+ } - - if (lane_id == 0) { - *scale_output = extract_required_scale_format(y_scale_inv); -@@ -373,19 +386,33 @@ __global__ void per_token_group_quant_8bit_kernel( - static_assert(sizeof(output_buf) == INPUT_PRIMARY_VEC_SIZE * sizeof(DST_DTYPE)); - - if constexpr (std::is_same_v) { -- const auto output_buf_ptr = reinterpret_cast<__nv_fp8x2_storage_t*>(&output_buf); -- static_assert(sizeof(output_buf) == INPUT_PRIMARY_VEC_SIZE / 2 * sizeof(__nv_fp8x2_storage_t)); -- static_assert(INPUT_PRIMARY_VEC_SIZE % 2 == 0); -+ if constexpr (SCALE_UE8M0) { -+ const auto output_buf_ptr = reinterpret_cast<__nv_fp8x2_storage_t*>(&output_buf); -+ static_assert(sizeof(output_buf) == INPUT_PRIMARY_VEC_SIZE / 2 * sizeof(__nv_fp8x2_storage_t)); -+ static_assert(INPUT_PRIMARY_VEC_SIZE % 2 == 0); - - #pragma unroll -- for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; j += 2) { -- float2 inputx2 = {static_cast(input_primary_vec[j]), static_cast(input_primary_vec[j + 1])}; -- float2 outputx2 = fmul2_rn(inputx2, y_scale_repeated); -+ for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; j += 2) { -+ float2 inputx2 = { -+ static_cast(input_primary_vec[j]), static_cast(input_primary_vec[j + 1])}; -+ float2 outputx2 = fmul2_rn(inputx2, {y_scale, y_scale}); - -- outputx2.x = fminf(fmaxf(outputx2.x, dst_dtype_info::MIN), dst_dtype_info::MAX); -- outputx2.y = fminf(fmaxf(outputx2.y, dst_dtype_info::MIN), dst_dtype_info::MAX); -+ outputx2.x = fminf(fmaxf(outputx2.x, dst_dtype_info::MIN), dst_dtype_info::MAX); -+ outputx2.y = fminf(fmaxf(outputx2.y, dst_dtype_info::MIN), dst_dtype_info::MAX); - -- output_buf_ptr[j / 2] = __nv_cvt_float2_to_fp8x2(outputx2, __NV_SATFINITE, __NV_E4M3); -+ output_buf_ptr[j / 2] = __nv_cvt_float2_to_fp8x2(outputx2, __NV_SATFINITE, __NV_E4M3); -+ } -+ } else { -+ // Match the legacy SM90 quantizer's scalar conversion exactly; -+ // packed conversion makes a different tie choice for rare values. -+ const auto output_buf_ptr = reinterpret_cast(&output_buf); -+#pragma unroll -+ for (uint32_t j = 0; j < INPUT_PRIMARY_VEC_SIZE; ++j) { -+ float val = static_cast(input_primary_vec[j]); -+ float q_val = -+ fminf(fmaxf(__fdiv_rn(val, y_scale_inv), dst_dtype_info::MIN), dst_dtype_info::MAX); -+ output_buf_ptr[j] = DST_DTYPE(q_val); -+ } - } - } else { - const auto output_buf_ptr = reinterpret_cast(&output_buf); -@@ -474,7 +501,8 @@ void sgl_per_token_group_quant_8bit_v2( - hidden_dim_num_groups, \ - scale_expert_stride, \ - scale_hidden_stride, \ -- num_tokens_per_expert); \ -+ num_tokens_per_expert, \ -+ static_cast(max_8bit)); \ - } while (0) - - #define LAUNCH_KERNEL(GROUP_SIZE, T, DST_DTYPE) \ -@@ -497,13 +525,51 @@ void sgl_per_token_group_quant_8bit_v2( - NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, true); \ - } \ - } else { \ -- LAUNCH_KERNEL_INNER(NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true); \ -+ if (masked_layout) { \ -+ LAUNCH_KERNEL_INNER( \ -+ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, false); \ -+ } else { \ -+ LAUNCH_KERNEL_INNER( \ -+ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, uint32_t, true, true, false); \ -+ } \ - } \ - } else { \ -- LAUNCH_KERNEL_INNER(NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true); \ -+ if (fuse_silu_and_mul) { \ -+ if (masked_layout) { \ -+ LAUNCH_KERNEL_INNER( \ -+ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, true); \ -+ } else { \ -+ LAUNCH_KERNEL_INNER( \ -+ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, true); \ -+ } \ -+ } else { \ -+ if (masked_layout) { \ -+ LAUNCH_KERNEL_INNER( \ -+ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, false); \ -+ } else { \ -+ LAUNCH_KERNEL_INNER( \ -+ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, true, false, false); \ -+ } \ -+ } \ - } \ - } else { \ -- LAUNCH_KERNEL_INNER(NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false); \ -+ if (fuse_silu_and_mul) { \ -+ if (masked_layout) { \ -+ LAUNCH_KERNEL_INNER( \ -+ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, true); \ -+ } else { \ -+ LAUNCH_KERNEL_INNER( \ -+ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, true); \ -+ } \ -+ } else { \ -+ if (masked_layout) { \ -+ LAUNCH_KERNEL_INNER( \ -+ MaskedLayoutScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, false); \ -+ } else { \ -+ LAUNCH_KERNEL_INNER( \ -+ NaiveScheduler, GROUP_SIZE, THREADS_PER_SUBWARP, T, DST_DTYPE, float, false, false, false); \ -+ } \ -+ } \ - } \ - } while (0) - -diff --git a/sgl-kernel/tests/test_per_token_group_quant_8bit.py b/sgl-kernel/tests/test_per_token_group_quant_8bit.py -index 24ebeb8778..9a3f986c12 100644 ---- a/sgl-kernel/tests/test_per_token_group_quant_8bit.py -+++ b/sgl-kernel/tests/test_per_token_group_quant_8bit.py -@@ -6,6 +6,7 @@ from pathlib import Path - - import pytest - import torch -+from sgl_kernel import silu_and_mul - from sgl_kernel.test_utils import ( - assert_all_close_or_tiny_diff, - create_per_token_group_quant_test_data, -@@ -22,6 +23,61 @@ from sglang.srt.utils import get_bool_env_var, is_hip - _is_hip = is_hip() - fp8_type_ = torch.float8_e4m3fnuz if _is_hip else torch.float8_e4m3fn - -+ -+def test_masked_fused_silu_quant_matches_contiguous_sm90(): -+ if torch.cuda.get_device_capability() != (9, 0): -+ pytest.skip("This test covers the Hopper train-rollout alignment path") -+ -+ torch.manual_seed(20260806) -+ num_experts, max_tokens, hidden_dim, group_size = 4, 128, 2048, 128 -+ masked_m = torch.tensor([1, 17, 64, 127], device="cuda", dtype=torch.int32) -+ gate_up = torch.randn( -+ num_experts, -+ max_tokens, -+ hidden_dim * 2, -+ device="cuda", -+ dtype=torch.bfloat16, -+ ) -+ gate_up[1, 3, :256] *= 7 -+ -+ fused_q, fused_s = sglang_per_token_group_quant_8bit( -+ x=gate_up, -+ group_size=group_size, -+ dst_dtype=torch.float8_e4m3fn, -+ column_major_scales=True, -+ scale_tma_aligned=True, -+ scale_ue8m0=False, -+ fuse_silu_and_mul=True, -+ masked_m=masked_m, -+ enable_v2=True, -+ ) -+ -+ for expert, count in enumerate(masked_m.cpu().tolist()): -+ activated = silu_and_mul(gate_up[expert, :count].contiguous()) -+ reference_q, reference_s = sglang_per_token_group_quant_8bit( -+ x=activated, -+ group_size=group_size, -+ dst_dtype=torch.float8_e4m3fn, -+ column_major_scales=True, -+ scale_tma_aligned=True, -+ scale_ue8m0=False, -+ fuse_silu_and_mul=False, -+ masked_m=None, -+ enable_v2=False, -+ ) -+ torch.testing.assert_close( -+ fused_q[expert, :count].view(torch.uint8), -+ reference_q.view(torch.uint8), -+ rtol=0, -+ atol=0, -+ ) -+ torch.testing.assert_close( -+ fused_s[expert, :count], -+ reference_s, -+ rtol=0, -+ atol=0, -+ ) -+ - configs = list( - itertools.product( - [1, 4, 16, 64, 127, 128, 512, 1024, 4096, 8192], # num_tokens -diff --git a/test/registered/debug_utils/test_tensor_dump_forward_hook.py b/test/registered/debug_utils/test_tensor_dump_forward_hook.py -index 4ebbfd17db..42e247e592 100644 ---- a/test/registered/debug_utils/test_tensor_dump_forward_hook.py -+++ b/test/registered/debug_utils/test_tensor_dump_forward_hook.py -@@ -1,4 +1,4 @@ --import unittest -+from pathlib import Path - - import torch - from torch import nn -@@ -6,103 +6,100 @@ from torch import nn - from sglang.srt.debug_utils.tensor_dump_forward_hook import ( - register_forward_hook_for_model, - ) --from sglang.srt.distributed.parallel_state import ( -- init_distributed_environment, -- initialize_model_parallel, --) --from sglang.srt.layers.layernorm import RMSNorm --from sglang.srt.layers.linear import LinearBase --from sglang.srt.models.qwen2 import Qwen2MLP --from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler --from sglang.srt.utils import add_prefix --from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -- --register_cuda_ci( -- est_time=9, -- stage="base-b", -- runner_config="1-gpu-small", -- disabled="Test uses pytest-style function without TestCase class - see #17145", --) --register_amd_ci( -- est_time=15, -- suite="stage-b-test-1-gpu-small-amd", -- disabled="Test uses pytest-style function without TestCase class - see #17145", --) -+from sglang.srt.model_executor.forward_batch_info import ForwardBatch - --TEST_HIDDEN_SIZE = 32 - -+class _Layer(nn.Module): -+ def __init__(self): -+ super().__init__() -+ self.experts = nn.Sequential(nn.Linear(2, 2, bias=False)) -+ self.shared_experts = nn.Sequential(nn.Linear(2, 2, bias=False)) -+ -+ def forward(self, value): -+ return self.experts(value) + self.shared_experts(value) - --class SimpleModel(nn.Module): - -- def __init__(self) -> None: -+class _InnerModel(nn.Module): -+ def __init__(self): - super().__init__() -- self.hidden_size = TEST_HIDDEN_SIZE -- self.rms_norm_eps = 1e-5 -- self.mlp = Qwen2MLP( -- hidden_size=self.hidden_size, -- intermediate_size=self.hidden_size, -- hidden_act="silu", -- quant_config=None, -- prefix=add_prefix("mlp", ""), -- ) -- self.layernorm = RMSNorm(self.hidden_size, eps=self.rms_norm_eps) -- -- @torch.no_grad() -- def forward( -- self, -- hidden_states: torch.Tensor, -- ) -> torch.Tensor: -- hidden_states = self.layernorm(hidden_states) -- hidden_states = self.mlp(hidden_states) -- return hidden_states -- -- --class MockCausalLM(nn.Module): -- def __init__(self) -> None: -+ self.layers = nn.ModuleList([_Layer(), _Layer()]) -+ -+ def forward(self, value, forward_batch): -+ del forward_batch -+ for layer in self.layers: -+ value = layer(value) -+ return value -+ -+ -+class _Model(nn.Module): -+ def __init__(self): - super().__init__() -- self.model = SimpleModel() -- -- @torch.no_grad() -- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: -- return self.model(hidden_states) -- -- --def init_weights(module): -- if isinstance(module, LinearBase): -- torch.nn.init.uniform_(module.weight) -- if module.bias is not None: -- torch.nn.init.zeros_(module.bias) -- elif isinstance(module, RMSNorm): -- torch.nn.init.ones_(module.weight) -- -- --def test_model_forward_dump(tmp_path): -- set_global_server_args_for_scheduler(ServerArgs(model_path="dummy")) -- init_distributed_environment( -- backend="nccl", -- world_size=1, -- rank=0, -- local_rank=0, -- distributed_init_method="tcp://127.0.0.1:2646", -- ) -- initialize_model_parallel() -- model = MockCausalLM() -- model.apply(init_weights) -- model = model.cuda().bfloat16() -+ self.model = _InnerModel() -+ -+ def forward(self, value, forward_batch): -+ return self.model(value, forward_batch) -+ -+ -+def _make_forward_batch() -> ForwardBatch: -+ batch = object.__new__(ForwardBatch) -+ batch.input_ids = torch.tensor([11, 12]) -+ batch.seq_lens = torch.tensor([2], dtype=torch.int32) -+ batch.positions = torch.tensor([0, 1]) -+ batch.req_pool_indices = torch.tensor([3], dtype=torch.int32) -+ batch.extend_seq_lens = torch.tensor([2], dtype=torch.int32) -+ batch.extend_prefix_lens = torch.tensor([0], dtype=torch.int32) -+ batch.rids = ["request-0"] -+ batch.forward_mode = "extend" -+ return batch -+ -+ -+def test_layer_outputs_only_dump(monkeypatch, tmp_path: Path): -+ monkeypatch.setenv("SGLANG_TENSOR_DUMP_LAYER_OUTPUTS_ONLY", "1") -+ monkeypatch.setenv("SGLANG_TENSOR_DUMP_CHUNK_SIZE", "16") -+ model = _Model() - dumper = register_forward_hook_for_model( -- model, tmp_path / "sglang_dump", [0], 0, 0, 0 -+ model, -+ str(tmp_path), -+ dump_layers=[0, 1], -+ tp_size=1, -+ tp_rank=0, -+ pp_rank=0, -+ ) -+ -+ model(torch.ones((2, 2)), _make_forward_batch()) -+ -+ (dump_file,) = list(Path(dumper.get_dump_dir()).glob("Chunk*.pt")) -+ (values,) = torch.load(dump_file, weights_only=False) -+ assert "model.layers.0" in values -+ assert "model.layers.1" in values -+ assert not any("experts" in name for name in values) -+ assert values["model.forward_batch_info.rids"] == ["request-0"] -+ torch.testing.assert_close( -+ values["model.forward_batch_info.positions"], torch.tensor([0, 1]) - ) - -- dir_path = dumper.get_dump_dir() -- inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01 -- result = model(inp.cuda()) -- data = torch.load(f"{dir_path}/Pass00000.pt") -- assert "model.layernorm" in data -- assert "model.mlp.down_proj" in data -- assert torch.allclose( -- data["model.mlp.down_proj"], result.cpu(), rtol=1e-5, atol=1e-5 -+ -+def test_layer_outputs_only_can_select_nested_modules(monkeypatch, tmp_path: Path): -+ monkeypatch.setenv("SGLANG_TENSOR_DUMP_LAYER_OUTPUTS_ONLY", "1") -+ monkeypatch.setenv( -+ "SGLANG_TENSOR_DUMP_MODULE_SUFFIXES", "mlp.experts,experts,shared_experts" -+ ) -+ model = _Model() -+ dumper = register_forward_hook_for_model( -+ model, -+ str(tmp_path), -+ dump_layers=[0], -+ tp_size=1, -+ tp_rank=0, -+ pp_rank=0, - ) - -+ model(torch.ones((2, 2)), _make_forward_batch()) - --if __name__ == "__main__": -- unittest.main() -+ data = torch.load( -+ Path(dumper.get_dump_dir()) / "Pass00000.pt", weights_only=False -+ ) -+ assert "model.layers.0" in data -+ assert "model.layers.0.experts" in data -+ assert "model.layers.0.shared_experts" in data -+ assert not any(name.startswith("model.layers.1") for name in data) -diff --git a/test/registered/moe/test_deepep_ll_prefill_staging.py b/test/registered/moe/test_deepep_ll_prefill_staging.py -new file mode 100644 -index 0000000000..e3c883a93d ---- /dev/null -+++ b/test/registered/moe/test_deepep_ll_prefill_staging.py -@@ -0,0 +1,47 @@ -+import pytest -+ -+from sglang.srt.layers.moe.fused_moe_triton.layer import ( -+ compute_deepep_ll_prefill_staging_slices, -+) -+ -+ -+@pytest.mark.parametrize( -+ ("local_num_tokens", "max_num_tokens", "capacity", "expected"), -+ [ -+ (0, 0, 64, []), -+ (63, 63, 64, [slice(0, 63)]), -+ (130, 130, 64, [slice(0, 64), slice(64, 128), slice(128, 130)]), -+ (17, 130, 64, [slice(0, 17), slice(17, 17), slice(17, 17)]), -+ ], -+) -+def test_compute_deepep_ll_prefill_staging_slices( -+ local_num_tokens: int, -+ max_num_tokens: int, -+ capacity: int, -+ expected: list[slice], -+): -+ assert ( -+ compute_deepep_ll_prefill_staging_slices( -+ local_num_tokens=local_num_tokens, -+ max_num_tokens=max_num_tokens, -+ capacity=capacity, -+ ) -+ == expected -+ ) -+ -+ -+@pytest.mark.parametrize( -+ ("local_num_tokens", "max_num_tokens", "capacity"), -+ [(1, 1, 0), (-1, 1, 64), (2, 1, 64)], -+) -+def test_compute_deepep_ll_prefill_staging_slices_rejects_invalid_inputs( -+ local_num_tokens: int, -+ max_num_tokens: int, -+ capacity: int, -+): -+ with pytest.raises(ValueError): -+ compute_deepep_ll_prefill_staging_slices( -+ local_num_tokens=local_num_tokens, -+ max_num_tokens=max_num_tokens, -+ capacity=capacity, -+ ) -diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py -index bcad2b1111..81cece1017 100644 ---- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py -+++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py -@@ -1228,16 +1228,14 @@ DSA_DECODE_IMPL_VARIANTS: tuple[str, ...] = ( - "aiter", - ) - --# Impls that accept an FP8-stored K cache. The flashmla *sparse* and FA3 --# kernels require BF16 K (`kv must have dtype torch::kBFloat16`), so they --# fall back to the inline-quantize-of-bf16 path that production *doesn't* --# take in FP8 deployments. The `flashmla_kv` decode kernel and *both* --# flashmla prefill kernels are the production-relevant FP8 paths. -+# Impls that support an FP8-stored K cache. FlashMLA sparse requires BF16 K, -+# so the backend gathers and dequantizes only the selected pages before calling -+# it. The `flashmla_kv` kernels consume the packed cache directly. - DSA_FP8_COMPATIBLE_PREFILL_IMPLS: frozenset[str] = frozenset( - {"flashmla_sparse", "flashmla_kv", "flashmla_auto"} - ) - DSA_FP8_COMPATIBLE_DECODE_IMPLS: frozenset[str] = frozenset( -- {"flashmla_kv", "flashmla_auto"} -+ {"flashmla_sparse", "flashmla_kv", "flashmla_auto"} - ) - - -@@ -1400,15 +1398,12 @@ def run_dsa_sparse_fp8_decode_case( - *, - dsa_decode_backend: str = "flashmla_kv", - ) -> None: -- """FP8-KV-cache decode. Only `flashmla_kv` (and `flashmla_auto` -- which resolves to it for FP8) accepts an FP8-stored K cache; -- `flashmla_sparse` and `fa3` decode kernels assert BF16 K and would -- fall back to the inline-quantize-of-bf16 path that production -- doesn't take in FP8 deployments.""" -+ """FP8-KV-cache decode. `flashmla_kv` reads the packed cache directly; -+ `flashmla_sparse` gathers and dequantizes the selected pages to BF16.""" - if dsa_decode_backend not in DSA_FP8_COMPATIBLE_DECODE_IMPLS: - testcase.skipTest( - f"DSA decode impl `{dsa_decode_backend}` does not support FP8 KV " -- f"cache (only `flashmla_kv` / `flashmla_auto` read FP8 K directly)." -+ "cache." - ) - if not case.forward_mode.is_decode(): - raise ValueError("run_dsa_sparse_fp8_decode_case expects a DECODE case.") diff --git a/docker/patch/latest/sglang-pull_weights.patch b/docker/patch/latest/sglang-pull_weights.patch index c15e64758b..267ac16435 100644 --- a/docker/patch/latest/sglang-pull_weights.patch +++ b/docker/patch/latest/sglang-pull_weights.patch @@ -1,8 +1,8 @@ diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py -index c74f7666a9972ba7ee5b0c94c805189e640fbbda..107c6eadf6a33e19f4f34d3668c22856ea221e1a 100644 +index 0189493c17a23efb9c81eca2ef3129063f73d40a..00335625112005b5af5ded168cc6d2eea9f69ea9 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py -@@ -129,6 +129,7 @@ from sglang.srt.managers.io_struct import ( +@@ -134,6 +134,7 @@ from sglang.srt.managers.io_struct import ( PauseGenerationReqInput, PostProcessWeightsReqInput, ProfileReq, @@ -10,7 +10,7 @@ index c74f7666a9972ba7ee5b0c94c805189e640fbbda..107c6eadf6a33e19f4f34d3668c22856 ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, SendWeightsToRemoteInstanceReqInput, -@@ -1344,6 +1345,21 @@ async def post_process_weights(req: PostProcessWeightsReqInput, request: Request +@@ -1431,6 +1432,21 @@ async def post_process_weights(req: PostProcessWeightsReqInput, request: Request ) @@ -33,11 +33,11 @@ index c74f7666a9972ba7ee5b0c94c805189e640fbbda..107c6eadf6a33e19f4f34d3668c22856 @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weight_version( diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index e13964675d0dc3bceffa0a7df63e11111c1e416d..9de60810205e1ded3c1a4975c86a887e8c942145 100644 +index 18b27b1b64652245d92107d0dc7e02d11d98a9cd..e14df7ca9cd0fddc488993da5f4ab07de7b411e0 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -1739,6 +1739,22 @@ class CheckWeightsReqOutput(BaseReq, kw_only=True): - payload: Optional[Dict[str, Any]] = None +@@ -1905,6 +1905,22 @@ class CheckWeightsReqOutput(BaseReq, kw_only=True): + payload: Optional[List[ChecksumInfo]] = None +class PullWeightsReqInput(BaseReq, kw_only=True): @@ -60,10 +60,10 @@ index e13964675d0dc3bceffa0a7df63e11111c1e416d..9de60810205e1ded3c1a4975c86a887e forward_sleep_time: Optional[float] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index d4169bebe0b58fc4f32d04e307e5132e66f364ff..e62eae2d4b613219a6227ba7080d9e89feac6207 100644 +index ecbd9bb4fe6ec40c84262addc1d0265ad8b281d3..947321d742879ff1849c376bd4734ccf3d654dc2 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py -@@ -128,6 +128,7 @@ from sglang.srt.managers.io_struct import ( +@@ -143,6 +143,7 @@ from sglang.srt.managers.io_struct import ( PauseGenerationReqInput, PostProcessWeightsReqInput, ProfileReq, @@ -71,7 +71,7 @@ index d4169bebe0b58fc4f32d04e307e5132e66f364ff..e62eae2d4b613219a6227ba7080d9e89 ReleaseMemoryOccupationReqInput, RemoveExternalCorpusReqInput, RemoveExternalCorpusReqOutput, -@@ -1437,6 +1438,10 @@ class Scheduler( +@@ -1560,6 +1561,10 @@ class Scheduler( CheckWeightsReqInput, self.weight_updater.check_weights, ), @@ -83,10 +83,10 @@ index d4169bebe0b58fc4f32d04e307e5132e66f364ff..e62eae2d4b613219a6227ba7080d9e89 ( ProfileReq, diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py -index 505bd2994fbf4812c5069d4391cb80197d42f80c..2ed3b21d75c748a567478686d335184920dcf26b 100644 +index e2aea439e360a29c997f2cbd0ed565db8a0e3167..100a3642deb0698544b696e1009bdcf520803a3c 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py -@@ -28,6 +28,8 @@ from sglang.srt.managers.io_struct import ( +@@ -30,6 +30,8 @@ from sglang.srt.managers.io_struct import ( InitWeightsUpdateGroupReqOutput, PostProcessWeightsReqInput, PostProcessWeightsReqOutput, @@ -95,7 +95,7 @@ index 505bd2994fbf4812c5069d4391cb80197d42f80c..2ed3b21d75c748a567478686d3351849 ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, ResumeMemoryOccupationReqInput, -@@ -309,6 +311,40 @@ class SchedulerWeightUpdaterManager: +@@ -333,6 +335,40 @@ class SchedulerWeightUpdaterManager: return ResumeMemoryOccupationReqOutput() @@ -137,7 +137,7 @@ index 505bd2994fbf4812c5069d4391cb80197d42f80c..2ed3b21d75c748a567478686d3351849 try: payload = self.tp_worker.model_runner.check_weights( diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py -index b733d20234ea1d3cb8b9bae77259229bea76a059..22495de1cc4f74db5332363dd990cd2b7acab1de 100644 +index efe7ba9d360d7ed5f08274269b0543f8d04af0ba..cb638169f156e96c416a87e0109d07a8c7ff3a3f 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -53,6 +53,8 @@ from sglang.srt.managers.io_struct import ( @@ -149,7 +149,7 @@ index b733d20234ea1d3cb8b9bae77259229bea76a059..22495de1cc4f74db5332363dd990cd2b ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, RemoveExternalCorpusReqInput, -@@ -102,6 +104,7 @@ _COMMUNICATOR_SPECS = [ +@@ -104,6 +106,7 @@ _COMMUNICATOR_SPECS = [ ("update_weights_from_tensor", UpdateWeightsFromTensorReqOutput), ("update_weights_from_ipc", UpdateWeightsFromIPCReqOutput), ("post_process_weights", PostProcessWeightsReqOutput), @@ -157,7 +157,7 @@ index b733d20234ea1d3cb8b9bae77259229bea76a059..22495de1cc4f74db5332363dd990cd2b ("get_weights_by_name", GetWeightsByNameReqOutput), ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), -@@ -760,6 +763,15 @@ class TokenizerControlMixin: +@@ -810,6 +813,15 @@ class TokenizerControlMixin: results = await self.post_process_weights_communicator(obj) return FanOutCommunicator.merge_results(results) @@ -174,20 +174,20 @@ index b733d20234ea1d3cb8b9bae77259229bea76a059..22495de1cc4f74db5332363dd990cd2b self: TokenizerManager, obj: CheckWeightsReqInput, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index 5af52950ff182997cf51badc33e53dc45475a93b..bd708d202f110af01d4a2ff583a8f153869a84c5 100644 +index 0361963fd2f5518bc436fc19741d7bd01f74cc80..9363e8c75f7892aed93e6ec8793ddbcda3dbcd7a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py -@@ -2376,6 +2376,10 @@ class ServerArgs: - nargs="*", +@@ -3096,6 +3096,10 @@ class ServerArgs: ), + NS("model"), ] = None + custom_pull_weights_pre_read_hook: A[ + Optional[str], + "Import path of a hook(source_dir, target_version) that /pull_weights calls before reading published weights.", + ] = None weight_loader_disable_mmap: A[ - bool, - "Disable mmap while loading weight using safetensors.", + bool, "Disable mmap while loading weight using safetensors.", NS("model") + ] = False diff --git a/python/sglang/srt/weight_sync/local_checkpoint.py b/python/sglang/srt/weight_sync/local_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..2fc0593931940342eebf849751ec518a26585f85 diff --git a/docker/patch/latest/sglang-release_hicache.patch b/docker/patch/latest/sglang-release_hicache.patch index c8f6f9473f..2e716f2f5a 100644 --- a/docker/patch/latest/sglang-release_hicache.patch +++ b/docker/patch/latest/sglang-release_hicache.patch @@ -1,11 +1,11 @@ diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py -index 40e50efa440f8ec48a1b8c30b412a8d3cf525ea3..505bd2994fbf4812c5069d4391cb80197d42f80c 100644 +index 2c876d270cb1a1cf5476b2ef6e79de6c9fb4431a..e2aea439e360a29c997f2cbd0ed565db8a0e3167 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py -@@ -206,10 +206,18 @@ class SchedulerWeightUpdaterManager: +@@ -224,10 +224,21 @@ class SchedulerWeightUpdaterManager: if tags is None or len(tags) == 0: tags = GPU_MEMORY_ALL_TYPES - + + tags_to_offload = [] for tag in tags: + if tag not in self.offload_tags and tag not in tags_to_offload: @@ -14,20 +14,19 @@ index 40e50efa440f8ec48a1b8c30b412a8d3cf525ea3..505bd2994fbf4812c5069d4391cb8019 + if not tags_to_offload: + return ReleaseMemoryOccupationReqOutput() + ++ if GPU_MEMORY_TYPE_WEIGHTS in tags_to_offload: ++ self._assert_weight_cache_inactive("release_memory_occupation") ++ + for tag in tags_to_offload: self.offload_tags.add(tag) - + - if GPU_MEMORY_TYPE_KV_CACHE in tags: + if GPU_MEMORY_TYPE_KV_CACHE in tags_to_offload: scheduler = self.scheduler if scheduler is not None: if scheduler.disaggregation_mode == DisaggregationMode.DECODE: -@@ -223,18 +231,24 @@ class SchedulerWeightUpdaterManager: - elif scheduler.disaggregation_mode == DisaggregationMode.PREFILL: - queue = getattr(scheduler, "disagg_prefill_bootstrap_queue", None) - if queue is not None: -- queue.release_memory_occupation() -+ queue.release_memory_occupation() +@@ -244,16 +255,21 @@ class SchedulerWeightUpdaterManager: + queue.release_memory_occupation() self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE) self.flush_cache() + if scheduler is not None and scheduler.server_args.release_hicache: @@ -36,24 +35,25 @@ index 40e50efa440f8ec48a1b8c30b412a8d3cf525ea3..505bd2994fbf4812c5069d4391cb8019 + ) + if release is not None: + release() - + - if GPU_MEMORY_TYPE_WEIGHTS in tags: +- self._assert_weight_cache_inactive("release_memory_occupation") + if GPU_MEMORY_TYPE_WEIGHTS in tags_to_offload: self.stashed_model_static_state = _export_static_state( self.tp_worker.model_runner.model ) torch.distributed.barrier(self.tp_cpu_group) self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_WEIGHTS) - + - if GPU_MEMORY_TYPE_CUDA_GRAPH in tags: + if GPU_MEMORY_TYPE_CUDA_GRAPH in tags_to_offload: self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_CUDA_GRAPH) - + torch.get_device_module().synchronize() -@@ -247,13 +261,21 @@ class SchedulerWeightUpdaterManager: +@@ -266,14 +282,24 @@ class SchedulerWeightUpdaterManager: if tags is None or len(tags) == 0: tags = GPU_MEMORY_ALL_TYPES - + + tags_to_resume = [] for tag in tags: + if tag in self.offload_tags and tag not in tags_to_resume: @@ -62,22 +62,26 @@ index 40e50efa440f8ec48a1b8c30b412a8d3cf525ea3..505bd2994fbf4812c5069d4391cb8019 + if not tags_to_resume: + return ResumeMemoryOccupationReqOutput() + ++ if GPU_MEMORY_TYPE_WEIGHTS in tags_to_resume: ++ self._assert_weight_cache_inactive("resume_memory_occupation") ++ + for tag in tags_to_resume: self.offload_tags.remove(tag) - + - if GPU_MEMORY_TYPE_CUDA_GRAPH in tags: + if GPU_MEMORY_TYPE_CUDA_GRAPH in tags_to_resume: self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_CUDA_GRAPH) - + - if GPU_MEMORY_TYPE_WEIGHTS in tags: +- self._assert_weight_cache_inactive("resume_memory_occupation") + if GPU_MEMORY_TYPE_WEIGHTS in tags_to_resume: self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS) torch.distributed.barrier(self.tp_cpu_group) _import_static_state( -@@ -262,10 +284,16 @@ class SchedulerWeightUpdaterManager: +@@ -282,10 +308,16 @@ class SchedulerWeightUpdaterManager: ) del self.stashed_model_static_state - + - if GPU_MEMORY_TYPE_KV_CACHE in tags: + if GPU_MEMORY_TYPE_KV_CACHE in tags_to_resume: self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_KV_CACHE) @@ -93,13 +97,13 @@ index 40e50efa440f8ec48a1b8c30b412a8d3cf525ea3..505bd2994fbf4812c5069d4391cb8019 for queue_name in ( "disagg_decode_transfer_queue", diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py -index 381af58be2a7e5922cb1e5d0d583dea182b1e56a..ecd143cc6e9d56bd9e5013f604449abc3b230496 100644 +index dbcec3547f5cf007cb820631f54229f3d38d0233..72e60d4e65860606567ba8db0c0ba639b81d28cb 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py -@@ -207,6 +207,19 @@ class HiMambaRadixCache(MambaRadixCache): - ) - super().reset() - +@@ -210,6 +210,19 @@ class HiMambaRadixCache(MambaRadixCache): + def release_host_resources(self) -> None: + self.host_pool_group.destroy() + + def release_memory_occupation(self) -> None: + if self.enable_storage: + logger.warning( @@ -117,13 +121,13 @@ index 381af58be2a7e5922cb1e5d0d583dea182b1e56a..ecd143cc6e9d56bd9e5013f604449abc # Backup invariant (for write-through mode): backed-up nodes must form a # contiguous prefix from root — no gaps. Skip if parent isn't backed diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py -index 378a29e936e6c1ad3213f710cbf5bd9380bd019f..7733c5a7583bfa55ea03b435976562b169c4d9a8 100644 +index e9c2d85daa46202d5f3e3b2ab7a61ac56faba9be..5c7ca9cf5edd0d3eb676aa862ff34902718e0e0e 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py -@@ -730,6 +730,19 @@ class HiRadixCache(RadixCache): - self.evictable_host_leaves.clear() - super().reset() - +@@ -786,6 +786,19 @@ class HiRadixCache(RadixCache): + if self.token_to_kv_pool_host is not None: + self.token_to_kv_pool_host.destroy() + + def release_memory_occupation(self) -> None: + if self.enable_storage: + logger.warning( @@ -141,63 +145,24 @@ index 378a29e936e6c1ad3213f710cbf5bd9380bd019f..7733c5a7583bfa55ea03b435976562b1 height = 0 while node != self.root_node: diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py -index 28c97f0cd1092e5478282a5cdf3961a5ca1b341d..110680f950fd7f7c7a0d3243523afd87db507313 100644 +index 2670de52c5260fce0aeb57aaff4f75b3b13c7c64..d907d07447f8827149a8853fb78566c27716f61b 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py -@@ -114,7 +114,9 @@ class MHATokenToKVPoolHost(HostKVCache): - self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel( - element_size=self.element_dim * self.dtype.itemsize - ) -+ self._post_init_host_buffers() - -+ def _post_init_host_buffers(self) -> None: - if self.layout == "page_first": - # Transpose [page, layer, ...] -> [layer, page, ...] to get per-layer views - # This swaps strides without copying data -@@ -686,6 +688,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache): - - self.init_kv_buffer() - self.lock = threading.RLock() -+ self._host_memory_released = False - self.clear() - - self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel( -@@ -696,6 +699,15 @@ class MHATokenToKOnlyPoolHost(HostKVCache): - dtype=torch.uint64, - device=self.device_pool.device, - ) -+ self._post_init_host_buffers() -+ -+ def _host_buffer_attr_names(self): -+ return ("k_buffer",) -+ -+ def _init_host_buffers(self) -> None: -+ self.init_kv_buffer() -+ -+ def _post_init_host_buffers(self) -> None: - if self.layout == "page_first": - transposed = self.k_buffer.transpose(0, 1) - self.k_data_refs = [transposed[i] for i in range(self.layer_num)] -@@ -1285,7 +1297,12 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): - self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel( - element_size=self.kv_cache_dim * self.dtype.itemsize - ) -+ self._post_init_host_buffers() - -+ def _host_buffer_attr_names(self): -+ return ("kv_buffer", "k_buffer", "v_buffer", "index_k_buffer") -+ -+ def _post_init_host_buffers(self) -> None: - if self.layout == "page_first": - # Transpose [page, layer, ...] -> [layer, page, ...] to get per-layer views - # This swaps strides without copying data -@@ -1801,8 +1818,18 @@ class MambaPoolHost(HostKVCache): +@@ -86,6 +86,7 @@ class MambaPoolHost(HostKVCache): + self.layout = layout + self.pin_memory = pin_memory + self.device = device ++ self.pool_label = "mamba" + self.allocator = get_allocator_from_storage(allocator_type) + self.num_mamba_layers = device_pool.num_mamba_layers + +@@ -157,8 +158,18 @@ class MambaPoolHost(HostKVCache): self.init_kv_buffer() self._init_write_back_staging_buffers() self.lock = threading.RLock() + self._host_memory_released = False self.clear() - + + def _host_buffer_attr_names(self): + return ("temporal_buffer", "conv_buffer") + @@ -208,47 +173,51 @@ index 28c97f0cd1092e5478282a5cdf3961a5ca1b341d..110680f950fd7f7c7a0d3243523afd87 + self._init_write_back_staging_buffers() + def init_kv_buffer(self): - alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] - -@@ -2361,6 +2388,12 @@ class LogicalHostPool: - def get_ksize_per_token(self): - return 0 - + _host_alloc = ALLOC_MEMORY_FUNCS[self.device_pool.device] + +@@ -714,6 +725,12 @@ class LogicalHostPool: + def get_dummy_flat_data_page(self): + return torch.empty(0, dtype=torch.uint8) + + def release_memory_occupation(self) -> None: + pass + + def resume_memory_occupation(self) -> None: + pass + - - class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): - """Host mirror for a DeepSeek V4 paged KV/indexer sub-pool.""" -@@ -3183,6 +3216,18 @@ class HostPoolGroup: + def set_from_flat_data_page(self, index, data_page): + pass + +@@ -1578,6 +1595,14 @@ class HostPoolGroup: for entry in self.entries: entry.host_pool.clear() - + + def release_memory_occupation(self) -> None: + for entry in self.entries: -+ release = getattr(entry.host_pool, "release_memory_occupation", None) -+ if release is not None: -+ release() ++ entry.host_pool.release_memory_occupation() + + def resume_memory_occupation(self) -> None: + for entry in self.entries: -+ resume = getattr(entry.host_pool, "resume_memory_occupation", None) -+ if resume is not None: -+ resume() -+ - def available_size(self): - return self.anchor_entry.host_pool.available_size() - -@@ -3330,8 +3375,26 @@ class DSAIndexerPoolHost(HostKVCache): ++ entry.host_pool.resume_memory_occupation() ++ + def destroy(self) -> None: + for entry in self.entries: + entry.host_pool.destroy() +@@ -1728,6 +1753,7 @@ class DSAIndexerPoolHost(HostKVCache): + self.layout = layout + self.pin_memory = pin_memory + self.device = device ++ self.pool_label = "dsa-indexer" + self.allocator = get_allocator_from_storage(allocator_type) + self.dtype = device_pool.store_dtype + self.start_layer = device_pool.start_layer +@@ -1773,8 +1799,26 @@ class DSAIndexerPoolHost(HostKVCache): self.can_use_write_back_jit = False self._init_write_back_staging_buffers() self.lock = threading.RLock() + self._host_memory_released = False self.clear() - + + def _requested_host_memory_bytes(self) -> int: + return ( + self.page_num @@ -270,12 +239,12 @@ index 28c97f0cd1092e5478282a5cdf3961a5ca1b341d..110680f950fd7f7c7a0d3243523afd87 return ( self.indexer_size_per_token * self.layer_num * self.indexer_dtype.itemsize diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py -index 39b809af2b378439d0868b52b9e1272e6dd05875..ade3cdc3965e62e8af1cd1c985d5d8a240f8cfc0 100644 +index 57dc8c1a2a2079f54938d841fa3c988f36c0b641..2255d05d5df75d050de778762d36a7c2eb4ca3f5 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/python/sglang/srt/mem_cache/pool_host/base.py @@ -1,6 +1,9 @@ from __future__ import annotations - + import abc +import ctypes +import ctypes.util @@ -283,10 +252,10 @@ index 39b809af2b378439d0868b52b9e1272e6dd05875..ade3cdc3965e62e8af1cd1c985d5d8a2 import logging import threading from functools import wraps -@@ -24,6 +27,27 @@ _is_hip = is_hip() - # Host RAM to leave free when sizing HiCache pools (OS, other processes). - HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3) - +@@ -26,6 +29,27 @@ HICACHE_HOST_MEMORY_RESERVE_BYTES: int = 10 * (1024**3) + + _WRITE_BACK_STAGING_PAGE_CHUNK = 64 + +try: + _libc = ctypes.CDLL(ctypes.util.find_library("c") or "libc.so.6") + _libc.malloc_trim.argtypes = [ctypes.c_size_t] @@ -308,13 +277,13 @@ index 39b809af2b378439d0868b52b9e1272e6dd05875..ade3cdc3965e62e8af1cd1c985d5d8a2 + if _libc is not None: + _libc.malloc_trim(0) + - + def sync_fixed_hicache_size(size: int, host_size: int) -> int: """Sync fixed-size HiCache token capacity across PP ranks. -@@ -120,9 +144,21 @@ class HostKVCache(abc.ABC): +@@ -138,9 +162,21 @@ class HostKVCache(abc.ABC): device_pool.size, ) - + - # Verify there is enough available host memory. + self._check_host_memory_available() + @@ -322,7 +291,7 @@ index 39b809af2b378439d0868b52b9e1272e6dd05875..ade3cdc3965e62e8af1cd1c985d5d8a2 + self._host_memory_released = False + self._init_host_buffers() + -+ # A lock for synchronized operations on memory allocation and state transitions. ++ # Initialize allocation state after all host buffers are ready. + self.clear() + + def _requested_host_memory_bytes(self) -> int: @@ -335,21 +304,28 @@ index 39b809af2b378439d0868b52b9e1272e6dd05875..ade3cdc3965e62e8af1cd1c985d5d8a2 available_bytes = host_mem.available - HICACHE_HOST_MEMORY_RESERVE_BYTES if requested_bytes > available_bytes: raise ValueError( -@@ -131,15 +167,73 @@ class HostKVCache(abc.ABC): +@@ -149,19 +185,80 @@ class HostKVCache(abc.ABC): f"{available_bytes / 1e9:.2f} GB free. Please reduce the " f"size of the hierarchical cache." ) - else: - logger.info( -- f"Allocating {requested_bytes / 1e9:.2f} GB host memory for hierarchical KV cache." +- "Allocating %s hierarchical KV host pool: %d tokens, %.2f GB host memory.", +- pool_label, +- self.size, +- requested_bytes / 1e9, - ) + logger.info( -+ f"Allocating {requested_bytes / 1e9:.2f} GB host memory for hierarchical KV cache." ++ "Allocating %s hierarchical KV host pool: %d tokens, %.2f GB host memory.", ++ self.pool_label, ++ self.size, ++ requested_bytes / 1e9, + ) - + + def _init_host_buffers(self) -> None: self.kv_buffer = self.init_kv_buffer() - + self.fd = getattr(self.allocator, "fd", None) + - # A lock for synchronized operations on memory allocation and state transitions. - self.lock = threading.RLock() + def _post_init_host_buffers(self) -> None: @@ -368,6 +344,11 @@ index 39b809af2b378439d0868b52b9e1272e6dd05875..ade3cdc3965e62e8af1cd1c985d5d8a2 + "data_ptrs", + "index_k_data_refs", + "index_k_data_ptrs", ++ "staging_k_buffer", ++ "staging_v_buffer", ++ "staging_buffer", ++ "temporal_staging_buffer", ++ "conv_staging_buffers", + ) + + def _release_host_buffers(self) -> None: @@ -387,14 +368,14 @@ index 39b809af2b378439d0868b52b9e1272e6dd05875..ade3cdc3965e62e8af1cd1c985d5d8a2 + + @synchronized + def release_memory_occupation(self) -> None: -+ if not hasattr(self, "_host_memory_released"): -+ return -+ if self._host_memory_released: ++ if not hasattr(self, "_host_memory_released") or self._host_memory_released: + return + self._release_host_buffers() + self.mem_state = torch.empty((0,), dtype=torch.uint8, device=self.device) + self.free_slots = torch.empty((0,), dtype=torch.int64) + self.slot_used = torch.empty((0,), dtype=torch.bool) ++ self.release_slots = [] ++ self.num_release_slots = 0 + self._host_memory_released = True + _trim_host_allocator() + logger.info( @@ -404,25 +385,102 @@ index 39b809af2b378439d0868b52b9e1272e6dd05875..ade3cdc3965e62e8af1cd1c985d5d8a2 + + @synchronized + def resume_memory_occupation(self) -> None: -+ if not hasattr(self, "_host_memory_released"): -+ return -+ if not self._host_memory_released: ++ if not hasattr(self, "_host_memory_released") or not self._host_memory_released: + return + self._check_host_memory_available() + self._init_host_buffers() + self._post_init_host_buffers() + self._host_memory_released = False self.clear() - + def destroy(self): +@@ -175,14 +272,7 @@ class HostKVCache(abc.ABC): + if getattr(self, "_destroyed", False): + return + self._destroyed = True +- buffers = getattr(self, "kv_buffer", None) +- if buffers is not None and self.pin_memory and (_is_cuda or _is_hip): +- if not isinstance(buffers, (list, tuple)): +- buffers = [buffers] +- for buf in buffers: +- if buf is not None: +- _cuda_host_unregister(buf) +- self.kv_buffer = None ++ self._release_host_buffers() + + @abc.abstractmethod + def get_size_per_token(self): +diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py +index 150c37b726285d9a5abcc7ac446e11882f46395b..c51b8e3daf0ac67d8a46da2152c2d8ed27a29cea 100644 +--- a/python/sglang/srt/mem_cache/pool_host/mha.py ++++ b/python/sglang/srt/mem_cache/pool_host/mha.py +@@ -101,7 +101,9 @@ class MHATokenToKVPoolHost(HostKVCache): + self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel( + element_size=self.element_dim * self.dtype.itemsize + ) ++ self._post_init_host_buffers() + ++ def _post_init_host_buffers(self) -> None: + if self.layout == "page_first": + # Transpose [page, layer, ...] -> [layer, page, ...] to get per-layer views + # This swaps strides without copying data +@@ -644,6 +646,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache): + self.layout = layout + self.pin_memory = pin_memory + self.device = device ++ self.pool_label = "index-k" + self.allocator = get_allocator_from_storage(allocator_type) + self.dtype = device_pool.store_dtype + self.start_layer = device_pool.start_layer +@@ -677,6 +680,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache): + + self.init_kv_buffer() + self.lock = threading.RLock() ++ self._host_memory_released = False + self.clear() + + self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel( +@@ -687,6 +691,15 @@ class MHATokenToKOnlyPoolHost(HostKVCache): + dtype=torch.uint64, + device=self.device_pool.device, + ) ++ self._post_init_host_buffers() ++ ++ def _host_buffer_attr_names(self): ++ return ("k_buffer",) ++ ++ def _init_host_buffers(self) -> None: ++ self.init_kv_buffer() ++ ++ def _post_init_host_buffers(self) -> None: + if self.layout == "page_first": + transposed = self.k_buffer.transpose(0, 1) + self.k_data_refs = [transposed[i] for i in range(self.layer_num)] +diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py +index e311520722dc7fc725c1c2b9c359df5edf1e687f..67b6f323a82096d00156161a6cce68e1774dbc6b 100644 +--- a/python/sglang/srt/mem_cache/pool_host/mla.py ++++ b/python/sglang/srt/mem_cache/pool_host/mla.py +@@ -88,7 +88,12 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): + self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel( + element_size=self.kv_cache_dim * self.dtype.itemsize + ) ++ self._post_init_host_buffers() + ++ def _host_buffer_attr_names(self): ++ return ("kv_buffer", "k_buffer", "v_buffer", "index_k_buffer") ++ ++ def _post_init_host_buffers(self) -> None: + if self.layout == "page_first": + # Transpose [page, layer, ...] -> [layer, page, ...] to get per-layer views + # This swaps strides without copying data diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py -index 5fd3c95ef4e534d08876f5afc7145fda25f02011..9c309b7d2b12c68f1f70f920aad5d4111ee1baa9 100644 +index c11a02e30f68e7b81834a41e9af5838da461e684..74a38bc9581e63ec50d8c5adefb8149c111428f6 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py -@@ -487,6 +487,21 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): - ) - self._record_all_cleared_event() - +@@ -305,6 +305,21 @@ class UnifiedRadixCache(BasePrefixCache): + + self.tree_core._record_all_cleared_event() + + def release_memory_occupation(self) -> None: + if self.enable_storage: + logger.warning( @@ -442,17 +500,17 @@ index 5fd3c95ef4e534d08876f5afc7145fda25f02011..9c309b7d2b12c68f1f70f920aad5d411 """Initialize HiCache infrastructure.""" from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index 3049e6d85b2ec11bf9258609886baa9b9b19abfb..5af52950ff182997cf51badc33e53dc45475a93b 100644 +index 18664959e959affb24db6b1669234e651535864d..0361963fd2f5518bc436fc19741d7bd01f74cc80 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py -@@ -2052,6 +2052,10 @@ class ServerArgs: - Optional[str], +@@ -2669,6 +2669,10 @@ class ServerArgs: "A dictionary in JSON string format, or a string starting with a leading '@' and a config file in JSON/YAML/TOML format, containing extra configuration for the storage backend.", + NS("memory"), ] = None + release_hicache: A[ + bool, + "Release hierarchical cache host memory during release_memory_occupation.", + ] = False - + # ------------------------------------------------------------------------- # Hierarchical sparse attention diff --git a/docker/patch/latest/sglang-top_p.patch b/docker/patch/latest/sglang-top_p.patch index ae82decdd8..6556fb9d94 100644 --- a/docker/patch/latest/sglang-top_p.patch +++ b/docker/patch/latest/sglang-top_p.patch @@ -1,49 +1,60 @@ diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py -index 50e13ffa5dbbe8cfc72e79edfc2881dad1357442..1a96a2378bb5aa5e55512ec0721dab642d8484e0 100644 +index 56bc3d42f6..3c0eb6fa84 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py -@@ -1550,6 +1550,8 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): +@@ -782,9 +782,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): + continue + + if poll == KVPoll.Bootstrapping: +- entry_time = ( +- decode_req.req.time_stats.decode_prealloc_queue_entry_time +- ) ++ entry_time = decode_req.req.time_stats.decode_prealloc_queue_entry_time + if entry_time > 0 and now - entry_time > bootstrap_timeout: + error_message = ( + f"Decode prealloc timeout for request rank={self.tp_rank} " +@@ -1831,6 +1829,8 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): output_token_logprobs_idx, output_top_logprobs_val, output_top_logprobs_idx, + output_top_p_token_ids_len, + output_top_p_token_ids, - output_topk_p, - output_topk_index, - output_hidden_states, -@@ -1653,6 +1655,11 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): + output_token_sampling_mask_len, + output_token_sampling_mask_idx, + output_token_sampling_logprobs, +@@ -1955,6 +1955,11 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): : decode_req.req.logprob.top_logprobs_num ].tolist() ) -+ top_p_token_ids_len = output_top_p_token_ids_len[0].item() ++ top_p_token_ids_len = int(output_top_p_token_ids_len[0].item()) + if top_p_token_ids_len > 0: + decode_req.req.logprob.output_top_p_token_ids.append( + output_top_p_token_ids[:top_p_token_ids_len].tolist() + ) - - decode_req.kv_receiver.clear() - decode_req.kv_receiver = None + if decode_req.req.return_sampling_mask: + assert ( + output_token_sampling_mask_idx is not None diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py -index e6ca3bd580ef8d295b5abb6919ff59b8f2a48083..1d4f85ef99e668bbe32b2c19915ebc0cf474d399 100644 +index 237f293ec5..d2e0604253 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -1,5 +1,6 @@ from __future__ import annotations - + +import logging import os import random from collections import deque -@@ -32,6 +33,8 @@ if TYPE_CHECKING: +@@ -46,6 +47,8 @@ if is_npu(): ######################### FAKE_BOOTSTRAP_HOST = "2.2.2.2" _IS_HIP = is_hip() +MAX_PD_TOP_P_TOKEN_IDS = 4096 +logger = logging.getLogger(__name__) - - - def get_dsa_seed_metadata_dim(hf_config) -> int: -@@ -276,6 +279,12 @@ class MetadataBuffers: + + + def poll_and_all_reduce_pp( +@@ -314,6 +317,12 @@ class MetadataBuffers: self.output_top_logprobs_idx = torch.zeros( (size, max_top_logprobs_num), dtype=torch.int32, device=device ) @@ -53,46 +64,28 @@ index e6ca3bd580ef8d295b5abb6919ff59b8f2a48083..1d4f85ef99e668bbe32b2c19915ebc0c + self.output_top_p_token_ids = torch.zeros( + (size, MAX_PD_TOP_P_TOKEN_IDS), dtype=torch.int32, device=device + ) - # For PD + spec decode - self.output_topk_p = torch.zeros( - (size, 16), dtype=torch.float32, device=device -@@ -308,6 +317,8 @@ class MetadataBuffers: - self.output_token_logprobs_idx.data_ptr(), - self.output_top_logprobs_val.data_ptr(), - self.output_top_logprobs_idx.data_ptr(), -+ self.output_top_p_token_ids_len.data_ptr(), -+ self.output_top_p_token_ids.data_ptr(), - self.output_topk_p.data_ptr(), - self.output_topk_index.data_ptr(), - self.output_hidden_states.data_ptr(), -@@ -320,6 +331,8 @@ class MetadataBuffers: - self.output_token_logprobs_idx.nbytes, - self.output_top_logprobs_val.nbytes, - self.output_top_logprobs_idx.nbytes, -+ self.output_top_p_token_ids_len.nbytes, -+ self.output_top_p_token_ids.nbytes, - self.output_topk_p.nbytes, - self.output_topk_index.nbytes, - self.output_hidden_states.nbytes, -@@ -332,6 +345,8 @@ class MetadataBuffers: - self.output_token_logprobs_idx[0].nbytes, - self.output_top_logprobs_val[0].nbytes, - self.output_top_logprobs_idx[0].nbytes, -+ self.output_top_p_token_ids_len[0].nbytes, -+ self.output_top_p_token_ids[0].nbytes, - self.output_topk_p[0].nbytes, - self.output_topk_index[0].nbytes, - self.output_hidden_states[0].nbytes, -@@ -351,6 +366,8 @@ class MetadataBuffers: + self.output_token_sampling_mask_len = None + self.output_token_sampling_mask_idx = None + self.output_token_sampling_logprobs = None +@@ -359,6 +368,8 @@ class MetadataBuffers: + self.output_token_logprobs_idx, + self.output_top_logprobs_val, + self.output_top_logprobs_idx, ++ self.output_top_p_token_ids_len, ++ self.output_top_p_token_ids, + ] + if self.enable_sampling_mask: + bufs.extend( +@@ -398,6 +409,8 @@ class MetadataBuffers: self.output_token_logprobs_idx[idx].clone(), self.output_top_logprobs_val[idx].clone(), self.output_top_logprobs_idx[idx].clone(), + self.output_top_p_token_ids_len[idx].clone(), + self.output_top_p_token_ids[idx].clone(), - self.output_topk_p[idx].clone(), - self.output_topk_index[idx].clone(), - self.output_hidden_states[idx].clone(), -@@ -383,6 +400,7 @@ class MetadataBuffers: + sampling_mask_len, + sampling_mask_idx, + sampling_logprobs, +@@ -433,6 +446,7 @@ class MetadataBuffers: self.cached_tokens[req.metadata_buffer_index][4] = image_t self.cached_tokens[req.metadata_buffer_index][5] = audio_t self.cached_tokens[req.metadata_buffer_index][6] = video_t @@ -100,225 +93,63 @@ index e6ca3bd580ef8d295b5abb6919ff59b8f2a48083..1d4f85ef99e668bbe32b2c19915ebc0c if req.return_logprob: if req.logprob.output_token_logprobs_val: # not none or empty list self.output_token_logprobs_val[req.metadata_buffer_index][0] = ( -@@ -409,6 +427,28 @@ class MetadataBuffers: +@@ -467,6 +481,27 @@ class MetadataBuffers: dtype=torch.int32, device="cpu", ) -+ if req.logprob.output_top_p_token_ids: # not none or empty list ++ if req.logprob.output_top_p_token_ids: + output_top_p_token_ids = req.logprob.output_top_p_token_ids[0] + if len(output_top_p_token_ids) > MAX_PD_TOP_P_TOKEN_IDS: + logger.warning( -+ "PD top-p token replay payload for the first output token " -+ "has %s ids, exceeding the metadata buffer cap %s. " -+ "Falling back to the sampled token only.", ++ "PD top-p replay payload has %s ids, exceeding cap %s; " ++ "falling back to the sampled token.", + len(output_top_p_token_ids), + MAX_PD_TOP_P_TOKEN_IDS, + ) + output_top_p_token_ids = [int(req.output_ids[0])] + top_p_len = len(output_top_p_token_ids) -+ self.output_top_p_token_ids_len[req.metadata_buffer_index][0] = ( -+ top_p_len -+ ) -+ self.output_top_p_token_ids[req.metadata_buffer_index][ -+ :top_p_len -+ ] = torch.tensor( -+ output_top_p_token_ids, -+ dtype=torch.int32, -+ device="cpu", ++ self.output_top_p_token_ids_len[req.metadata_buffer_index][ ++ 0 ++ ] = top_p_len ++ self.output_top_p_token_ids[req.metadata_buffer_index][:top_p_len] = ( ++ torch.tensor( ++ output_top_p_token_ids, ++ dtype=torch.int32, ++ device="cpu", ++ ) + ) - # For PD + spec decode - if req.hidden_states_tensor is not None: - # speculative_eagle_topk should not be greater than 16 currently + if req.return_sampling_mask: + if not self.enable_sampling_mask: + raise RuntimeError( diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py -index 45dc18463eed570a7c09029b966513eee0e584e7..ec107713dd87b5698fbdd0edeb8795b63f9c00b8 100644 +index 480ed81a5e..de5b0bcc8d 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py -@@ -173,6 +173,8 @@ class LogitsProcessorOutput: +@@ -176,6 +176,8 @@ class LogitsProcessorOutput: List[Union[List[float], torch.Tensor]] ] = None next_token_token_ids_logprobs_idx: Optional[List] = None -+ # The kept token ids used by rollout top-p replay. One tensor/list per sampled row. ++ # Kept token ids used by the legacy rollout top-p replay protocol. + next_token_top_p_token_ids: Optional[List[Optional[torch.Tensor]]] = None - - ## Part 3: Prefill-only. This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor - # The logprobs of input tokens. shape: [#token] -diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py -index e83e4157faf6a9d316995ac721be5d6bb22a2d34..2276a73b2e812063b6b2018fda1be606b5655878 100644 ---- a/python/sglang/srt/layers/sampler.py -+++ b/python/sglang/srt/layers/sampler.py -@@ -12,7 +12,12 @@ from sglang.srt.layers.dp_attention import ( - ) - from sglang.srt.layers.logits_processor import LogitsProcessorOutput - from sglang.srt.layers.utils.hash import murmur_hash32 --from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logprobs -+from sglang.srt.layers.utils.logprob import ( -+ get_token_ids_logprobs, -+ get_top_logprobs, -+ get_top_p_token_ids_from_probs, -+ renorm_logprob_over_top_p, -+) - from sglang.srt.runtime_context import get_flags - from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo - from sglang.srt.sampling.sampling_params import TOP_K_ALL -@@ -127,6 +132,9 @@ class Sampler(nn.Module): - _aiter_greedy_sample(batch_next_token_ids, logits) - else: - batch_next_token_ids = torch.argmax(logits, -1) -+ self._attach_greedy_top_p_token_ids_to_output( -+ logits_output, batch_next_token_ids, sampling_info -+ ) - if return_logprob: - original_logprobs = logprobs = torch.nn.functional.log_softmax( - logits, dim=-1 -@@ -157,6 +165,7 @@ class Sampler(nn.Module): - if self.use_ascend_backend: - # Ascend backend: sample from logits directly. - batch_next_token_ids, logprobs = self._forward_ascend_backend( -+ logits_output, - logits, - sampling_info, - simple_sampling_case, -@@ -184,15 +193,45 @@ class Sampler(nn.Module): - logits[:] = torch.softmax(logits, dim=-1) - probs = logits - -+ self._attach_top_p_token_ids_to_output( -+ logits_output, -+ probs, -+ sampling_info, -+ simple_sampling_case, -+ ) - batch_next_token_ids = self._sample_from_probs( - probs, sampling_info, positions, simple_sampling_case - ) - if return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB: -- logprobs = ( -- logprobs_via_logsoftmax_kernel -- if logprobs_via_logsoftmax_kernel is not None -- else torch.log(probs) -- ) -+ top_p_logprobs = None -+ if sampling_info.need_return_top_p_token_ids: -+ # Force-keep the sampled token in the renorm denominator: -+ # SGLang samples with the flashinfer kernel but -+ # renorm_logprob_over_top_p computes the nucleus with a -+ # torch reimplementation. The two can disagree at the -+ # nucleus boundary, so a sampled token may fall outside -+ # the torch nucleus and get a -inf renormalized logprob -+ # (-> NaN downstream). Force-keeping it makes the -+ # denominator ``nucleus ∪ {sampled}``, matching the -+ # trainer which also force-keeps the target token. -+ top_p_logprobs = renorm_logprob_over_top_p( -+ probs=probs, -+ top_ks=sampling_info.top_ks, -+ top_ps=sampling_info.top_ps, -+ min_ps=sampling_info.min_ps, -+ need_top_p_sampling=sampling_info.need_top_p_sampling, -+ need_min_p_sampling=sampling_info.need_min_p_sampling, -+ request_mask=sampling_info.return_top_p_token_ids, -+ force_keep_token_ids=batch_next_token_ids, -+ ) -+ if top_p_logprobs is not None: -+ logprobs = top_p_logprobs -+ else: -+ logprobs = ( -+ logprobs_via_logsoftmax_kernel -+ if logprobs_via_logsoftmax_kernel is not None -+ else torch.log(probs) -+ ) - del probs - - # Attach logprobs to logits_output (in-place modification) -@@ -321,6 +360,7 @@ class Sampler(nn.Module): - - def _forward_ascend_backend( - self, -+ logits_output: LogitsProcessorOutput, - logits: torch.Tensor, - sampling_info: SamplingBatchInfo, - simple_sampling_case: bool, -@@ -337,6 +377,15 @@ class Sampler(nn.Module): - when return_logprob is False or SGLANG_RETURN_ORIGINAL_LOGPROB is set. - """ - logits.div_(sampling_info.temperatures) -+ if sampling_info.need_return_top_p_token_ids and not simple_sampling_case: -+ probs = torch.softmax(logits, dim=-1) -+ self._attach_top_p_token_ids_to_output( -+ logits_output, -+ probs, -+ sampling_info, -+ simple_sampling_case, -+ ) -+ del probs - batch_next_token_ids = self._sample_from_logits( - logits, sampling_info, simple_sampling_case, positions - ) -@@ -345,6 +394,49 @@ class Sampler(nn.Module): - logprobs = torch.log_softmax(logits, dim=-1) - return batch_next_token_ids, logprobs - -+ def _attach_greedy_top_p_token_ids_to_output( -+ self, -+ logits_output: LogitsProcessorOutput, -+ batch_next_token_ids: torch.Tensor, -+ sampling_info: SamplingBatchInfo, -+ ) -> None: -+ if not sampling_info.need_return_top_p_token_ids: -+ return -+ -+ request_mask = sampling_info.return_top_p_token_ids -+ logits_output.next_token_top_p_token_ids = [ -+ batch_next_token_ids[i : i + 1].to(torch.int32) -+ if bool(request_mask[i].item()) -+ else None -+ for i in range(len(batch_next_token_ids)) -+ ] -+ -+ def _attach_top_p_token_ids_to_output( -+ self, -+ logits_output: LogitsProcessorOutput, -+ probs: torch.Tensor, -+ sampling_info: SamplingBatchInfo, -+ simple_sampling_case: bool, -+ ) -> None: -+ if ( -+ not sampling_info.need_return_top_p_token_ids -+ or sampling_info.return_top_p_token_ids is None -+ or simple_sampling_case -+ ): -+ return -+ -+ top_p_token_ids = get_top_p_token_ids_from_probs( -+ probs=probs, -+ top_ks=sampling_info.top_ks, -+ top_ps=sampling_info.top_ps, -+ min_ps=sampling_info.min_ps, -+ need_top_p_sampling=sampling_info.need_top_p_sampling, -+ need_min_p_sampling=sampling_info.need_min_p_sampling, -+ request_mask=sampling_info.return_top_p_token_ids, -+ ) -+ if top_p_token_ids is not None: -+ logits_output.next_token_top_p_token_ids = top_p_token_ids -+ - def _attach_logprobs_to_output( - self, - logits_output: LogitsProcessorOutput, -diff --git a/python/sglang/srt/layers/utils/logprob.py b/python/sglang/srt/layers/utils/logprob.py -index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758af1b28503 100644 ---- a/python/sglang/srt/layers/utils/logprob.py -+++ b/python/sglang/srt/layers/utils/logprob.py -@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, List, Optional + # Sparse top-k/top-p/min-p support ids and selected-token logprob after + # truncation/renormalization. Only populated when requested. + next_token_sampling_mask_idx: Optional[List[Optional[List[int]]]] = None +diff --git a/python/sglang/srt/layers/logprob_processor.py b/python/sglang/srt/layers/logprob_processor.py +index 5ea86a6a5b..98b1a283ad 100644 +--- a/python/sglang/srt/layers/logprob_processor.py ++++ b/python/sglang/srt/layers/logprob_processor.py +@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Tuple import torch - + from sglang.srt.environ import envs +from sglang.srt.sampling.sampling_params import TOP_K_ALL - + if TYPE_CHECKING: - from sglang.srt.layers.logits_processor import LogitsMetadata -@@ -85,6 +86,110 @@ def get_top_logprobs( - ) - - + from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput +@@ -301,11 +302,120 @@ def get_token_ids_logprobs_chunk( + return next_split_pruned_len + + +def _top_p_filter_rows( + top_ks: torch.Tensor, + top_ps: torch.Tensor, @@ -327,12 +158,11 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a + need_min_p_sampling: bool, + request_mask: torch.Tensor, +) -> torch.Tensor: -+ """Rows that were requested AND actually have a top-k/top-p/min-p filter.""" + row_has_filter = top_ks != TOP_K_ALL + if need_top_p_sampling: -+ row_has_filter = row_has_filter | (top_ps != 1.0) ++ row_has_filter |= top_ps != 1.0 + if need_min_p_sampling: -+ row_has_filter = row_has_filter | (min_ps > 0) ++ row_has_filter |= min_ps > 0 + return request_mask & row_has_filter + + @@ -344,13 +174,8 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a + need_top_p_sampling: bool, + need_min_p_sampling: bool, +) -> tuple[torch.Tensor, torch.Tensor]: -+ """Boolean nucleus keep-mask in descending-prob order, plus the sort indices. -+ -+ Reproduces SGLang's sampler truncation (rank < top_k, cumulative prob within -+ top_p, prob >= top1 * min_p) so replay sees the exact set the sampler keeps. -+ """ + probs_sort, probs_idx = probs.sort(dim=-1, descending=True) -+ ranks = torch.arange(probs_sort.shape[-1], device=probs_sort.device).view(1, -1) ++ ranks = torch.arange(probs_sort.shape[-1], device=probs.device).view(1, -1) + keep = ranks < top_ks.view(-1, 1) + if need_top_p_sampling: + keep &= (torch.cumsum(probs_sort, dim=-1) - probs_sort) <= top_ps.view(-1, 1) @@ -370,32 +195,32 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a + force_keep_token_ids: Optional[torch.Tensor] = None, +) -> Optional[torch.Tensor]: + rows = _top_p_filter_rows( -+ top_ks, top_ps, min_ps, need_top_p_sampling, need_min_p_sampling, request_mask ++ top_ks, ++ top_ps, ++ min_ps, ++ need_top_p_sampling, ++ need_min_p_sampling, ++ request_mask, + ) + if not bool(rows.any().item()): + return None + + keep, probs_idx = _top_p_keep_mask_sorted( -+ probs, top_ks, top_ps, min_ps, need_top_p_sampling, need_min_p_sampling ++ probs, ++ top_ks, ++ top_ps, ++ min_ps, ++ need_top_p_sampling, ++ need_min_p_sampling, + ) -+ # Scatter the keep-mask back to vocab order so we renormalize directly over -+ # vocab ids (and can force-keep specific token ids). + keep_vocab = torch.empty_like(keep) + keep_vocab.scatter_(-1, probs_idx, keep) -+ + if force_keep_token_ids is not None: -+ # Force-keep the sampled/accepted token so its renormalized logprob is -+ # finite even when SGLang's sampling kernel (e.g. flashinfer) keeps a -+ # boundary token that this torch nucleus drops. This matches the trainer, -+ # which also force-keeps the target token before renormalizing, so the -+ # rollout and training denominators are both ``nucleus ∪ {token}``. -+ # Non-filter rows are overwritten by the ``torch.where`` below, so -+ # force-keeping every row is harmless and avoids a row gather. -+ row_idx = torch.arange(keep_vocab.shape[0], device=keep_vocab.device) -+ keep_vocab[row_idx, force_keep_token_ids] = True ++ rows_idx = torch.arange(keep_vocab.shape[0], device=probs.device) ++ keep_vocab[rows_idx, force_keep_token_ids] = True + + kept_probs = probs * keep_vocab -+ kept_probs = kept_probs / kept_probs.sum(dim=-1, keepdim=True).clamp_min(1e-12) ++ kept_probs /= kept_probs.sum(dim=-1, keepdim=True).clamp_min(1e-12) + return torch.where(rows.view(-1, 1), torch.log(kept_probs), torch.log(probs)) + + @@ -409,13 +234,23 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a + request_mask: torch.Tensor, +) -> Optional[List[Optional[torch.Tensor]]]: + rows = _top_p_filter_rows( -+ top_ks, top_ps, min_ps, need_top_p_sampling, need_min_p_sampling, request_mask ++ top_ks, ++ top_ps, ++ min_ps, ++ need_top_p_sampling, ++ need_min_p_sampling, ++ request_mask, + ) + if not bool(rows.any().item()): + return None + + keep, probs_idx = _top_p_keep_mask_sorted( -+ probs, top_ks, top_ps, min_ps, need_top_p_sampling, need_min_p_sampling ++ probs, ++ top_ks, ++ top_ps, ++ min_ps, ++ need_top_p_sampling, ++ need_min_p_sampling, + ) + return [ + probs_idx[i][keep[i]].to(torch.int32) if bool(rows[i].item()) else None @@ -423,10 +258,8 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a + ] + + - def get_token_ids_logprobs_raw( - logprobs: torch.Tensor, - token_ids_logprobs_list: List[Optional[List[int]]], -@@ -297,13 +402,14 @@ def compute_spec_v2_logprobs( + def compute_spec_v2_logprobs( + batch, logits_output, predict: torch.Tensor, accept_index: torch.Tensor, @@ -434,27 +267,18 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a speculative_num_steps: int, ): """Compute logprobs for accepted tokens after spec v2 verify sampling. - - Gathers logits at accepted positions, applies log_softmax (temperature-scaled -- if not greedy), and populates logits_output.next_token_logprobs (plus optional -- top-k / token-ids logprobs) so they flow through copy_to_cpu(). -+ if not greedy), and populates logits_output.next_token_logprobs plus optional -+ top-k / token-ids / top-p replay metadata so they flow through copy_to_cpu(). - """ - bs = len(batch.seq_lens) - max_accept = speculative_num_steps + 1 -@@ -311,6 +417,7 @@ def compute_spec_v2_logprobs( - +@@ -320,6 +430,7 @@ def compute_spec_v2_logprobs( + flat_accept_idx = accept_index.long().reshape(-1) gathered_logits = logits_output.next_token_logits[flat_accept_idx] + temperatures = None - + if batch.sampling_info.is_all_greedy or envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get(): gathered_logprobs = torch.nn.functional.log_softmax(gathered_logits, dim=-1) -@@ -332,6 +439,80 @@ def compute_spec_v2_logprobs( +@@ -341,6 +452,81 @@ def compute_spec_v2_logprobs( ] logits_output.next_token_logprobs = token_logprobs.reshape(bs, max_accept) - + + if batch.sampling_info.need_return_top_p_token_ids: + valid_accept_mask = ( + torch.arange(max_accept, device=device).view(1, -1) @@ -462,16 +286,19 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a + ).reshape(-1) + request_mask = ( + torch.repeat_interleave( -+ batch.sampling_info.return_top_p_token_ids, max_accept ++ batch.sampling_info.return_top_p_token_ids, ++ max_accept, + ) + & valid_accept_mask + ) + + if batch.sampling_info.is_all_greedy: + logits_output.next_token_top_p_token_ids = [ -+ accepted_token_ids[i : i + 1].to(torch.int32) -+ if bool(request_mask[i].item()) -+ else None ++ ( ++ accepted_token_ids[i : i + 1].to(torch.int32) ++ if bool(request_mask[i].item()) ++ else None ++ ) + for i in range(bs * max_accept) + ] + else: @@ -483,15 +310,18 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a + ) + probs = torch.softmax(gathered_logits / temperatures, dim=-1) + expanded_top_ks = torch.repeat_interleave( -+ batch.sampling_info.top_ks, max_accept ++ batch.sampling_info.top_ks, ++ max_accept, + ) + expanded_top_ps = torch.repeat_interleave( -+ batch.sampling_info.top_ps, max_accept ++ batch.sampling_info.top_ps, ++ max_accept, + ) + expanded_min_ps = torch.repeat_interleave( -+ batch.sampling_info.min_ps, max_accept ++ batch.sampling_info.min_ps, ++ max_accept, + ) -+ top_p_token_ids = get_top_p_token_ids_from_probs( ++ logits_output.next_token_top_p_token_ids = get_top_p_token_ids_from_probs( + probs=probs, + top_ks=expanded_top_ks, + top_ps=expanded_top_ps, @@ -500,9 +330,6 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a + need_min_p_sampling=False, + request_mask=request_mask, + ) -+ if top_p_token_ids is not None: -+ logits_output.next_token_top_p_token_ids = top_p_token_ids -+ + renorm_logprobs = renorm_logprob_over_top_p( + probs=probs, + top_ks=expanded_top_ks, @@ -511,44 +338,194 @@ index fde20e211a80b89735f56629d7ba1a959b4b42d6..ab1b26362e104a56407a42b559fb758a + need_top_p_sampling=batch.sampling_info.need_top_p_sampling, + need_min_p_sampling=False, + request_mask=request_mask, -+ # Force-keep the accepted token: a small fraction of -+ # speculatively accepted tokens land outside their own top-p -+ # nucleus, which would give a -inf renormalized logprob. -+ # Force-keeping makes the denominator ``nucleus ∪ {accepted}``, -+ # matching the trainer which also force-keeps the target token, -+ # so these tokens stay finite and on-policy. + force_keep_token_ids=accepted_token_ids.long(), + ) + if renorm_logprobs is not None: + idx = torch.arange(bs * max_accept, device=device) -+ renorm_token_logprobs = renorm_logprobs[idx, accepted_token_ids.long()] ++ renorm_token_logprobs = renorm_logprobs[ ++ idx, ++ accepted_token_ids.long(), ++ ] + renorm_token_logprobs.clamp_( + min=torch.finfo(renorm_token_logprobs.dtype).min + ) + logits_output.next_token_logprobs = renorm_token_logprobs.reshape( -+ bs, max_accept ++ bs, ++ max_accept, + ) + if batch.top_logprobs_nums and any(x > 0 for x in batch.top_logprobs_nums): top_logprobs_nums_expanded = [ num for num in batch.top_logprobs_nums for _ in range(max_accept) +diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py +index dc9e5128d3..d96f534859 100644 +--- a/python/sglang/srt/layers/sampler.py ++++ b/python/sglang/srt/layers/sampler.py +@@ -14,6 +14,8 @@ from sglang.srt.layers.dp_attention import ( + from sglang.srt.layers.logits_processor import LogitsProcessorOutput + from sglang.srt.layers.logprob_processor import ( + OutputLogprobProcessor, ++ get_top_p_token_ids_from_probs, ++ renorm_logprob_over_top_p, + ) + from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args + from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo +@@ -131,6 +133,11 @@ class Sampler(nn.Module): + _aiter_greedy_sample(batch_next_token_ids, logits) + else: + batch_next_token_ids = torch.argmax(logits, -1) ++ self._attach_greedy_top_p_token_ids_to_output( ++ logits_output, ++ batch_next_token_ids, ++ sampling_info, ++ ) + if return_sampling_mask: + self._attach_greedy_sampling_mask_to_output( + logits_output, sampling_info, batch_next_token_ids +@@ -165,6 +172,7 @@ class Sampler(nn.Module): + if self.use_ascend_backend: + # Ascend backend: sample from logits directly. + batch_next_token_ids, logprobs = self._forward_ascend_backend( ++ logits_output, + logits, + sampling_info, + simple_sampling_case, +@@ -207,6 +215,12 @@ class Sampler(nn.Module): + logits[:] = torch.softmax(logits, dim=-1) + probs = logits + ++ self._attach_top_p_token_ids_to_output( ++ logits_output, ++ probs, ++ sampling_info, ++ simple_sampling_case, ++ ) + batch_next_token_ids = self._sample_from_probs( + probs, sampling_info, positions, simple_sampling_case + ) +@@ -221,10 +235,26 @@ class Sampler(nn.Module): + sampling_mask_data, + ) + if return_logprob and not SGLANG_RETURN_ORIGINAL_LOGPROB: ++ top_p_logprobs = None ++ if sampling_info.need_return_top_p_token_ids: ++ top_p_logprobs = renorm_logprob_over_top_p( ++ probs=probs, ++ top_ks=sampling_info.top_ks, ++ top_ps=sampling_info.top_ps, ++ min_ps=sampling_info.min_ps, ++ need_top_p_sampling=sampling_info.need_top_p_sampling, ++ need_min_p_sampling=sampling_info.need_min_p_sampling, ++ request_mask=sampling_info.return_top_p_token_ids, ++ force_keep_token_ids=batch_next_token_ids, ++ ) + logprobs = ( +- logprobs_via_logsoftmax_kernel +- if logprobs_via_logsoftmax_kernel is not None +- else torch.log(probs) ++ top_p_logprobs ++ if top_p_logprobs is not None ++ else ( ++ logprobs_via_logsoftmax_kernel ++ if logprobs_via_logsoftmax_kernel is not None ++ else torch.log(probs) ++ ) + ) + del probs + +@@ -295,6 +325,51 @@ class Sampler(nn.Module): + raise ValueError(f"Invalid sampling backend: {backend}") + return batch_next_token_ids + ++ def _attach_greedy_top_p_token_ids_to_output( ++ self, ++ logits_output: LogitsProcessorOutput, ++ batch_next_token_ids: torch.Tensor, ++ sampling_info: SamplingBatchInfo, ++ ) -> None: ++ if not sampling_info.need_return_top_p_token_ids: ++ return ++ ++ request_mask = sampling_info.return_top_p_token_ids ++ logits_output.next_token_top_p_token_ids = [ ++ ( ++ batch_next_token_ids[i : i + 1].to(torch.int32) ++ if bool(request_mask[i].item()) ++ else None ++ ) ++ for i in range(len(batch_next_token_ids)) ++ ] ++ ++ def _attach_top_p_token_ids_to_output( ++ self, ++ logits_output: LogitsProcessorOutput, ++ probs: torch.Tensor, ++ sampling_info: SamplingBatchInfo, ++ simple_sampling_case: bool, ++ ) -> None: ++ if ( ++ not sampling_info.need_return_top_p_token_ids ++ or sampling_info.return_top_p_token_ids is None ++ or simple_sampling_case ++ ): ++ return ++ ++ top_p_token_ids = get_top_p_token_ids_from_probs( ++ probs=probs, ++ top_ks=sampling_info.top_ks, ++ top_ps=sampling_info.top_ps, ++ min_ps=sampling_info.min_ps, ++ need_top_p_sampling=sampling_info.need_top_p_sampling, ++ need_min_p_sampling=sampling_info.need_min_p_sampling, ++ request_mask=sampling_info.return_top_p_token_ids, ++ ) ++ if top_p_token_ids is not None: ++ logits_output.next_token_top_p_token_ids = top_p_token_ids ++ + def _compute_sampling_mask_from_probs( + self, probs: torch.Tensor, sampling_info: SamplingBatchInfo + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: +@@ -466,6 +541,7 @@ class Sampler(nn.Module): + + def _forward_ascend_backend( + self, ++ logits_output: LogitsProcessorOutput, + logits: torch.Tensor, + sampling_info: SamplingBatchInfo, + simple_sampling_case: bool, +@@ -482,6 +558,15 @@ class Sampler(nn.Module): + when return_logprob is False or SGLANG_RETURN_ORIGINAL_LOGPROB is set. + """ + logits.div_(sampling_info.temperatures) ++ if sampling_info.need_return_top_p_token_ids and not simple_sampling_case: ++ probs = torch.softmax(logits, dim=-1) ++ self._attach_top_p_token_ids_to_output( ++ logits_output, ++ probs, ++ sampling_info, ++ simple_sampling_case, ++ ) ++ del probs + batch_next_token_ids = self._sample_from_logits( + logits, sampling_info, simple_sampling_case, positions + ) diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py -index fa9a1d262f90faa9e6c76ecf7e7d3d89451b35f5..9f90d58036262179d32bdff0eea414f76674a735 100644 +index 19a7f4938a..96e40f6ab6 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py -@@ -441,6 +441,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): +@@ -471,6 +471,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): input_token_ids_logprobs_idx=recv_obj.input_token_ids_logprobs_idx, output_token_ids_logprobs_val=recv_obj.output_token_ids_logprobs_val, output_token_ids_logprobs_idx=recv_obj.output_token_ids_logprobs_idx, + output_top_p_token_ids=recv_obj.output_top_p_token_ids, output_token_entropy_val=recv_obj.output_token_entropy_val, - output_hidden_states=recv_obj.output_hidden_states, - routed_experts=routed_experts, + output_token_sampling_mask=recv_obj.output_token_sampling_mask, + output_token_sampling_logprobs=recv_obj.output_token_sampling_logprobs, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index 17aab81733e7b0688c0ad00408588d58b02c0cbc..e13964675d0dc3bceffa0a7df63e11111c1e416d 100644 +index e14df7ca9c..aebf916e3c 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -1189,6 +1189,7 @@ TopLogprobValues = Optional[List[Optional[List[Optional[List[float]]]]]] +@@ -1262,6 +1262,7 @@ TopLogprobValues = Optional[List[Optional[List[Optional[List[float]]]]]] TopLogprobIndices = Optional[List[Optional[List[Optional[List[int]]]]]] TokenIdsLogprobValues = Optional[List[Optional[List[Optional[List[float]]]]]] TokenIdsLogprobIndices = Optional[List[Optional[List[Optional[List[int]]]]]] @@ -556,27 +533,27 @@ index 17aab81733e7b0688c0ad00408588d58b02c0cbc..e13964675d0dc3bceffa0a7df63e1111 HiddenStateChunk = List[Optional[Union[float, List[float]]]] OutputHiddenStates = Optional[List[Optional[List[HiddenStateChunk]]]] CachedTokensDetails = Dict[str, Union[int, str]] -@@ -1229,6 +1230,7 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True): +@@ -1335,6 +1336,7 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True): input_token_ids_logprobs_idx: TokenIdsLogprobIndices output_token_ids_logprobs_val: TokenIdsLogprobValues output_token_ids_logprobs_idx: TokenIdsLogprobIndices + output_top_p_token_ids: TopPTokenIds output_token_entropy_val: Optional[List[Optional[float]]] - - # Hidden states -@@ -1305,6 +1307,7 @@ class BatchStrOutput(BaseBatchReq, kw_only=True): + # Per-request chunks of output-token sampling supports. None when no request + # in the batch asks for return_sampling_mask. +@@ -1429,6 +1431,7 @@ class BatchStrOutput(BaseBatchReq, kw_only=True): input_token_ids_logprobs_idx: TokenIdsLogprobIndices output_token_ids_logprobs_val: TokenIdsLogprobValues output_token_ids_logprobs_idx: TokenIdsLogprobIndices + output_top_p_token_ids: TopPTokenIds output_token_entropy_val: Optional[List[Optional[float]]] - - # Hidden states + # Detokenizer pass-through for BatchTokenIDOutput.output_token_sampling_*. + # None when sampling masks are not returned. diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py -index 94f5895171c229b6b0aa23f0c091a3b963de52f1..d284d33d3241d2e19ad3b5f4dcecdcbf7e6d1e72 100644 +index 0a2cd2b572..c8f062e4b7 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py -@@ -220,6 +220,9 @@ def _handle_output_by_index(output, i): +@@ -238,6 +238,9 @@ def _handle_output_by_index(output, i): output_token_ids_logprobs_idx=_extract_field_by_index( output, "output_token_ids_logprobs_idx", i, check_length=False ), @@ -586,7 +563,7 @@ index 94f5895171c229b6b0aa23f0c091a3b963de52f1..d284d33d3241d2e19ad3b5f4dcecdcbf output_token_entropy_val=_extract_field_by_index( output, "output_token_entropy_val", i, check_length=False ), -@@ -313,6 +316,9 @@ def _handle_output_by_index(output, i): +@@ -355,6 +358,9 @@ def _handle_output_by_index(output, i): output_token_ids_logprobs_idx=_extract_field_by_index( output, "output_token_ids_logprobs_idx", i, check_length=False ), @@ -597,30 +574,41 @@ index 94f5895171c229b6b0aa23f0c091a3b963de52f1..d284d33d3241d2e19ad3b5f4dcecdcbf output, "output_token_entropy_val", i, check_length=False ), diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py -index 99a45d9cab48d10d2d6e6bac54f33a58ef738664..986480f0c1ff8a573e6b37f3a8e8f72357087ea6 100755 +index 0334e3530d..8515785846 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py -@@ -661,6 +661,7 @@ class ReqLogprob: +@@ -755,6 +755,7 @@ class ReqLogprob: None ) output_token_ids_logprobs_idx: Optional[list] = None + output_top_p_token_ids: Optional[list] = None - - - class Req(ReqDllmMixin): -@@ -916,6 +917,7 @@ class Req(ReqDllmMixin): + + + @dataclasses.dataclass(slots=True, kw_only=True) +@@ -1030,6 +1031,7 @@ class Req(ReqDllmMixin): # Can contain either lists or GPU tensors (delayed copy optimization for prefill-only scoring) self.logprob.output_token_ids_logprobs_val = [] self.logprob.output_token_ids_logprobs_idx = [] + self.logprob.output_top_p_token_ids = [] - self.hidden_states: List[List[float]] = [] - self.hidden_states_tensor = None # Note: use tensor instead of list to transfer hidden_states when PD + MTP - self.output_topk_p = None + if return_sampling_mask: + self.output_token_sampling_mask = [] + self.output_token_sampling_logprobs = [] +@@ -2740,9 +2742,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): + + retracted_reqs = [] + first_iter = True +- num_minimum_reqs = ( +- 0 if server_args.disaggregation_mode == "decode" else 1 +- ) ++ num_minimum_reqs = 0 if server_args.disaggregation_mode == "decode" else 1 + while first_iter or ( + not self.check_decode_mem(selected_indices=sorted_indices) + ): diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py -index 6e20e6532852eb5d46ff770dc0e13a0eae0514c3..0dd6ff0051aa4fc9d0a83aa05ff8b713394146a9 100644 +index dbf83eea96..63ddf8c14e 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py -@@ -386,6 +386,11 @@ class SchedulerBatchResultProcessor: +@@ -417,6 +417,11 @@ class SchedulerBatchResultProcessor: logits_output.next_token_token_ids_logprobs_val = [ v.tolist() for v in logits_output.next_token_token_ids_logprobs_val ] @@ -629,10 +617,10 @@ index 6e20e6532852eb5d46ff770dc0e13a0eae0514c3..0dd6ff0051aa4fc9d0a83aa05ff8b713 + v.tolist() if torch.is_tensor(v) else v + for v in logits_output.next_token_top_p_token_ids + ] - + def _apply_prefill_logprobs( self, -@@ -766,6 +771,11 @@ class SchedulerBatchResultProcessor: +@@ -948,6 +953,11 @@ class SchedulerBatchResultProcessor: logits_output.next_token_token_ids_logprobs_val = [ v.tolist() for v in logits_output.next_token_token_ids_logprobs_val ] @@ -642,9 +630,9 @@ index 6e20e6532852eb5d46ff770dc0e13a0eae0514c3..0dd6ff0051aa4fc9d0a83aa05ff8b713 + for v in logits_output.next_token_top_p_token_ids + ] return next_token_ids, next_token_logprobs - + def _apply_decode_logprobs( -@@ -808,6 +818,14 @@ class SchedulerBatchResultProcessor: +@@ -990,6 +1000,14 @@ class SchedulerBatchResultProcessor: req.logprob.output_token_ids_logprobs_idx.append( logits_output.next_token_token_ids_logprobs_idx[flat_idx] ) @@ -656,17 +644,17 @@ index 6e20e6532852eb5d46ff770dc0e13a0eae0514c3..0dd6ff0051aa4fc9d0a83aa05ff8b713 + if torch.is_tensor(row_top_p_token_ids): + row_top_p_token_ids = row_top_p_token_ids.tolist() + req.logprob.output_top_p_token_ids.append(row_top_p_token_ids) - - def _handle_finish_state_updated_req( + + def add_sampling_mask_return_values( self, diff --git a/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py b/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py -index d97d9ae801e130546165f53f2f209d0a233d10c9..a9f9372ef73fc73e1876eded7e05ce1f9e0c0a42 100644 +index 379501f0e3..420c559b2d 100644 --- a/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/logprob_result_processor.py -@@ -314,6 +314,13 @@ class SchedulerLogprobResultProcessor: +@@ -348,6 +348,13 @@ class SchedulerLogprobResultProcessor: output.next_token_token_ids_logprobs_idx[i] ) - + + if output.next_token_top_p_token_ids: + row_top_p_token_ids = output.next_token_top_p_token_ids[i] + if row_top_p_token_ids is not None: @@ -675,29 +663,29 @@ index d97d9ae801e130546165f53f2f209d0a233d10c9..a9f9372ef73fc73e1876eded7e05ce1f + req.logprob.output_top_p_token_ids.append(row_top_p_token_ids) + return num_input_logprobs - + def _initialize_empty_logprob_containers(self, req: Req) -> None: diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py -index 6b691d84eb707525c98f2b88ca3b30348c6c9869..7ddd361722759675fff1eaf776f19507175d4b56 100644 +index 18a42407af..63867fe79c 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/python/sglang/srt/managers/scheduler_components/output_streamer.py -@@ -297,6 +297,7 @@ class _GenerationStreamAccumulator: +@@ -324,6 +324,7 @@ class _GenerationStreamAccumulator: input_token_ids_logprobs_idx: Optional[list] = None output_token_ids_logprobs_val: Optional[list] = None output_token_ids_logprobs_idx: Optional[list] = None + output_top_p_token_ids: Optional[list] = None - - def __post_init__(self) -> None: - if self.return_hidden_states: -@@ -319,6 +320,7 @@ class _GenerationStreamAccumulator: + output_token_sampling_mask: Optional[list] = None + output_token_sampling_logprobs: Optional[list] = None + # Rust server mode: the Rust detokenizer reconstructs text/ids from the raw +@@ -355,6 +356,7 @@ class _GenerationStreamAccumulator: self.input_token_ids_logprobs_idx = [] self.output_token_ids_logprobs_val = [] self.output_token_ids_logprobs_idx = [] + self.output_top_p_token_ids = [] - - def accept(self, *, req: Req) -> None: - if req.finished(): -@@ -472,6 +474,11 @@ class _GenerationStreamAccumulator: + if self.return_sampling_mask: + self.output_token_sampling_mask = [] + self.output_token_sampling_logprobs = [] +@@ -534,6 +536,11 @@ class _GenerationStreamAccumulator: send_output_token_logprobs_offset:logprob_end ] ) @@ -709,38 +697,30 @@ index 6b691d84eb707525c98f2b88ca3b30348c6c9869..7ddd361722759675fff1eaf776f19507 req.send_output_token_logprobs_offset = logprob_end else: self.output_token_logprobs_val.append([]) -@@ -480,6 +487,7 @@ class _GenerationStreamAccumulator: +@@ -542,6 +549,7 @@ class _GenerationStreamAccumulator: self.output_top_logprobs_idx.append([]) self.output_token_ids_logprobs_val.append([]) self.output_token_ids_logprobs_idx.append([]) + self.output_top_p_token_ids.append([]) - - if self.return_hidden_states: - if req.return_hidden_states: -@@ -558,6 +566,7 @@ class _GenerationStreamAccumulator: + + if self.return_sampling_mask: + if req.return_sampling_mask: +@@ -667,6 +675,7 @@ class _GenerationStreamAccumulator: input_token_ids_logprobs_idx=self.input_token_ids_logprobs_idx, output_token_ids_logprobs_val=self.output_token_ids_logprobs_val, output_token_ids_logprobs_idx=self.output_token_ids_logprobs_idx, + output_top_p_token_ids=self.output_top_p_token_ids, output_token_entropy_val=None, - output_hidden_states=self.output_hidden_states, - routed_experts=self.routed_experts, + output_token_sampling_mask=self.output_token_sampling_mask, + output_token_sampling_logprobs=self.output_token_sampling_logprobs, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index f4331b64b42a081aa31e63370efc82a96e416e40..44358b387775266e79c2355f80e5f18ab7eac05d 100644 +index 7472d4b9f9..e2567e998b 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py -@@ -212,6 +212,7 @@ class ReqState: - output_top_logprobs: List[Any] = dataclasses.field(default_factory=list) - input_token_ids_logprobs: List[Any] = dataclasses.field(default_factory=list) - output_token_ids_logprobs: List[Any] = dataclasses.field(default_factory=list) -+ output_top_p_token_ids: List[List[int]] = dataclasses.field(default_factory=list) - customized_info_accumulated: Dict[str, List[Any]] = dataclasses.field( - default_factory=dict - ) -@@ -233,6 +234,21 @@ def _slice_streaming_output_meta_info( - meta_info[key] = meta_info[key][last_output_offset:] - - +@@ -195,6 +195,21 @@ _INCREMENTAL_STREAMING_META_INFO_KEYS = ( + ) + + +def _b64_encode_int32(values: List[int]) -> str: + int32_values = array("i", values) + assert int32_values.itemsize == 4 @@ -756,13 +736,21 @@ index f4331b64b42a081aa31e63370efc82a96e416e40..44358b387775266e79c2355f80e5f18a + return _b64_encode_int32(token_ids), _b64_encode_int32(offsets) + + - class InputFormat(Enum): - """Input format types for tokenization handling.""" - -@@ -2223,6 +2239,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): + class RequestAbortedError(ValueError): + status_code = 499 + +@@ -278,6 +293,7 @@ class ReqState: + output_top_logprobs: List[Any] = dataclasses.field(default_factory=list) + input_token_ids_logprobs: List[Any] = dataclasses.field(default_factory=list) + output_token_ids_logprobs: List[Any] = dataclasses.field(default_factory=list) ++ output_top_p_token_ids: List[List[int]] = dataclasses.field(default_factory=list) + customized_info_accumulated: Dict[str, List[Any]] = dataclasses.field( + default_factory=dict + ) +@@ -2619,6 +2635,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): meta_info["input_token_ids_logprobs"] = state.input_token_ids_logprobs meta_info["output_token_ids_logprobs"] = state.output_token_ids_logprobs - + + if state.output_top_p_token_ids and meta_info.get("finish_reason") is not None: + token_ids, offsets = _encode_top_p_token_ids(state.output_top_p_token_ids) + meta_info["top_p_token_ids"] = token_ids @@ -771,10 +759,10 @@ index f4331b64b42a081aa31e63370efc82a96e416e40..44358b387775266e79c2355f80e5f18a def convert_logprob_style( self, meta_info: dict, -@@ -2283,6 +2304,13 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): +@@ -2688,6 +2709,13 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): recv_obj.output_token_ids_logprobs_idx[recv_obj_index] ) - + + output_top_p_token_ids = getattr(recv_obj, "output_top_p_token_ids", None) + if ( + output_top_p_token_ids is not None @@ -786,10 +774,10 @@ index f4331b64b42a081aa31e63370efc82a96e416e40..44358b387775266e79c2355f80e5f18a meta_info, state, diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py -index 87fbb35ded8d306360cb65333984bfba1600e4dd..6123941282e1a7209ca1a6a91b4c379a923553d0 100644 +index fe883c264d..ddaf20cbec 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py -@@ -121,6 +121,11 @@ class GenerationBatchResult: +@@ -144,6 +144,11 @@ class GenerationBatchResult: _async_d2h(v) if torch.is_tensor(v) else v for v in self.logits_output.next_token_token_ids_logprobs_val ] @@ -801,41 +789,41 @@ index 87fbb35ded8d306360cb65333984bfba1600e4dd..6123941282e1a7209ca1a6a91b4c379a if return_hidden_states and self.logits_output.hidden_states is not None: self.logits_output.hidden_states = _async_d2h( self.logits_output.hidden_states -@@ -207,6 +212,7 @@ def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict: +@@ -236,6 +241,7 @@ def get_logprob_dict_from_result(result: GenerationBatchResult) -> dict: "next_token_top_logprobs_idx": result.logits_output.next_token_top_logprobs_idx, "next_token_token_ids_logprobs_val": result.logits_output.next_token_token_ids_logprobs_val, "next_token_token_ids_logprobs_idx": result.logits_output.next_token_token_ids_logprobs_idx, + "next_token_top_p_token_ids": result.logits_output.next_token_top_p_token_ids, + "next_token_sampling_mask_idx": result.logits_output.next_token_sampling_mask_idx, + "next_token_sampling_logprobs": result.logits_output.next_token_sampling_logprobs, "input_token_logprobs": result.logits_output.input_token_logprobs, - "input_top_logprobs_val": result.logits_output.input_top_logprobs_val, - "input_top_logprobs_idx": result.logits_output.input_top_logprobs_idx, -@@ -231,6 +237,9 @@ def get_logprob_from_pp_outputs( +@@ -262,6 +268,10 @@ def get_logprob_from_pp_outputs( next_token_token_ids_logprobs_idx=next_pp_outputs[ "next_token_token_ids_logprobs_idx" ], + next_token_top_p_token_ids=next_pp_outputs.tensors.get( -+ "next_token_top_p_token_ids", None ++ "next_token_top_p_token_ids", ++ None, + ), + next_token_sampling_mask_idx=next_pp_outputs["next_token_sampling_mask_idx"], + next_token_sampling_logprobs=next_pp_outputs["next_token_sampling_logprobs"], input_token_logprobs=next_pp_outputs["input_token_logprobs"], - input_top_logprobs_val=next_pp_outputs["input_top_logprobs_val"], - input_top_logprobs_idx=next_pp_outputs["input_top_logprobs_idx"], diff --git a/python/sglang/srt/sampling/sampling_batch_info.py b/python/sglang/srt/sampling/sampling_batch_info.py -index cfb22d84194a7affcefa01933f0a4bafa9b4b0b6..567bbe5d5d39dc00183975a7d4e7335a9b4fe806 100644 +index a85faafd24..8019d13dc6 100644 --- a/python/sglang/srt/sampling/sampling_batch_info.py +++ b/python/sglang/srt/sampling/sampling_batch_info.py -@@ -42,6 +42,11 @@ class SamplingBatchInfo: - - # Masking tensors for grammar-guided structured outputs - vocab_size: int -+ -+ # Whether to return kept token ids from top-p replay sampling. +@@ -77,6 +77,10 @@ class SamplingBatchInfo: + return_sampling_masks: Optional[List[bool]] = None + sampling_mask_max_top_k: int = 0 + ++ # Legacy slime top-p replay request, carried in sampling custom_params. + return_top_p_token_ids: Optional[torch.Tensor] = None + need_return_top_p_token_ids: bool = False + - grammars: Optional[List] = None - rids_int: Optional[torch.Tensor] = None - bootstrap_room_ids_int: Optional[torch.Tensor] = None -@@ -105,6 +110,16 @@ class SamplingBatchInfo: + # Device + device: str = "cuda" + +@@ -115,6 +119,16 @@ class SamplingBatchInfo: dtype=torch.float, pin_memory=_pin, ).to(device, non_blocking=True) @@ -852,7 +840,7 @@ index cfb22d84194a7affcefa01933f0a4bafa9b4b0b6..567bbe5d5d39dc00183975a7d4e7335a sampling_seed = ( torch.tensor( [ -@@ -191,6 +206,8 @@ class SamplingBatchInfo: +@@ -207,6 +221,8 @@ class SamplingBatchInfo: need_top_p_sampling=any(r.sampling_params.top_p != 1.0 for r in reqs), need_top_k_sampling=any(r.sampling_params.top_k != TOP_K_ALL for r in reqs), need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs), @@ -861,7 +849,7 @@ index cfb22d84194a7affcefa01933f0a4bafa9b4b0b6..567bbe5d5d39dc00183975a7d4e7335a vocab_size=vocab_size, penalizer_orchestrator=penalizer_orchestrator, has_custom_logit_processor=has_custom_logit_processor, -@@ -293,6 +310,7 @@ class SamplingBatchInfo: +@@ -311,6 +327,7 @@ class SamplingBatchInfo: "top_ps", "top_ks", "min_ps", @@ -869,7 +857,7 @@ index cfb22d84194a7affcefa01933f0a4bafa9b4b0b6..567bbe5d5d39dc00183975a7d4e7335a "sampling_seed", ]: value = getattr(self, item, None) -@@ -401,6 +419,7 @@ class SamplingBatchInfo: +@@ -436,6 +453,7 @@ class SamplingBatchInfo: "top_ps", "top_ks", "min_ps", @@ -877,59 +865,43 @@ index cfb22d84194a7affcefa01933f0a4bafa9b4b0b6..567bbe5d5d39dc00183975a7d4e7335a "sampling_seed", ]: self_val = getattr(self, item, None) -@@ -412,6 +431,7 @@ class SamplingBatchInfo: +@@ -448,6 +466,7 @@ class SamplingBatchInfo: self.need_top_p_sampling |= other.need_top_p_sampling self.need_top_k_sampling |= other.need_top_k_sampling self.need_min_p_sampling |= other.need_min_p_sampling + self.need_return_top_p_token_ids |= other.need_return_top_p_token_ids - + self.adjusted_merge_batch(other) - -diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py -index baa3674964975b1c35535c1f98bdb2565304f432..b6479805724112af05209e126baebbe51f616c13 100644 ---- a/python/sglang/srt/speculative/eagle_worker_v2.py -+++ b/python/sglang/srt/speculative/eagle_worker_v2.py -@@ -1678,7 +1678,12 @@ class EAGLEWorkerV2(BaseSpecWorker): - - if batch.return_logprob and not batch.forward_mode.is_idle(): - compute_spec_v2_logprobs( -- batch, logits_output, predict, accept_index, self.speculative_num_steps -+ batch, -+ logits_output, -+ predict, -+ accept_index, -+ accept_lens, -+ self.speculative_num_steps, - ) - - if not batch.forward_mode.is_idle() and self.topk > 1: -diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py -index 5361342c414b5711e640b6954b8602d6e6abf8cb..7d07cf33b966c46fc7e72ecca91b527c947150b1 100644 ---- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py -+++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py -@@ -837,7 +837,12 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): - - if batch.return_logprob and not batch.forward_mode.is_idle(): - compute_spec_v2_logprobs( -- batch, logits_output, predict, accept_index, self.speculative_num_steps -+ batch, -+ logits_output, -+ predict, -+ accept_index, -+ accept_lens, -+ self.speculative_num_steps, - ) - - next_draft_input = EagleDraftInput(bonus_tokens=bonus_tokens) + +diff --git a/python/sglang/srt/speculative/eagle_worker_common.py b/python/sglang/srt/speculative/eagle_worker_common.py +index 8e0c018264..806b114197 100644 +--- a/python/sglang/srt/speculative/eagle_worker_common.py ++++ b/python/sglang/srt/speculative/eagle_worker_common.py +@@ -625,7 +625,14 @@ def run_eagle_verify( + bonus_tokens = torch.empty((0,), device=device, dtype=torch.int32) + + if batch.return_logprob and not batch.forward_mode.is_idle(): +- compute_spec_v2_logprobs(batch, logits_output, predict, accept_index, num_steps) ++ compute_spec_v2_logprobs( ++ batch, ++ logits_output, ++ predict, ++ accept_index, ++ accept_lens, ++ num_steps, ++ ) + + if finalize_tree_path and not batch.forward_mode.is_idle() and topk > 1: + # topk == 1 needs nothing here: the accepted path is already the front diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py -index 4cbcb7f0ef1f14bea2ae4cafe841afc786ef15f9..acda994c854c34d7e3b5dcaea2bef685bff1e3e0 100644 +index 84cdecdde6..3b61fcfc0d 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py -@@ -463,6 +463,7 @@ class NGRAMWorker(BaseSpecWorker): +@@ -487,6 +487,7 @@ class NGRAMWorker(BaseSpecWorker): logits_output, predict, accept_index, + accept_lens, self.draft_token_num - 1, ) - + diff --git a/docker/patch/latest/sglang.patch b/docker/patch/latest/sglang.patch index 967585629a..7ae199967d 100644 --- a/docker/patch/latest/sglang.patch +++ b/docker/patch/latest/sglang.patch @@ -1,28 +1,30 @@ diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py -index 2164cade4909b0712654791f05d35dfe137796db..50e13ffa5dbbe8cfc72e79edfc2881dad1357442 100644 +index 240e20deffb725382bc3bd7db277f4722605b295..56bc3d42f67fc9da48eee231b6851df630041ea8 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -21,6 +21,7 @@ Life cycle of a request in the decode server from __future__ import annotations - + import logging +import os import time from collections import deque from dataclasses import dataclass -@@ -659,12 +660,33 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): - [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group - ) - +@@ -769,6 +770,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): + [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group + ) + + bootstrap_timeout = float( + os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") + ) + now = time.perf_counter() + - for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): - if rids_to_check is not None and decode_req.req.rid not in rids_to_check: + for decode_req, poll in zip(self.queue, polls): + if poll is None: continue - +@@ -776,7 +782,23 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): + continue + if poll == KVPoll.Bootstrapping: - pass + entry_time = ( @@ -45,10 +47,10 @@ index 2164cade4909b0712654791f05d35dfe137796db..50e13ffa5dbbe8cfc72e79edfc2881da elif poll == KVPoll.WaitingForInput: decode_req.waiting_for_input = True decode_req.req.time_stats.set_bootstrap_done_time() -@@ -1469,6 +1491,17 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): +@@ -1639,6 +1661,17 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): return host_indices return kv_loc - + + def release_memory_occupation(self): + self.queue.clear() + self.retracted_queue.clear() @@ -60,13 +62,13 @@ index 2164cade4909b0712654791f05d35dfe137796db..50e13ffa5dbbe8cfc72e79edfc2881da + if hasattr(self.kv_manager, "register_buffer_to_engine"): + self.kv_manager.register_buffer_to_engine() + - - class DecodeTransferQueue(DecodeHiCacheTransferMixin): - """ -@@ -1678,6 +1711,11 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): + + def alloc_for_decode_prealloc_hisparse( + allocator: BaseTokenToKVPoolAllocator, +@@ -1995,6 +2028,11 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): else: polls = self._poll_with_metadata_gate() - + + transfer_timeout = float( + os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") + ) @@ -75,7 +77,7 @@ index 2164cade4909b0712654791f05d35dfe137796db..50e13ffa5dbbe8cfc72e79edfc2881da transferred_reqs = [] indices_to_remove = set() for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): -@@ -1754,7 +1792,17 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): +@@ -2071,7 +2109,17 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): KVPoll.WaitingForInput, KVPoll.Transferring, ]: @@ -93,18 +95,18 @@ index 2164cade4909b0712654791f05d35dfe137796db..50e13ffa5dbbe8cfc72e79edfc2881da + decode_req.kv_receiver.abort() else: raise ValueError(f"Unexpected poll case: {poll}") - -@@ -1780,6 +1828,9 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): - + +@@ -2097,6 +2145,9 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): + def release_memory_occupation(self): """Clean up in-flight transfers before releasing GPU memory.""" + for decode_req in self.queue: + if decode_req.kv_receiver is not None: + decode_req.kv_receiver.abort() self.queue.clear() - + def resume_memory_occupation(self): -@@ -1988,6 +2039,11 @@ class SchedulerDisaggregationDecodeMixin: +@@ -2320,6 +2371,11 @@ class SchedulerDisaggregationDecodeMixin: resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs() self.waiting_queue.extend(resumed_reqs) if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0: @@ -115,21 +117,21 @@ index 2164cade4909b0712654791f05d35dfe137796db..50e13ffa5dbbe8cfc72e79edfc2881da + self.waiting_queue.extend(transferred_reqs) # if there are still retracted requests, we do not allocate new requests return - + diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py -index 3e8153e9699568331f6ac876ef868dfa42ac9764..aadb70ee753829034dda77859678f3131fc554ee 100644 +index 1907e0ee69f8d3a5b9b9c425c148201e3046c89d..14048ad0c54fbe7b57681ecd3d43a3a60229bb22 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py -@@ -851,7 +851,7 @@ class MooncakeKVManager(CommonKVManager): +@@ -1050,7 +1050,7 @@ class MooncakeKVManager(CommonKVManager): for i, dst_aux_ptr in enumerate(dst_aux_ptrs): length = prefill_aux_item_lens[i] src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index - dst_addr = dst_aux_ptrs[i] + length * req.dst_aux_index + dst_addr = dst_aux_ptr + length * req.dst_aux_index transfer_blocks.append((src_addr, dst_addr, length)) - + return self._transfer_data(req.mooncake_session_id, transfer_blocks) -@@ -1354,12 +1354,6 @@ class MooncakeKVManager(CommonKVManager): +@@ -1714,12 +1714,6 @@ class MooncakeKVManager(CommonKVManager): if ret != 0: with self.session_lock: self.session_failures[req.mooncake_session_id] += 1 @@ -142,41 +144,12 @@ index 3e8153e9699568331f6ac876ef868dfa42ac9764..aadb70ee753829034dda77859678f313 self.record_failure( kv_chunk.room, f"Failed to send kv chunk of {kv_chunk.room} to " -@@ -1377,12 +1371,27 @@ class MooncakeKVManager(CommonKVManager): - - if kv_chunk.is_last_chunk: - if kv_chunk.state_indices: -- self.maybe_send_extra( -+ ret = self.maybe_send_extra( - req, - kv_chunk.state_indices, - executor, - target_rank_registration_info, - ) -+ if ret != 0: -+ self.record_failure( -+ kv_chunk.room, -+ f"Failed to send extra state of {kv_chunk.room} to " -+ f"{NetworkAddress(req.endpoint, req.dst_port).to_host_port_str()}", -+ ) -+ self.update_status(kv_chunk.room, KVPoll.Failed) -+ self.sync_status_to_decode_endpoint( -+ req.endpoint, -+ req.dst_port, -+ req.room, -+ KVPoll.Failed, -+ prefill_unique_rank, -+ ) -+ break - - # Only the last chunk we need to send the aux data - ret = self.send_aux( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py -index 499a9fff12898f9244731fa85f87c2d38f764fb6..b9c672d174390847a94b1a4f598a83af3658c91d 100644 +index 5642c64efaa480cf2ec3cb01ede685f0da290ae3..c32d65034fc9bb7d87a5b092e625f4de007332b3 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -21,6 +21,8 @@ from __future__ import annotations - + import hashlib import logging +import os @@ -184,19 +157,19 @@ index 499a9fff12898f9244731fa85f87c2d38f764fb6..b9c672d174390847a94b1a4f598a83af from array import array from collections import deque from http import HTTPStatus -@@ -335,6 +337,11 @@ class PrefillBootstrapQueue: - self.scheduler.attn_tp_cpu_group, - ) - +@@ -416,6 +418,11 @@ class PrefillBootstrapQueue: + self.scheduler.attn_tp_cpu_group, + ) + + bootstrap_timeout = float( + os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") + ) + now = time.perf_counter() + for i, (req, poll) in enumerate(zip(self.queue, polls)): - if ( - rids_to_check is not None -@@ -351,6 +358,27 @@ class PrefillBootstrapQueue: + if poll is None: + continue +@@ -425,6 +432,27 @@ class PrefillBootstrapQueue: indices_to_remove.add(i) failed_reqs.append(req) elif poll == KVPoll.Bootstrapping: @@ -222,12 +195,12 @@ index 499a9fff12898f9244731fa85f87c2d38f764fb6..b9c672d174390847a94b1a4f598a83af + self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() + continue if ( - req.time_stats.prefill_retry_count - < self.scheduler.server_args.optimistic_prefill_retries -@@ -741,6 +769,11 @@ class SchedulerDisaggregationPrefillMixin: + req.prefill_attempt_count + < self.scheduler.server_args.optimistic_prefill_attempts +@@ -832,6 +860,11 @@ class SchedulerDisaggregationPrefillMixin: self.attn_tp_cpu_group, ) - + + transfer_timeout = float( + os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") + ) @@ -236,10 +209,11 @@ index 499a9fff12898f9244731fa85f87c2d38f764fb6..b9c672d174390847a94b1a4f598a83af undone_reqs: List[Req] = [] # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue for req, poll in zip(self.disagg_prefill_inflight_queue, polls): -@@ -765,7 +798,27 @@ class SchedulerDisaggregationPrefillMixin: - continue - +@@ -865,8 +898,27 @@ class SchedulerDisaggregationPrefillMixin: + continue + if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: +- # todo: set Transferring correctly in backend - undone_reqs.append(req) + entry_time = req.time_stats.prefill_transfer_queue_entry_time + if entry_time > 0 and now - entry_time > transfer_timeout: @@ -263,13 +237,13 @@ index 499a9fff12898f9244731fa85f87c2d38f764fb6..b9c672d174390847a94b1a4f598a83af + else: + undone_reqs.append(req) elif poll == KVPoll.Success: # transfer done - release_kv_cache(req, self.tree_cache) # unlock the tree - req.finished_reason = FINISH_LENGTH(length=0) + if not isinstance(req.finished_reason, FINISH_ABORT): + req.finished_reason = FINISH_LENGTH(length=0) diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py -index f58dc2784fbb824f431a53ffc5c1091dc4e18947..4a51574b3487dbc1e474ce92ba56b5be7be3b9d3 100644 +index 6eb0ad1af3a5e7c377d783a2b5ae1eaa7209c714..a8a1e701aad5635deff8ae369a13dd4099733f5e 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py -@@ -69,6 +69,7 @@ from sglang.srt.managers.io_struct import ( +@@ -73,6 +73,7 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterReqInput, MultimodalDataInputFormat, OpenSessionReqInput, @@ -277,10 +251,10 @@ index f58dc2784fbb824f431a53ffc5c1091dc4e18947..4a51574b3487dbc1e474ce92ba56b5be ProfileReq, ProfileReqType, ReleaseMemoryOccupationReqInput, -@@ -1127,6 +1128,20 @@ class Engine(EngineScoreMixin, EngineBase): +@@ -1429,6 +1430,20 @@ class Engine(EngineScoreMixin, EngineBase): self.tokenizer_manager.update_weights_from_ipc(obj, None) ) - + + def post_process_weights( + self, + restore_weights_before_load: bool = False, @@ -299,10 +273,10 @@ index f58dc2784fbb824f431a53ffc5c1091dc4e18947..4a51574b3487dbc1e474ce92ba56b5be """Get weights by parameter name.""" obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py -index bd9d7eafa1c9721a01d98c397deebf17e95a06dc..c74f7666a9972ba7ee5b0c94c805189e640fbbda 100644 +index ac38b4668c8119dbd1a17c9eabff4e5fb0430faa..0189493c17a23efb9c81eca2ef3129063f73d40a 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py -@@ -127,6 +127,7 @@ from sglang.srt.managers.io_struct import ( +@@ -132,6 +132,7 @@ from sglang.srt.managers.io_struct import ( OpenSessionReqInput, ParseFunctionCallReq, PauseGenerationReqInput, @@ -310,7 +284,7 @@ index bd9d7eafa1c9721a01d98c397deebf17e95a06dc..c74f7666a9972ba7ee5b0c94c805189e ProfileReq, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, -@@ -686,10 +687,8 @@ async def model_info(): +@@ -762,10 +763,8 @@ async def model_info(): @app.get("/weight_version") async def weight_version(): """Get the current weight version.""" @@ -320,10 +294,10 @@ index bd9d7eafa1c9721a01d98c397deebf17e95a06dc..c74f7666a9972ba7ee5b0c94c805189e - ) + result = await model_info() + return {"weight_version": result.get("weight_version", None)} - - + + @app.get("/get_server_info") -@@ -706,9 +705,18 @@ async def get_server_info(): +@@ -782,9 +781,18 @@ async def get_server_info(): async def server_info(): """Get the server information.""" # Returns internal states per DP. @@ -342,13 +316,13 @@ index bd9d7eafa1c9721a01d98c397deebf17e95a06dc..c74f7666a9972ba7ee5b0c94c805189e + server_info_timeout, + ) + internal_states = [] - + server_args = _global_state.tokenizer_manager.server_args - -@@ -1322,6 +1330,20 @@ async def update_weights_from_ipc( + +@@ -1409,6 +1417,20 @@ async def update_weights_from_ipc( return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) - - + + +@app.post("/post_process_weights") +@auth_level(AuthLevel.ADMIN_OPTIONAL) +async def post_process_weights(req: PostProcessWeightsReqInput, request: Request): @@ -367,46 +341,29 @@ index bd9d7eafa1c9721a01d98c397deebf17e95a06dc..c74f7666a9972ba7ee5b0c94c805189e @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weight_version( diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py -index 59b2256815db726ce9f0104842b2acb86385a065..5e5f7aa2a5b64eb0d20d4af14aecc4be9c508c42 100644 +index 79136182bd2d80709984f0ae88946aaf8f477836..a74aba3631680af328dab31431ba83dc60d12741 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -2,6 +2,7 @@ from __future__ import annotations - + import contextlib import logging +import os - from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union - -@@ -168,6 +169,15 @@ def _uses_dsa_attention_backend(forward_batch: ForwardBatch) -> bool: - return backend_name in ("dsa", "nsa") - - -+def _match_head_gate_q_scale( -+ weights: torch.Tensor, q_scale: torch.Tensor -+) -> torch.Tensor: -+ if weights.shape[1] < q_scale.shape[1]: -+ assert q_scale.shape[1] % weights.shape[1] == 0 -+ weights = weights.repeat_interleave(q_scale.shape[1] // weights.shape[1], dim=1) -+ return weights -+ -+ - if _is_cuda: - from sglang.jit_kernel.dsv4 import fused_q_indexer_rope_first_quant - from sglang.jit_kernel.dsv32 import ( -@@ -223,6 +233,7 @@ if _is_cuda: - ) -> torch.Tensor: - out = torch.mm(x, weight.t(), out_dtype=torch.float32) - weights = out * n_heads_inv_sqrt -+ weights = _match_head_gate_q_scale(weights, q_scale) - weights = weights.unsqueeze(-1) * q_scale * softmax_scale - return weights - -@@ -440,6 +451,15 @@ class Indexer(MultiPlatformOp): - self.k_norm = LayerNorm( - self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32 - ) -+ server_args = get_global_server_args() + + import torch +@@ -18,6 +19,7 @@ from sglang.srt.layers.attention.dsa.dsa_indexer_metadata import BaseIndexerMeta + from sglang.srt.layers.attention.dsa.dsa_npu_indexer import DSANPUIndexerMixin + from sglang.srt.layers.attention.dsa.dsa_prefill_cuda_graph import ( + GRAPH_WEIGHTS_PROJ_LORA_ERROR, ++ _match_head_gate_q_scale, + _is_in_piecewise_or_breakable_cuda_graph, + bcg_dsa_indexer_prefill_split, + pcg_dsa_indexer_prefill_split, +@@ -235,6 +237,13 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + self.index_topk = index_topk + self.q_lora_rank = q_lora_rank + self.layer_id = layer_id + env_neox_style = os.environ.get("INDEXER_ROPE_NEOX_STYLE") + if env_neox_style is not None: + if env_neox_style not in ("0", "1"): @@ -414,39 +371,29 @@ index 59b2256815db726ce9f0104842b2acb86385a065..5e5f7aa2a5b64eb0d20d4af14aecc4be + "INDEXER_ROPE_NEOX_STYLE must be either '0' or '1' when set." + ) + is_neox_style = env_neox_style == "1" -+ - self.rotary_emb = get_rope_wrapper( - rope_head_dim, - rotary_dim=rope_head_dim, -@@ -447,7 +467,7 @@ class Indexer(MultiPlatformOp): - base=rope_theta, # type: ignore - rope_scaling=rope_scaling, - is_neox_style=is_neox_style, -- device=get_global_server_args().device, -+ device=server_args.device, - ) - self.block_size = block_size - self.scale_fmt = scale_fmt -@@ -507,6 +527,7 @@ class Indexer(MultiPlatformOp): + self.use_dsa_indexer_fusion = ( + _is_cuda + and not envs.SGLANG_DISABLE_DSA_INDEXER_FUSION.get() +@@ -365,6 +374,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): ): weights = self._weights_proj_bf16_in_fp32_out(x) weights = weights * self.n_heads**-0.5 + weights = _match_head_gate_q_scale(weights, q_scale) weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale return weights - -@@ -514,6 +535,7 @@ class Indexer(MultiPlatformOp): + +@@ -372,6 +382,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): def _apply_q_scale_and_softmax_scale( self, weights: torch.Tensor, q_scale: torch.Tensor ): + weights = _match_head_gate_q_scale(weights, q_scale) return weights.unsqueeze(-1) * q_scale * self.softmax_scale - + @torch.compile(dynamic=True) -@@ -539,6 +561,12 @@ class Indexer(MultiPlatformOp): +@@ -397,6 +408,12 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): return max_kv_len <= self.index_topk return False - + + def _maybe_repeat_query_heads(self, query: torch.Tensor) -> torch.Tensor: + if query.shape[1] < 32: + assert 32 % query.shape[1] == 0 @@ -456,7 +403,7 @@ index 59b2256815db726ce9f0104842b2acb86385a065..5e5f7aa2a5b64eb0d20d4af14aecc4be def _get_q_k_bf16( self, q_lora: torch.Tensor, -@@ -1770,6 +1798,7 @@ class Indexer(MultiPlatformOp): +@@ -1604,6 +1621,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): query, key, weights_raw = self._get_q_k_bf16( q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch ) @@ -464,22 +411,59 @@ index 59b2256815db726ce9f0104842b2acb86385a065..5e5f7aa2a5b64eb0d20d4af14aecc4be q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) with torch.cuda.stream(self.alt_stream): self._store_index_k_cache( -@@ -1787,6 +1816,7 @@ class Indexer(MultiPlatformOp): - query, key, weights_raw = self._get_q_k_bf16( - q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch +@@ -1625,6 +1643,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + enable_dual_stream, + forward_batch=forward_batch, ) + query = self._maybe_repeat_query_heads(query) - + if enable_dual_stream: current_stream = torch.cuda.current_stream() +diff --git a/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py b/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py +index f90e7ba95b4a1e08fa8743fd59555b728ba561e3..5f24b16483eabcc81c23fdb50ae3ffcf312901e2 100644 +--- a/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py ++++ b/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py +@@ -19,6 +19,16 @@ from sglang.srt.utils.custom_op import register_custom_op + + _is_cuda = is_cuda() + ++ ++def _match_head_gate_q_scale( ++ weights: torch.Tensor, q_scale: torch.Tensor ++) -> torch.Tensor: ++ if weights.shape[1] < q_scale.shape[1]: ++ assert q_scale.shape[1] % weights.shape[1] == 0 ++ weights = weights.repeat_interleave(q_scale.shape[1] // weights.shape[1], dim=1) ++ return weights ++ ++ + GRAPH_WEIGHTS_PROJ_LORA_ERROR = ( + "DSA indexer weights_proj LoRA is incompatible with " + "piecewise/breakable CUDA graph; remove the explicit " +@@ -55,6 +65,7 @@ if _is_cuda: + q_scale: torch.Tensor, + ) -> torch.Tensor: + weights = weights_raw * n_heads_inv_sqrt ++ weights = _match_head_gate_q_scale(weights, q_scale) + return weights.unsqueeze(-1) * q_scale * softmax_scale + + def _logits_head_gate_graph_fake_impl( +@@ -81,6 +92,7 @@ if _is_cuda: + ) -> torch.Tensor: + out = torch.mm(x, weight.t(), out_dtype=torch.float32) + weights = out * n_heads_inv_sqrt ++ weights = _match_head_gate_q_scale(weights, q_scale) + weights = weights.unsqueeze(-1) * q_scale * softmax_scale + return weights + diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py -index 0495e951ff746e44dbeebe844615418589d796fc..07b75b5c8c76f36f9ec7382364cfb84e577bc2cd 100644 +index 3490fc6a184eed9263f1647b8cc4960e8a9d40d3..b448595feb72b0c32aa50fdf1c1baf274e977ab5 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py -@@ -931,6 +931,10 @@ class CompressedTensorsLinearMethod(LinearMethodBase): +@@ -1021,6 +1021,10 @@ class CompressedTensorsLinearMethod(LinearMethodBase): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.scheme.process_weights_after_loading(layer) - + + def restore_weights_before_loading(self, layer: torch.nn.Module) -> None: + if hasattr(layer.scheme, "restore_weights_before_loading"): + layer.scheme.restore_weights_before_loading(layer) @@ -487,10 +471,10 @@ index 0495e951ff746e44dbeebe844615418589d796fc..07b75b5c8c76f36f9ec7382364cfb84e def create_weights( self, layer: torch.nn.Module, -@@ -985,6 +989,10 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase): +@@ -1075,6 +1079,10 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.scheme.process_weights_after_loading(layer) - + + def restore_weights_before_loading(self, layer: torch.nn.Module) -> None: + if hasattr(layer.scheme, "restore_weights_before_loading"): + layer.scheme.restore_weights_before_loading(layer) @@ -499,13 +483,13 @@ index 0495e951ff746e44dbeebe844615418589d796fc..07b75b5c8c76f36f9ec7382364cfb84e self, layer: torch.nn.Module, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index 8435da43cfea6a7f3c007cf9aa6c8ce5f6084136..17aab81733e7b0688c0ad00408588d58b02c0cbc 100644 +index 928348eb52ad811cb8055cb99d44d1eaa3b0e4e6..18b27b1b64652245d92107d0dc7e02d11d98a9cd 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -1715,6 +1715,16 @@ class ResumeMemoryOccupationReqOutput(BaseReq, kw_only=True): +@@ -1852,6 +1852,16 @@ class ReleaseMemoryOccupationReqOutput(BaseReq, kw_only=True): pass - - + + +class PostProcessWeightsReqInput(BaseReq, kw_only=True): + restore_weights_before_load: bool = False + post_process_quantization: bool = False @@ -516,33 +500,14 @@ index 8435da43cfea6a7f3c007cf9aa6c8ce5f6084136..17aab81733e7b0688c0ad00408588d58 + message: str + + - class CheckWeightsReqInput(BaseReq, kw_only=True): - action: str = "checksum" - allow_quant_error: bool = False -@@ -2030,7 +2040,7 @@ class GetLoadsReqInput(BaseReq, kw_only=True): - """Request for /v1/loads endpoint.""" - - VALID_SECTIONS = frozenset( -- {"core", "memory", "spec", "lora", "disagg", "queues", "all"} -+ {"core", "memory", "spec", "lora", "disagg", "queues", "inflight", "all"} - ) - - include: List[str] = msgspec.field(default_factory=lambda: ["all"]) -@@ -2074,6 +2084,9 @@ class GetLoadsReqOutput(BaseReq, kw_only=True): - lora: Optional[LoRAMetrics] = None - disaggregation: Optional[DisaggregationMetrics] = None - queues: Optional[QueueMetrics] = None -+ # Per-request breakdown of every queue, only populated when "inflight" or -+ # "all" is requested. -+ inflight: Optional[List[Dict[str, Any]]] = None - - - class SetInjectDumpMetadataReqInput(BaseReq, kw_only=True): + class ResumeMemoryOccupationReqInput(BaseReq, kw_only=True): + # Optional tags to identify the memory region, which is primarily used for RL + # Currently we only support `weights` and `kv_cache` diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py -index 55c88afa1978c189ade9e46a953661c2a69e9d29..99a45d9cab48d10d2d6e6bac54f33a58ef738664 100755 +index a45d122f0ff5831135ab841be0a3b2ea747ac0b2..0334e3530d716723d9c6a640654b3a30f9c1f44f 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py -@@ -977,6 +977,7 @@ class Req(ReqDllmMixin): +@@ -1103,6 +1103,7 @@ class Req(ReqDllmMixin): self.metrics_collector = metrics_collector if time_stats is not None: self.time_stats = SchedulerReqTimeStats.new_from_obj(time_stats) @@ -550,8 +515,8 @@ index 55c88afa1978c189ade9e46a953661c2a69e9d29..99a45d9cab48d10d2d6e6bac54f33a58 else: self.time_stats = SchedulerReqTimeStats(disagg_mode=disagg_mode) self.time_stats.set_metrics_collector(metrics_collector) -@@ -2490,11 +2491,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): - +@@ -2739,11 +2740,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): + retracted_reqs = [] first_iter = True + num_minimum_reqs = ( @@ -565,11 +530,11 @@ index 55c88afa1978c189ade9e46a953661c2a69e9d29..99a45d9cab48d10d2d6e6bac54f33a58 + if len(sorted_indices) <= num_minimum_reqs: + # Unified mode keeps one request; decode disaggregation may retract all. break - + first_iter = False -@@ -2505,7 +2509,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): +@@ -2754,7 +2758,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.release_req(idx, len(sorted_indices), server_args) - + reqs_to_abort: List[Req] = [] - if len(sorted_indices) <= 1 and not self.check_decode_mem( + if len(sorted_indices) <= num_minimum_reqs and not self.check_decode_mem( @@ -577,10 +542,10 @@ index 55c88afa1978c189ade9e46a953661c2a69e9d29..99a45d9cab48d10d2d6e6bac54f33a58 ): # Even the last remaining request cannot fit in memory. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index 0630f942398044341766dd9c0e67cb738f4e3cc7..d4169bebe0b58fc4f32d04e307e5132e66f364ff 100644 +index a9615a14c512345b2d63323ff3ebd2fe3d9a0dca..ecbd9bb4fe6ec40c84262addc1d0265ad8b281d3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py -@@ -126,6 +126,7 @@ from sglang.srt.managers.io_struct import ( +@@ -141,6 +141,7 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterReqOutput, OpenSessionReqInput, PauseGenerationReqInput, @@ -588,7 +553,7 @@ index 0630f942398044341766dd9c0e67cb738f4e3cc7..d4169bebe0b58fc4f32d04e307e5132e ProfileReq, ReleaseMemoryOccupationReqInput, RemoveExternalCorpusReqInput, -@@ -1416,6 +1417,10 @@ class Scheduler( +@@ -1539,6 +1540,10 @@ class Scheduler( UpdateWeightsFromIPCReqInput, self.weight_updater.update_weights_from_ipc, ), @@ -599,7 +564,7 @@ index 0630f942398044341766dd9c0e67cb738f4e3cc7..d4169bebe0b58fc4f32d04e307e5132e ( GetWeightsByNameReqInput, self.weight_updater.get_weights_by_name, -@@ -3988,6 +3993,12 @@ class Scheduler( +@@ -4457,6 +4462,12 @@ class Scheduler( # The request will still run one decode forward pass. # Then we reuse all existing code to clean up the KV cache allocation. logger.debug(f"Abort running request. {req.rid=}") @@ -610,112 +575,13 @@ index 0630f942398044341766dd9c0e67cb738f4e3cc7..d4169bebe0b58fc4f32d04e307e5132e + if sender is not None and hasattr(sender, "abort"): + sender.abort() req.to_finish = FINISH_ABORT() - + def _pause_engine(self) -> Tuple[List[Req], int]: -diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py -index b49409fc7c15cf9b5ba8c07306b7d8ff2cf46980..aa818c0a3b88a0f6bb97cc8bdc749e9af36f1c8e 100644 ---- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py -+++ b/python/sglang/srt/managers/scheduler_components/load_inquirer.py -@@ -212,6 +212,88 @@ class SchedulerLoadInquirer: - retracted=self.get_stats().num_retracted_reqs, - ) - -+ inflight = None -+ if include_all or "inflight" in include: -+ now_perf = time.perf_counter() -+ inflight_queues = [("running", self.get_running_batch().reqs, None)] -+ if self.disaggregation_mode == DisaggregationMode.PREFILL: -+ inflight_queues += [ -+ ("waiting", self.get_waiting_queue(), "wait_queue_entry_time"), -+ ( -+ "bootstrap", -+ self.get_disagg_prefill_bootstrap_queue().queue, -+ "prefill_bootstrap_queue_entry_time", -+ ), -+ ( -+ "prefill_inflight", -+ self.get_disagg_prefill_inflight_queue(), -+ "prefill_transfer_queue_entry_time", -+ ), -+ ] -+ elif self.disaggregation_mode == DisaggregationMode.DECODE: -+ inflight_queues += [ -+ ("waiting", self.get_waiting_queue(), "wait_queue_entry_time"), -+ ( -+ "prealloc", -+ self.get_disagg_decode_prealloc_queue().queue, -+ "decode_prealloc_queue_entry_time", -+ ), -+ ( -+ "transfer", -+ self.get_disagg_decode_transfer_queue().queue, -+ "decode_transfer_queue_entry_time", -+ ), -+ ( -+ "retracted", -+ self.get_disagg_decode_prealloc_queue().retracted_queue, -+ "decode_prealloc_queue_entry_time", -+ ), -+ ] -+ else: -+ inflight_queues.append( -+ ("waiting", self.get_waiting_queue(), "wait_queue_entry_time") -+ ) -+ -+ def describe_req(entry, stage, entry_time_field): -+ req = getattr(entry, "req", entry) -+ info = { -+ "rid": getattr(req, "rid", None), -+ "bootstrap_room": getattr(req, "bootstrap_room", None), -+ "seqlen": getattr(entry, "seqlen", None), -+ "stage": stage, -+ } -+ if entry_time_field is not None: -+ time_stats = getattr(req, "time_stats", None) -+ entry_time = ( -+ getattr(time_stats, entry_time_field, 0.0) -+ if time_stats -+ else 0.0 -+ ) -+ info["age_s"] = ( -+ round(now_perf - entry_time, 3) if entry_time else None -+ ) -+ if entry is not req: -+ info["waiting_for_input"] = getattr( -+ entry, "waiting_for_input", None -+ ) -+ info["timeout_cancel_issued"] = getattr( -+ entry, "timeout_cancel_issued", None -+ ) -+ return info -+ -+ inflight = [] -+ for name, queue, entry_time_field in inflight_queues: -+ inflight.append( -+ { -+ "name": name, -+ "num_reqs": len(queue), -+ "reqs": [ -+ describe_req(entry, name, entry_time_field) -+ for entry in queue -+ ], -+ } -+ ) -+ - return GetLoadsReqOutput( - dp_rank=self.ps.dp_rank, - timestamp=time.time(), -@@ -231,4 +313,5 @@ class SchedulerLoadInquirer: - lora=lora, - disaggregation=disaggregation, - queues=queues, -+ inflight=inflight, - ) diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py -index 90945b82d63a8cf5116bf8b7563abd947ecebb1e..a34ee065157090033dd07a0ef7c64225e81de4ab 100644 +index 39eeec2e71ca3628c3139911f3e6ee6f4e771b50..12ba5133d890c6956b5b97ed961593530e03b911 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/python/sglang/srt/managers/scheduler_components/profiler_manager.py -@@ -384,7 +384,7 @@ class SchedulerProfilerManager: +@@ -397,7 +397,7 @@ class SchedulerProfilerManager: if self.profiler_prefill_ct > self.profiler_target_prefill_ct: if self.profile_in_progress: self._stop_profile(stage=ForwardMode.EXTEND) @@ -725,10 +591,10 @@ index 90945b82d63a8cf5116bf8b7563abd947ecebb1e..a34ee065157090033dd07a0ef7c64225 if self.profile_in_progress: # force trace flush diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py -index 8c135d12f7196e9e858db77995233af1c8cc4fbb..40e50efa440f8ec48a1b8c30b412a8d3cf525ea3 100644 +index 653b28c4186823a01801e14a624b1828eac51538..2c876d270cb1a1cf5476b2ef6e79de6c9fb4431a 100644 --- a/python/sglang/srt/managers/scheduler_components/weight_updater.py +++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py -@@ -26,6 +26,8 @@ from sglang.srt.managers.io_struct import ( +@@ -28,6 +28,8 @@ from sglang.srt.managers.io_struct import ( GetWeightsByNameReqOutput, InitWeightsUpdateGroupReqInput, InitWeightsUpdateGroupReqOutput, @@ -737,10 +603,10 @@ index 8c135d12f7196e9e858db77995233af1c8cc4fbb..40e50efa440f8ec48a1b8c30b412a8d3 ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, ResumeMemoryOccupationReqInput, -@@ -181,6 +183,19 @@ class SchedulerWeightUpdaterManager: - parameter = self.tp_worker.get_weights_by_name(recv_req) - return GetWeightsByNameReqOutput(parameter=parameter) - +@@ -199,6 +201,19 @@ class SchedulerWeightUpdaterManager: + f"Restart with --weight-cache-mode off to use this operation." + ) + + def post_process_weights(self, recv_req: PostProcessWeightsReqInput): + success, message = self.tp_worker.post_process_weights(recv_req) + if ( @@ -758,7 +624,7 @@ index 8c135d12f7196e9e858db77995233af1c8cc4fbb..40e50efa440f8ec48a1b8c30b412a8d3 assert ( self.is_fully_idle() diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py -index 53b376705a4fbf6b2cb22330d62ef78655d128be..b733d20234ea1d3cb8b9bae77259229bea76a059 100644 +index 7cfd98b8f318af4373b86dd5d25603ba1a5fc664..efe7ba9d360d7ed5f08274269b0543f8d04af0ba 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -48,6 +48,8 @@ from sglang.srt.managers.io_struct import ( @@ -770,7 +636,7 @@ index 53b376705a4fbf6b2cb22330d62ef78655d128be..b733d20234ea1d3cb8b9bae77259229b ProfileReq, ProfileReqOutput, ProfileReqType, -@@ -99,6 +101,7 @@ _COMMUNICATOR_SPECS = [ +@@ -101,6 +103,7 @@ _COMMUNICATOR_SPECS = [ ("send_weights_to_remote_instance", SendWeightsToRemoteInstanceReqOutput), ("update_weights_from_tensor", UpdateWeightsFromTensorReqOutput), ("update_weights_from_ipc", UpdateWeightsFromIPCReqOutput), @@ -778,10 +644,10 @@ index 53b376705a4fbf6b2cb22330d62ef78655d128be..b733d20234ea1d3cb8b9bae77259229b ("get_weights_by_name", GetWeightsByNameReqOutput), ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), -@@ -747,6 +750,16 @@ class TokenizerControlMixin: +@@ -797,6 +800,16 @@ class TokenizerControlMixin: self.auto_create_handle_loop() await self.resume_memory_occupation_communicator(obj) - + + async def post_process_weights( + self: TokenizerManager, + obj: PostProcessWeightsReqInput, @@ -796,10 +662,10 @@ index 53b376705a4fbf6b2cb22330d62ef78655d128be..b733d20234ea1d3cb8b9bae77259229b self: TokenizerManager, obj: CheckWeightsReqInput, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index 74688606f99a0960f0abb602021250e7fabd2128..f4331b64b42a081aa31e63370efc82a96e416e40 100644 +index 80351f330430e0129d5730dece42fd82913a99a0..7472d4b9f95b7f8e59c05953484325a2e36e9a9f 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py -@@ -2408,25 +2408,23 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): +@@ -2840,27 +2840,25 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): priority = getattr(state.obj, "priority", None) if priority is not None: labels["priority"] = str(priority) @@ -811,11 +677,15 @@ index 74688606f99a0960f0abb602021250e7fabd2128..f4331b64b42a081aa31e63370efc82a9 state.ttft_observed = True state.last_completion_tokens = completion_tokens - self.metrics_collector.observe_time_to_first_token( -- labels, state.time_stats.get_first_token_latency() +- labels, +- state.time_stats.get_first_token_latency(), +- stream=getattr(state.obj, "stream", False), - ) + if self.disaggregation_mode != DisaggregationMode.PREFILL: + self.metrics_collector.observe_time_to_first_token( -+ labels, state.time_stats.get_first_token_latency() ++ labels, ++ state.time_stats.get_first_token_latency(), ++ stream=getattr(state.obj, "stream", False), + ) else: num_new_tokens = completion_tokens - state.last_completion_tokens @@ -829,14 +699,14 @@ index 74688606f99a0960f0abb602021250e7fabd2128..f4331b64b42a081aa31e63370efc82a9 state.time_stats.set_last_time() - state.last_completion_tokens = completion_tokens + state.last_completion_tokens = completion_tokens - + if state.finished: # Get detailed cache breakdown if available diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py -index 3db49daf82c030abd34ab43859e9e53958b777c0..6c361995008e0e2178d7c3c8bf7abd8fd55d5416 100644 +index 0fdf5b96657f12efcb568724112b85f5d3155155..8f30c4e1994c66705965f19b3125c8790a931aaa 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py -@@ -29,6 +29,7 @@ from sglang.srt.managers.io_struct import ( +@@ -30,6 +30,7 @@ from sglang.srt.managers.io_struct import ( InitWeightsUpdateGroupReqInput, LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, @@ -844,25 +714,24 @@ index 3db49daf82c030abd34ab43859e9e53958b777c0..6c361995008e0e2178d7c3c8bf7abd8f SendWeightsToRemoteInstanceReqInput, UnloadLoRAAdapterReqInput, UpdateWeightFromDiskReqInput, -@@ -162,6 +163,13 @@ class BaseTpWorker(ABC): +@@ -218,6 +219,12 @@ class BaseTpWorker(ABC): ) return success, message - + + def post_process_weights(self, recv_req: PostProcessWeightsReqInput): -+ success, message = self.model_runner.post_process_weights( ++ return self.model_runner.post_process_weights( + restore_weights_before_load=recv_req.restore_weights_before_load, + post_process_quantization=recv_req.post_process_quantization, + ) -+ return success, message + - def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): - - monkey_patch_torch_reductions() + def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput): + parameter = self.model_runner.weight_exporter.get_weights_by_name( + recv_req.name, recv_req.truncate_size diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py -index c6ce33af36ca98e175ced37c6e76378f209918b5..378a29e936e6c1ad3213f710cbf5bd9380bd019f 100644 +index 5ebac56f30eb01f42e027d1892f6defd026c3034..e9c2d85daa46202d5f3e3b2ab7a61ac56faba9be 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py -@@ -1034,9 +1034,7 @@ class HiRadixCache(RadixCache): +@@ -1137,9 +1137,7 @@ class HiRadixCache(RadixCache): self._update_leaf_status(node) self._update_host_leaf_status(node) if node.parent is None: @@ -872,22 +741,22 @@ index c6ce33af36ca98e175ced37c6e76378f209918b5..378a29e936e6c1ad3213f710cbf5bd93 + break node = node.parent return DecLockRefResult(delta=delta) - -@@ -1142,6 +1140,7 @@ class HiRadixCache(RadixCache): + +@@ -1245,6 +1243,7 @@ class HiRadixCache(RadixCache): self._update_host_leaf_status(node) # update leaf status for the parent because the node is evicted self._update_leaf_status(node.parent) + self._update_host_leaf_status(node.parent) return num_evicted - + def _evict_backuped(self, node: TreeNode): diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py -index 87e2638c1bc4a1afb6cbef92bdef2fa4ad1f3b77..25f2539857a083347f84993b8361519a08ed6fb5 100644 +index a9b2b8b994a03e401b8658b5bdcef6e0328769ed..27b1c48e8bcc92d17408f07dfcb7c6bce21c9546 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py -@@ -3065,9 +3065,12 @@ class DSATokenToKVPool(MLATokenToKVPool): - else: - assert self.page_size == 64 +@@ -4391,9 +4391,12 @@ class DSATokenToKVPool(MLATokenToKVPool): + def _create_index_buffers(self): + num_pages = (self.index_buf_size + self.page_size + 1) // self.page_size with ( - torch.cuda.use_mem_pool(self.custom_mem_pool) - if self.custom_mem_pool @@ -902,12 +771,12 @@ index 87e2638c1bc4a1afb6cbef92bdef2fa4ad1f3b77..25f2539857a083347f84993b8361519a self.index_k_with_scale_buffer = [ torch.zeros( diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py -index 6ebf09939ce2d92260716167fe0f7ed0cdb712cc..0a0555a28b63088442545633eeb37ce8debb1c2c 100644 +index 2916f8c6474f3697c7efbd510310fc96437dcab3..32b81c490e60cb87087d3ae47eb0ec51a534b4e0 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py -@@ -491,6 +491,9 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache): +@@ -491,6 +491,9 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): return - + token_ids = req.get_fill_ids() + kv_committed_len = getattr(req, "kv_committed_len", None) + if kv_committed_len is not None and len(token_ids) > kv_committed_len: @@ -915,7 +784,7 @@ index 6ebf09939ce2d92260716167fe0f7ed0cdb712cc..0a0555a28b63088442545633eeb37ce8 kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : len(token_ids) ] -@@ -619,9 +622,7 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache): +@@ -619,9 +622,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): node.lock_ref -= 1 self._update_leaf_status(node) if node.parent is None: @@ -925,54 +794,38 @@ index 6ebf09939ce2d92260716167fe0f7ed0cdb712cc..0a0555a28b63088442545633eeb37ce8 + break node = node.parent return DecLockRefResult(delta=delta) - + diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index ea17a25587f432f224d965abc2f79e23bc185fa2..159ef8fc5a3ddcc30366c79c12ed4b89371d14a3 100644 +index 8e957f66156bef42145acbca8de8f0407abaeb20..e2afff41c8a609b62baa5dd82f023adfc28c799b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py -@@ -19,6 +19,7 @@ import contextlib - import datetime - import gc - import inspect -+import json - import logging - import os - import socket -@@ -558,7 +559,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -376,9 +376,12 @@ class ModelRunner: self.war_fastpath_read_done_event: Optional[torch.cuda.Event] = None - + # CPU offload -- set_offloader(create_offloader_from_server_args(server_args, dp_rank=dp_rank)) +- set_offloader( +- create_offloader_from_server_args(server_args, dp_rank=self.ps.dp_rank) +- ) + if not is_draft_worker: + set_offloader( -+ create_offloader_from_server_args(server_args, dp_rank=dp_rank) ++ create_offloader_from_server_args( ++ server_args, dp_rank=self.ps.dp_rank ++ ) + ) - - self._weight_checker = WeightChecker(model_runner=self) - -@@ -1842,7 +1846,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): - weight_name_filter: Optional[Callable[[str], bool]] = None, - recapture_cuda_graph: bool = False, - ) -> tuple[bool, str]: -- """Update engine weights in-place from the disk.""" -+ """Update engine weights in-place from disk.""" - logger.info( - f"Update engine weights online from disk begin. " - f"avail mem={get_available_gpu_memory(self.device, self.gpu_id, empty_cache=False):.2f} GB" -@@ -3058,6 +3062,12 @@ class ModelRunner(ModelRunnerKVCacheMixin): + + self._weight_checker = WeightChecker(get_model=lambda: self.model, ps=self.ps) + +@@ -1441,6 +1444,9 @@ class ModelRunner: output.expert_distribution_metrics = recorder_outputs.get("metrics") - - no_copy_to_cpu = not self.server_args.disable_overlap_schedule -+ cuda_graph_num_tokens = None -+ decode_graph_runner = getattr(self, "decode_cuda_graph_runner", None) -+ if getattr(decode_graph_runner, "bs", None): -+ cuda_graph_num_tokens = decode_graph_runner.bs * getattr( -+ decode_graph_runner, "num_tokens_per_bs", 1 -+ ) + + no_copy_to_cpu = not get_schedule().disable_overlap_schedule ++ cuda_graph_num_tokens = getattr(self.decode_cuda_graph_runner, "bs", None) ++ if cuda_graph_num_tokens is not None: ++ cuda_graph_num_tokens *= self.decode_cuda_graph_runner.captured_req_width if ( not self.is_draft_worker and (experts_capturer := get_global_experts_capturer()) is not None -@@ -3065,7 +3075,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -1448,7 +1454,7 @@ class ModelRunner: output.routed_experts_output = experts_capturer.on_forward_end( forward_batch=forward_batch, can_run_graph=output.can_run_graph, @@ -980,8 +833,8 @@ index ea17a25587f432f224d965abc2f79e23bc185fa2..159ef8fc5a3ddcc30366c79c12ed4b89 + cuda_graph_batch=cuda_graph_num_tokens, no_copy_to_cpu=no_copy_to_cpu, ) - -@@ -3073,7 +3083,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): + +@@ -1456,7 +1462,7 @@ class ModelRunner: output.indexer_topk_output = indexer_capturer.on_forward_end( forward_batch=forward_batch, can_run_graph=output.can_run_graph, @@ -989,11 +842,11 @@ index ea17a25587f432f224d965abc2f79e23bc185fa2..159ef8fc5a3ddcc30366c79c12ed4b89 + cuda_graph_batch=cuda_graph_num_tokens, no_copy_to_cpu=no_copy_to_cpu, ) - -@@ -3359,6 +3369,39 @@ class ModelRunner(ModelRunnerKVCacheMixin): - logger.error(f"IPC weight update failed: {e}") - return False, str(e) - + +@@ -1712,6 +1718,39 @@ class ModelRunner: + action=action, allow_quant_error=allow_quant_error + ) + + def post_process_weights( + self, + restore_weights_before_load: bool = False, @@ -1008,7 +861,7 @@ index ea17a25587f432f224d965abc2f79e23bc185fa2..159ef8fc5a3ddcc30366c79c12ed4b89 + target_device = torch.device(self.device) + + if restore_weights_before_load: -+ for _, module in self.model.named_modules(): ++ for module in self.model.modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is not None and hasattr( + quant_method, "restore_weights_before_loading" @@ -1017,7 +870,7 @@ index ea17a25587f432f224d965abc2f79e23bc185fa2..159ef8fc5a3ddcc30366c79c12ed4b89 + quant_method.restore_weights_before_loading(module) + + if post_process_quantization: -+ for _, module in self.model.named_modules(): ++ for module in self.model.modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is not None and hasattr( + quant_method, "process_weights_after_loading" @@ -1027,40 +880,26 @@ index ea17a25587f432f224d965abc2f79e23bc185fa2..159ef8fc5a3ddcc30366c79c12ed4b89 + + return True, "Success" + - def prealloc_symmetric_memory_pool(self): - # PyTorch mempools never de-fragment memory in OOM scenarios, so we need to pre-allocate a large chunk of memory to limit fragmentation. - if ( -@@ -3408,6 +3451,13 @@ class ModelRunner(ModelRunnerKVCacheMixin): - return output - - -+def _resolve_torch_dtype(dtype: Union[str, torch.dtype]) -> torch.dtype: -+ if isinstance(dtype, torch.dtype): -+ return dtype -+ dtype_name = dtype.split(".", 1)[1] if dtype.startswith("torch.") else dtype -+ return getattr(torch, dtype_name) -+ -+ - def _model_load_weights_direct(model, named_tensors: List[Tuple[str, torch.Tensor]]): - params_dict = dict(model.named_parameters()) - for name, tensor in named_tensors: + def _expand_eplb_metadata_for_scale( + self, + from_ep_size: int, diff --git a/python/sglang/srt/models/glm4v_moe.py b/python/sglang/srt/models/glm4v_moe.py -index 83d55414efca159a2e7bf8b20d46346a91fcf41b..025570d76a720603506add59d45040a71ebe3a31 100644 +index 0b5e21f587a49b9c9df8af42172e41e087b81bb5..49e349e9f6eab6bd7d0740df99903b8d1c88f7a3 100644 --- a/python/sglang/srt/models/glm4v_moe.py +++ b/python/sglang/srt/models/glm4v_moe.py -@@ -42,6 +42,8 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): - +@@ -41,6 +41,8 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): + self.pp_group = get_pp_group() self.config = config + self.config.encoder_only = getattr(config, "encoder_only", False) + self.config.language_only = getattr(config, "language_only", False) - self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder + self.use_data_parallel = get_mm().mm_enable_dp_encoder vision_utils.update_vit_attn_dummy_heads_config(self.config) self.tp_size = get_parallel().tp_size -@@ -49,11 +51,30 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): +@@ -48,11 +50,30 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): self.num_fused_shared_experts = 0 self.determine_num_fused_shared_experts() - + - self.model = Glm4MoeModel( - config, - quant_config, @@ -1082,7 +921,7 @@ index 83d55414efca159a2e7bf8b20d46346a91fcf41b..025570d76a720603506add59d45040a7 + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), -+ use_attn_tp_group=get_global_server_args().enable_dp_lm_head, ++ use_attn_tp_group=get_parallel().enable_dp_lm_head, + ) + else: + # ranks other than the last rank will have a placeholder layer @@ -1093,10 +932,10 @@ index 83d55414efca159a2e7bf8b20d46346a91fcf41b..025570d76a720603506add59d45040a7 self.visual = Glm4vVisionModel( config.vision_config, quant_config=quant_config, -@@ -61,24 +82,14 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): +@@ -60,24 +81,14 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): use_data_parallel=self.use_data_parallel, ) - + - if self.pp_group.is_last_rank: - if self.pp_group.world_size == 1 and self.config.tie_word_embeddings: - self.lm_head = self.model.embed_tokens @@ -1106,7 +945,7 @@ index 83d55414efca159a2e7bf8b20d46346a91fcf41b..025570d76a720603506add59d45040a7 - config.hidden_size, - quant_config=quant_config, - prefix=add_prefix("lm_head", prefix), -- use_attn_tp_group=get_flags().enable_dp_lm_head, +- use_attn_tp_group=get_parallel().enable_dp_lm_head, - ) - else: - # ranks other than the last rank will have a placeholder layer @@ -1121,10 +960,10 @@ index 83d55414efca159a2e7bf8b20d46346a91fcf41b..025570d76a720603506add59d45040a7 + or {} + ) + self.is_mrope_enabled = "mrope_section" in rope_config - + # For EAGLE3 support self.capture_aux_hidden_states = False -@@ -221,6 +232,10 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): +@@ -220,6 +231,10 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue @@ -1134,17 +973,17 @@ index 83d55414efca159a2e7bf8b20d46346a91fcf41b..025570d76a720603506add59d45040a7 + continue if name not in params_dict: continue - -@@ -236,6 +251,8 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): + +@@ -235,6 +250,8 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue + if "visual" in name or self.config.encoder_only: + continue - + # Mark as expert weight regardless of whether we can process it is_expert_weight = True -@@ -267,6 +284,10 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): +@@ -266,6 +283,10 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue @@ -1154,12 +993,12 @@ index 83d55414efca159a2e7bf8b20d46346a91fcf41b..025570d76a720603506add59d45040a7 + continue if name not in params_dict: continue - + diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py -index 75c807a5d7f9ff2a31a4b3ba1fa5e63c12909df0..56675ef0ade24c0e12d32bed7c3fc754dd6d9744 100644 +index 83697f140a9f0bb0683d90e50f168248478fba8c..aae8b5c80071e53980099e6c5052ddf109c61f89 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py -@@ -1169,9 +1169,14 @@ class Qwen3LLMModel(Qwen3Model): +@@ -1171,9 +1171,14 @@ class Qwen3LLMModel(Qwen3Model): # To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack # The order matters because addition with different tensors is not associative in practice. # Deepstack for prev_layer is applied at the start of current layer via post_residual_addition. @@ -1183,7 +1022,7 @@ index db684259d2f6aea80ead030556042c7333384fc3..17d2cb6958afe6311a6a61e93625c289 +++ b/python/sglang/srt/multimodal/processors/glm4v.py @@ -1,7 +1,13 @@ from typing import List, Union - + +import torch + from sglang.srt.layers.rotary_embedding import MRotaryEmbedding @@ -1202,13 +1041,13 @@ index db684259d2f6aea80ead030556042c7333384fc3..17d2cb6958afe6311a6a61e93625c289 self.VIDEO_END_TOKEN_ID = hf_config.video_end_token_id + self.IM_START_TOKEN_ID = self.IMAGE_START_TOKEN_ID + self.IM_END_TOKEN_ID = self.IMAGE_END_TOKEN_ID - + # Vision config self.IMAGE_FACTOR = 28 @@ -60,6 +68,39 @@ class Glm4vImageProcessor(SGLangBaseProcessor): video_token_id=self.IM_TOKEN_ID, ).build(_processor) - + + def get_mm_data(self, prompt, embeddings, img_grid_thw): + input_ids, offsets, _ = self.build_input_ids(prompt, img_grid_thw=img_grid_thw) + image_embeddings = ( @@ -1246,10 +1085,10 @@ index db684259d2f6aea80ead030556042c7333384fc3..17d2cb6958afe6311a6a61e93625c289 image_grid_thw = None video_grid_thw = None diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py -index b8774ebade56891fe0aef2f54aa26e39b0c63b23..fa01537b2010c6ee002de514281c0d36a2c3bad2 100644 +index a53e93730d45c3aadc7b4ac6969f2b02bdd2f4b6..1ae2c6f4f6566052d16f6d40a4b66389b50b9b18 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py +++ b/python/sglang/srt/multimodal/processors/qwen_vl.py -@@ -678,7 +678,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor): +@@ -723,7 +723,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor): **kwargs, ): entry_time = time.perf_counter() @@ -1259,22 +1098,32 @@ index b8774ebade56891fe0aef2f54aa26e39b0c63b23..fa01537b2010c6ee002de514281c0d36 image_data=image_data, video_data=request_obj.video_data, diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py -index 27734d988bf416ce94e3cecc72690b9096d2ba75..8c4c5f0edccfba0e60bb0d90caa240b6caff03f2 100644 +index 3bcb6b3629e0f77d68ee9b3c83278a5b300ba27d..2c31446fc4b50b9bf04290c1cee001088c624e60 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -346,7 +346,7 @@ class ReqTimeStatsBase: state["trace_ctx"] = TraceNullContext() - + for key in state.keys(): - if key.endswith("time"): + if key.endswith("time") and state[key] > 0.0: state[key] = convert_time_cross_thread( state[key], state["diff_realtime_monotonic"], -@@ -632,9 +632,20 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): - return {} - +@@ -623,6 +623,9 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + transfer_speed_gb_s: float = 0.0 + transfer_total_mb: float = 0.0 + ++ # Number of prefill retries for this request ++ prefill_retry_count: int = 0 ++ + has_timing_data: bool = False + + def __getstate__(self) -> object: +@@ -632,9 +635,20 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + state = { + "has_timing_data": True, + "enable_metrics": self.enable_metrics, + "disagg_mode": self.disagg_mode, "wait_queue_entry_time": self.wait_queue_entry_time, @@ -1292,8 +1141,8 @@ index 27734d988bf416ce94e3cecc72690b9096d2ba75..8c4c5f0edccfba0e60bb0d90caa240b6 "diff_realtime_monotonic": global_diff_realtime_monotonic, } return state -@@ -1150,6 +1161,13 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): - +@@ -1149,6 +1163,13 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + def convert_to_output_meta_info(self): meta_data = {} + @@ -1306,7 +1155,7 @@ index 27734d988bf416ce94e3cecc72690b9096d2ba75..8c4c5f0edccfba0e60bb0d90caa240b6 if self.forward_entry_time > 0.0: meta_data["forward_entry_time"] = convert_time_to_realtime( self.forward_entry_time -@@ -1163,6 +1181,67 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): +@@ -1162,6 +1183,67 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): "queue_time": self.get_queueing_time(), } ) @@ -1372,13 +1221,13 @@ index 27734d988bf416ce94e3cecc72690b9096d2ba75..8c4c5f0edccfba0e60bb0d90caa240b6 + if self.transfer_total_mb > 0.0: + meta_data["pd_transfer_total_mb"] = self.transfer_total_mb return meta_data - + def format_duration(self, duration: float) -> str: diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py -index 7cf12912ce7ac82ce4f0d2cb9046ec52b8df022f..d48cffa15e64a2ff8ed1a7e20a46e44df0dfa8e8 100644 +index 13371f43259d797a3a6c068d14dc43fbd73a2fc4..1bd918564e3cb94eab8d5c7d9eb6bffce2d2e30d 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py -@@ -554,8 +554,10 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): +@@ -566,8 +566,10 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): forward_batch.seq_lens, forward_batch.out_cache_loc, forward_batch.positions, @@ -1392,14 +1241,14 @@ index 7cf12912ce7ac82ce4f0d2cb9046ec52b8df022f..d48cffa15e64a2ff8ed1a7e20a46e44d ] if buffers.rids_int is not None and forward_batch.rids_int is not None: diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py -index 82db90da3bfb071542366f2d902e71364f0ea2cf..6a1fe354b2b2b2c59309deb1b91af06ada759d9b 100644 +index 6636d06abe7d89db7cb7e6e79cb551366bb41175..ead53179e41abb7b8862b4965ba673980d30328b 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py -@@ -2382,7 +2382,13 @@ def _get_fastapi_request_path(request) -> Tuple[str, bool]: +@@ -2482,7 +2482,13 @@ def _get_fastapi_request_path(request) -> Tuple[str, bool]: for route in request.app.routes: match, child_scope = route.matches(request.scope) if match == Match.FULL: -- return route.path, True +- return getattr(route, "path", request.url.path), True + path = getattr(route, "path", None) or getattr(route, "path_format", None) + if path is None: + child_route = child_scope.get("route") @@ -1407,14 +1256,14 @@ index 82db90da3bfb071542366f2d902e71364f0ea2cf..6a1fe354b2b2b2c59309deb1b91af06a + child_route, "path_format", None + ) + return path or request.url.path, True - + return request.url.path, False - -@@ -2733,6 +2739,7 @@ class SafeUnpickler(pickle.Unpickler): + +@@ -2834,6 +2840,7 @@ class SafeUnpickler(pickle.Unpickler): "sglang.srt.utils.", "sglang.srt.disaggregation.", "sglang.srt.managers.", + "slime.", "torch_npu.", } - + diff --git a/slime/backends/megatron_utils/alignment/deepgemm_moe_forward.py b/slime/backends/megatron_utils/alignment/deepgemm_moe_forward.py index 30e9640e17..ec43943ac9 100644 --- a/slime/backends/megatron_utils/alignment/deepgemm_moe_forward.py +++ b/slime/backends/megatron_utils/alignment/deepgemm_moe_forward.py @@ -34,6 +34,7 @@ _should_log_deepgemm_summary, _sum_to_parameter_dtype, ) +from slime.backends.sglang_utils.compat import import_sglang_module logger = logging.getLogger(__name__) @@ -1125,11 +1126,18 @@ def _load_deepgemm_ops() -> _DeepGEMMOps: import deep_gemm from sgl_kernel import silu_and_mul from sglang.srt.layers import deep_gemm_wrapper - from sglang.srt.layers.moe.ep_moe.kernels import tma_align_input_scale - from sglang.srt.layers.quantization.fp8_kernel import sglang_per_token_group_quant_fp8 from slime.backends.megatron_utils.kernels.fp8_kernel import blockwise_cast_to_fp8_triton + ep_moe_kernels = import_sglang_module( + "sglang.kernels.ops.moe.ep_moe_kernels", + "sglang.srt.layers.moe.ep_moe.kernels", + ) + fp8_kernel = import_sglang_module( + "sglang.kernels.ops.quantization.fp8_kernel", + "sglang.srt.layers.quantization.fp8_kernel", + ) + _configure_batch_invariant(deep_gemm, deep_gemm_wrapper) scale_ue8m0 = bool(deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0) m_alignment = int(deep_gemm.get_mk_alignment_for_contiguous_layout()) @@ -1143,8 +1151,8 @@ def _load_deepgemm_ops() -> _DeepGEMMOps: # and activation/weight scales TMA-aligned as a separate step. Unchanged. return _DeepGEMMOps( quantize_weight=blockwise_cast_to_fp8_triton, - quantize_activation=sglang_per_token_group_quant_fp8, - align_input_scale=tma_align_input_scale, + quantize_activation=fp8_kernel.sglang_per_token_group_quant_fp8, + align_input_scale=ep_moe_kernels.tma_align_input_scale, grouped_gemm=deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig, silu_and_mul=silu_and_mul, scale_ue8m0=False, @@ -1166,8 +1174,8 @@ def _quantize_weight_ue8m0(weight: torch.Tensor, block: tuple[int, int]): return _DeepGEMMOps( quantize_weight=_quantize_weight_ue8m0, - quantize_activation=sglang_per_token_group_quant_fp8, - align_input_scale=tma_align_input_scale, + quantize_activation=fp8_kernel.sglang_per_token_group_quant_fp8, + align_input_scale=ep_moe_kernels.tma_align_input_scale, grouped_gemm=deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig, silu_and_mul=silu_and_mul, scale_ue8m0=True, @@ -2448,7 +2456,10 @@ def forward( ) if topk_indices.shape != topk_weights.shape or topk_indices.shape != output_index.shape: raise ValueError("DeepEP gather IDs, weights, and output indices must align") - from sglang.srt.layers.moe.ep_moe.kernels import ep_gather + ep_moe_kernels = import_sglang_module( + "sglang.kernels.ops.moe.ep_moe_kernels", + "sglang.srt.layers.moe.ep_moe.kernels", + ) output_shape = (topk_indices.shape[0], hidden_states.shape[1]) output = torch.empty( @@ -2456,7 +2467,7 @@ def forward( device=hidden_states.device, dtype=hidden_states.dtype, ) - ep_gather( + ep_moe_kernels.ep_gather( hidden_states, topk_indices, topk_weights, diff --git a/slime/backends/sglang_utils/compat.py b/slime/backends/sglang_utils/compat.py new file mode 100644 index 0000000000..7d56cd5b75 --- /dev/null +++ b/slime/backends/sglang_utils/compat.py @@ -0,0 +1,15 @@ +"""Compatibility helpers for SGLang modules moved between supported releases.""" + +from importlib import import_module +from types import ModuleType + + +def import_sglang_module(current_path: str, legacy_path: str) -> ModuleType: + """Import an SGLang 0.5.17 module, falling back to its 0.5.15 path.""" + try: + return import_module(current_path) + except ModuleNotFoundError as error: + missing_path = error.name + if missing_path is None or not (current_path == missing_path or current_path.startswith(f"{missing_path}.")): + raise + return import_module(legacy_path) diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index 3e562a7e38..08fed2616f 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -65,7 +65,6 @@ def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process: from sglang.srt.entrypoints.http_server import launch_server multiprocessing.set_start_method("spawn", force=True) - server_args.host = server_args.host.strip("[]") p = multiprocessing.Process(target=launch_server, args=(server_args,)) p.start() @@ -188,7 +187,11 @@ def _sanity_check_server_args(actual_server_args, expect_server_args): def _init_normal(self, server_args_dict): logger.info(f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}") - self.process = launch_server_process(ServerArgs(**server_args_dict)) + local_server_args = { + **server_args_dict, + "host": server_args_dict["host"].strip("[]"), + } + self.process = launch_server_process(ServerArgs(**local_server_args)) self._register_to_router(server_args_dict) def _register_to_router(self, server_args_dict): diff --git a/slime/utils/routing_replay.py b/slime/utils/routing_replay.py index c26a9e192e..3d94136da5 100644 --- a/slime/utils/routing_replay.py +++ b/slime/utils/routing_replay.py @@ -1,6 +1,9 @@ import os + import torch +from slime.backends.sglang_utils.compat import import_sglang_module + ROUTING_REPLAY = None ORDERED_TOPK_CAPTURE_ROUTER = None @@ -27,6 +30,87 @@ def consume_ordered_topk(module): return module.__dict__.pop("_slime_ordered_topk_indices", None) +def _run_sglang_fused_sigmoid_topk(logits, expert_bias, topk): + moe_fused_gate = import_sglang_module( + "sglang.kernels.ops.moe.moe_fused_gate", + "sglang.jit_kernel.moe_fused_gate", + ).moe_fused_gate + + return moe_fused_gate( + logits, + expert_bias, + topk, + scoring_func="sigmoid", + renormalize=True, + ) + + +class _SGLangFusedSigmoidTopK(torch.autograd.Function): + """SGLang's fused router forward with the corresponding analytic backward.""" + + @staticmethod + def forward(ctx, logits, expert_bias, topk): + weights, indices = _run_sglang_fused_sigmoid_topk(logits, expert_bias, topk) + ctx.save_for_backward(logits, weights, indices) + ctx.mark_non_differentiable(indices) + return weights, indices + + @staticmethod + def backward(ctx, grad_weights, _grad_indices): + if not ctx.needs_input_grad[0]: + return None, None, None + + logits, weights, indices = ctx.saved_tensors + if grad_weights is None: + return torch.zeros_like(logits), None, None + + indices_long = indices.to(dtype=torch.long) + selected_logits = logits.gather(1, indices_long).float() + selected_scores = torch.sigmoid(selected_logits) + score_sum = selected_scores.sum(dim=-1, keepdim=True) + centered_grad = grad_weights.float() - (grad_weights.float() * weights).sum(dim=-1, keepdim=True) + grad_selected = centered_grad * selected_scores * (1.0 - selected_scores) / score_sum + grad_logits = torch.zeros_like(logits) + grad_logits.scatter_(1, indices_long, grad_selected.to(dtype=logits.dtype)) + return grad_logits, None, None + + +def maybe_sglang_fused_sigmoid_topk( + logits, + topk, + *, + use_pre_softmax=False, + num_groups=None, + group_topk=None, + scaling_factor=None, + score_function="softmax", + expert_bias=None, + replay_active=False, +): + """Use SGLang's batch-invariant GLM router only for an aligned router.""" + if ORDERED_TOPK_CAPTURE_ROUTER is None or replay_active: + return None + if ( + score_function != "sigmoid" + or group_topk + or num_groups + or expert_bias is None + or topk <= 1 + or logits.dtype != torch.float32 + or (scaling_factor is not None and float(scaling_factor) != 1.0) + ): + return None + + # Megatron's ``use_pre_softmax`` only changes the softmax branch. The + # sigmoid path always selects from sigmoid(logits) + expert_bias and then + # renormalizes the selected unbiased sigmoid scores, matching SGLang. + del use_pre_softmax + + weights, indices = _SGLangFusedSigmoidTopK.apply(logits, expert_bias, topk) + _capture_ordered_topk(indices) + return weights, indices + + def register_ordered_topk_capture(module): """Capture one forward's SGLang-compatible top-k order without R3.""" if getattr(module, "_slime_ordered_topk_capture_registered", False): diff --git a/slime_plugins/models/glm5/glm5.py b/slime_plugins/models/glm5/glm5.py index 656ed024f4..de671071f3 100644 --- a/slime_plugins/models/glm5/glm5.py +++ b/slime_plugins/models/glm5/glm5.py @@ -27,6 +27,8 @@ from megatron.core.transformer.transformer_config import MLATransformerConfig from transformers import AutoConfig +from slime.backends.sglang_utils.compat import import_sglang_module + from .ops.indexer import generate_varlen_mask_params, lighting_indexer from .ops.sparse_mla import SGLangSparseMLA, SparseMLA @@ -170,7 +172,10 @@ def _apply_sglang_rope_forward( cos_sin_cache: torch.Tensor, positions: torch.Tensor, ) -> torch.Tensor: - from sglang.jit_kernel.rope import apply_rope_with_cos_sin_cache_inplace + apply_rope_with_cos_sin_cache_inplace = import_sglang_module( + "sglang.kernels.ops.attention.rope", + "sglang.jit_kernel.rope", + ).apply_rope_with_cos_sin_cache_inplace output = torch.empty_strided(value.size(), value.stride(), dtype=value.dtype, device=value.device) output.copy_(value) @@ -235,8 +240,14 @@ def _get_sglang_rope_cache( class _DSAKVFP8QAT(torch.autograd.Function): @staticmethod def forward(ctx, kv: torch.Tensor): - from sglang.srt.layers.attention.dsa.dequant_k_cache import dequantize_k_cache - from sglang.srt.layers.attention.dsa.quant_k_cache import quantize_k_cache + dequantize_k_cache = import_sglang_module( + "sglang.kernels.ops.attention.dsa.dequant_k_cache", + "sglang.srt.layers.attention.dsa.dequant_k_cache", + ).dequantize_k_cache + quantize_k_cache = import_sglang_module( + "sglang.kernels.ops.attention.dsa.quant_k_cache", + "sglang.srt.layers.attention.dsa.quant_k_cache", + ).quantize_k_cache if kv.dtype != torch.bfloat16 or kv.shape[-2:] != (1, 576): raise ValueError( diff --git a/slime_plugins/models/glm5/ops/indexer.py b/slime_plugins/models/glm5/ops/indexer.py index 8dcac9053c..7fe1f6f8aa 100644 --- a/slime_plugins/models/glm5/ops/indexer.py +++ b/slime_plugins/models/glm5/ops/indexer.py @@ -2,6 +2,8 @@ import torch +from slime.backends.sglang_utils.compat import import_sglang_module + # Bind flashinfer.comm's module-level `cudart = CudaRTLibrary()` to the *real* # libcudart before tilelang loads its `libcudart_stub.so`. On a cu12 box whose # flashinfer pulled in cu13 deps, importing flashinfer.comm *after* the stub is @@ -31,9 +33,16 @@ def _sglang_fp8_indexer_logits( if index_q.shape[-1] != 128: raise ValueError("SGLang FP8 indexer alignment requires head_dim=128, " f"got {index_q.shape[-1]}") import deep_gemm - from sglang.jit_kernel.fused_store_index_cache import fused_store_index_k_cache from sglang.srt.layers.attention.dsa.dsa_indexer import rotate_activation - from sglang.srt.layers.attention.dsa.triton_kernel import act_quant + + fused_store_index_k_cache = import_sglang_module( + "sglang.kernels.ops.attention.fused_store_index_cache", + "sglang.jit_kernel.fused_store_index_cache", + ).fused_store_index_k_cache + act_quant = import_sglang_module( + "sglang.kernels.ops.attention.dsa.triton_kernel", + "sglang.srt.layers.attention.dsa.triton_kernel", + ).act_quant q_rotated = rotate_activation(index_q.contiguous()) if index_k.ndim == 3: @@ -150,14 +159,31 @@ def forward( clean_logits=True, ) if topk_indices is None: - index_score, topk_indices = pytorch_topk_with_invalid_padding(logits, topk) - if os.getenv("MEGATRON_USE_SGLANG_FP8_INDEXER", "0") == "1": - invalid_sort_key = torch.iinfo(torch.int32).max - topk_indices = torch.sort( - topk_indices.masked_fill(topk_indices < 0, invalid_sort_key), - dim=-1, - ).values - topk_indices = topk_indices.masked_fill(topk_indices == invalid_sort_key, -1) + use_sglang_fp8_indexer = os.getenv("MEGATRON_USE_SGLANG_FP8_INDEXER", "0") == "1" + if use_sglang_fp8_indexer and logits.shape[-1] >= topk: + from sglang.srt.layers.attention.dsa.dsa_topk_backend import deterministic_flashinfer_topk + + starts = cu_seqlen_ks.to(torch.int32) + local_indices = deterministic_flashinfer_topk( + logits, + (cu_seqlen_ke - cu_seqlen_ks).to(torch.int32), + topk, + row_starts=starts, + ) + topk_indices = torch.where( + local_indices >= 0, + local_indices + starts.unsqueeze(-1), + local_indices, + ) + else: + _, topk_indices = pytorch_topk_with_invalid_padding(logits, topk) + if use_sglang_fp8_indexer: + invalid_sort_key = torch.iinfo(torch.int32).max + topk_indices = torch.sort( + topk_indices.masked_fill(topk_indices < 0, invalid_sort_key), + dim=-1, + ).values + topk_indices = topk_indices.masked_fill(topk_indices == invalid_sort_key, -1) index_score = pytorch_extract_topk_scores(logits, topk_indices) diff --git a/tests/test_deepgemm_moe_forward.py b/tests/test_deepgemm_moe_forward.py index f82b8e713a..9e1591b15c 100644 --- a/tests/test_deepgemm_moe_forward.py +++ b/tests/test_deepgemm_moe_forward.py @@ -7,10 +7,55 @@ import torch.nn.functional as F from slime.backends.megatron_utils.alignment import deepgemm_moe_forward +from slime.backends.sglang_utils import compat as sglang_compat NUM_GPUS = 1 +def test_sglang_compat_prefers_current_module(monkeypatch): + current_module = object() + calls = [] + + def fake_import(path): + calls.append(path) + return current_module + + monkeypatch.setattr(sglang_compat, "import_module", fake_import) + + assert sglang_compat.import_sglang_module("sglang.current", "sglang.legacy") is current_module + assert calls == ["sglang.current"] + + +def test_sglang_compat_falls_back_for_moved_namespace(monkeypatch): + legacy_module = object() + calls = [] + + def fake_import(path): + calls.append(path) + if path == "sglang.current.module": + raise ModuleNotFoundError(name="sglang.current") + return legacy_module + + monkeypatch.setattr(sglang_compat, "import_module", fake_import) + + assert sglang_compat.import_sglang_module("sglang.current.module", "sglang.legacy.module") is legacy_module + assert calls == ["sglang.current.module", "sglang.legacy.module"] + + +def test_sglang_compat_does_not_mask_missing_dependency(monkeypatch): + missing_dependency = ModuleNotFoundError(name="triton") + + def fake_import(path): + raise missing_dependency + + monkeypatch.setattr(sglang_compat, "import_module", fake_import) + + with pytest.raises(ModuleNotFoundError) as exc_info: + sglang_compat.import_sglang_module("sglang.current.module", "sglang.legacy.module") + + assert exc_info.value is missing_dependency + + @pytest.mark.parametrize("value", ["1", "true", "yes", "on"]) def test_configures_batch_invariant_in_megatron_actor(monkeypatch, value): state = {"enabled": False} @@ -601,7 +646,12 @@ def test_cuda_fused_static_route_gradient_is_bitwise(monkeypatch): def test_ordered_ep_gather_backward_matches_weighted_route_sum(monkeypatch): - import sglang.srt.layers.moe.ep_moe.kernels as kernels + kernels = SimpleNamespace(ep_gather=None) + monkeypatch.setattr( + deepgemm_moe_forward, + "import_sglang_module", + lambda current_path, legacy_path: kernels, + ) def fake_ep_gather(input_tensor, recv_ids, recv_weights, input_index, output): output.zero_() @@ -836,6 +886,89 @@ def megatron_topk(values, topk, num_groups=None, group_topk=None): torch.testing.assert_close(actual, expected, rtol=0, atol=0) +def test_sglang_fused_sigmoid_topk_forward_and_backward(monkeypatch): + from slime.utils import routing_replay + + router = SimpleNamespace() + monkeypatch.setattr(routing_replay, "ORDERED_TOPK_CAPTURE_ROUTER", router) + + def fake_fused_topk(logits, expert_bias, topk): + scores = torch.sigmoid(logits) + indices = torch.topk(scores + expert_bias, topk, dim=1, sorted=False).indices + selected = scores.gather(1, indices) + return selected / selected.sum(dim=-1, keepdim=True), indices.to(torch.int32) + + monkeypatch.setattr(routing_replay, "_run_sglang_fused_sigmoid_topk", fake_fused_topk) + logits = torch.tensor( + [[-1.5, 0.25, 1.75, -0.5], [0.75, -2.0, 0.5, 1.25]], + dtype=torch.float32, + requires_grad=True, + ) + bias = torch.tensor([0.1, -0.2, 0.05, 0.3], dtype=torch.float32) + grad_weights = torch.tensor([[0.5, -1.25], [1.5, 0.75]], dtype=torch.float32) + + weights, indices = routing_replay.maybe_sglang_fused_sigmoid_topk( + logits, + 2, + use_pre_softmax=True, + score_function="sigmoid", + expert_bias=bias, + ) + (weights * grad_weights).sum().backward() + + reference_logits = logits.detach().clone().requires_grad_(True) + reference_scores = torch.sigmoid(reference_logits) + reference_selected = reference_scores.gather(1, indices.long()) + reference_weights = reference_selected / reference_selected.sum(dim=-1, keepdim=True) + (reference_weights * grad_weights).sum().backward() + + torch.testing.assert_close(weights, reference_weights, rtol=0, atol=0) + torch.testing.assert_close(logits.grad, reference_logits.grad, rtol=1e-6, atol=1e-7) + torch.testing.assert_close( + routing_replay.consume_ordered_topk(router), + indices, + rtol=0, + atol=0, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_sglang_fused_sigmoid_topk_is_bitwise_batch_invariant(): + try: + moe_fused_gate = sglang_compat.import_sglang_module( + "sglang.kernels.ops.moe.moe_fused_gate", + "sglang.jit_kernel.moe_fused_gate", + ).moe_fused_gate + except (ImportError, RuntimeError) as exc: + pytest.skip(f"SGLang fused router is unavailable: {exc}") + + generator = torch.Generator(device="cuda").manual_seed(1234) + logits = torch.randn((73, 256), device="cuda", dtype=torch.float32, generator=generator) + bias = torch.randn((256,), device="cuda", dtype=torch.float32, generator=generator) * 0.125 + expected_weights, expected_indices = moe_fused_gate( + logits, + bias, + 8, + scoring_func="sigmoid", + renormalize=True, + ) + + for chunk_size in (1, 3, 17, 64): + actual_weights = torch.empty_like(expected_weights) + actual_indices = torch.empty_like(expected_indices) + for start in range(0, logits.shape[0], chunk_size): + stop = min(start + chunk_size, logits.shape[0]) + actual_weights[start:stop], actual_indices[start:stop] = moe_fused_gate( + logits[start:stop], + bias, + 8, + scoring_func="sigmoid", + renormalize=True, + ) + torch.testing.assert_close(actual_indices, expected_indices, rtol=0, atol=0) + torch.testing.assert_close(actual_weights, expected_weights, rtol=0, atol=0) + + def test_static_combine_backward_matches_dynamic_path(): route_values = _small_bf16((8, 128), offset=3) weights = torch.tensor( @@ -960,7 +1093,12 @@ def test_static_combine_backward_handles_padding_with_reused_storage(): def test_low_latency_ep_gather_preserves_order_and_backward(monkeypatch): - import sglang.srt.layers.moe.ep_moe.kernels as kernels + kernels = SimpleNamespace(ep_gather=None) + monkeypatch.setattr( + deepgemm_moe_forward, + "import_sglang_module", + lambda current_path, legacy_path: kernels, + ) calls = [] @@ -1959,13 +2097,19 @@ def _require_cuda_deepgemm(monkeypatch): for module_name in ( "deep_gemm", "sglang.srt.layers.deep_gemm_wrapper", - "sglang.srt.layers.quantization.fp8_kernel", "sgl_kernel", ): try: importlib.import_module(module_name) except Exception as exc: pytest.skip(f"{module_name} is unavailable: {exc}") + try: + sglang_compat.import_sglang_module( + "sglang.kernels.ops.quantization.fp8_kernel", + "sglang.srt.layers.quantization.fp8_kernel", + ) + except Exception as exc: + pytest.skip(f"SGLang FP8 kernel is unavailable: {exc}") def test_cuda_deepgemm_moe_forward_diff_against_bf16_reference_with_expert_m_padding( diff --git a/tests/test_glm52_6layer_deterministic_e2e.py b/tests/test_glm52_6layer_deterministic_e2e.py index 1f94cc0732..de3dc9fe8e 100644 --- a/tests/test_glm52_6layer_deterministic_e2e.py +++ b/tests/test_glm52_6layer_deterministic_e2e.py @@ -64,6 +64,9 @@ assert "align_fp8_quantization" in inspect.signature(Buffer.low_latency_dispatch).parameters, ( "DeepEP lacks align_fp8_quantization" ) +assert "fp8_quantization_mode" in inspect.signature(Buffer.low_latency_dispatch).parameters, ( + "DeepEP lacks fp8_quantization_mode" +) assert "enable_fp32_moe_router" in ServerArgs.__dataclass_fields__, "SGLang lacks enable_fp32_moe_router" """ @@ -231,7 +234,7 @@ def _train_args( "--sglang-mem-fraction-static 0.70 --sglang-enable-dp-attention --sglang-enable-dp-lm-head " f"--sglang-ep-size {NUM_GPUS} --sglang-dp-size {NUM_GPUS} --sglang-moe-dp-size 1 " "--sglang-moe-dense-tp-size 1 --sglang-moe-a2a-backend deepep --sglang-deepep-mode low_latency " - "--sglang-moe-runner-backend deep_gemm --sglang-fp8-gemm-runner-backend deep_gemm " + "--sglang-moe-runner-backend deep_gemm --sglang-fp8-gemm-backend deep_gemm " f"--sglang-page-size 64 --sglang-kv-cache-dtype {kv_cache_dtype} --sglang-attention-backend dsa " "--sglang-dsa-prefill-backend flashmla_sparse --sglang-dsa-decode-backend flashmla_sparse " "--sglang-dsa-topk-backend torch --sglang-chunked-prefill-size 4096 --sglang-context-length 8192 " @@ -407,7 +410,10 @@ def _run(cmd, env=None, cwd=None, check=True, stream=False): sys.stdout.write(line) sys.stdout.flush() proc.wait() - return proc.returncode, "".join(lines) + output = "".join(lines) + if check and proc.returncode != 0: + raise RuntimeError(f"{cmd} failed ({proc.returncode}):\n{output}") + return proc.returncode, output def test_glm52_6layer_deterministic_train_rollout_alignment(): diff --git a/tests/utils/test_sglang_arguments.py b/tests/utils/test_sglang_arguments.py index 48c98ec6f6..df6132b07a 100644 --- a/tests/utils/test_sglang_arguments.py +++ b/tests/utils/test_sglang_arguments.py @@ -1,8 +1,8 @@ -from argparse import Namespace +from argparse import ArgumentParser, Namespace import pytest -from slime.backends.sglang_utils.arguments import validate_args +from slime.backends.sglang_utils.arguments import add_sglang_arguments, validate_args NUM_GPUS = 0 @@ -30,5 +30,14 @@ def test_validate_args_canonicalizes_moe_data_parallel_size(): assert args.sglang_moe_dp_size == 1 +@pytest.mark.unit +def test_sglang_0517_accepts_fp8_gemm_backend_option(): + parser = add_sglang_arguments(ArgumentParser()) + + args = parser.parse_args(["--sglang-fp8-gemm-backend", "deep_gemm"]) + + assert args.sglang_fp8_gemm_runner_backend == "deep_gemm" + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__]))