Coverage for src/flag_gems/fused/FLA/chunk_o.py: 22%

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 

6# ruff: noqa: E501 

7 

8import torch 

9import triton 

10import triton.language as tl 

11 

12from flag_gems.fused.FLA.index import prepare_chunk_indices 

13from flag_gems.fused.FLA.triton_ops_helper import exp 

14from flag_gems.fused.FLA.utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper 

15from flag_gems.utils import libentry, libtuner 

16 

17BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] 

18NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8] 

19 

20 

21@libentry() 

22@triton.heuristics( 

23 { 

24 "USE_G": lambda args: args["g"] is not None, 

25 "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, 

26 } 

27) 

28@libtuner( 

29 configs=[ 

30 triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) 

31 for BK in BKV_LIST 

32 for BV in BKV_LIST 

33 for num_warps in NUM_WARPS 

34 for num_stages in [2, 3, 4] 

35 ], 

36 key=["H", "K", "V", "BT"], 

37) 

38@triton.jit(do_not_specialize=["T"]) 

39def chunk_fwd_kernel_o( 

40 q, 

41 k, 

42 v, 

43 h, 

44 g, 

45 o, 

46 cu_seqlens, 

47 chunk_indices, 

48 scale, 

49 T, 

50 H: tl.constexpr, 

51 Hg: tl.constexpr, 

52 K: tl.constexpr, 

53 V: tl.constexpr, 

54 BT: tl.constexpr, 

55 BK: tl.constexpr, 

56 BV: tl.constexpr, 

57 USE_G: tl.constexpr, 

58 IS_VARLEN: tl.constexpr, 

59): 

60 i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) 

61 i_b, i_h = i_bh // H, i_bh % H 

62 

63 if IS_VARLEN: 

64 i_tg = i_t 

65 i_n, i_t = ( 

66 tl.load(chunk_indices + i_t * 2).to(tl.int32), 

67 tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), 

68 ) 

69 bos, eos = ( 

70 tl.load(cu_seqlens + i_n).to(tl.int32), 

71 tl.load(cu_seqlens + i_n + 1).to(tl.int32), 

72 ) 

73 T = eos - bos 

74 NT = tl.cdiv(T, BT) 

75 else: 

76 NT = tl.cdiv(T, BT) 

77 i_tg = i_b * NT + i_t 

78 bos, eos = i_b * T, i_b * T + T 

79 

80 # offset calculation 

81 q += (bos * Hg + i_h // (H // Hg)) * K 

82 k += (bos * Hg + i_h // (H // Hg)) * K 

83 v += (bos * H + i_h) * V 

84 o += (bos * H + i_h) * V 

85 h += (i_tg * H + i_h).to(tl.int64) * K * V 

86 

87 b_o = tl.zeros([BT, BV], dtype=tl.float32) 

88 b_A = tl.zeros([BT, BT], dtype=tl.float32) 

89 

90 for i_k in range(tl.cdiv(K, BK)): 

91 p_q = tl.make_block_ptr( 

92 q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) 

93 ) 

94 p_k = tl.make_block_ptr( 

95 k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) 

96 ) 

97 p_h = tl.make_block_ptr( 

98 h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0) 

99 ) 

100 # [BT, BK] 

101 b_q = tl.load(p_q, boundary_check=(0, 1)) 

102 # [BK, BT] 

103 b_k = tl.load(p_k, boundary_check=(0, 1)) 

104 # [BK, BV] 

105 b_h = tl.load(p_h, boundary_check=(0, 1)) 

106 

107 # [BT, BK] @ [BK, BV] -> [BT, BV] 

108 b_o += tl.dot(b_q, b_h) 

109 # [BT, BK] @ [BK, BT] -> [BT, BT] 

110 b_A += tl.dot(b_q, b_k) 

111 

112 if USE_G: 

113 g += bos * H + i_h 

114 p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) 

115 b_g = tl.load(p_g, boundary_check=(0,)) 

116 b_o = b_o * exp(b_g)[:, None] 

117 b_A = b_A * exp(b_g[:, None] - b_g[None, :]) 

118 

119 o_t = i_t * BT + tl.arange(0, BT) 

120 m_t = o_t < T 

121 m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) 

122 b_A = tl.where(m_A, b_A, 0) 

123 

124 p_v = tl.make_block_ptr( 

125 v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) 

126 ) 

127 p_o = tl.make_block_ptr( 

128 o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) 

129 ) 

130 b_v = tl.load(p_v, boundary_check=(0, 1)) 

131 

132 # to fix mma -> mma layout conversion 

133 # already solved by triton v3.2 or higher 

134 b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale 

135 tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) 

136 

137 

138def chunk_fwd_o( 

139 q: torch.Tensor, 

140 k: torch.Tensor, 

141 v: torch.Tensor, 

142 h: torch.Tensor, 

143 g: torch.Tensor | None = None, # cumsum of log decay 

144 scale: float | None = None, 

145 cu_seqlens: torch.LongTensor | None = None, 

146 chunk_size: int = 64, 

147) -> torch.Tensor: 

148 B, T, Hg, K, V = *q.shape, v.shape[-1] 

149 H = v.shape[-2] 

150 BT = 64 if FLA_GDN_FIX_BT else min(chunk_size, max(16, triton.next_power_of_2(T))) 

151 chunk_indices = ( 

152 prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None 

153 ) 

154 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) 

155 if scale is None: 

156 scale = k.shape[-1] ** -0.5 

157 

158 o = torch.empty_like(v) 

159 

160 def grid(meta): 

161 return (triton.cdiv(V, meta["BV"]), NT, B * H) 

162 

163 chunk_fwd_kernel_o[grid]( 

164 q, 

165 k, 

166 v, 

167 h, 

168 g, 

169 o, 

170 cu_seqlens, 

171 chunk_indices, 

172 scale, 

173 T=T, 

174 H=H, 

175 Hg=Hg, 

176 K=K, 

177 V=V, 

178 BT=BT, 

179 ) 

180 return o