| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.platforms import current_platform |
| 7 | from vllm.triton_utils import tl, triton |
| 8 | |
| 9 | |
| 10 | @triton.jit |
| 11 | def _gather_initial_states_kernel( |
| 12 | state_ptr, |
| 13 | indices_ptr, |
| 14 | has_initial_state_ptr, |
| 15 | output_ptr, |
| 16 | stride_state_batch, |
| 17 | stride_indices, |
| 18 | stride_has_initial_state, |
| 19 | row_size: tl.constexpr, |
| 20 | BLOCK_SIZE: tl.constexpr, |
| 21 | launch_pdl: tl.constexpr, |
| 22 | ): |
| 23 | block_idx = tl.program_id(0) |
| 24 | batch_idx = tl.program_id(1) |
| 25 | offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 26 | mask = offsets < row_size |
| 27 | |
| 28 | if launch_pdl: |
| 29 | tl.extra.cuda.gdc_wait() |
| 30 | tl.extra.cuda.gdc_launch_dependents() |
| 31 | |
| 32 | has_initial_state = tl.load( |
| 33 | has_initial_state_ptr + batch_idx * stride_has_initial_state |
| 34 | ).to(tl.int1) |
| 35 | state_idx = tl.load(indices_ptr + batch_idx * stride_indices).to(tl.int64) |
| 36 | state_idx = tl.where(has_initial_state, state_idx, 0) |
| 37 | values = tl.load( |
| 38 | state_ptr + state_idx * stride_state_batch + offsets, |
| 39 | mask=mask & has_initial_state, |
| 40 | other=0.0, |
| 41 | ) |
| 42 | tl.store(output_ptr + batch_idx * row_size + offsets, values, mask=mask) |
| 43 | |
| 44 | |
| 45 | def gather_initial_states( |
| 46 | state: torch.Tensor, |
| 47 | indices: torch.Tensor, |
| 48 | has_initial_state: torch.Tensor, |
| 49 | ) -> torch.Tensor: |
| 50 | """Gather dense state rows, replacing uninitialized rows with zeros.""" |
| 51 | assert state.ndim >= 2 |
| 52 | assert state.is_cuda or state.is_xpu |
| 53 | assert indices.ndim == 1 and has_initial_state.ndim == 1 |
| 54 | assert indices.shape == has_initial_state.shape |
| 55 | assert indices.device == state.device |
| 56 | assert has_initial_state.device == state.device |
| 57 | assert indices.dtype in (torch.int32, torch.int64) |
| 58 | assert has_initial_state.dtype == torch.bool |
| 59 | |
| 60 | row_size = state[0].numel() |
| 61 | # Mamba pages may pad stride(0), but each state row remains dense. |
| 62 | assert state[0].is_contiguous() |
| 63 | output = torch.empty( |
| 64 | (indices.numel(), *state.shape[1:]), |
| 65 | dtype=state.dtype, |
| 66 | device=state.device, |
| 67 | ) |
| 68 | block_size = min(triton.next_power_of_2(row_size), 1024) |
| 69 | grid = (triton.cdiv(row_size, block_size), indices.numel()) |
| 70 | _gather_initial_states_kernel[grid]( |
| 71 | state, |
| 72 | indices, |
| 73 | has_initial_state, |
| 74 | output, |
| 75 | state.stride(0), |
| 76 | indices.stride(0), |
| 77 | has_initial_state.stride(0), |
| 78 | row_size=row_size, |
| 79 | BLOCK_SIZE=block_size, |
| 80 | num_warps=8, |
| 81 | launch_pdl=current_platform.is_arch_support_pdl(), |
| 82 | ) |
| 83 | return output |