From 4ea3642250b378a9497ccdd88c7c0ff01d899125 Mon Sep 17 00:00:00 2001
From: liupeng374 <782420244@qq.com>
Date: Sun, 14 Dec 2025 16:38:12 +0800
Subject: [PATCH] [NPU] bug fix for multi stream (#15048)

---
 python/sglang/srt/environ.py               | 1 +
 python/sglang/srt/models/deepseek_nextn.py | 7 ++++++-
 python/sglang/srt/models/deepseek_v2.py    | 6 +++++-
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index 3524a0128..d3eb3470d 100644
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -239,6 +239,7 @@ class Envs:
 
     # NPU
     SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False)
+    SGLANG_NPU_USE_MULTI_STREAM = EnvBool(False)
 
     # Quantization
     SGLANG_INT4_WEIGHT = EnvBool(False)
diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py
index 2aa062e5d..bbd4af816 100644
--- a/python/sglang/srt/models/deepseek_nextn.py
+++ b/python/sglang/srt/models/deepseek_nextn.py
@@ -22,6 +22,7 @@ from transformers import PretrainedConfig
 
 from sglang.srt.configs.model_config import is_deepseek_nsa
 from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers.attention.nsa.utils import (
     can_cp_split,
@@ -97,7 +98,11 @@ class DeepseekModelNextN(nn.Module):
 
         self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
 
-        self.alt_stream = torch.cuda.Stream() if _is_cuda or _is_npu else None
+        self.alt_stream = (
+            torch.cuda.Stream()
+            if _is_cuda or envs.SGLANG_NPU_USE_MULTI_STREAM.get()
+            else None
+        )
 
         layer_name = "decoder"
         if _is_npu and (
diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
index c9895f63c..1f5ce486a 100644
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -3020,7 +3020,11 @@ class DeepseekV2Model(nn.Module):
         else:
             self.embed_tokens = PPMissingLayer()
 
-        self.alt_stream = torch.cuda.Stream() if _is_cuda or _is_npu else None
+        self.alt_stream = (
+            torch.cuda.Stream()
+            if _is_cuda or envs.SGLANG_NPU_USE_MULTI_STREAM.get()
+            else None
+        )
 
         self.layers, self.start_layer, self.end_layer = make_layers(
             config.num_hidden_layers,