Refactor prefill_forward_llama_vision

MohammedTaherMcW · MohammedTaherMcW · commit 99b75aedd8a7 · 2025-08-11T15:23:50.000Z
diff --git a/models/tt_transformers/tt/generator.py b/models/tt_transformers/tt/generator.py
@@ -624,8 +624,6 @@ def prefill_forward_llama_vision(
                 kv_cache=model_kv_cache,
                 cross_page_table=user_cross_page_table,
                 model_id=model_id,
-                image_grid_thw=kwargs["image_grid_thw"][idx] if kwargs.get("image_grid_thw") else None,
-                input_ids=kwargs["input_ids"][idx] if kwargs.get("input_ids") else None,
             )
 
             if xattn_caches is not None:

Original file line number	Diff line number	Diff line change
`@@ -624,8 +624,6 @@ def prefill_forward_llama_vision(`
`624`	`624`	`kv_cache=model_kv_cache,`
`625`	`625`	`cross_page_table=user_cross_page_table,`
`626`	`626`	`model_id=model_id,`
`627`		`- image_grid_thw=kwargs["image_grid_thw"][idx] if kwargs.get("image_grid_thw") else None,`
`628`		`- input_ids=kwargs["input_ids"][idx] if kwargs.get("input_ids") else None,`
`629`	`627`	`)`
`630`	`628`
`631`	`629`	`if xattn_caches is not None:`