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

Fix Bug on dsv3.2 (#18553)

Browse Source 
This PR affects only the NPU. If any issues arise, please contact iforgetmyname.
This commit is contained in:
BourneSun0527
2026-02-11 14:39:01 +08:00
committed by GitHub
parent d84d2063d3
commit 2cc235e795
2 changed files with 16 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+12 -6
View File
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
import torch import torch
from einops import rearrange from einops import rearrange
from sglang.srt.environ import envs
from sglang.srt.layers.layernorm import LayerNorm from sglang.srt.layers.layernorm import LayerNorm
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.layers.utils import MultiPlatformOp
@@ -1190,13 +1191,17 @@ class Indexer(MultiPlatformOp):
) # [bs, n, d] ) # [bs, n, d]
q = torch.cat([q_pe, q_nope], dim=-1) q = torch.cat([q_pe, q_nope], dim=-1)
indexer_weight_stream = get_indexer_weight_stream() if envs.SGLANG_NPU_USE_MULTI_STREAM.get():
indexer_weight_stream.wait_stream(torch.npu.current_stream()) indexer_weight_stream = get_indexer_weight_stream()
with torch.npu.stream(indexer_weight_stream): indexer_weight_stream.wait_stream(torch.npu.current_stream())
with torch.npu.stream(indexer_weight_stream):
x = x.view(-1, self.hidden_size)
weights = self.weights_proj(x.float())[0].to(torch.bfloat16)
weights.record_stream(indexer_weight_stream)
weights_event = indexer_weight_stream.record_event()
else:
x = x.view(-1, self.hidden_size) x = x.view(-1, self.hidden_size)
weights = self.weights_proj(x.float())[0].to(torch.bfloat16) weights = self.weights_proj(x.float())[0].to(torch.bfloat16)
weights.record_stream(indexer_weight_stream)
weights_event = indexer_weight_stream.record_event()
k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128]
k = self.k_norm(k_proj) k = self.k_norm(k_proj)
@@ -1278,7 +1283,8 @@ class Indexer(MultiPlatformOp):
if self.alt_stream is not None: if self.alt_stream is not None:
torch.npu.current_stream().wait_event(q_rope_event) torch.npu.current_stream().wait_event(q_rope_event)
torch.npu.current_stream().wait_event(weights_event) if envs.SGLANG_NPU_USE_MULTI_STREAM.get():
torch.npu.current_stream().wait_event(weights_event)
block_table = forward_batch.attn_backend.forward_metadata.block_tables block_table = forward_batch.attn_backend.forward_metadata.block_tables
if ( if (
python/sglang/srt/managers/overlap_utils.py
+4 -2
View File
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional
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 get_compiler_backend from sglang.srt.utils import get_compiler_backend, is_npu
if TYPE_CHECKING: if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import ModelWorkerBatch from sglang.srt.managers.schedule_batch import ModelWorkerBatch
@@ -14,8 +14,10 @@ if TYPE_CHECKING:
from sglang.srt.speculative.eagle_info import EagleDraftInput from sglang.srt.speculative.eagle_info import EagleDraftInput
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
_is_npu = is_npu()
@torch.compile(dynamic=True, backend=get_compiler_backend())
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
def _resolve_future_token_ids(input_ids, future_token_ids_map): def _resolve_future_token_ids(input_ids, future_token_ids_map):
input_ids[:] = torch.where( input_ids[:] = torch.where(
input_ids < 0, input_ids < 0,