[Perf][GDN] Build cudagraph-capture metadata without a device sync (#55404)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Roger Wang <rogerw@inferact.ai> Co-authored-by: Kimi Code <noreply@moonshot.cn>
Nick Hill committed
Sep 4, 2026 at 17:40 UTC
6cbb3c154ef1449d2b3c9131a237f36faa695734
3 files changed
+67
-2
tests/models/kimi_k3/test_kda_metadata.py
+62
@@ -631,6 +631,68 @@ def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn():
631
_assert_matches_shared_gdn(reference, actual)
632
633
634
+@pytest.mark.parametrize(
635
+ "builder_cls",
636
+ [
637
+ GDNAttentionMetadataBuilder,
638
+ # The Kimi-K3 builder stages spec-decode state indices on the device.
639
+ pytest.param(
640
+ KimiK3KDAMetadataBuilder,
641
+ marks=pytest.mark.skipif(
642
+ not torch.cuda.is_available(), reason="requires CUDA"
643
+ ),
644
+ ),
645
+ ],
646
+)
647
+def test_cudagraph_capture_metadata_avoids_device_to_host_copy(
648
+ builder_cls: type[AttentionMetadataBuilder], monkeypatch: pytest.MonkeyPatch
649
+):
650
+ """`build_for_cudagraph_capture` is not capture-only: a DP rank with nothing
651
+ scheduled re-stages its FULL-graph metadata through it on every dummy step,
652
+ so it must not synchronize the device. The host-side draft counts have to
653
+ come from `query_start_loc_cpu` and match the device-derived values."""
654
+ num_speculative_tokens = 2
655
+ batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3])
656
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
657
+ common_attn_metadata = create_common_attn_metadata(
658
+ batch, BLOCK_SIZE, device
659
+ ).replace(is_prefilling=torch.zeros(batch.batch_size, dtype=torch.bool))
660
+ builder = _make_builder(
661
+ builder_cls, num_speculative_tokens, full_cuda_graph=True, device=device
662
+ )
663
+
664
+ device_to_host_copies: list[torch.Size] = []
665
+ original_cpu = torch.Tensor.cpu
666
+
667
+ def counting_cpu(self: torch.Tensor, *args, **kwargs):
668
+ device_to_host_copies.append(self.shape)
669
+ return original_cpu(self, *args, **kwargs)
670
+
671
+ monkeypatch.setattr(torch.Tensor, "cpu", counting_cpu)
672
+ actual = builder.build_for_cudagraph_capture(common_attn_metadata)
673
+ monkeypatch.undo()
674
+
675
+ assert not device_to_host_copies, device_to_host_copies
676
+
677
+ num_accepted_tokens = torch.diff(common_attn_metadata.query_start_loc)
678
+ reference = builder.build(
679
+ 0,
680
+ common_attn_metadata,
681
+ num_accepted_tokens,
682
+ (num_accepted_tokens - 1).cpu(),
683
+ )
684
+ assert actual.num_spec_decodes == batch.batch_size
685
+ assert actual.num_decodes == 0
686
+ assert actual.num_prefills == 0
687
+ for field in fields(type(actual)):
688
+ actual_value = getattr(actual, field.name)
689
+ reference_value = getattr(reference, field.name)
690
+ if isinstance(actual_value, torch.Tensor):
691
+ torch.testing.assert_close(actual_value, reference_value)
692
+ else:
693
+ assert actual_value == reference_value, field.name
694
+
695
+
696
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
697
def test_recoverssm_spec_cudagraph_stages_one_checkpoint_per_request():
698
device = torch.device("cuda")
vllm/v1/attention/backends/gdn_attn.py
+2
-1
@@ -542,6 +542,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
542
)
543
544
num_accepted_tokens = torch.diff(m.query_start_loc)
545
- num_decode_draft_tokens_cpu = (num_accepted_tokens - 1).cpu()
545
+ num_decode_draft_tokens_cpu = torch.diff(m.query_start_loc_cpu).sub_(1)
546
+ assert num_decode_draft_tokens_cpu.shape == num_accepted_tokens.shape
547
548
return self.build(0, m, num_accepted_tokens, num_decode_draft_tokens_cpu)
vllm/v1/attention/backends/short_conv_attn.py
+3
-1
@@ -546,7 +546,9 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
546
547
if self.use_spec_decode:
548
num_accepted_tokens = torch.diff(m.query_start_loc)
549
- num_decode_draft_tokens_cpu = (num_accepted_tokens - 1).cpu()
549
+
550
+ num_decode_draft_tokens_cpu = torch.diff(m.query_start_loc_cpu).sub_(1)
551
+ assert num_decode_draft_tokens_cpu.shape == num_accepted_tokens.shape
552
return self.build(
553
0,
554
m,