| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import Any |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.triton_utils import tl, triton |
| 9 | from vllm.v1.attention.backends.recoverssm_metadata import RecoverSSMMetadata |
| 10 | from vllm.v1.worker.utils import AttentionGroup |
| 11 | |
| 12 | |
| 13 | class RecoverSSMState: |
| 14 | """Coordinates RecoverSSM metadata between attention and postprocessing.""" |
| 15 | |
| 16 | def __init__(self) -> None: |
| 17 | self._step: tuple[RecoverSSMMetadata, ...] | None = None |
| 18 | |
| 19 | def record_step( |
| 20 | self, |
| 21 | attn_metadata: dict[str, Any], |
| 22 | attn_groups: list[list[AttentionGroup]], |
| 23 | *, |
| 24 | for_capture: bool, |
| 25 | ) -> None: |
| 26 | if for_capture: |
| 27 | self._step = None |
| 28 | return |
| 29 | |
| 30 | step: list[RecoverSSMMetadata] = [] |
| 31 | for group_list in attn_groups: |
| 32 | for group in group_list: |
| 33 | metadata = attn_metadata[group.layer_names[0]] |
| 34 | if isinstance(metadata, RecoverSSMMetadata): |
| 35 | step.append(metadata) |
| 36 | self._step = tuple(step) |
| 37 | |
| 38 | def commit_step( |
| 39 | self, |
| 40 | num_sampled: torch.Tensor | int, |
| 41 | idx_mapping: torch.Tensor, |
| 42 | *, |
| 43 | state_indices: torch.Tensor | None, |
| 44 | num_accepted_tokens: torch.Tensor, |
| 45 | ) -> None: |
| 46 | step = self._step |
| 47 | self._step = None |
| 48 | if isinstance(num_sampled, int) or step is None: |
| 49 | return |
| 50 | |
| 51 | for metadata in step: |
| 52 | postprocess_meta = metadata.commit_recoverssm_state(num_sampled) |
| 53 | if postprocess_meta is None: |
| 54 | continue |
| 55 | assert state_indices is not None |
| 56 | # RecoverSSM already restored the accepted state. Update its running |
| 57 | # column and reset the next-step copy bias to the neutral value. |
| 58 | _postprocess_recoverssm_align_kernel[(postprocess_meta.num_spec_decodes,)]( |
| 59 | idx_mapping, |
| 60 | num_sampled, |
| 61 | postprocess_meta.request_indices, |
| 62 | postprocess_meta.num_computed_tokens, |
| 63 | state_indices, |
| 64 | num_accepted_tokens, |
| 65 | MAMBA_BLOCK_SIZE=postprocess_meta.block_size, |
| 66 | BLOCK_TABLE_WIDTH=postprocess_meta.block_table.shape[1], |
| 67 | ) |
| 68 | |
| 69 | |
| 70 | @triton.heuristics( |
| 71 | {"HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None} |
| 72 | ) |
| 73 | @triton.jit |
| 74 | def _postprocess_recoverssm_align_kernel( |
| 75 | idx_mapping_ptr, |
| 76 | num_sampled_ptr, |
| 77 | request_indices_ptr, |
| 78 | num_computed_ptr, |
| 79 | state_idx_ptr, |
| 80 | num_accepted_ptr, |
| 81 | HAS_REQUEST_INDICES: tl.constexpr, |
| 82 | MAMBA_BLOCK_SIZE: tl.constexpr, |
| 83 | BLOCK_TABLE_WIDTH: tl.constexpr, |
| 84 | ): |
| 85 | spec_idx = tl.program_id(0) |
| 86 | batch_idx = spec_idx |
| 87 | if HAS_REQUEST_INDICES: |
| 88 | batch_idx = tl.load(request_indices_ptr + spec_idx) |
| 89 | req_state_idx = tl.load(idx_mapping_ptr + batch_idx) |
| 90 | if req_state_idx < 0: |
| 91 | return |
| 92 | num_sampled = tl.load(num_sampled_ptr + batch_idx) |
| 93 | num_computed = tl.load(num_computed_ptr + batch_idx) |
| 94 | tl.store( |
| 95 | state_idx_ptr + req_state_idx, |
| 96 | tl.minimum( |
| 97 | (num_computed + num_sampled) // MAMBA_BLOCK_SIZE, |
| 98 | BLOCK_TABLE_WIDTH - 1, |
| 99 | ), |
| 100 | ) |
| 101 | tl.store(num_accepted_ptr + req_state_idx, 1) |