| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # SPDX-FileCopyrightText: Songlin Yang, Yu Zhang |
| 4 | # |
| 5 | # This file contains code copied from the flash-linear-attention project. |
| 6 | # The original source code was licensed under the MIT license and included |
| 7 | # the following copyright notice: |
| 8 | # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang |
| 9 | # ruff: noqa: E501 |
| 10 | import torch |
| 11 | |
| 12 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 13 | from vllm.utils.torch_utils import async_tensor_h2d |
| 14 | from vllm.triton_utils import triton |
| 15 | |
| 16 | from .utils import tensor_cache |
| 17 | |
| 18 | |
| 19 | @tensor_cache |
| 20 | def prepare_lens(cu_seqlens: torch.Tensor) -> torch.Tensor: |
| 21 | return cu_seqlens[1:] - cu_seqlens[:-1] |
| 22 | |
| 23 | |
| 24 | @tensor_cache |
| 25 | def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor: |
| 26 | # This will be fixed by https://github.com/vllm-project/vllm/pull/51540. |
| 27 | with gpu_sync_allowed(): |
| 28 | chunk_counts = triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() |
| 29 | indices = torch.cat([torch.arange(n) for n in chunk_counts]) |
| 30 | chunk_indices = torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1) |
| 31 | return async_tensor_h2d( |
| 32 | chunk_indices, device=cu_seqlens.device, dtype=cu_seqlens.dtype |
| 33 | ) |
| 34 | |
| 35 | |
| 36 | @tensor_cache |
| 37 | def prepare_chunk_offsets(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor: |
| 38 | return torch.cat( |
| 39 | [cu_seqlens.new_zeros(1), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)] |
| 40 | ).cumsum(-1) |