chore: optimize deepstack buffer handling for MM Qwen3 models (#1643)

Signed-off-by: AlpinDale <alpindale@gmail.com>

Author: AlpinDale

aphrodite/model_executor/models/qwen3_omni_moe_thinker.py

@@ -1643,6 +1643,9 @@ def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = ""):
                     )
                     for _ in range(self.deepstack_num_level)
                 ]
+                # Tracks the valid token span currently stored in the buffer.
+                # Zero means there is no active deepstack payload to consume.
+                self.deepstack_input_embeds_num_tokens = 0
 
         with self._mark_language_model(aphrodite_config):
             self.language_model = Qwen3MoeLLMForCausalLM(
@@ -1661,6 +1664,13 @@ def _get_deepstack_input_embeds(
     ) -> IntermediateTensors | None:
         if not getattr(self, "deepstack_input_embeds", None):
             return None  # If vision tower is skipped
+        if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
+            return None
+        if num_tokens > self.deepstack_input_embeds_num_tokens:
+            raise ValueError(
+                "Requested more deepstack tokens than available in buffer: "
+                f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
+            )
 
         # get deepstack_input_embeds from buffer, and clear the buffer
         return IntermediateTensors(
@@ -1689,14 +1699,25 @@ def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> N
         for idx in range(self.deepstack_num_level):
             self.deepstack_input_embeds[idx][:num_tokens].copy_(deepstack_input_embeds[idx])
 
+        self.deepstack_input_embeds_num_tokens = num_tokens
+
     def _clear_deepstack_input_embeds(self, num_tokens: int) -> None:
         if not getattr(self, "deepstack_input_embeds", None):
             return
+        if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
+            return
 
         # clear deepstack_input_embeds in buffer
         if num_tokens > 0:
+            if num_tokens > self.deepstack_input_embeds_num_tokens:
+                raise ValueError(
+                    "Requested to clear more deepstack tokens than available in "
+                    "buffer: "
+                    f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
+                )
             for idx in range(self.deepstack_num_level):
                 self.deepstack_input_embeds[idx][:num_tokens].zero_()
+            self.deepstack_input_embeds_num_tokens = 0
 
     def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:
         mm_input_by_modality = {}

aphrodite/model_executor/models/qwen3_vl.py

@@ -1592,6 +1592,9 @@ def __init__(self, *, aphrodite_config: AphroditeConfig, prefix: str = "model"):
                     )
                     for _ in range(self.deepstack_num_level)
                 ]
+                # Tracks the valid token span currently stored in the buffer.
+                # Zero means there is no active deepstack payload to consume.
+                self.deepstack_input_embeds_num_tokens = 0
 
         with self._mark_language_model(aphrodite_config):
             self.language_model = Qwen3LLMForCausalLM(
@@ -1612,6 +1615,13 @@ def _get_deepstack_input_embeds(
     ) -> IntermediateTensors | None:
         if not getattr(self, "deepstack_input_embeds", None):
             return None  # If vision tower is skipped
+        if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
+            return None
+        if num_tokens > self.deepstack_input_embeds_num_tokens:
+            raise ValueError(
+                "Requested more deepstack tokens than available in buffer: "
+                f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
+            )
 
         # get deepstack_input_embeds from buffer, and clear the buffer
         return IntermediateTensors(
@@ -1639,15 +1649,25 @@ def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> N
             ]
         for idx in range(self.deepstack_num_level):
             self.deepstack_input_embeds[idx][:num_tokens].copy_(deepstack_input_embeds[idx])
+        self.deepstack_input_embeds_num_tokens = num_tokens
 
     def _clear_deepstack_input_embeds(self, num_tokens: int) -> None:
         if not getattr(self, "deepstack_input_embeds", None):
             return
+        if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0:
+            return
 
         # clear deepstack_input_embeds in buffer
         if num_tokens > 0:
+            if num_tokens > self.deepstack_input_embeds_num_tokens:
+                raise ValueError(
+                    "Requested to clear more deepstack tokens than available in "
+                    "buffer: "
+                    f"{num_tokens=} > {self.deepstack_input_embeds_num_tokens=}"
+                )
             for idx in range(self.deepstack_num_level):
                 self.deepstack_input_embeds[idx][:num_tokens].zero_()
+            self.deepstack_input_embeds_num_tokens = 0
 
     # -- SupportsEncoderCudaGraph protocol methods --