Pass per-forward overrides to ForwardBatch.init_new as explicit arguments (#30670)

2026-07-15 14:25:59 +08:00
parent 861d97d24d
commit e77d95c3d5
16 changed files with 147 additions and 55 deletions
@@ -507,7 +507,11 @@ def extend(reqs, model_runner):
        )
        batch.prefill_input_ids_cpu = None

-    forward_batch = ForwardBatch.init_new(batch, model_runner)
+    forward_batch = ForwardBatch.init_new(
+        batch,
+        model_runner,
+        return_hidden_states_before_norm=False,
+    )
    logits_output = model_runner.forward(forward_batch).logits_output
    next_token_ids = model_runner.sample(logits_output, forward_batch)
    return next_token_ids, logits_output.next_token_logits, batch
@@ -518,7 +522,11 @@ def decode(input_token_ids, batch, model_runner):
    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)
+    forward_batch = ForwardBatch.init_new(
+        batch,
+        model_runner,
+        return_hidden_states_before_norm=False,
+    )
    logits_output = model_runner.forward(forward_batch).logits_output
    next_token_ids = model_runner.sample(logits_output, forward_batch)
    return next_token_ids, logits_output.next_token_logits