[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
@@ -59,7 +59,7 @@ class TestMlpSyncPadUnpad(CustomTestCase):
batch = SimpleNamespace(
global_num_tokens=[2, 0, 3],
global_num_tokens_for_logprob=[2, 0, 3],
can_run_dp_cuda_graph=True,
can_run_decode_cuda_graph=True,
)
fb.init_mlp_sync_metadata(batch, torch.device("cpu"))
@@ -71,7 +71,7 @@ class TestMlpSyncPadUnpad(CustomTestCase):
torch.testing.assert_close(
fb.global_num_tokens_for_logprob_gpu, torch.tensor([4, 0, 6])
)
self.assertTrue(fb.can_run_dp_cuda_graph)
self.assertTrue(fb.can_run_decode_cuda_graph)
def test_draft_input_without_hidden_states_can_be_padded(self):
spec_info = SimpleNamespace(