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, 15:00 UTC
[Opt] Optimize deepstack buffer handling for multimodal Qwen3 models (#40145)
Signed-off-by: xiaoming <1259730330@qq.com>
labAxiaoming committed
Apr 25, 2026 at 21:04 UTC
6646c0c7e0c921709c9b194e3988dfaabda5ee15
2 files changed
+40
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+20
@@ -1753,6 +1753,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1753
)
1754
for _ in range(self.deepstack_num_level)
1755
]
1756
+ # Tracks the valid token span currently stored in the buffer.
1757
+ # Zero means there is no active deepstack payload to consume.
1758
+ self.deepstack_input_embeds_num_tokens = 0
1759
1760
with self._mark_language_model(vllm_config):
1761
self.language_model = Qwen3MoeLLMForCausalLM(
@@ -1773,6 +1776,13 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1776
) -> IntermediateTensors | None:
1777
if not getattr(self, "deepstack_input_embeds", None):
1778
return None # If vision tower is skipped
1779
+ if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
1780
+ return None
1781
+ if num_tokens > self.deepstack_input_embeds_num_tokens:
1782
+ raise ValueError(
1783
+ "Requested more deepstack tokens than available in buffer: "
1784
+ f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
1785
+ )
1786
1787
# get deepstack_input_embeds from buffer, and clear the buffer
1788
return IntermediateTensors(
@@ -1804,15 +1814,25 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1814
self.deepstack_input_embeds[idx][:num_tokens].copy_(
1815
deepstack_input_embeds[idx]
1816
)
1817
+ self.deepstack_input_embeds_num_tokens = num_tokens
1818
1819
def _clear_deepstack_input_embeds(self, num_tokens: int) -> None:
1820
if not getattr(self, "deepstack_input_embeds", None):
1821
return
1822
+ if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
1823
+ return
1824
1825
# clear deepstack_input_embeds in buffer
1826
if num_tokens > 0:
1827
+ if num_tokens > self.deepstack_input_embeds_num_tokens:
1828
+ raise ValueError(
1829
+ "Requested to clear more deepstack tokens than available in "
1830
+ "buffer: "
1831
+ f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
1832
+ )
1833
for idx in range(self.deepstack_num_level):
1834
self.deepstack_input_embeds[idx][:num_tokens].zero_()
1835
+ self.deepstack_input_embeds_num_tokens = 0
1836
1837
def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
1838
mm_input_by_modality = {}
vllm/model_executor/models/qwen3_vl.py
+20
@@ -1675,6 +1675,9 @@ class Qwen3VLForConditionalGeneration(
1675
)
1676
for _ in range(self.deepstack_num_level)
1677
]
1678
+ # Tracks the valid token span currently stored in the buffer.
1679
+ # Zero means there is no active deepstack payload to consume.
1680
+ self.deepstack_input_embeds_num_tokens = 0
1681
1682
with self._mark_language_model(vllm_config):
1683
self.language_model = Qwen3LLMForCausalLM(
@@ -1702,6 +1705,13 @@ class Qwen3VLForConditionalGeneration(
1705
) -> IntermediateTensors | None:
1706
if not getattr(self, "deepstack_input_embeds", None):
1707
return None # If vision tower is skipped
1708
+ if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
1709
+ return None
1710
+ if num_tokens > self.deepstack_input_embeds_num_tokens:
1711
+ raise ValueError(
1712
+ "Requested more deepstack tokens than available in buffer: "
1713
+ f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
1714
+ )
1715
1716
# get deepstack_input_embeds from buffer, and clear the buffer
1717
return IntermediateTensors(
@@ -1733,15 +1743,25 @@ class Qwen3VLForConditionalGeneration(
1743
self.deepstack_input_embeds[idx][:num_tokens].copy_(
1744
deepstack_input_embeds[idx]
1745
)
1746
+ self.deepstack_input_embeds_num_tokens = num_tokens
1747
1748
def _clear_deepstack_input_embeds(self, num_tokens: int) -> None:
1749
if not getattr(self, "deepstack_input_embeds", None):
1750
return
1751
+ if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
1752
+ return
1753
1754
# clear deepstack_input_embeds in buffer
1755
if num_tokens > 0:
1756
+ if num_tokens > self.deepstack_input_embeds_num_tokens:
1757
+ raise ValueError(
1758
+ "Requested to clear more deepstack tokens than available in "
1759
+ "buffer: "
1760
+ f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
1761
+ )
1762
for idx in range(self.deepstack_num_level):
1763
self.deepstack_input_embeds[idx][:num_tokens].zero_()
1764
+ self.deepstack_input_embeds_num_tokens = 0
1765
1766
# -- SupportsEncoderCudaGraph protocol methods --
1767