[Spec][test] fix(kv_canary): assert draft-extend-v2 oracle tokens in token_oracle test (#28325)

2026-06-15 19:14:09 -07:00
parent 448af67a98
commit c6d9d73fd6
1 changed files with 11 additions and 1 deletions
@@ -46,8 +46,18 @@ class TestTokenOracleManager(CustomTestCase):
            expected_inputs_out=expected_inputs,
        )

+        # draft-extend-v2 is not an extend mode, so fill_expected_inputs records
+        # deterministic oracle tokens (not input_ids). The req row [3, 7] expands
+        # to num_tokens_per_req=4 draft tokens each -> one row per draft token.
+        expected_generalized_req_ids = torch.tensor(
+            [3, 3, 3, 3, 7, 7, 7, 7], dtype=torch.int64, device=self.device
+        )
+        expected_tokens = manager.oracle.expected_tokens(
+            generalized_req_ids=expected_generalized_req_ids,
+            positions=forward_batch.positions.to(torch.int64),
+        )
        self.assertTrue(
-            torch.equal(expected_inputs.tokens[:8], forward_batch.input_ids)
+            torch.equal(expected_inputs.tokens[:8], expected_tokens.to(torch.int64))
        )
        self.assertTrue(
            torch.equal(expected_inputs.positions[:8], forward_batch.positions)