38 lines
1.0 KiB
Python
38 lines
1.0 KiB
Python
# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/utils/index.py
|
|
# -*- coding: utf-8 -*-
|
|
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import triton
|
|
import triton.language as tl
|
|
|
|
from sglang.srt.layers.attention.fla.utils import tensor_cache
|
|
|
|
|
|
@tensor_cache
|
|
def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor:
|
|
return cu_seqlens[1:] - cu_seqlens[:-1]
|
|
|
|
|
|
@tensor_cache
|
|
def prepare_chunk_indices(
|
|
cu_seqlens: torch.LongTensor, chunk_size: int
|
|
) -> torch.LongTensor:
|
|
indices = torch.cat(
|
|
[
|
|
torch.arange(n)
|
|
for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()
|
|
]
|
|
)
|
|
return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens)
|
|
|
|
|
|
@tensor_cache
|
|
def prepare_chunk_offsets(
|
|
cu_seqlens: torch.LongTensor, chunk_size: int
|
|
) -> torch.LongTensor:
|
|
return torch.cat(
|
|
[cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]
|
|
).cumsum(-1)
|