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

Disable torch.compile for NPU in speculative overlap utils (#26403)

Browse Source
This commit is contained in:
xdtbynd
2026-05-27 12:30:53 +08:00
committed by GitHub
parent 0c34fc5ace
commit 21d0e74aff
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/overlap_utils.py
+4 -3
View File
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional, Union
import torch import torch
from sglang.srt.speculative.spec_utils import spec_need_hidden_states from sglang.srt.speculative.spec_utils import spec_need_hidden_states
from sglang.srt.utils import is_cuda, is_hip from sglang.srt.utils import is_cuda, is_hip, is_npu
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -16,6 +16,7 @@ if TYPE_CHECKING:
_is_cuda = is_cuda() _is_cuda = is_cuda()
_is_hip = is_hip() _is_hip = is_hip()
_is_npu = is_npu()
# Token-buf consume tracking: init to -1, assert non-negative on gather, # Token-buf consume tracking: init to -1, assert non-negative on gather,
# write -1 back. Catches "gather without intermediate stash" bugs. CI enables # write -1 back. Catches "gather without intermediate stash" bugs. CI enables
@@ -23,7 +24,7 @@ _is_hip = is_hip()
_DEBUG_ASSERT = os.getenv("SGLANG_IS_IN_CI", "").lower() == "true" _DEBUG_ASSERT = os.getenv("SGLANG_IS_IN_CI", "").lower() == "true"
@torch.compile(dynamic=True) @torch.compile(dynamic=True, disable=_is_npu)
def _assert_nonneg_and_invalidate( def _assert_nonneg_and_invalidate(
values: torch.Tensor, buf: torch.Tensor, indices: torch.Tensor values: torch.Tensor, buf: torch.Tensor, indices: torch.Tensor
) -> None: ) -> None:
@@ -33,7 +34,7 @@ def _assert_nonneg_and_invalidate(
buf[indices] = -1 buf[indices] = -1
@torch.compile(dynamic=True) @torch.compile(dynamic=True, disable=_is_npu)
def _gather_spec_extras( def _gather_spec_extras(
indices: torch.Tensor, indices: torch.Tensor,
topk_p_buf: torch.Tensor, topk_p_buf: torch.Tensor,