Coverage for src/flag_gems/fused/FLA/index.py: 100%

1# This file contains code copied from the flash-linear-attention project. 

2# The original source code was licensed under the MIT license and included 

3# the following copyright notice: 

4# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang 

5# ruff: noqa: E501 

6import torch 

7import triton 

8 

9from flag_gems.fused.FLA.utils import tensor_cache 

10 

11 

12@tensor_cache 

13def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: 

14 return cu_seqlens[1:] - cu_seqlens[:-1] 

15 

16 

17@tensor_cache 

18def prepare_chunk_indices( 

19 cu_seqlens: torch.LongTensor, chunk_size: int 

20) -> torch.LongTensor: 

21 chunk_counts = triton.cdiv(prepare_lens(cu_seqlens), chunk_size) 

22 chunk_offsets = torch.cat([cu_seqlens.new_tensor([0]), chunk_counts]).cumsum(-1) 

23 chunk_arange = torch.arange(chunk_offsets[-1], device=cu_seqlens.device) 

24 seq_ids = torch.repeat_interleave( 

25 torch.arange(chunk_counts.numel(), device=cu_seqlens.device), chunk_counts 

26 ) 

27 chunk_ids = chunk_arange - torch.repeat_interleave(chunk_offsets[:-1], chunk_counts) 

28 return torch.stack([seq_ids, chunk_ids], 1).to(cu_seqlens) 

29 

30 

31@tensor_cache 

32def prepare_chunk_offsets( 

33 cu_seqlens: torch.LongTensor, chunk_size: int 

34) -> torch.LongTensor: 

35 return torch.cat( 

36 [cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)] 

37 ).cumsum(-1)