Fix post schedule and simplify embedding cache

Author: gty111

gllm/model_runner.py

@@ -35,8 +35,6 @@ class EmbeddingInfo:
     embedding: torch.Tensor = None
     prompt_positions: torch.Tensor = None
     mrope_position_delta: torch.Tensor = None
-    stale: bool = False
-
 
 class ModelRunner:
     def __init__(
@@ -257,7 +255,6 @@ def mm_prepare_inputs(self, seqs: List[Sequence]):
             position = None
             if seq.computed_prompt:
                 embedding_info = self.embedding_cache[seq.seq_id]
-                assert embedding_info.stale
                 embedding = self.model.embed_input_ids(
                     torch.tensor(seq.to_compute_tokens)
                 )
@@ -269,10 +266,7 @@ def mm_prepare_inputs(self, seqs: List[Sequence]):
                 position = torch.tensor(position, device="cpu")
             else:
                 embedding_info = None
-                if (
-                    seq.seq_id not in self.embedding_cache
-                    or self.embedding_cache[seq.seq_id].stale
-                ):
+                if seq.seq_id not in self.embedding_cache:
                     mm_embeddings = None
                     image_grid_thw: torch.Tensor = None
                     video_grid_thw: torch.Tensor = None
@@ -338,7 +332,6 @@ def mm_prepare_inputs(self, seqs: List[Sequence]):
                     ]
                 if seq.seq_len == seq.prompt_len:
                     # invalidate embedding_cache
-                    embedding_info.stale = True
                     embedding_info.embedding = None
             batch_embeddings.append(embedding)
             batch_positions.append(position)
@@ -482,3 +475,5 @@ def step_once(self):
 
     def free(self, seq: Sequence):
         self.memory_manager.free(seq)
+        if self.use_mm:
+            self.embedding_cache.pop(seq.seq_id)

gllm/scheduler.py

@@ -160,10 +160,16 @@ def post_schedule(self, schedule_seqs: List[Sequence]):
         for seq in schedule_seqs:
             if seq.has_schedule:
                 post_schedule_seq = copy.copy(seq)
-                post_schedule_seq.to_compute_tokens = seq[
-                    seq.computed_token_num : seq.seq_len
-                ]
-                post_schedule_seq.token_ids = None
+                # MM prefill may still need full prompt token_ids to
+                # build (or rebuild) cached multimodal embeddings/positions.
+                # Keep token_ids for unfinished MM prefills; drop otherwise to
+                # reduce IPC payload.
+                keep_full_token_ids = self.model_runner.use_mm and (not seq.computed_prompt)
+                post_schedule_seq.token_ids = seq.token_ids if keep_full_token_ids else None
+                if not keep_full_token_ids:
+                    post_schedule_seq.to_compute_tokens = seq[
+                        seq.computed_token_num : seq.seq_len
+                    ]
                 post_schedule_seqs.append(post_schedule_seq)
             else:
                 post_schedule_seqs.append(seq)