drop output ids (#25774)

2026-05-19 17:50:47 -07:00
parent 425dffbde3
commit 7f154ba449
6 changed files with 32 additions and 31 deletions
@@ -467,7 +467,7 @@ def extend(reqs, model_runner):

@torch.no_grad
 def decode(input_token_ids, batch, model_runner):
-    batch.output_ids = input_token_ids
+    batch.input_ids = input_token_ids.to(torch.int64)
    batch.prepare_for_decode()
    _maybe_prepare_mlp_sync_batch(batch, model_runner)
    forward_batch = ForwardBatch.init_new(batch, model_runner)