[Feature] Coordinate FullCG prefill across DP-attention ranks (#35640)
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
This commit is contained in:
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user