Fix mm embedding to be 3d tensor

muskansh-google · muskansh-google · commit b41fe2c130cf · 2026-03-25T03:50:41.000Z
diff --git a/tpu_inference/runner/compilation_manager.py b/tpu_inference/runner/compilation_manager.py
@@ -129,7 +129,7 @@ def _precompile_input_embeddings_merger(self) -> None:
             )
             sharding = NamedSharding(self.runner.mesh, PartitionSpec())
             dummy_multimodal_embeddings = self._create_dummy_tensor(
-                (num_tokens, hidden_size),
+                (1, num_tokens, hidden_size),
                 self.runner.vllm_config.model_config.dtype,
                 sharding=sharding)
             dummy_input_ids = self._create_dummy_tensor((num_tokens, ),

Original file line number	Diff line number	Diff line change
`@@ -129,7 +129,7 @@ def _precompile_input_embeddings_merger(self) -> None:`
`129`	`129`	`)`
`130`	`130`	`sharding = NamedSharding(self.runner.mesh, PartitionSpec())`
`131`	`131`	`dummy_multimodal_embeddings = self._create_dummy_tensor(`
`132`		`- (num_tokens, hidden_size),`
	`132`	`+ (1, num_tokens, hidden_size),`
`133`	`133`	`self.runner.vllm_config.model_config.dtype,`
`134`	`134`	`sharding=sharding)`
`135`	`135`	`dummy_input_ids = self._create_dummy_tensor((num_tokens, ),`