From e5e541446f950f26e83552fb75bfcb51485861fd Mon Sep 17 00:00:00 2001 From: MOVZX Date: Wed, 30 Sep 2026 07:11:03 +0200 Subject: [PATCH] recoverssm: FULL cudagraphs for uniform spec-decode batches (port of upstream #58863, 3aa3d5b586) Port of the upstream vLLM change #58863 (Apache-2.0). The GDN and PLE RecoverSSM metadata builders declare UNIFORM_BATCH, so --use-replayssm resolves FULL_AND_PIECEWISE; graph padding rows get a null state slot and write nothing. Co-Authored-By: Hinata Signed-off-by: MOVZX --- vllm/v1/attention/backends/gdn_recoverssm.py | 29 ++++++++++++++++---- vllm/v1/attention/backends/ple_recoverssm.py | 23 +++++++++++++--- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/vllm/v1/attention/backends/gdn_recoverssm.py b/vllm/v1/attention/backends/gdn_recoverssm.py index a45142ab3c..8b6349a772 100644 --- a/vllm/v1/attention/backends/gdn_recoverssm.py +++ b/vllm/v1/attention/backends/gdn_recoverssm.py @@ -76,18 +76,19 @@ class GDNRecoverSSMMetadata(GDNAttentionMetadata, RecoverSSMMetadata): class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): - # The verify/commit metadata is not staged for FULL graph replay yet; declaring - # NEVER makes the runner fall back to PIECEWISE (attention compiled piecewise). - _cudagraph_support = AttentionCGSupport.NEVER + # Uniform spec-decode batches replay as FULL graphs: the base builder stages the + # state indices, query offsets and accepted counts into persistent buffers and pads + # the batch with null-state, zero-length rows, which the verify kernel skips. The + # commit runs after sampling, outside the graph, on the staged tensors. + _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) - # The builder may be created before the runner resolves the CUDA graph mode. - self.use_full_cuda_graph = False max_reqs = vllm_config.scheduler_config.max_num_seqs self._ones = torch.ones(max_reqs, dtype=torch.int32, device=device) self._context: Any = None self._logged = False + self._capturing = False def _get_context(self): if self._context is None: @@ -120,6 +121,9 @@ class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): m = common_attn_metadata rows = None if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: + if self._capturing and m.is_prefilling is None: + # CUDA graph capture: dummy decode batches, every row a spec window. + m = m.replace(is_prefilling=torch.zeros(m.num_reqs, dtype=torch.bool)) if m.is_prefilling is None or m.is_prefilling.device.type != "cpu": raise ValueError("GDN RecoverSSM needs the CPU is_prefilling mask") # Positive counts only defeat the base builder's "no drafts -> regular @@ -144,7 +148,11 @@ class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): n = meta.num_spec_decodes assert meta.spec_state_indices_tensor is not None base["spec_state_indices_tensor"] = meta.spec_state_indices_tensor[:, :1] - base["num_accepted_tokens"] = self._ones[:n] + # One accepted token per row: the checkpoint already holds the accepted + # state. Sized like the (possibly graph-padded) state indices. + base["num_accepted_tokens"] = self._ones[ + : base["spec_state_indices_tensor"].shape[0] + ] if rows is None: raise ValueError("GDN RecoverSSM: spec rows without draft counts") request_indices = recoverssm_request_indices( @@ -170,6 +178,15 @@ class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder): recoverssm_context=self._get_context() if commit is not None else None, ) + def build_for_cudagraph_capture( + self, common_attn_metadata: CommonAttentionMetadata + ): + self._capturing = True + try: + return super().build_for_cudagraph_capture(common_attn_metadata) + finally: + self._capturing = False + class GDNRecoverSSMAttentionBackend(GDNAttentionBackend): @staticmethod diff --git a/vllm/v1/attention/backends/ple_recoverssm.py b/vllm/v1/attention/backends/ple_recoverssm.py index 7402001fc0..9b43a6fd92 100644 --- a/vllm/v1/attention/backends/ple_recoverssm.py +++ b/vllm/v1/attention/backends/ple_recoverssm.py @@ -230,17 +230,19 @@ class PleRecoverSSMMetadata(PleShortConvAttentionMetadata, RecoverSSMMetadata): class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): - # Same as the GDN RecoverSSM builder: PIECEWISE only for now. - _cudagraph_support = AttentionCGSupport.NEVER + # Same as the GDN RecoverSSM builder: uniform spec-decode batches replay as FULL + # graphs on the base builder's staged, null-padded buffers; the commit runs after + # sampling, outside the graph. + _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) - self.use_full_cuda_graph = False self._ones = torch.ones( vllm_config.scheduler_config.max_num_seqs, dtype=torch.int32, device=device ) self._ctx: _PleConvCommit | None = None self._logged = False + self._capturing = False def _get_ctx(self): if self._ctx is None: @@ -270,6 +272,9 @@ class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): m = common_attn_metadata rows = None if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: + if self._capturing and m.is_prefilling is None: + # CUDA graph capture: dummy decode batches, every row a spec window. + m = m.replace(is_prefilling=torch.zeros(m.num_reqs, dtype=torch.bool)) if m.is_prefilling is None or m.is_prefilling.device.type != "cpu": raise ValueError("PLE RecoverSSM needs the CPU is_prefilling mask") # the same classification as the GDN builder, so both commits map the same @@ -295,7 +300,10 @@ class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): return meta base = {f.name: getattr(meta, f.name) for f in fields(meta)} n = meta.num_spec_decodes - base["num_accepted_tokens"] = self._ones[:n] + # sized like the (possibly graph-padded) state indices + assert meta.spec_state_indices_tensor is not None + num_rows = meta.spec_state_indices_tensor.shape[0] + base["num_accepted_tokens"] = self._ones[:num_rows] if rows is None: raise ValueError("PLE RecoverSSM: spec rows without draft counts") req_idx = recoverssm_request_indices(rows, n, m.query_start_loc.device) @@ -315,6 +323,13 @@ class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder): **base, recoverssm_commit=commit, recoverssm_context=self._get_ctx() ) + def build_for_cudagraph_capture(self, common_attn_metadata): + self._capturing = True + try: + return super().build_for_cudagraph_capture(common_attn_metadata) + finally: + self._capturing = False + class PleRecoverSSMAttentionBackend(PleShortConvAttentionBackend): @staticmethod