A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 13:45 UTC
[DSV4] Align aux stream API with DeepseekV4DecoderLayer (#41171)
Signed-off-by: zixi-qi <zixi@inferact.ai>
qizixi committed
Apr 28, 2026 at 17:22 UTC
6fb3f7b46b12ea63265afbe6d53d6f15a5de7b3a
1 file changed
+7
-5
vllm/model_executor/models/deepseek_v4_mtp.py
+7
-5
@@ -35,7 +35,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
35
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36
from vllm.platforms import current_platform
37
from vllm.sequence import IntermediateTensors
38
-from vllm.utils.multi_stream_utils import AuxStreamType
38
39
from .deepseek_mtp import SharedHead
40
from .deepseek_v2 import get_spec_layer_idx_from_weight_name
@@ -65,6 +64,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
64
vllm_config: VllmConfig,
65
topk_indices_buffer: torch.Tensor,
66
prefix: str,
67
+ aux_stream_list: list[torch.cuda.Stream] | None = None,
68
) -> None:
69
super().__init__()
70
@@ -112,14 +112,11 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
112
self.shared_head = SharedHead(
113
config=config, prefix=prefix, quant_config=quant_config
114
)
115
- self.aux_stream_dict = {
116
- AuxStreamType.Attention: torch.cuda.Stream(),
117
- }
115
self.mtp_block = DeepseekV4DecoderLayer(
116
vllm_config,
117
prefix,
118
topk_indices_buffer=topk_indices_buffer,
122
- aux_stream_dict=self.aux_stream_dict,
119
+ aux_stream_list=aux_stream_list,
120
)
121
122
def forward(
@@ -169,6 +166,10 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
166
device=self.device,
167
)
168
169
+ # Three aux streams shared across all MTP layers, mirroring
170
+ # DeepseekV4Model.
171
+ aux_stream_list = [torch.cuda.Stream() for _ in range(3)]
172
+
173
# to map the exact layer index from weights
174
self.layers = torch.nn.ModuleDict(
175
{
@@ -176,6 +177,7 @@ class DeepSeekV4MultiTokenPredictor(nn.Module):
177
vllm_config,
178
self.topk_indices_buffer,
179
f"{prefix}.layers.{idx}",
180
+ aux_stream_list=aux_stream_list,
181
)
182
for idx in range(
183
self.mtp_start_layer_idx,