| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass, replace |
| 4 | from typing import Any |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.config import VllmConfig |
| 9 | from vllm.utils.torch_utils import async_tensor_h2d |
| 10 | from vllm.v1.attention.backend import ( |
| 11 | AttentionBackend, |
| 12 | AttentionCGSupport, |
| 13 | CommonAttentionMetadata, |
| 14 | ) |
| 15 | from vllm.v1.attention.backends.mamba_attn import ( |
| 16 | BaseMambaAttentionMetadata, |
| 17 | BaseMambaAttentionMetadataBuilder, |
| 18 | ) |
| 19 | from vllm.v1.attention.backends.utils import ( |
| 20 | NULL_BLOCK_ID, |
| 21 | mamba_get_block_table_tensor, |
| 22 | ) |
| 23 | from vllm.v1.kv_cache_interface import MambaSpec |
| 24 | |
| 25 | |
| 26 | class ShortConvAttentionBackend(AttentionBackend): |
| 27 | @staticmethod |
| 28 | def get_name() -> str: |
| 29 | return "SHORT_CONV_ATTN" |
| 30 | |
| 31 | @staticmethod |
| 32 | def get_builder_cls() -> type["ShortConvAttentionMetadataBuilder"]: |
| 33 | return ShortConvAttentionMetadataBuilder |
| 34 | |
| 35 | @classmethod |
| 36 | def is_ssm(cls) -> bool: |
| 37 | return True |
| 38 | |
| 39 | |
| 40 | @dataclass |
| 41 | class ShortConvAttentionMetadata(BaseMambaAttentionMetadata): |
| 42 | pass |
| 43 | |
| 44 | |
| 45 | class ShortConvAttentionMetadataBuilder( |
| 46 | BaseMambaAttentionMetadataBuilder[ShortConvAttentionMetadata] |
| 47 | ): |
| 48 | metadata_cls = ShortConvAttentionMetadata |
| 49 | |
| 50 | |
| 51 | @dataclass |
| 52 | class PleShortConvAttentionMetadata(ShortConvAttentionMetadata): |
| 53 | # Number of speculative-decode (multi-query / MTP) requests and the total |
| 54 | # number of tokens they contribute. These are 0 when spec-decode is off. |
| 55 | num_spec_decodes: int = 0 |
| 56 | num_spec_decode_tokens: int = 0 |
| 57 | num_actual_tokens: int = 0 |
| 58 | |
| 59 | # Max query length among spec-decode requests (== num_speculative_tokens + 1). |
| 60 | # Used as ``max_query_len`` for the varlen spec causal_conv1d_update. |
| 61 | spec_query_len: int = 1 |
| 62 | |
| 63 | # Max query length among the non-spec *prefill* requests, precomputed |
| 64 | # CPU-side in the builder. The dilated PLE short-conv uses it to size its |
| 65 | # packing buffer without a device->host sync (``lengths.max().item()``). |
| 66 | # 0 when there are no prefill requests. |
| 67 | max_prefill_query_len: int = 0 |
| 68 | query_start_loc: torch.Tensor | None = None |
| 69 | |
| 70 | # ``state_indices_tensor`` keeps the historical (non-spec) layout used by |
| 71 | # all existing short-conv consumers: the conv-state slot for each regular |
| 72 | # decode followed by each prefill request. When spec-decode is active this |
| 73 | # only covers the non-spec requests. |
| 74 | state_indices_tensor: torch.Tensor | None = None |
| 75 | has_initial_states_d: torch.Tensor | None = None |
| 76 | |
| 77 | # ``non_spec_query_start_loc`` is the varlen cumulative token offset over |
| 78 | # the non-spec requests only (decodes then prefills). It equals |
| 79 | # ``query_start_loc`` when there are no spec-decode requests. |
| 80 | non_spec_query_start_loc: torch.Tensor | None = None |
| 81 | |
| 82 | # Speculative-decode (MTP) conv metadata. Only column 0 of the block table |
| 83 | # is needed for the convolution state, so these tensors are 1-D over the |
| 84 | # spec-decode requests. |
| 85 | spec_query_start_loc: torch.Tensor | None = None # [num_spec_decodes + 1] |
| 86 | spec_state_indices_tensor: torch.Tensor | None = None # [num_spec_decodes] |
| 87 | spec_sequence_masks: torch.Tensor | None = None # [batch] |
| 88 | spec_token_indx: torch.Tensor | None = None |
| 89 | non_spec_token_indx: torch.Tensor | None = None |
| 90 | num_decode_draft_tokens_cpu: torch.Tensor | None = None |
| 91 | |
| 92 | |
| 93 | class PleShortConvAttentionBackend(ShortConvAttentionBackend): |
| 94 | @staticmethod |
| 95 | def get_name() -> str: |
| 96 | return "PLE_SHORT_CONV_ATTN" |
| 97 | |
| 98 | @staticmethod |
| 99 | def get_builder_cls() -> type["PleShortConvAttentionMetadataBuilder"]: |
| 100 | return PleShortConvAttentionMetadataBuilder |
| 101 | |
| 102 | |
| 103 | class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder): |
| 104 | metadata_cls = PleShortConvAttentionMetadata |
| 105 | # Spec-decode requires a uniform (multi-token) decode batch for full |
| 106 | # CUDA graph capture, matching the GDN backend. |
| 107 | _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH |
| 108 | reorder_batch_threshold: int = 1 |
| 109 | supports_update_block_table = False |
| 110 | needs_causal_conv1d_metadata = False |
| 111 | |
| 112 | def __init__( |
| 113 | self, |
| 114 | kv_cache_spec: MambaSpec, |
| 115 | layer_names: list[str], |
| 116 | vllm_config: VllmConfig, |
| 117 | device: torch.device, |
| 118 | ) -> None: |
| 119 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 120 | self.num_spec = self.num_spec_tokens |
| 121 | self.use_full_cuda_graph = ( |
| 122 | self.compilation_config.cudagraph_mode.has_full_cudagraphs() |
| 123 | ) |
| 124 | |
| 125 | max_num_seqs = vllm_config.scheduler_config.max_num_seqs |
| 126 | max_capture_size = self.compilation_config.max_cudagraph_capture_size |
| 127 | self.decode_cudagraph_max_bs = max_num_seqs |
| 128 | self.decode_cudagraph_max_tokens = max_num_seqs * (self.num_spec + 1) |
| 129 | if max_capture_size is not None: |
| 130 | self.decode_cudagraph_max_bs = min( |
| 131 | self.decode_cudagraph_max_bs, max_capture_size |
| 132 | ) |
| 133 | self.decode_cudagraph_max_tokens = min( |
| 134 | self.decode_cudagraph_max_tokens, max_capture_size |
| 135 | ) |
| 136 | |
| 137 | # Persistent buffers reused during full CUDA graph capture and replay. |
| 138 | self.spec_state_indices_tensor = torch.empty( |
| 139 | (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device |
| 140 | ) |
| 141 | self.spec_sequence_masks = torch.empty( |
| 142 | (self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device |
| 143 | ) |
| 144 | self.spec_query_start_loc = torch.empty( |
| 145 | (self.decode_cudagraph_max_bs + 1,), dtype=torch.int32, device=device |
| 146 | ) |
| 147 | self.num_accepted_tokens = torch.empty( |
| 148 | (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device |
| 149 | ) |
| 150 | self.has_initial_states_d = torch.empty( |
| 151 | (self.decode_cudagraph_max_bs,), dtype=torch.bool, device=device |
| 152 | ) |
| 153 | |
| 154 | def _build_non_spec_metadata( |
| 155 | self, |
| 156 | common_prefix_len: int, |
| 157 | common_attn_metadata: CommonAttentionMetadata, |
| 158 | fast_build: bool, |
| 159 | num_decode_draft_tokens_cpu: torch.Tensor | None, |
| 160 | **kwargs: Any, |
| 161 | ) -> PleShortConvAttentionMetadata: |
| 162 | metadata = super().build( |
| 163 | common_prefix_len, |
| 164 | common_attn_metadata, |
| 165 | fast_build, |
| 166 | num_accepted_tokens=None, |
| 167 | **kwargs, |
| 168 | ) |
| 169 | assert isinstance(metadata, PleShortConvAttentionMetadata) |
| 170 | |
| 171 | state_indices_d = metadata.state_indices_tensor_d |
| 172 | if state_indices_d is not None and state_indices_d.dim() > 1: |
| 173 | state_indices_d = state_indices_d[:, 0] |
| 174 | state_indices_p = metadata.state_indices_tensor_p |
| 175 | if metadata.num_prefills == 0: |
| 176 | assert state_indices_d is not None |
| 177 | # BaseMambaAttentionMetadataBuilder pads decode state indices into |
| 178 | # a persistent tensor for full CUDA graphs. Keep those rows so the |
| 179 | # PLE decode receives one cache slot per graph-padded token. |
| 180 | state_indices_tensor = state_indices_d |
| 181 | elif metadata.num_decodes == 0: |
| 182 | assert state_indices_p is not None |
| 183 | state_indices_tensor = state_indices_p[: metadata.num_prefills] |
| 184 | else: |
| 185 | assert state_indices_d is not None |
| 186 | assert state_indices_p is not None |
| 187 | state_indices_tensor = torch.cat( |
| 188 | (state_indices_d, state_indices_p[: metadata.num_prefills]) |
| 189 | ) |
| 190 | |
| 191 | has_initial_states_d = None |
| 192 | if metadata.num_decodes > 0: |
| 193 | num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() |
| 194 | has_initial_states_d = num_computed_tokens[: metadata.num_decodes] > 0 |
| 195 | if ( |
| 196 | self.use_full_cuda_graph |
| 197 | and metadata.num_prefills == 0 |
| 198 | and metadata.num_decodes <= self.decode_cudagraph_max_bs |
| 199 | ): |
| 200 | assert state_indices_d is not None |
| 201 | # Prepare tensors for CUDA graph replay. Padded rows have no |
| 202 | # initial state and use NULL_BLOCK_ID in state_indices_d. |
| 203 | num_decode_rows = state_indices_d.numel() |
| 204 | self.has_initial_states_d[: metadata.num_decodes].copy_( |
| 205 | has_initial_states_d, non_blocking=True |
| 206 | ) |
| 207 | self.has_initial_states_d[metadata.num_decodes : num_decode_rows].fill_( |
| 208 | False |
| 209 | ) |
| 210 | has_initial_states_d = self.has_initial_states_d[:num_decode_rows] |
| 211 | |
| 212 | max_prefill_query_len = 0 |
| 213 | if metadata.num_prefills > 0: |
| 214 | query_lens_cpu = torch.diff(common_attn_metadata.query_start_loc_cpu) |
| 215 | max_prefill_query_len = int( |
| 216 | query_lens_cpu[ |
| 217 | metadata.num_decodes : ( |
| 218 | metadata.num_decodes + metadata.num_prefills |
| 219 | ) |
| 220 | ] |
| 221 | .max() |
| 222 | .item() |
| 223 | ) |
| 224 | |
| 225 | return replace( |
| 226 | metadata, |
| 227 | num_actual_tokens=common_attn_metadata.num_actual_tokens, |
| 228 | spec_query_len=self.num_spec + 1, |
| 229 | max_prefill_query_len=max_prefill_query_len, |
| 230 | query_start_loc=common_attn_metadata.query_start_loc, |
| 231 | state_indices_tensor=state_indices_tensor, |
| 232 | has_initial_states_d=has_initial_states_d, |
| 233 | non_spec_query_start_loc=common_attn_metadata.query_start_loc, |
| 234 | num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, |
| 235 | ) |
| 236 | |
| 237 | def build( # type: ignore[override] |
| 238 | self, |
| 239 | common_prefix_len: int, |
| 240 | common_attn_metadata: CommonAttentionMetadata, |
| 241 | fast_build: bool = False, |
| 242 | *, |
| 243 | num_accepted_tokens: torch.Tensor | None = None, |
| 244 | num_decode_draft_tokens_cpu: torch.Tensor | None = None, |
| 245 | **kwargs: Any, |
| 246 | ) -> PleShortConvAttentionMetadata: |
| 247 | m = common_attn_metadata |
| 248 | spec_sequence_masks_cpu: torch.Tensor | None = None |
| 249 | # Detect speculative-decode requests. We use -1 to mark prefill and |
| 250 | # plain-decode requests, so any value >= 0 is a (multi-query) |
| 251 | # spec-decode request. |
| 252 | if self.use_spec_decode and num_decode_draft_tokens_cpu is not None: |
| 253 | candidate_mask = num_decode_draft_tokens_cpu[: m.num_reqs] >= 0 |
| 254 | if bool(candidate_mask.any().item()): |
| 255 | spec_sequence_masks_cpu = candidate_mask |
| 256 | |
| 257 | if spec_sequence_masks_cpu is None: |
| 258 | return self._build_non_spec_metadata( |
| 259 | common_prefix_len, |
| 260 | common_attn_metadata, |
| 261 | fast_build, |
| 262 | num_decode_draft_tokens_cpu, |
| 263 | **kwargs, |
| 264 | ) |
| 265 | |
| 266 | del common_prefix_len, fast_build, kwargs |
| 267 | query_start_loc = m.query_start_loc |
| 268 | query_start_loc_cpu = m.query_start_loc_cpu |
| 269 | query_lens_cpu = torch.diff(query_start_loc_cpu) |
| 270 | block_table_tensor = mamba_get_block_table_tensor( |
| 271 | m.block_table_tensor, |
| 272 | m.seq_lens, |
| 273 | self.kv_cache_spec, |
| 274 | self.vllm_config.cache_config.mamba_cache_mode, |
| 275 | ) |
| 276 | |
| 277 | if query_start_loc.device.type == "cpu": |
| 278 | spec_sequence_masks = spec_sequence_masks_cpu |
| 279 | else: |
| 280 | spec_sequence_masks = async_tensor_h2d( |
| 281 | spec_sequence_masks_cpu, device=query_start_loc.device |
| 282 | ) |
| 283 | |
| 284 | has_initial_states_p = None |
| 285 | has_initial_states_d = None |
| 286 | |
| 287 | # Per-request classification by mask, NOT by position. With |
| 288 | # spec-decode, the front decode group can contain both spec-decode |
| 289 | # requests and plain non-spec single-token decodes. Spec requests are |
| 290 | # therefore not guaranteed to occupy the first num_spec_decodes slots. |
| 291 | non_spec_mask_cpu = ~spec_sequence_masks_cpu |
| 292 | decode_mask_cpu = non_spec_mask_cpu & (query_lens_cpu == 1) |
| 293 | prefill_mask_cpu = non_spec_mask_cpu & (query_lens_cpu > 1) |
| 294 | |
| 295 | num_spec_decodes = int(spec_sequence_masks_cpu.sum().item()) |
| 296 | num_decodes = int(decode_mask_cpu.sum().item()) |
| 297 | num_prefills = int(prefill_mask_cpu.sum().item()) |
| 298 | num_decode_tokens = num_decodes |
| 299 | num_prefill_tokens = int(query_lens_cpu[prefill_mask_cpu].sum().item()) |
| 300 | num_spec_decode_tokens = int( |
| 301 | query_lens_cpu[spec_sequence_masks_cpu].sum().item() |
| 302 | ) |
| 303 | # Max prefill query length, CPU-side (no device-to-host sync) for the |
| 304 | # PLE dilated short-conv packing buffer. |
| 305 | max_prefill_query_len = ( |
| 306 | int(query_lens_cpu[prefill_mask_cpu].max().item()) |
| 307 | if num_prefills > 0 |
| 308 | else 0 |
| 309 | ) |
| 310 | |
| 311 | assert num_accepted_tokens is not None |
| 312 | non_spec_req_idx: torch.Tensor | None = None |
| 313 | |
| 314 | if num_decodes == 0 and num_prefills == 0: |
| 315 | # All real requests are speculative; zero-length padding is trailing. |
| 316 | spec_token_indx = None |
| 317 | non_spec_token_indx = None |
| 318 | spec_state_indices_tensor = block_table_tensor[:num_spec_decodes, 0] |
| 319 | num_accepted_tokens = num_accepted_tokens[:num_spec_decodes] |
| 320 | non_spec_state_indices_tensor = None |
| 321 | spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] |
| 322 | non_spec_query_start_loc = None |
| 323 | else: |
| 324 | # Mixed batch: build a per-token group key consistent with the |
| 325 | # request grouping above (spec=0 | decode=1 | prefill=2) and a |
| 326 | # stable sort, so tokens of each request stay contiguous and in |
| 327 | # request order. This yields spec tokens first, then the non-spec |
| 328 | # [decode, prefill] tokens. |
| 329 | query_lens = torch.diff(query_start_loc) |
| 330 | spec_req_idx_cpu = spec_sequence_masks_cpu.nonzero(as_tuple=True)[0] |
| 331 | decode_req_idx_cpu = decode_mask_cpu.nonzero(as_tuple=True)[0] |
| 332 | prefill_req_idx_cpu = prefill_mask_cpu.nonzero(as_tuple=True)[0] |
| 333 | non_spec_req_idx_cpu = torch.cat((decode_req_idx_cpu, prefill_req_idx_cpu)) |
| 334 | spec_req_idx = async_tensor_h2d( |
| 335 | spec_req_idx_cpu, device=query_start_loc.device |
| 336 | ) |
| 337 | non_spec_req_idx = async_tensor_h2d( |
| 338 | non_spec_req_idx_cpu, device=query_start_loc.device |
| 339 | ) |
| 340 | decode_req_idx = async_tensor_h2d( |
| 341 | decode_req_idx_cpu, device=query_start_loc.device |
| 342 | ) |
| 343 | req_group = torch.full( |
| 344 | (m.num_reqs,), |
| 345 | 2, |
| 346 | dtype=torch.int64, |
| 347 | device=query_start_loc.device, |
| 348 | ) |
| 349 | req_group.index_fill_(0, spec_req_idx, 0) |
| 350 | req_group.index_fill_(0, decode_req_idx, 1) |
| 351 | token_group = torch.repeat_interleave( |
| 352 | req_group, query_lens, output_size=int(query_start_loc_cpu[-1]) |
| 353 | ) |
| 354 | token_perm = torch.argsort(token_group, stable=True) |
| 355 | spec_token_indx = token_perm[:num_spec_decode_tokens] |
| 356 | non_spec_token_indx = token_perm[num_spec_decode_tokens:] |
| 357 | |
| 358 | spec_state_indices_tensor = block_table_tensor[spec_req_idx, 0] |
| 359 | num_accepted_tokens = num_accepted_tokens[spec_req_idx] |
| 360 | non_spec_state_indices_tensor = block_table_tensor[non_spec_req_idx, 0] |
| 361 | spec_query_start_loc = torch.zeros( |
| 362 | num_spec_decodes + 1, |
| 363 | dtype=torch.int32, |
| 364 | device=query_start_loc.device, |
| 365 | ) |
| 366 | torch.cumsum(query_lens[spec_req_idx], dim=0, out=spec_query_start_loc[1:]) |
| 367 | non_spec_query_start_loc = torch.zeros( |
| 368 | num_decodes + num_prefills + 1, |
| 369 | dtype=torch.int32, |
| 370 | device=query_start_loc.device, |
| 371 | ) |
| 372 | torch.cumsum( |
| 373 | query_lens[non_spec_req_idx], |
| 374 | dim=0, |
| 375 | out=non_spec_query_start_loc[1:], |
| 376 | ) |
| 377 | |
| 378 | # Compute the conv-state slots for the non-spec decode/prefill split. |
| 379 | if non_spec_state_indices_tensor is None: |
| 380 | state_indices_tensor = block_table_tensor[:0, 0] |
| 381 | else: |
| 382 | state_indices_tensor = non_spec_state_indices_tensor |
| 383 | |
| 384 | # Build the regular decode/prefill state metadata inherited from the |
| 385 | # generic short-conv metadata contract. |
| 386 | query_start_loc_p = None |
| 387 | query_start_loc_d = None |
| 388 | state_indices_tensor_p = None |
| 389 | state_indices_tensor_d = None |
| 390 | if num_decodes > 0 or num_prefills > 0: |
| 391 | num_computed_tokens = m.compute_num_computed_tokens() |
| 392 | assert non_spec_req_idx is not None |
| 393 | num_computed_tokens = num_computed_tokens[non_spec_req_idx] |
| 394 | |
| 395 | state_indices_tensor_d = state_indices_tensor[:num_decodes] |
| 396 | state_indices_tensor_p = state_indices_tensor[ |
| 397 | num_decodes : num_decodes + num_prefills |
| 398 | ] |
| 399 | if num_decodes > 0: |
| 400 | has_initial_states_d = num_computed_tokens[:num_decodes] > 0 |
| 401 | assert non_spec_query_start_loc is not None |
| 402 | query_start_loc_d = non_spec_query_start_loc[: num_decodes + 1] |
| 403 | if num_prefills > 0: |
| 404 | has_initial_states_p = ( |
| 405 | num_computed_tokens[num_decodes : num_decodes + num_prefills] > 0 |
| 406 | ) |
| 407 | assert non_spec_query_start_loc is not None |
| 408 | query_start_loc_p = ( |
| 409 | non_spec_query_start_loc[num_decodes:] - num_decode_tokens |
| 410 | ) |
| 411 | |
| 412 | # Prepare persistent tensors for CUDA graph capture and replay. |
| 413 | # ``m.num_actual_tokens`` is already padded by the model runner. |
| 414 | # Request-level buffers use ``m.num_reqs`` while token-level buffers |
| 415 | # use their independently bounded token count. |
| 416 | batch_size = m.num_reqs |
| 417 | if ( |
| 418 | self.use_full_cuda_graph |
| 419 | and num_prefills == 0 |
| 420 | and num_decodes == 0 |
| 421 | and spec_sequence_masks is not None |
| 422 | and num_spec_decodes <= self.decode_cudagraph_max_bs |
| 423 | and num_spec_decode_tokens <= self.decode_cudagraph_max_tokens |
| 424 | ): |
| 425 | assert spec_state_indices_tensor is not None |
| 426 | self.spec_state_indices_tensor[:num_spec_decodes].copy_( |
| 427 | spec_state_indices_tensor, non_blocking=True |
| 428 | ) |
| 429 | spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size] |
| 430 | spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID) |
| 431 | |
| 432 | self.spec_sequence_masks[:batch_size].copy_( |
| 433 | spec_sequence_masks[:batch_size], non_blocking=True |
| 434 | ) |
| 435 | spec_sequence_masks = self.spec_sequence_masks[:batch_size] |
| 436 | |
| 437 | assert spec_query_start_loc is not None |
| 438 | self.spec_query_start_loc[: num_spec_decodes + 1].copy_( |
| 439 | spec_query_start_loc, non_blocking=True |
| 440 | ) |
| 441 | spec_num_query_tokens = spec_query_start_loc[-1] |
| 442 | spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1] |
| 443 | spec_query_start_loc[num_spec_decodes + 1 :].fill_(spec_num_query_tokens) |
| 444 | |
| 445 | assert num_accepted_tokens is not None |
| 446 | self.num_accepted_tokens[:num_spec_decodes].copy_( |
| 447 | num_accepted_tokens, non_blocking=True |
| 448 | ) |
| 449 | num_accepted_tokens = self.num_accepted_tokens[:batch_size] |
| 450 | num_accepted_tokens[num_spec_decodes:].fill_(1) |
| 451 | |
| 452 | return PleShortConvAttentionMetadata( |
| 453 | num_prefills=num_prefills, |
| 454 | num_prefill_tokens=num_prefill_tokens, |
| 455 | num_decodes=num_decodes, |
| 456 | num_decode_tokens=num_decode_tokens, |
| 457 | num_reqs=m.num_reqs, |
| 458 | num_spec_decodes=num_spec_decodes, |
| 459 | num_spec_decode_tokens=num_spec_decode_tokens, |
| 460 | num_actual_tokens=m.num_actual_tokens, |
| 461 | spec_query_len=self.num_spec + 1, |
| 462 | max_prefill_query_len=max_prefill_query_len, |
| 463 | query_start_loc=query_start_loc, |
| 464 | state_indices_tensor=state_indices_tensor, |
| 465 | has_initial_states_p=has_initial_states_p, |
| 466 | has_initial_states_d=has_initial_states_d, |
| 467 | non_spec_query_start_loc=non_spec_query_start_loc, |
| 468 | spec_query_start_loc=spec_query_start_loc, |
| 469 | spec_state_indices_tensor=spec_state_indices_tensor, |
| 470 | spec_sequence_masks=spec_sequence_masks, |
| 471 | spec_token_indx=spec_token_indx, |
| 472 | non_spec_token_indx=non_spec_token_indx, |
| 473 | num_accepted_tokens=num_accepted_tokens, |
| 474 | num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, |
| 475 | query_start_loc_p=query_start_loc_p, |
| 476 | query_start_loc_d=query_start_loc_d, |
| 477 | state_indices_tensor_p=state_indices_tensor_p, |
| 478 | state_indices_tensor_d=state_indices_tensor_d, |
| 479 | seq_lens=m.seq_lens, |
| 480 | ) |
| 481 | |
| 482 | def build_for_cudagraph_capture( |
| 483 | self, common_attn_metadata: CommonAttentionMetadata |
| 484 | ) -> PleShortConvAttentionMetadata: |
| 485 | """Build metadata for full CUDA graph capture. |
| 486 | |
| 487 | Currently, only decode is supported for full CUDA graphs with |
| 488 | short-conv. |
| 489 | """ |
| 490 | m = common_attn_metadata |
| 491 | assert ( |
| 492 | m.num_reqs <= self.decode_cudagraph_max_bs |
| 493 | and m.num_actual_tokens <= self.decode_cudagraph_max_tokens |
| 494 | ), ( |
| 495 | "ShortConv only supports decode-only full CUDAGraph capture. " |
| 496 | f"Make sure batch size ({m.num_reqs}) <= " |
| 497 | f"cudagraph capture size ({self.decode_cudagraph_max_bs}) and " |
| 498 | f"number of tokens ({m.num_actual_tokens}) <= " |
| 499 | f"token capture size ({self.decode_cudagraph_max_tokens})." |
| 500 | ) |
| 501 | |
| 502 | if self.use_spec_decode: |
| 503 | num_accepted_tokens = torch.diff(m.query_start_loc) |
| 504 | |
| 505 | num_decode_draft_tokens_cpu = torch.diff(m.query_start_loc_cpu).sub_(1) |
| 506 | assert num_decode_draft_tokens_cpu.shape == num_accepted_tokens.shape |
| 507 | return self.build( |
| 508 | 0, |
| 509 | m, |
| 510 | num_accepted_tokens=num_accepted_tokens, |
| 511 | num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, |
| 512 | ) |
| 513 | return self.build(0, m) |