From bde0efdbb78a57dc10375e8d0686cf862332192c Mon Sep 17 00:00:00 2001
From: artem-spector <artem_spector@yahoo.com>
Date: Tue, 28 Apr 2026 10:43:30 +0300
Subject: [PATCH] [Bugfix][Granite4Vision] Fix deepstack buffer causing decode
 slowdown in compiled mode (#40917)

Signed-off-by: artemspector <artems@il.ibm.com>
Co-authored-by: artemspector <artems@il.ibm.com>
---
 vllm/model_executor/models/granite4_vision.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py
index 147f02eced9..710fc94ee5f 100644
--- a/vllm/model_executor/models/granite4_vision.py
+++ b/vllm/model_executor/models/granite4_vision.py
@@ -887,9 +887,10 @@ class Granite4VisionForConditionalGeneration(
             and get_pp_group().is_first_rank
             and self._ds_layer_indices
         ):
+            n = inputs_embeds.size(0)
             ds: IntermediateTensors | None = IntermediateTensors(
                 {
-                    f"ds_{llm_layer}": self._ds_buffers[lvl]
+                    f"ds_{llm_layer}": self._ds_buffers[lvl][:n]
                     for lvl, llm_layer in enumerate(self._ds_layer_indices)
                 }
             )