[Feature] Coordinate FullCG prefill across DP-attention ranks (#35640)

Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
This commit is contained in:
Aurick Qiao
2026-08-26 02:16:02 -07:00
committed by GitHub
co-authored by Yuwei An
parent 2511743bd7
commit 58ecbba0bd
15 changed files with 155 additions and 57 deletions
@@ -67,7 +67,7 @@ class TestDraftDpSyncMetadata(CustomTestCase):
batch = SimpleNamespace(
global_num_tokens=[1, 3, 0, 2],
global_num_tokens_for_logprob=[1, 3, 0, 2],
can_run_dp_cuda_graph=True,
can_run_decode_cuda_graph=True,
)
with patch(
@@ -84,7 +84,7 @@ class TestDraftDpSyncMetadata(CustomTestCase):
self.assertEqual(forward_batch.num_token_non_padded.item(), 6)
self.assertEqual(forward_batch.num_token_non_padded.dtype, torch.int32)
self.assertEqual(forward_batch.num_token_non_padded_cpu, 6)
self.assertTrue(forward_batch.can_run_dp_cuda_graph)
self.assertTrue(forward_batch.can_run_decode_cuda_graph)
class TestBusyIdleGraphKeyIdentity(CustomTestCase):