[Speculate Decoding] Fix reasoning_phase_token_constraint call args in SpeculativeSampler (#7402)

2026-04-23 00:17:25 +08:00 · 2026-04-15 12:45:23 +08:00
parent 8995a38fa4
commit f7a2418ce2
1 changed files with 2 additions and 1 deletions
@@ -996,7 +996,8 @@ class SpeculativeSampler(nn.Layer):
        if self.enf_gen_phase_tag:
            reasoning_phase_token_constraint(
                logits,
-                sampling_metadata.pre_token_ids,
+                token_ids_all,
+                prompt_lens,
                share_inputs["stop_flags"],
                share_inputs["seq_lens_this_time"],
                share_inputs["seq_lens_encoder"],