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

[spec-overlap] bugfix for pd disaggregation and npu (#14088)

Browse Source 
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
This commit is contained in:
liupeng374
2025-12-01 22:58:20 +08:00
committed by GitHub
parent e55731b6e8
commit 2e8f54e61e
9 changed files with 32 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/release-docker-npu-nightly.yml vendored
View File

@@ -73,6 +73,6 @@ jobs:
push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }}
provenance: false
build-args: |
SGLANG_KERNEL_NPU_TAG=20251120
SGLANG_KERNEL_NPU_TAG=20251128
CANN_VERSION=${{ matrix.cann_version }}
DEVICE_TYPE=${{ matrix.device_type }}

2
.github/workflows/release-docker-npu.yml vendored
View File

@@ -70,6 +70,6 @@ jobs:
push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }}
provenance: false
build-args: |
SGLANG_KERNEL_NPU_TAG=20251120
SGLANG_KERNEL_NPU_TAG=20251128
CANN_VERSION=${{ matrix.cann_version }}
DEVICE_TYPE=${{ matrix.device_type }}

10
python/sglang/srt/managers/scheduler.py
View File

@@ -824,8 +824,10 @@ class Scheduler(
draft_token_to_kv_pool = (
self.draft_worker.draft_worker.draft_runner.token_to_kv_pool
)
model_config = self.draft_worker.draft_worker.draft_runner.model_config
else:
draft_token_to_kv_pool = self.draft_worker.model_runner.token_to_kv_pool
model_config = self.draft_worker.model_config
if (
self.disaggregation_mode == DisaggregationMode.DECODE
@@ -837,12 +839,12 @@ class Scheduler(
self.disagg_metadata_buffers = MetadataBuffers(
buffer_size,
hidden_size=(
self.draft_worker.model_config.hidden_size
model_config.hidden_size
if self.spec_algorithm.is_eagle()
else 16 # minimal padding size for RDMA
),
hidden_states_dtype=(
self.draft_worker.model_config.dtype
model_config.dtype
if self.spec_algorithm.is_eagle()
else torch.float32
),
@@ -890,12 +892,12 @@ class Scheduler(
self.disagg_metadata_buffers = MetadataBuffers(
buffer_size,
hidden_size=(
self.draft_worker.model_config.hidden_size
model_config.hidden_size
if self.spec_algorithm.is_eagle()
else 16 # minimal padding size for RDMA
),
hidden_states_dtype=(
self.draft_worker.model_config.dtype
model_config.dtype
if self.spec_algorithm.is_eagle()
else torch.float32
),

19
python/sglang/srt/mem_cache/allocator_ascend.py
View File

@@ -8,7 +8,7 @@ if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import KVCache
from sglang.srt.mem_cache.allocator import PagedTokenToKVPoolAllocator
from sglang.srt.utils import get_num_new_pages
from sglang.srt.utils import get_num_new_pages, next_power_of_2
def alloc_extend_kernel_ascend(
@@ -104,21 +104,24 @@ class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
return None
if num_new_pages_item < 200:
import sgl_kernel_npu # noqa: F401
from sgl_kernel_npu.mem_cache.allocator import alloc_extend_kernel
out_indices = torch.empty(
(extend_num_tokens,),
dtype=torch.int64,
device=self.device,
)
torch.ops.npu.alloc_extend(
prefix_lens.to(torch.int64),
seq_lens.to(torch.int64),
last_loc.to(torch.int64),
max_num_extend_tokens = next_power_of_2(extend_num_tokens)
bs = prefix_lens.shape[0]
alloc_extend_kernel[(bs,)](
prefix_lens,
seq_lens,
last_loc,
self.free_pages,
self.page_size,
out_indices,
num_new_pages,
next_power_of_2(bs),
self.page_size,
max_num_extend_tokens,
)
else:

4
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -386,8 +386,8 @@ class EagleDraftWorker(BaseDraftWorker):
spec_info.hidden_states = hidden_states
# Run forward
logits_output = self.draft_runner.model.forward(
forward_batch.input_ids, forward_batch.positions, forward_batch
logits_output, _ = self.draft_runner.forward(
forward_batch, skip_attn_backend_init=True
)
if self.server_args.enable_nan_detection:
detect_nan(logits_output)

29
python/sglang/srt/speculative/spec_utils.py
View File

@@ -117,26 +117,15 @@ def assign_req_to_token_pool_func(
out_cache_loc: torch.Tensor,
batch_size: int,
):
if _is_cuda or _is_hip:
assign_req_to_token_pool[(batch_size,)](
req_pool_indices,
req_to_token,
start_offset,
end_offset,
out_cache_loc,
req_to_token.shape[1],
next_power_of_2(batch_size),
)
elif _is_npu:
import sgl_kernel_npu # noqa: F401
torch.ops.npu.cache_loc_assign(
req_pool_indices,
req_to_token,
start_offset.to(torch.int64),
end_offset.to(torch.int64),
out_cache_loc,
)
assign_req_to_token_pool[(batch_size,)](
req_pool_indices,
req_to_token,
start_offset,
end_offset,
out_cache_loc,
req_to_token.shape[1],
next_power_of_2(batch_size),
)
@triton.jit

2
scripts/ci/npu_ci_install_dependency.sh
View File

@@ -49,7 +49,7 @@ wget -O "${BISHENG_NAME}" "${BISHENG_URL}" && chmod a+x "${BISHENG_NAME}" && "./
### Install sgl-kernel-npu
SGL_KERNEL_NPU_TAG="20251120"
SGL_KERNEL_NPU_TAG="20251128"
git clone --depth 1 https://github.com/sgl-project/sgl-kernel-npu.git --branch ${SGL_KERNEL_NPU_TAG}
(cd sgl-kernel-npu && bash ./build.sh && ${PIP_INSTALL} output/deep_ep*.whl output/sgl_kernel_npu*.whl && cd "$(python3 -m pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so)

0
test/manual/ascend/test_ascend_deepseek_mtp.py → test/srt/ascend/test_ascend_deepseek_mtp.py
View File

2
test/srt/run_suite.py
View File

@@ -375,7 +375,7 @@ suite_ascend = {
],
"per-commit-16-npu-a3": [
TestFile("ascend/test_ascend_deepep.py", 400),
# TestFile("ascend/test_ascend_deepseek_mtp.py", 400),
TestFile("ascend/test_ascend_deepseek_mtp.py", 400),
],
}