From 755f97c6223d4e70b11cc3aac934bde6c5fa2438 Mon Sep 17 00:00:00 2001 From: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:16:00 +0800 Subject: [PATCH] [CI] Fix the DSpark dp-tier unit test fixture after #34919 (#38314) --- test/registered/spec/dspark/test_dspark_dp_tier.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/test/registered/spec/dspark/test_dspark_dp_tier.py b/test/registered/spec/dspark/test_dspark_dp_tier.py index 5ea44346e..95f3a1ca8 100644 --- a/test/registered/spec/dspark/test_dspark_dp_tier.py +++ b/test/registered/spec/dspark/test_dspark_dp_tier.py @@ -62,6 +62,7 @@ class TestDraftDpSyncMetadata(CustomTestCase): num_tokens_for_logprob_per_req=1, ) proposer.draft_model_runner = SimpleNamespace(device="cpu") + proposer._num_token_non_padded = torch.empty((1,), dtype=torch.int32) forward_batch = SimpleNamespace(input_ids=torch.arange(6)) batch = SimpleNamespace( @@ -81,8 +82,13 @@ class TestDraftDpSyncMetadata(CustomTestCase): [1, 3, 0, 2], ) self.assertEqual(forward_batch.global_num_tokens_cpu, [6, 18, 0, 12]) - # Metadata fill sets only the invariant GLOBAL count; the LOCAL - # num_token_non_padded is derived later when the draft forward localizes. + # The LOCAL count reuses the proposer's persistent buffer; the GLOBAL + # count is the invariant every DP rank agrees on. + self.assertIs( + forward_batch.num_token_non_padded, proposer._num_token_non_padded + ) + self.assertEqual(forward_batch.num_token_non_padded.item(), 6) + self.assertEqual(forward_batch.num_token_non_padded_cpu, 6) self.assertEqual(forward_batch.global_num_token_non_padded.item(), 6) self.assertEqual(forward_batch.global_num_token_non_padded.dtype, torch.int32) self.assertEqual(forward_batch.global_num_token_non_padded_cpu, 6)