| 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.model_executor.warmup.jit_warmup import WarmupIntRange |
| 9 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 10 | DispatchSpec, |
| 11 | TritonWarmupTensor, |
| 12 | triton_kernel_dispatcher_with_warmup, |
| 13 | ) |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.triton_utils import tl, triton |
| 16 | |
| 17 | |
| 18 | @triton.jit |
| 19 | def _scatter_states_kernel( |
| 20 | state_ptr, |
| 21 | src_ptr, |
| 22 | indices_ptr, |
| 23 | stride_state_batch, |
| 24 | stride_src_batch, |
| 25 | stride_indices, |
| 26 | row_size: tl.constexpr, |
| 27 | BLOCK_SIZE: tl.constexpr, |
| 28 | launch_pdl: tl.constexpr, |
| 29 | ): |
| 30 | block_idx = tl.program_id(0) |
| 31 | batch_idx = tl.program_id(1) |
| 32 | offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 33 | mask = offsets < row_size |
| 34 | |
| 35 | if launch_pdl: |
| 36 | tl.extra.cuda.gdc_wait() |
| 37 | tl.extra.cuda.gdc_launch_dependents() |
| 38 | |
| 39 | state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64) |
| 40 | values = tl.load(src_ptr + batch_idx * stride_src_batch + offsets, mask=mask) |
| 41 | tl.store(state_ptr + state_idx * stride_state_batch + offsets, values, mask=mask) |
| 42 | |
| 43 | |
| 44 | def _scatter_states_warmup_inputs( |
| 45 | *, |
| 46 | state_shape: tuple[int, ...], |
| 47 | state_dtype: torch.dtype, |
| 48 | indices_dtype: torch.dtype, |
| 49 | max_num_tokens: int, |
| 50 | ) -> dict[str, object]: |
| 51 | num_tokens: Any = WarmupIntRange(1, max_num_tokens + 1) |
| 52 | return dict( |
| 53 | state=TritonWarmupTensor( |
| 54 | state_dtype, |
| 55 | shape=(1,) + state_shape, |
| 56 | ), |
| 57 | src=TritonWarmupTensor( |
| 58 | state_dtype, |
| 59 | shape=(num_tokens,) + state_shape, |
| 60 | ), |
| 61 | indices=TritonWarmupTensor( |
| 62 | indices_dtype, |
| 63 | shape=(num_tokens,), |
| 64 | ), |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | @triton_kernel_dispatcher_with_warmup( |
| 69 | kernel=_scatter_states_kernel, |
| 70 | warmup_inputs=_scatter_states_warmup_inputs, |
| 71 | ) |
| 72 | def scatter_states( |
| 73 | state: torch.Tensor, |
| 74 | src: torch.Tensor, |
| 75 | indices: torch.Tensor, |
| 76 | ) -> DispatchSpec: |
| 77 | """Scatter ``src`` rows into ``state`` at ``indices`` (in place). |
| 78 | |
| 79 | Equivalent to ``state[indices] = src`` but non-atomic and bandwidth-bound, |
| 80 | since mamba cache slots are unique per sequence. ``gather_initial_states`` |
| 81 | is the read-side counterpart. |
| 82 | """ |
| 83 | assert state.ndim >= 2 |
| 84 | assert state.is_cuda |
| 85 | assert src.ndim == state.ndim |
| 86 | assert indices.ndim == 1 |
| 87 | assert indices.device == state.device |
| 88 | assert src.shape[1:] == state.shape[1:] |
| 89 | assert src.shape[0] == indices.shape[0] |
| 90 | assert indices.dtype in (torch.int32, torch.int64) |
| 91 | |
| 92 | row_size = state[0].numel() |
| 93 | assert state[0].is_contiguous() |
| 94 | assert src[0].is_contiguous() |
| 95 | block_size = min(triton.next_power_of_2(row_size), 1024) |
| 96 | grid = (triton.cdiv(row_size, block_size), indices.numel()) |
| 97 | return grid, dict( |
| 98 | stride_state_batch=state.stride(0), |
| 99 | stride_src_batch=src.stride(0), |
| 100 | stride_indices=indices.stride(0), |
| 101 | row_size=row_size, |
| 102 | BLOCK_SIZE=block_size, |
| 103 | num_warps=8, |
| 104 | launch_pdl=current_platform.is_arch_support_pdl(), |
| 105 | ) |