@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1aa427fdc1

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Kernels] Add Flash Linear Attention Kernels (#24518)

Signed-off-by: youkaichao <youkaichao@gmail.com>

youkaichao committed Sep 10, 2025 at 00:04 UTC 1aa427fdc1d4ddb7a0c966f7c46badd1c3952c14

17 files changed +2671 -2

tools/mypy.sh

+1 -1

@@ -29,7 +29,7 @@ run_mypy vllm/engine

29 run_mypy vllm/executor

30 run_mypy vllm/inputs

31 run_mypy vllm/lora

32 -run_mypy vllm/model_executor

32 +run_mypy --exclude 'vllm/model_executor/layers/fla/ops' vllm/model_executor

33 run_mypy vllm/plugins

34 run_mypy vllm/worker

35 run_mypy vllm/v1

vllm/model_executor/layers/fla/__init__.py new

@@ -0,0 +1,8 @@

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:

vllm/model_executor/layers/fla/ops/__init__.py new

+17

@@ -0,0 +1,17 @@

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:

9 +from .chunk import chunk_gated_delta_rule

10 +from .fused_recurrent import fused_recurrent_gated_delta_rule

11 +from .layernorm_guard import RMSNormGated

12 +

13 +__all__ = [

14 + "RMSNormGated",

15 + "chunk_gated_delta_rule",

16 + "fused_recurrent_gated_delta_rule",

17 +]

vllm/model_executor/layers/fla/ops/chunk.py new

+225

@@ -0,0 +1,225 @@

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:

9 +# ruff: noqa: E501

10 +import warnings

11 +from typing import Optional

12 +

13 +import torch

14 +from einops import rearrange

15 +

16 +from .chunk_delta_h import chunk_gated_delta_rule_fwd_h

17 +from .chunk_o import chunk_fwd_o

18 +from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd

19 +from .cumsum import chunk_local_cumsum

20 +from .l2norm import l2norm_fwd

21 +from .solve_tril import solve_tril

22 +from .utils import SUPPRESS_LEVEL, input_guard

23 +from .wy_fast import recompute_w_u_fwd

24 +

25 +

26 +def chunk_gated_delta_rule_fwd(q: torch.Tensor,

27 + k: torch.Tensor,

28 + v: torch.Tensor,

29 + g: torch.Tensor,

30 + beta: torch.Tensor,

31 + scale: float,

32 + initial_state: torch.Tensor,

33 + output_final_state: bool,

34 + cu_seqlens: Optional[torch.LongTensor] = None):

35 + g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens)

36 + # obtain WY representation. u is actually the new v.

37 + A = chunk_scaled_dot_kkt_fwd(k=k,

38 + beta=beta,

39 + g_cumsum=g,

40 + cu_seqlens=cu_seqlens,

41 + output_dtype=torch.float32)

42 + A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)

43 + w, u = recompute_w_u_fwd(

44 + k=k,

45 + v=v,

46 + beta=beta,

47 + A=A,

48 + g_cumsum=g,

49 + cu_seqlens=cu_seqlens,

50 + )

51 + h, v_new, final_state = chunk_gated_delta_rule_fwd_h(

52 + k=k,

53 + w=w,

54 + u=u,

55 + g=g,

56 + initial_state=initial_state,

57 + output_final_state=output_final_state,

58 + cu_seqlens=cu_seqlens,

59 + )

60 + o = chunk_fwd_o(

61 + q=q,

62 + k=k,

63 + v=v_new,

64 + h=h,

65 + g=g,

66 + scale=scale,

67 + cu_seqlens=cu_seqlens,

68 + )

69 + if SUPPRESS_LEVEL < 3:

70 + return g, o, A, final_state, None, None, None

71 + elif SUPPRESS_LEVEL >= 3:

72 + return g, o, A, final_state, w, h, v_new

73 +

74 +

75 +class ChunkGatedDeltaRuleFunction(torch.autograd.Function):

76 +

77 + @staticmethod

78 + @input_guard

79 + @torch.amp.custom_fwd(device_type='cuda')

80 + def forward(ctx,

81 + q: torch.Tensor,

82 + k: torch.Tensor,

83 + v: torch.Tensor,

84 + g: torch.Tensor,

85 + beta: torch.Tensor,

86 + scale: float,

87 + initial_state: torch.Tensor,

88 + output_final_state: bool,

89 + cu_seqlens: Optional[torch.LongTensor] = None,

90 + use_qk_l2norm_in_kernel: bool = False):

91 + if use_qk_l2norm_in_kernel:

92 + q = l2norm_fwd(q)

93 + k = l2norm_fwd(k)

94 +

95 + g, o, A, final_state, w, h, v_new = chunk_gated_delta_rule_fwd(

96 + q=q,

97 + k=k,

98 + v=v,

99 + g=g,

100 + beta=beta,

101 + scale=scale,

102 + initial_state=initial_state,

103 + output_final_state=output_final_state,

104 + cu_seqlens=cu_seqlens,

105 + )

106 + ctx.scale = scale

107 + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel

108 + return o.to(q.dtype), final_state

109 +

110 +

111 +@torch.compiler.disable

112 +def chunk_gated_delta_rule(q: torch.Tensor,

113 + k: torch.Tensor,

114 + v: torch.Tensor,

115 + g: torch.Tensor,

116 + beta: torch.Tensor,

117 + scale: float = None,

118 + initial_state: torch.Tensor = None,

119 + output_final_state: bool = False,

120 + cu_seqlens: Optional[torch.LongTensor] = None,

121 + head_first: bool = False,

122 + use_qk_l2norm_in_kernel: bool = False):

123 + r"""

124 + Args:

125 + q (torch.Tensor):

126 + queries of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.

127 + k (torch.Tensor):

128 + keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.

129 + v (torch.Tensor):

130 + values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.

131 + g (torch.Tensor):

132 + (forget) gating tensor (in log space!) of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.

133 + beta (torch.Tensor):

134 + betas of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.

135 + scale (Optional[int]):

136 + Scale factor for the RetNet attention scores.

137 + If not provided, it will default to `1 / sqrt(K)`. Default: `None`.

138 + initial_state (Optional[torch.Tensor]):

139 + Initial state of shape `[N, H, K, V]` for `N` input sequences.

140 + For equal-length input sequences, `N` equals the batch size `B`.

141 + Default: `None`.

142 + output_final_state (Optional[bool]):

143 + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`.

144 + cu_seqlens (torch.LongTensor):

145 + Cumulative sequence lengths of shape `[N+1]` used for variable-length training,

146 + consistent with the FlashAttention API.

147 + head_first (Optional[bool]):

148 + Whether the inputs are in the head-first format, which is not supported for variable-length inputs.

149 + Default: `False`.

150 +

151 + Returns:

152 + o (torch.Tensor):

153 + Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.

154 + final_state (torch.Tensor):

155 + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`.

156 +

157 + Examples::

158 + >>> import torch

159 + >>> import torch.nn.functional as F

160 + >>> from einops import rearrange

161 + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule

162 + # inputs with equal lengths

163 + >>> B, T, H, K, V = 4, 2048, 4, 512, 512

164 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')

165 + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)

166 + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')

167 + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid()

168 + >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda'))

169 + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda')

170 + >>> o, ht = chunk_gated_delta_rule(

171 + q, k, v, g, beta,

172 + initial_state=h0,

173 + output_final_state=True

174 + )

175 + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required

176 + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))

177 + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected

178 + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)

179 + >>> o_var, ht_var = chunk_gated_delta_rule(

180 + q, k, v, g, beta,

181 + initial_state=h0,

182 + output_final_state=True,

183 + cu_seqlens=cu_seqlens

184 + )

185 + """

186 + assert q.dtype == k.dtype == v.dtype

187 + assert q.dtype != torch.float32, "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."

188 + assert len(

189 + beta.shape

190 + ) == 3, "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."

191 +

192 + if head_first:

193 + raise DeprecationWarning(

194 + "head_first is deprecated and will be removed in a future version. "

195 + "Please use head_first=False for now instead.",

196 + stacklevel=2)

197 + q, k, v, beta, g = map(

198 + lambda x: rearrange(x, 'b h t ... -> b t h ...'),

199 + (q, k, v, beta, g))

200 + if not head_first and q.shape[1] < q.shape[2]:

201 + warnings.warn(

202 + f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "

203 + "This may indicate the inputs were passed in head-first format [B, H, T, ...] "

204 + "when head_first=False was specified. "

205 + "Please verify your input tensor format matches the expected shape [B, T, H, ...].",

206 + stacklevel=2)

207 + if cu_seqlens is not None:

208 + if q.shape[0] != 1:

209 + raise ValueError(

210 + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."

211 + f"Please flatten variable-length inputs before processing.")

212 + if initial_state is not None and initial_state.shape[0] != len(

213 + cu_seqlens) - 1:

214 + raise ValueError(

215 + f"The number of initial states is expected to be equal to the number of input sequences, "

216 + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}."

217 + )

218 + if scale is None:

219 + scale = k.shape[-1]**-0.5

220 + o, final_state = ChunkGatedDeltaRuleFunction.apply(

221 + q, k, v, g, beta, scale, initial_state, output_final_state, cu_seqlens,

222 + use_qk_l2norm_in_kernel)

223 + if head_first:

224 + o = rearrange(o, 'b t h ... -> b h t ...')

225 + return o, final_state

vllm/model_executor/layers/fla/ops/chunk_delta_h.py new

+289

@@ -0,0 +1,289 @@

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:

9 +# ruff: noqa: E501

10 +from typing import Optional

11 +

12 +import torch

13 +

14 +from vllm.triton_utils import tl, triton

15 +

16 +from .index import prepare_chunk_indices, prepare_chunk_offsets

17 +from .op import exp, safe_exp

18 +from .utils import is_nvidia_hopper, use_cuda_graph

19 +

20 +NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16]

21 +

22 +

23 +@triton.heuristics({

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

25 + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None,

26 + 'STORE_FINAL_STATE': lambda args: args['ht'] is not None,

27 + 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None,

28 + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,

29 +})

30 +@triton.autotune(

31 + configs=[

32 + triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages)

33 + for num_warps in [2, 4] for num_stages in [2, 3, 4] for BV in [32, 64]

34 + ],

35 + key=['H', 'K', 'V', 'BT', 'USE_G'],

36 + use_cuda_graph=use_cuda_graph,

37 +)

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

39 +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(

40 + k,

41 + v,

42 + w,

43 + v_new,

44 + g,

45 + h,

46 + h0,

47 + ht,

48 + cu_seqlens,

49 + chunk_offsets,

50 + T,

51 + H: tl.constexpr,

52 + Hg: tl.constexpr,

53 + K: tl.constexpr,

54 + V: tl.constexpr,

55 + BT: tl.constexpr,

56 + BV: tl.constexpr,

57 + USE_G: tl.constexpr,

58 + USE_INITIAL_STATE: tl.constexpr,

59 + STORE_FINAL_STATE: tl.constexpr,

60 + SAVE_NEW_VALUE: tl.constexpr,

61 + IS_VARLEN: tl.constexpr,

62 +):

63 + i_v, i_nh = tl.program_id(0), tl.program_id(1)

64 + i_n, i_h = i_nh // H, i_nh % H

65 + if IS_VARLEN:

66 + bos, eos = tl.load(cu_seqlens + i_n).to(

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

68 + T = eos - bos

69 + NT = tl.cdiv(T, BT)

70 + boh = tl.load(chunk_offsets + i_n).to(tl.int32)

71 + else:

72 + bos, eos = i_n * T, i_n * T + T

73 + NT = tl.cdiv(T, BT)

74 + boh = i_n * NT

75 +

76 + # [BK, BV]

77 + b_h1 = tl.zeros([64, BV], dtype=tl.float32)

78 + if K > 64:

79 + b_h2 = tl.zeros([64, BV], dtype=tl.float32)

80 + if K > 128:

81 + b_h3 = tl.zeros([64, BV], dtype=tl.float32)

82 + if K > 192:

83 + b_h4 = tl.zeros([64, BV], dtype=tl.float32)

84 +

85 + # calculate offset

86 + h += (boh * H + i_h) * K * V

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

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

89 + w += (bos * H + i_h) * K

90 + if SAVE_NEW_VALUE:

91 + v_new += (bos * H + i_h) * V

92 + stride_v = H * V

93 + stride_h = H * K * V

94 + stride_k = Hg * K

95 + stride_w = H * K

96 + if USE_INITIAL_STATE:

97 + h0 = h0 + i_nh * K * V

98 + if STORE_FINAL_STATE:

99 + ht = ht + i_nh * K * V

100 +

101 + # load initial state

102 + if USE_INITIAL_STATE:

103 + p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV),

104 + (1, 0))

105 + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32)

106 + if K > 64:

107 + p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV),

108 + (64, BV), (1, 0))

109 + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32)

110 + if K > 128:

111 + p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV),

112 + (64, BV), (1, 0))

113 + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32)

114 + if K > 192:

115 + p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV),

116 + (64, BV), (1, 0))

117 + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32)

118 +

119 + # main recurrence

120 + for i_t in range(NT):

121 + p_h1 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1),

122 + (0, i_v * BV), (64, BV), (1, 0))

123 + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1))

124 + if K > 64:

125 + p_h2 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1),

126 + (64, i_v * BV), (64, BV), (1, 0))

127 + tl.store(p_h2,

128 + b_h2.to(p_h2.dtype.element_ty),

129 + boundary_check=(0, 1))

130 + if K > 128:

131 + p_h3 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1),

132 + (128, i_v * BV), (64, BV), (1, 0))

133 + tl.store(p_h3,

134 + b_h3.to(p_h3.dtype.element_ty),

135 + boundary_check=(0, 1))

136 + if K > 192:

137 + p_h4 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1),

138 + (192, i_v * BV), (64, BV), (1, 0))

139 + tl.store(p_h4,

140 + b_h4.to(p_h4.dtype.element_ty),

141 + boundary_check=(0, 1))

142 +

143 + p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV),

144 + (BT, BV), (1, 0))

145 + p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_v, 1),

146 + (i_t * BT, i_v * BV), (BT, BV),

147 + (1, 0)) if SAVE_NEW_VALUE else None

148 + b_v_new = tl.zeros([BT, BV], dtype=tl.float32)

149 + p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 0),

150 + (BT, 64), (1, 0))

151 + b_w = tl.load(p_w, boundary_check=(0, 1))

152 + b_v_new += tl.dot(b_w, b_h1.to(b_w.dtype))

153 + if K > 64:

154 + p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 64),

155 + (BT, 64), (1, 0))

156 + b_w = tl.load(p_w, boundary_check=(0, 1))

157 + b_v_new += tl.dot(b_w, b_h2.to(b_w.dtype))

158 + if K > 128:

159 + p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 128),

160 + (BT, 64), (1, 0))

161 + b_w = tl.load(p_w, boundary_check=(0, 1))

162 + b_v_new += tl.dot(b_w, b_h3.to(b_w.dtype))

163 + if K > 192:

164 + p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 192),

165 + (BT, 64), (1, 0))

166 + b_w = tl.load(p_w, boundary_check=(0, 1))

167 + b_v_new += tl.dot(b_w, b_h4.to(b_w.dtype))

168 + b_v_new = -b_v_new + tl.load(p_v, boundary_check=(0, 1))

169 +

170 + if SAVE_NEW_VALUE:

171 + p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_v, 1),

172 + (i_t * BT, i_v * BV), (BT, BV), (1, 0))

173 + tl.store(p_v_new,

174 + b_v_new.to(p_v_new.dtype.element_ty),

175 + boundary_check=(0, 1))

176 +

177 + if USE_G:

178 + last_idx = min((i_t + 1) * BT, T) - 1

179 + b_g_last = tl.load(g + bos * H + last_idx * H + i_h)

180 + p_g = tl.make_block_ptr(g + bos * H + i_h, (T, ), (H, ),

181 + (i_t * BT, ), (BT, ), (0, ))

182 + b_g = tl.load(p_g, boundary_check=(0, ))

183 + b_v_new = b_v_new * safe_exp(b_g_last - b_g)[:, None]

184 + b_g_last = exp(b_g_last)

185 + b_h1 = b_h1 * b_g_last

186 + if K > 64:

187 + b_h2 = b_h2 * b_g_last

188 + if K > 128:

189 + b_h3 = b_h3 * b_g_last

190 + if K > 192:

191 + b_h4 = b_h4 * b_g_last

192 + b_v_new = b_v_new.to(k.dtype.element_ty)

193 + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT),

194 + (64, BT), (0, 1))

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

196 + b_h1 += tl.dot(b_k, b_v_new)

197 + if K > 64:

198 + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT),

199 + (64, BT), (0, 1))

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

201 + b_h2 += tl.dot(b_k, b_v_new)

202 + if K > 128:

203 + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT),

204 + (64, BT), (0, 1))

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

206 + b_h3 += tl.dot(b_k, b_v_new)

207 + if K > 192:

208 + p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT),

209 + (64, BT), (0, 1))

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

211 + b_h4 += tl.dot(b_k, b_v_new)

212 +

213 + # epilogue

214 + if STORE_FINAL_STATE:

215 + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV),

216 + (1, 0))

217 + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1))

218 + if K > 64:

219 + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV),

220 + (64, BV), (1, 0))

221 + tl.store(p_ht,

222 + b_h2.to(p_ht.dtype.element_ty),

223 + boundary_check=(0, 1))

224 + if K > 128:

225 + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV),

226 + (64, BV), (1, 0))

227 + tl.store(p_ht,

228 + b_h3.to(p_ht.dtype.element_ty),

229 + boundary_check=(0, 1))

230 + if K > 192:

231 + p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV),

232 + (64, BV), (1, 0))

233 + tl.store(p_ht,

234 + b_h4.to(p_ht.dtype.element_ty),

235 + boundary_check=(0, 1))

236 +

237 +

238 +def chunk_gated_delta_rule_fwd_h(

239 + k: torch.Tensor,

240 + w: torch.Tensor,

241 + u: torch.Tensor,

242 + g: Optional[torch.Tensor] = None,

243 + initial_state: Optional[torch.Tensor] = None,

244 + output_final_state: bool = False,

245 + chunk_size: int = 64, # SY: remove this argument and force chunk size 64?

246 + save_new_value: bool = True,

247 + cu_seqlens: Optional[torch.LongTensor] = None,

248 +) -> tuple[torch.Tensor, torch.Tensor]:

249 + B, T, Hg, K, V = *k.shape, u.shape[-1]

250 + H = u.shape[-2]

251 + BT = chunk_size

252 +

253 + chunk_indices = prepare_chunk_indices(

254 + cu_seqlens, chunk_size) if cu_seqlens is not None else None

255 + # N: the actual number of sequences in the batch with either equal or variable lengths

256 + if cu_seqlens is None:

257 + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None

258 + else:

259 + N, NT, chunk_offsets = len(cu_seqlens) - 1, len(

260 + chunk_indices), prepare_chunk_offsets(cu_seqlens, BT)

261 + assert K <= 256, "current kernel does not support head dimension larger than 256."

262 +

263 + h = k.new_empty(B, NT, H, K, V)

264 + final_state = k.new_empty(

265 + N, H, K, V, dtype=torch.float32) if output_final_state else None

266 +

267 + v_new = torch.empty_like(u) if save_new_value else None

268 +

269 + def grid(meta):

270 + return (triton.cdiv(V, meta['BV']), N * H)

271 +

272 + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid](

273 + k=k,

274 + v=u,

275 + w=w,

276 + v_new=v_new,

277 + g=g,

278 + h=h,

279 + h0=initial_state,

280 + ht=final_state,

281 + cu_seqlens=cu_seqlens,

282 + chunk_offsets=chunk_offsets,

283 + T=T,

284 + H=H,

285 + Hg=Hg,

286 + K=K,

287 + V=V,

288 + BT=BT)

289 + return h, v_new, final_state

vllm/model_executor/layers/fla/ops/chunk_o.py new

+176

@@ -0,0 +1,176 @@

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:

9 +

10 +# ruff: noqa: E501

11 +

12 +from typing import Optional

13 +

14 +import torch

15 +

16 +from vllm.triton_utils import tl, triton

17 +

18 +from .index import prepare_chunk_indices

19 +from .op import exp, safe_exp

20 +from .utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper

21 +

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

23 +NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8]

24 +

25 +

26 +@triton.heuristics({

27 + 'USE_G': lambda args: args['g'] is not None,

28 + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None

29 +})

30 +@triton.autotune(

31 + configs=[

32 + triton.Config({

33 + 'BK': BK,

34 + 'BV': BV

35 + },

36 + num_warps=num_warps,

37 + num_stages=num_stages) for BK in BKV_LIST

38 + for BV in BKV_LIST for num_warps in NUM_WARPS

39 + for num_stages in [2, 3, 4]

40 + ],

41 + key=['H', 'K', 'V', 'BT'],

42 +)

43 +@triton.jit(do_not_specialize=['T'])

44 +def chunk_fwd_kernel_o(

45 + q,

46 + k,

47 + v,

48 + h,

49 + g,

50 + o,

51 + cu_seqlens,

52 + chunk_indices,

53 + scale,

54 + T,

55 + H: tl.constexpr,

56 + Hg: tl.constexpr,

57 + K: tl.constexpr,

58 + V: tl.constexpr,

59 + BT: tl.constexpr,

60 + BK: tl.constexpr,

61 + BV: tl.constexpr,

62 + USE_G: tl.constexpr,

63 + IS_VARLEN: tl.constexpr,

64 +):

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

66 + i_b, i_h = i_bh // H, i_bh % H

67 +

68 + if IS_VARLEN:

69 + i_tg = i_t

70 + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(

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

72 + bos, eos = tl.load(cu_seqlens + i_n).to(

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

74 + T = eos - bos

75 + NT = tl.cdiv(T, BT)

76 + else:

77 + NT = tl.cdiv(T, BT)

78 + i_tg = i_b * NT + i_t

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

80 +

81 + # offset calculation

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

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

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

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

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

87 +

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

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

90 +

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

92 + p_q = tl.make_block_ptr(q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK),

93 + (BT, BK), (1, 0))

94 + p_k = tl.make_block_ptr(k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT),

95 + (BK, BT), (0, 1))

96 + p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV),

97 + (BK, BV), (1, 0))

98 + # [BT, BK]

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

100 + # [BK, BT]

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

102 + # [BK, BV]

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

104 +

105 + # [BT, BK] @ [BK, BV] -> [BT, BV]

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

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

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

109 +

110 + if USE_G:

111 + g += bos * H + i_h

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

113 + b_g = tl.load(p_g, boundary_check=(0, ))

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

115 + b_A = b_A * safe_exp(b_g[:, None] - b_g[None, :])

116 +

117 + o_i = tl.arange(0, BT)

118 + m_A = o_i[:, None] >= o_i[None, :]

119 + b_A = tl.where(m_A, b_A, 0)

120 +

121 + p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV),

122 + (BT, BV), (1, 0))

123 + p_o = tl.make_block_ptr(o, (T, V), (H * V, 1), (i_t * BT, i_v * BV),

124 + (BT, BV), (1, 0))

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

126 +

127 + # to fix mma -> mma layout conversion

128 + # already solved by triton v3.2 or higher

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

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

131 +

132 +

133 +def chunk_fwd_o(

134 + q: torch.Tensor,

135 + k: torch.Tensor,

136 + v: torch.Tensor,

137 + h: torch.Tensor,

138 + g: Optional[torch.Tensor] = None, # cumsum of log decay

139 + scale: Optional[float] = None,

140 + cu_seqlens: Optional[torch.LongTensor] = None,

141 + chunk_size: int = 64) -> torch.Tensor:

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

143 + H = v.shape[-2]

144 + if FLA_GDN_FIX_BT:

145 + BT = 64

146 + else:

147 + BT = min(chunk_size, max(16, triton.next_power_of_2(T)))

148 + chunk_indices = prepare_chunk_indices(

149 + cu_seqlens, BT) if cu_seqlens is not None else None

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

151 + if scale is None:

152 + scale = k.shape[-1]**-0.5

153 +

154 + o = torch.empty_like(v)

155 +

156 + def grid(meta):

157 + return (triton.cdiv(V, meta['BV']), NT, B * H)

158 +

159 + chunk_fwd_kernel_o[grid](

160 + q,

161 + k,

162 + v,

163 + h,

164 + g,

165 + o,

166 + cu_seqlens,

167 + chunk_indices,

168 + scale,

169 + T=T,

170 + H=H,

171 + Hg=Hg,

172 + K=K,

173 + V=V,

174 + BT=BT,

175 + )

176 + return o

vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py new

+138

@@ -0,0 +1,138 @@

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:

9 +# ruff: noqa: E501

10 +from typing import Optional

11 +

12 +import torch

13 +

14 +from vllm.triton_utils import tl, triton

15 +

16 +from .index import prepare_chunk_indices

17 +from .op import safe_exp

18 +

19 +

20 +@triton.heuristics({

21 + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,

22 + 'USE_G': lambda args: args['g_cumsum'] is not None

23 +})

24 +@triton.autotune(

25 + configs=[

26 + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages)

27 + for BK in [32, 64, 128] for num_warps in [2, 4, 8]

28 + for num_stages in [2, 3, 4]

29 + ],

30 + key=['H', 'K', 'BT', 'IS_VARLEN'],

31 +)

32 +@triton.jit(do_not_specialize=['T'])

33 +def chunk_scaled_dot_kkt_fwd_kernel(

34 + k,

35 + beta,

36 + g_cumsum,

37 + A,

38 + cu_seqlens,

39 + chunk_indices,

40 + T,

41 + H: tl.constexpr,

42 + Hg: tl.constexpr,

43 + K: tl.constexpr,

44 + BT: tl.constexpr,

45 + BK: tl.constexpr,

46 + IS_VARLEN: tl.constexpr,

47 + USE_G: tl.constexpr,

48 +):

49 + i_t, i_bh = tl.program_id(0), tl.program_id(1)

50 + i_b, i_h = i_bh // H, i_bh % H

51 + if IS_VARLEN:

52 + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(

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

54 + bos, eos = tl.load(cu_seqlens + i_n).to(

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

56 + T = eos - bos

57 + else:

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

59 + o_t = tl.arange(0, BT)

60 +

61 + p_beta = tl.make_block_ptr(beta + bos * H + i_h, (T, ), (H, ),

62 + (i_t * BT, ), (BT, ), (0, ))

63 + b_beta = tl.load(p_beta, boundary_check=(0, ))

64 +

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

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

67 + p_k = tl.make_block_ptr(k + (bos * Hg + i_h // (H // Hg)) * K, (T, K),

68 + (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK),

69 + (1, 0))

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

71 + b_kb = b_k * b_beta[:, None]

72 + b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k))

73 +

74 + if USE_G:

75 + p_g = tl.make_block_ptr(g_cumsum + bos * H + i_h, (T, ), (H, ),

76 + (i_t * BT, ), (BT, ), (0, ))

77 + b_g = tl.load(p_g, boundary_check=(0, ))

78 + b_g_diff = b_g[:, None] - b_g[None, :]

79 + b_A = b_A * safe_exp(b_g_diff)

80 +

81 + b_A = tl.where(o_t[:, None] > o_t[None, :], b_A, 0)

82 + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1),

83 + (i_t * BT, 0), (BT, BT), (1, 0))

84 + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))

85 +

86 +

87 +def chunk_scaled_dot_kkt_fwd(

88 + k: torch.Tensor,

89 + beta: torch.Tensor,

90 + g_cumsum: Optional[torch.Tensor] = None,

91 + cu_seqlens: Optional[torch.LongTensor] = None,

92 + chunk_size: int = 64,

93 + output_dtype: torch.dtype = torch.float32) -> torch.Tensor:

94 + r"""

95 + Compute beta * K * K^T.

96 +

97 + Args:

98 + k (torch.Tensor):

99 + The key tensor of shape `[B, T, H, K]`.

100 + beta (torch.Tensor):

101 + The beta tensor of shape `[B, T, H]`.

102 + g_cumsum (torch.Tensor):

103 + The cumulative sum of the gate tensor of shape `[B, T, H]`.

104 + Default: None

105 + cu_seqlens (torch.LongTensor):

106 + The cumulative sequence lengths of the input tensor.

107 + Default: None

108 + chunk_size (int):

109 + The chunk size. Default: 64.

110 + output_dtype (torch.dtype):

111 + The dtype of the output tensor. Default: `torch.float32`

112 +

113 + Returns:

114 + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.

115 + """

116 +

117 + B, T, Hg, K = k.shape

118 +

119 + H = beta.shape[-1]

120 + BT = chunk_size

121 + chunk_indices = prepare_chunk_indices(

122 + cu_seqlens, BT) if cu_seqlens is not None else None

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

124 + A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)

125 + chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)](

126 + k=k,

127 + beta=beta,

128 + g_cumsum=g_cumsum,

129 + A=A,

130 + cu_seqlens=cu_seqlens,

131 + chunk_indices=chunk_indices,

132 + T=T,

133 + H=H,

134 + Hg=Hg,

135 + K=K,

136 + BT=BT,

137 + )

138 + return A

vllm/model_executor/layers/fla/ops/cumsum.py new

+226

@@ -0,0 +1,226 @@

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:

9 +# ruff: noqa: E501

10 +import warnings

11 +from typing import Optional

12 +

13 +import torch

14 +

15 +from vllm.triton_utils import tl, triton

16 +

17 +from .index import prepare_chunk_indices

18 +from .utils import check_shared_mem, input_guard

19 +

20 +BS_LIST = [32, 64] if check_shared_mem() else [16, 32]

21 +

22 +

23 +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None})

24 +@triton.autotune(configs=[

25 + triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]

26 +],

27 + key=['B', 'H', 'BT', 'IS_VARLEN', 'REVERSE'])

28 +@triton.jit(do_not_specialize=['T'])

29 +def chunk_local_cumsum_scalar_kernel(

30 + s,

31 + o,

32 + cu_seqlens,

33 + chunk_indices,

34 + T,

35 + B: tl.constexpr,

36 + H: tl.constexpr,

37 + BT: tl.constexpr,

38 + REVERSE: tl.constexpr,

39 + IS_VARLEN: tl.constexpr,

40 + HEAD_FIRST: tl.constexpr,

41 +):

42 + i_t, i_bh = tl.program_id(0), tl.program_id(1)

43 + i_b, i_h = i_bh // H, i_bh % H

44 + if IS_VARLEN:

45 + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(

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

47 + bos, eos = tl.load(cu_seqlens + i_n).to(

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

49 + T = eos - bos

50 + else:

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

52 +

53 + if HEAD_FIRST:

54 + p_s = tl.make_block_ptr(s + bos * H + i_h * T, (T, ), (1, ),

55 + (i_t * BT, ), (BT, ), (0, ))

56 + p_o = tl.make_block_ptr(o + bos * H + i_h * T, (T, ), (1, ),

57 + (i_t * BT, ), (BT, ), (0, ))

58 + else:

59 + p_s = tl.make_block_ptr(s + bos * H + i_h, (T, ), (H, ), (i_t * BT, ),

60 + (BT, ), (0, ))

61 + p_o = tl.make_block_ptr(o + bos * H + i_h, (T, ), (H, ), (i_t * BT, ),

62 + (BT, ), (0, ))

63 + # [BT]

64 + b_s = tl.load(p_s, boundary_check=(0, )).to(tl.float32)

65 + b_o = tl.cumsum(b_s, axis=0)

66 + if REVERSE:

67 + b_z = tl.sum(b_s, axis=0)

68 + b_o = -b_o + b_z[None] + b_s

69 + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, ))

70 +

71 +

72 +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None})

73 +@triton.autotune(configs=[

74 + triton.Config({'BS': BS}, num_warps=num_warps) for BS in BS_LIST

75 + for num_warps in [2, 4, 8]

76 +],

77 + key=['B', 'H', 'S', 'BT', 'IS_VARLEN', 'REVERSE'])

78 +@triton.jit(do_not_specialize=['T'])

79 +def chunk_local_cumsum_vector_kernel(

80 + s,

81 + o,

82 + cu_seqlens,

83 + chunk_indices,

84 + T,

85 + B: tl.constexpr,

86 + H: tl.constexpr,

87 + S: tl.constexpr,

88 + BT: tl.constexpr,

89 + BS: tl.constexpr,

90 + REVERSE: tl.constexpr,

91 + IS_VARLEN: tl.constexpr,

92 + HEAD_FIRST: tl.constexpr,

93 +):

94 + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)

95 + i_b, i_h = i_bh // H, i_bh % H

96 + if IS_VARLEN:

97 + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(

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

99 + bos, eos = tl.load(cu_seqlens + i_n).to(

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

101 + T = eos - bos

102 + else:

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

104 +

105 + o_i = tl.arange(0, BT)

106 + if REVERSE:

107 + m_s = tl.where(o_i[:, None] <= o_i[None, :], 1., 0.)

108 + else:

109 + m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.)

110 +

111 + if HEAD_FIRST:

112 + p_s = tl.make_block_ptr(s + (bos * H + i_h * T) * S, (T, S), (S, 1),

113 + (i_t * BT, i_s * BS), (BT, BS), (1, 0))

114 + p_o = tl.make_block_ptr(o + (bos * H + i_h * T) * S, (T, S), (S, 1),

115 + (i_t * BT, i_s * BS), (BT, BS), (1, 0))

116 + else:

117 + p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H * S, 1),

118 + (i_t * BT, i_s * BS), (BT, BS), (1, 0))

119 + p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H * S, 1),

120 + (i_t * BT, i_s * BS), (BT, BS), (1, 0))

121 + # [BT, BS]

122 + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32)

123 + b_o = tl.dot(m_s, b_s, allow_tf32=False)

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

125 +

126 +

127 +def chunk_local_cumsum_scalar(

128 + g: torch.Tensor,

129 + chunk_size: int,

130 + reverse: bool = False,

131 + cu_seqlens: Optional[torch.Tensor] = None,

132 + head_first: bool = False,

133 + output_dtype: Optional[torch.dtype] = torch.float) -> torch.Tensor:

134 + if head_first:

135 + B, H, T = g.shape

136 + else:

137 + B, T, H = g.shape

138 + assert chunk_size == 2**(chunk_size.bit_length() -

139 + 1), "chunk_size must be a power of 2"

140 + BT = chunk_size

141 + chunk_indices = prepare_chunk_indices(

142 + cu_seqlens, BT) if cu_seqlens is not None else None

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

144 + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)

145 + grid = (NT, B * H)

146 + chunk_local_cumsum_scalar_kernel[grid](g_org,

147 + g,

148 + cu_seqlens,

149 + chunk_indices,

150 + T=T,

151 + B=B,

152 + H=H,

153 + BT=BT,

154 + HEAD_FIRST=head_first,

155 + REVERSE=reverse)

156 + return g

157 +

158 +

159 +def chunk_local_cumsum_vector(

160 + g: torch.Tensor,

161 + chunk_size: int,

162 + reverse: bool = False,

163 + cu_seqlens: Optional[torch.Tensor] = None,

164 + head_first: bool = False,

165 + output_dtype: Optional[torch.dtype] = torch.float) -> torch.Tensor:

166 + if head_first:

167 + B, H, T, S = g.shape

168 + else:

169 + B, T, H, S = g.shape

170 + BT = chunk_size

171 + chunk_indices = prepare_chunk_indices(

172 + cu_seqlens, chunk_size) if cu_seqlens is not None else None

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

174 + assert chunk_size == 2**(chunk_size.bit_length() -

175 + 1), "chunk_size must be a power of 2"

176 +

177 + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)

178 +

179 + def grid(meta):

180 + return (triton.cdiv(meta['S'], meta['BS']), NT, B * H)

181 +

182 + # keep cummulative normalizer in fp32

183 + # this kernel is equivalent to

184 + # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1)

185 + chunk_local_cumsum_vector_kernel[grid](g_org,

186 + g,

187 + cu_seqlens,

188 + chunk_indices,

189 + T=T,

190 + B=B,

191 + H=H,

192 + S=S,

193 + BT=BT,

194 + HEAD_FIRST=head_first,

195 + REVERSE=reverse)

196 + return g

197 +

198 +

199 +@input_guard

200 +def chunk_local_cumsum(g: torch.Tensor,

201 + chunk_size: int,

202 + reverse: bool = False,

203 + cu_seqlens: Optional[torch.Tensor] = None,

204 + head_first: bool = False,

205 + output_dtype: Optional[torch.dtype] = torch.float,

206 + **kwargs) -> torch.Tensor:

207 + if not head_first and g.shape[1] < g.shape[2]:

208 + warnings.warn(

209 + f"Input tensor shape suggests potential format mismatch: seq_len ({g.shape[1]}) < num_heads ({g.shape[2]}). "

210 + "This may indicate the inputs were passed in head-first format [B, H, T, ...] "

211 + "when head_first=False was specified. "

212 + "Please verify your input tensor format matches the expected shape [B, T, H, ...].",

213 + stacklevel=2)

214 + if cu_seqlens is not None:

215 + assert g.shape[

216 + 0] == 1, "Only batch size 1 is supported when cu_seqlens are provided"

217 + if len(g.shape) == 3:

218 + return chunk_local_cumsum_scalar(g, chunk_size, reverse, cu_seqlens,

219 + head_first, output_dtype)

220 + elif len(g.shape) == 4:

221 + return chunk_local_cumsum_vector(g, chunk_size, reverse, cu_seqlens,

222 + head_first, output_dtype)

223 + else:

224 + raise ValueError(f"Unsupported input shape {g.shape}. "

225 + f"which should be (B, T, H, D) if `head_first=False` "

226 + f"or (B, H, T, D) otherwise")

vllm/model_executor/layers/fla/ops/fused_recurrent.py new

+366

@@ -0,0 +1,366 @@

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:

9 +# ruff: noqa: E501

10 +from typing import Optional

11 +

12 +import torch

13 +

14 +from vllm.triton_utils import tl, triton

15 +

16 +from .op import exp

17 +

18 +

19 +@triton.heuristics({

20 + 'USE_INITIAL_STATE':

21 + lambda args: args['h0'] is not None,

22 + 'IS_VARLEN':

23 + lambda args: args['cu_seqlens'] is not None,

24 + "IS_CONTINUOUS_BATCHING":

25 + lambda args: args['ssm_state_indices'] is not None,

26 + "IS_SPEC_DECODING":

27 + lambda args: args['num_accepted_tokens'] is not None,

28 +})

29 +@triton.jit(do_not_specialize=['N', 'T'])

30 +def fused_recurrent_gated_delta_rule_fwd_kernel(

31 + q,

32 + k,

33 + v,

34 + g,

35 + beta,

36 + o,

37 + h0,

38 + ht,

39 + cu_seqlens,

40 + ssm_state_indices,

41 + num_accepted_tokens,

42 + scale,

43 + N: tl.constexpr, # num of sequences

44 + T: tl.constexpr, # num of tokens

45 + B: tl.constexpr,

46 + H: tl.constexpr,

47 + HV: tl.constexpr,

48 + K: tl.constexpr,

49 + V: tl.constexpr,

50 + BK: tl.constexpr,

51 + BV: tl.constexpr,

52 + stride_init_state_token: tl.constexpr,

53 + stride_final_state_token: tl.constexpr,

54 + stride_indices_seq: tl.constexpr,

55 + stride_indices_tok: tl.constexpr,

56 + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state

57 + INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace

58 + IS_BETA_HEADWISE: tl.

59 + constexpr, # whether beta is headwise vector or scalar,

60 + USE_QK_L2NORM_IN_KERNEL: tl.constexpr,

61 + IS_VARLEN: tl.constexpr,

62 + IS_CONTINUOUS_BATCHING: tl.constexpr,

63 + IS_SPEC_DECODING: tl.constexpr,

64 +):

65 + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)

66 + i_n, i_hv = i_nh // HV, i_nh % HV

67 + i_h = i_hv // (HV // H)

68 + if IS_VARLEN:

69 + bos, eos = tl.load(cu_seqlens + i_n).to(

70 + tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)

71 + all = T

72 + T = eos - bos

73 + else:

74 + bos, eos = i_n * T, i_n * T + T

75 + all = B * T

76 +

77 + if T == 0:

78 + # no tokens to process for this sequence

79 + return

80 +

81 + o_k = i_k * BK + tl.arange(0, BK)

82 + o_v = i_v * BV + tl.arange(0, BV)

83 +

84 + p_q = q + (bos * H + i_h) * K + o_k

85 + p_k = k + (bos * H + i_h) * K + o_k

86 + p_v = v + (bos * HV + i_hv) * V + o_v

87 + if IS_BETA_HEADWISE:

88 + p_beta = beta + (bos * HV + i_hv) * V + o_v

89 + else:

90 + p_beta = beta + bos * HV + i_hv

91 + p_g = g + bos * HV + i_hv

92 + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v

93 +

94 + mask_k = o_k < K

95 + mask_v = o_v < V

96 + mask_h = mask_k[:, None] & mask_v[None, :]

97 +

98 + b_h = tl.zeros([BK, BV], dtype=tl.float32)

99 + if USE_INITIAL_STATE:

100 + if IS_CONTINUOUS_BATCHING:

101 + if IS_SPEC_DECODING:

102 + i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1

103 + else:

104 + i_t = 0

105 + p_h0 = h0 + tl.load(ssm_state_indices + i_n * stride_indices_seq +

106 + i_t).to(tl.int64) * stride_init_state_token

107 + else:

108 + p_h0 = h0 + bos * HV * K * V

109 + p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :]

110 + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)

111 +

112 + for i_t in range(0, T):

113 + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)

114 + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)

115 + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)

116 + b_g = tl.load(p_g).to(tl.float32)

117 +

118 + if USE_QK_L2NORM_IN_KERNEL:

119 + b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6)

120 + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6)

121 + b_q = b_q * scale

122 + # [BK, BV]

123 + b_h *= exp(b_g)

124 + # [BV]

125 + b_v -= tl.sum(b_h * b_k[:, None], 0)

126 + if IS_BETA_HEADWISE:

127 + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)

128 + else:

129 + b_beta = tl.load(p_beta).to(tl.float32)

130 + b_v *= b_beta

131 + # [BK, BV]

132 + b_h += b_k[:, None] * b_v[None, :]

133 + # [BV]

134 + b_o = tl.sum(b_h * b_q[:, None], 0)

135 + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)

136 +

137 + # keep the states for multi-query tokens

138 + if INPLACE_FINAL_STATE:

139 + p_ht = ht + tl.load(ssm_state_indices + i_n * stride_indices_seq +

140 + i_t).to(tl.int64) * stride_final_state_token

141 + else:

142 + p_ht = ht + (bos + i_t) * stride_final_state_token

143 + p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :]

144 + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)

145 +

146 + p_q += H * K

147 + p_k += H * K

148 + p_o += HV * V

149 + p_v += HV * V

150 + p_g += HV

151 + p_beta += HV * (V if IS_BETA_HEADWISE else 1)

152 +

153 +

154 +def fused_recurrent_gated_delta_rule_fwd(

155 + q: torch.Tensor,

156 + k: torch.Tensor,

157 + v: torch.Tensor,

158 + g: torch.Tensor,

159 + beta: torch.Tensor,

160 + scale: float,

161 + initial_state: torch.Tensor,

162 + inplace_final_state: bool = True,

163 + cu_seqlens: Optional[torch.LongTensor] = None,

164 + ssm_state_indices: Optional[torch.Tensor] = None,

165 + num_accepted_tokens: Optional[torch.Tensor] = None,

166 + use_qk_l2norm_in_kernel: bool = False,

167 +) -> tuple[torch.Tensor, torch.Tensor]:

168 + B, T, H, K, V = *k.shape, v.shape[-1]

169 + HV = v.shape[2]

170 + N = B if cu_seqlens is None else len(cu_seqlens) - 1

171 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8)

172 + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)

173 + assert NK == 1, "NK > 1 is not supported yet"

174 + num_stages = 3

175 + num_warps = 1

176 +

177 + o = q.new_empty(NK, *v.shape)

178 + if inplace_final_state:

179 + final_state = initial_state

180 + else:

181 + final_state = q.new_empty(T, HV, K, V, dtype=initial_state.dtype)

182 +

183 + stride_init_state_token = initial_state.stride(0)

184 + stride_final_state_token = final_state.stride(0)

185 +

186 + if ssm_state_indices is None:

187 + stride_indices_seq, stride_indices_tok = 1, 1

188 + elif ssm_state_indices.ndim == 1:

189 + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1

190 + else:

191 + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()

192 +

193 + grid = (NK, NV, N * HV)

194 + fused_recurrent_gated_delta_rule_fwd_kernel[grid](

195 + q=q,

196 + k=k,

197 + v=v,

198 + g=g,

199 + beta=beta,

200 + o=o,

201 + h0=initial_state,

202 + ht=final_state,

203 + cu_seqlens=cu_seqlens,

204 + ssm_state_indices=ssm_state_indices,

205 + num_accepted_tokens=num_accepted_tokens,

206 + scale=scale,

207 + N=N,

208 + T=T,

209 + B=B,

210 + H=H,

211 + HV=HV,

212 + K=K,

213 + V=V,

214 + BK=BK,

215 + BV=BV,

216 + stride_init_state_token=stride_init_state_token,

217 + stride_final_state_token=stride_final_state_token,

218 + stride_indices_seq=stride_indices_seq,

219 + stride_indices_tok=stride_indices_tok,

220 + IS_BETA_HEADWISE=beta.ndim == v.ndim,

221 + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,

222 + INPLACE_FINAL_STATE=inplace_final_state,

223 + num_warps=num_warps,

224 + num_stages=num_stages,

225 + )

226 + o = o.squeeze(0)

227 + return o, final_state

228 +

229 +

230 +class FusedRecurrentFunction(torch.autograd.Function):

231 +

232 + @staticmethod

233 + def forward(ctx,

234 + q: torch.Tensor,

235 + k: torch.Tensor,

236 + v: torch.Tensor,

237 + g: torch.Tensor,

238 + beta: torch.Tensor,

239 + scale: float,

240 + initial_state: torch.Tensor,

241 + inplace_final_state: bool = True,

242 + cu_seqlens: Optional[torch.LongTensor] = None,

243 + ssm_state_indices: Optional[torch.Tensor] = None,

244 + num_accepted_tokens: Optional[torch.Tensor] = None,

245 + use_qk_l2norm_in_kernel: bool = False):

246 + o, final_state = fused_recurrent_gated_delta_rule_fwd(

247 + q=q.contiguous(),

248 + k=k.contiguous(),

249 + v=v.contiguous(),

250 + g=g.contiguous(),

251 + beta=beta.contiguous(),

252 + scale=scale,

253 + initial_state=initial_state,

254 + inplace_final_state=inplace_final_state,

255 + cu_seqlens=cu_seqlens,

256 + ssm_state_indices=ssm_state_indices,

257 + num_accepted_tokens=num_accepted_tokens,

258 + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,

259 + )

260 +

261 + return o, final_state

262 +

263 +

264 +def fused_recurrent_gated_delta_rule(

265 + q: torch.Tensor,

266 + k: torch.Tensor,

267 + v: torch.Tensor,

268 + g: torch.Tensor,

269 + beta: torch.Tensor = None,

270 + scale: float = None,

271 + initial_state: torch.Tensor = None,

272 + inplace_final_state: bool = True,

273 + cu_seqlens: Optional[torch.LongTensor] = None,

274 + ssm_state_indices: Optional[torch.Tensor] = None,

275 + num_accepted_tokens: Optional[torch.Tensor] = None,

276 + use_qk_l2norm_in_kernel: bool = False,

277 +) -> tuple[torch.Tensor, torch.Tensor]:

278 + r"""

279 + Args:

280 + q (torch.Tensor):

281 + queries of shape `[B, T, H, K]`.

282 + k (torch.Tensor):

283 + keys of shape `[B, T, H, K]`.

284 + v (torch.Tensor):

285 + values of shape `[B, T, HV, V]`.

286 + GVA is applied if `HV > H`.

287 + g (torch.Tensor):

288 + g (decays) of shape `[B, T, HV]`.

289 + beta (torch.Tensor):

290 + betas of shape `[B, T, HV]`.

291 + scale (Optional[int]):

292 + Scale factor for the RetNet attention scores.

293 + If not provided, it will default to `1 / sqrt(K)`. Default: `None`.

294 + initial_state (Optional[torch.Tensor]):

295 + Initial state of shape `[N, HV, K, V]` for `N` input sequences.

296 + For equal-length input sequences, `N` equals the batch size `B`.

297 + Default: `None`.

298 + inplace_final_state: bool:

299 + Whether to store the final state in-place to save memory.

300 + Default: `True`.

301 + cu_seqlens (torch.LongTensor):

302 + Cumulative sequence lengths of shape `[N+1]` used for variable-length training,

303 + consistent with the FlashAttention API.

304 + ssm_state_indices (Optional[torch.Tensor]):

305 + Indices to map the input sequences to the initial/final states.

306 + num_accepted_tokens (Optional[torch.Tensor]):

307 + Number of accepted tokens for each sequence during decoding.

308 +

309 + Returns:

310 + o (torch.Tensor):

311 + Outputs of shape `[B, T, HV, V]`.

312 + final_state (torch.Tensor):

313 + Final state of shape `[N, HV, K, V]`.

314 +

315 + Examples::

316 + >>> import torch

317 + >>> import torch.nn.functional as F

318 + >>> from einops import rearrange

319 + >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule

320 + # inputs with equal lengths

321 + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512

322 + >>> q = torch.randn(B, T, H, K, device='cuda')

323 + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)

324 + >>> v = torch.randn(B, T, HV, V, device='cuda')

325 + >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))

326 + >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()

327 + >>> h0 = torch.randn(B, HV, K, V, device='cuda')

328 + >>> o, ht = fused_gated_recurrent_delta_rule(

329 + q, k, v, g, beta,

330 + initial_state=h0,

331 + )

332 + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required

333 + >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))

334 + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected

335 + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)

336 + >>> o_var, ht_var = fused_gated_recurrent_delta_rule(

337 + q, k, v, g, beta,

338 + initial_state=h0,

339 + cu_seqlens=cu_seqlens

340 + )

341 + """

342 + if cu_seqlens is not None and q.shape[0] != 1:

343 + raise ValueError(

344 + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."

345 + f"Please flatten variable-length inputs before processing.")

346 + if scale is None:

347 + scale = k.shape[-1]**-0.5

348 + else:

349 + assert scale > 0, "scale must be positive"

350 + if beta is None:

351 + beta = torch.ones_like(q[..., 0])

352 + o, final_state = FusedRecurrentFunction.apply(

353 + q,

354 + k,

355 + v,

356 + g,

357 + beta,

358 + scale,

359 + initial_state,

360 + inplace_final_state,

361 + cu_seqlens,

362 + ssm_state_indices,

363 + num_accepted_tokens,

364 + use_qk_l2norm_in_kernel,

365 + )

366 + return o, final_state

vllm/model_executor/layers/fla/ops/index.py new

+39

@@ -0,0 +1,39 @@

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:

9 +# ruff: noqa: E501

10 +import torch

11 +

12 +from vllm.triton_utils import triton

13 +

14 +from .utils import tensor_cache

15 +

16 +

17 +@tensor_cache

18 +def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor:

19 + return cu_seqlens[1:] - cu_seqlens[:-1]

20 +

21 +

22 +@tensor_cache

23 +def prepare_chunk_indices(cu_seqlens: torch.LongTensor,

24 + chunk_size: int) -> torch.LongTensor:

25 + indices = torch.cat([

26 + torch.arange(n)

27 + for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()

28 + ])

29 + return torch.stack([indices.eq(0).cumsum(0) - 1, indices],

30 + 1).to(cu_seqlens)

31 +

32 +

33 +@tensor_cache

34 +def prepare_chunk_offsets(cu_seqlens: torch.LongTensor,

35 + chunk_size: int) -> torch.LongTensor:

36 + return torch.cat([

37 + cu_seqlens.new_tensor([0]),

38 + triton.cdiv(prepare_lens(cu_seqlens), chunk_size)

39 + ]).cumsum(-1)

vllm/model_executor/layers/fla/ops/l2norm.py new

+143

@@ -0,0 +1,143 @@

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:

9 +

10 +import os

11 +from typing import Optional

12 +

13 +import torch

14 +

15 +from vllm.triton_utils import tl, triton

16 +

17 +BT_LIST = [8, 16, 32, 64, 128]

18 +

19 +USE_DEFAULT_FLA_NORM = int(os.getenv("USE_DEFAULT_FLA_NORM", "0"))

20 +

21 +

22 +@triton.autotune(configs=[

23 + triton.Config({}, num_warps=num_warps)

24 + for num_warps in [1, 2, 4, 8, 16, 32]

25 +],

26 + key=['D'])

27 +@triton.jit

28 +def l2norm_fwd_kernel1(

29 + x,

30 + y,

31 + D,

32 + BD: tl.constexpr,

33 + eps,

34 +):

35 + i_t = tl.program_id(0)

36 + x += i_t * D

37 + y += i_t * D

38 + # Compute mean and variance

39 + cols = tl.arange(0, BD)

40 + mask = cols < D

41 + b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32)

42 + b_var = tl.sum(b_x * b_x, axis=0)

43 + b_rstd = 1 / tl.sqrt(b_var + eps)

44 + # tl.store(Rstd + i_t, rstd)

45 + # Normalize and apply linear transformation

46 + b_y = b_x * b_rstd

47 + tl.store(y + cols, b_y, mask=mask)

48 +

49 +

50 +@triton.autotune(configs=[

51 + triton.Config({'BT': BT}, num_warps=num_warps)

52 + for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST

53 +],

54 + key=['D'])

55 +@triton.jit(do_not_specialize=["NB"])

56 +def l2norm_fwd_kernel(

57 + x,

58 + y,

59 + eps,

60 + NB,

61 + T,

62 + D: tl.constexpr,

63 + BT: tl.constexpr,

64 + BD: tl.constexpr,

65 +):

66 + i_t = tl.program_id(0)

67 + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))

68 + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)

69 + b_var = tl.sum(b_x * b_x, axis=1)

70 + b_y = b_x / tl.sqrt(b_var + eps)[:, None]

71 + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))

72 + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))

73 +

74 +

75 +@triton.jit

76 +def l2norm_fwd_kernel2(X, Y, eps, M, N: tl.constexpr, MBLOCK: tl.constexpr):

77 + xoffset = tl.program_id(0) * MBLOCK

78 + row_idx = xoffset + tl.arange(0, MBLOCK)[:, None]

79 + xmask = row_idx < M

80 + rindex = tl.arange(0, N)[None, :]

81 + xs = tl.load(X + (rindex + N * row_idx), None).to(tl.float32)

82 + square = tl.broadcast_to(xs * xs, [MBLOCK, N])

83 + square_sum = tl.sum(tl.where(xmask, square, 0), 1)[:, None]

84 + rsqrt = tl.rsqrt(square_sum + eps)

85 + tl.store(Y + (rindex + N * row_idx), xs * rsqrt, xmask)

86 +

87 +

88 +def l2norm_fwd(x: torch.Tensor,

89 + eps: float = 1e-6,

90 + output_dtype: Optional[torch.dtype] = None):

91 + x_shape_og = x.shape

92 + x = x.view(-1, x.shape[-1])

93 + # allocate output

94 + if output_dtype is None:

95 + y = torch.empty_like(x)

96 + else:

97 + y = torch.empty_like(x, dtype=output_dtype)

98 + assert y.stride(-1) == 1

99 + T, D = x.shape[0], x.shape[-1]

100 + # rstd = torch.empty((T,), dtype=torch.float32, device=x.device)

101 + # Less than 64KB per feature: enqueue fused kernel

102 + MAX_FUSED_SIZE = 65536 // x.element_size()

103 + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))

104 + if D > BD:

105 + raise RuntimeError("This layer doesn't support feature dim >= 64KB.")

106 +

107 + if not USE_DEFAULT_FLA_NORM:

108 + MBLOCK = 32

109 + # M, N = x.shape

110 + l2norm_fwd_kernel2[(triton.cdiv(T, MBLOCK), )](

111 + x,

112 + y,

113 + eps,

114 + T,

115 + D,

116 + MBLOCK,

117 + )

118 + else:

119 + if D <= 512:

120 + NB = triton.cdiv(T, 2048)

121 +

122 + def grid(meta):

123 + return (triton.cdiv(T, meta['BT']), )

124 +

125 + l2norm_fwd_kernel[grid](

126 + x,

127 + y,

128 + eps,

129 + NB=NB,

130 + T=T,

131 + D=D,

132 + BD=BD,

133 + )

134 + else:

135 + l2norm_fwd_kernel1[(T, )](

136 + x,

137 + y,

138 + eps=eps,

139 + D=D,

140 + BD=BD,

141 + )

142 +

143 + return y.view(x_shape_og)

vllm/model_executor/layers/fla/ops/layernorm_guard.py new

+337

@@ -0,0 +1,337 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +# SPDX-FileCopyrightText: Tri Dao

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:

9 +

10 +# ruff: noqa: E501

11 +# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html

12 +# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate.

13 +# This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling.

14 +# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine.

15 +

16 +from typing import Optional

17 +

18 +import torch

19 +import torch.nn as nn

20 +import torch.nn.functional as F

21 +from einops import rearrange

22 +

23 +from vllm.triton_utils import tl, triton

24 +

25 +from .utils import input_guard

26 +

27 +

28 +def rms_norm_ref(x,

29 + weight,

30 + bias,

31 + z=None,

32 + eps=1e-6,

33 + group_size=None,

34 + norm_before_gate=True,

35 + upcast=True):

36 + dtype = x.dtype

37 + weight = weight.float()

38 + bias = bias.float() if bias is not None else None

39 + if upcast:

40 + x = x.float()

41 + z = z.float() if z is not None else z

42 + if z is not None and not norm_before_gate:

43 + x = x * F.silu(z)

44 + if group_size is None:

45 + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps)

46 + out = (x * rstd * weight) + bias if bias is not None else (x * rstd *

47 + weight)

48 + else:

49 + x_group = rearrange(x, "... (g d) -> ... g d", d=group_size)

50 + rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) +

51 + eps)

52 + out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight

53 + if bias is not None:

54 + out = out + bias

55 + if z is not None and norm_before_gate:

56 + out *= F.silu(z)

57 + return out.to(dtype)

58 +

59 +

60 +@triton.heuristics({

61 + "HAS_BIAS": lambda args: args["B"] is not None,

62 + "HAS_Z": lambda args: args["Z"] is not None,

63 +})

64 +@triton.jit

65 +def layer_norm_fwd_kernel(

66 + X, # pointer to the input

67 + Y, # pointer to the output

68 + W, # pointer to the weights

69 + B, # pointer to the biases

70 + Z, # pointer to the other branch

71 + Mean, # pointer to the mean

72 + Rstd, # pointer to the 1/std

73 + stride_x_row, # how much to increase the pointer when moving by 1 row

74 + stride_y_row,

75 + stride_z_row,

76 + M, # number of rows in X

77 + N, # number of columns in X

78 + eps, # epsilon to avoid division by zero

79 + BLOCK_N: tl.constexpr,

80 + HAS_BIAS: tl.constexpr,

81 + HAS_Z: tl.constexpr,

82 + NORM_BEFORE_GATE: tl.constexpr,

83 + IS_RMS_NORM: tl.constexpr,

84 +):

85 + # Map the program id to the row of X and Y it should compute.

86 + row = tl.program_id(0)

87 + group = tl.program_id(1)

88 + X += row * stride_x_row + group * N

89 + Y += row * stride_y_row + group * N

90 + if HAS_Z:

91 + Z += row * stride_z_row + group * N

92 + if not IS_RMS_NORM:

93 + Mean += group * M

94 + Rstd += group * M

95 + W += group * N

96 + if HAS_BIAS:

97 + B += group * N

98 + # Compute mean and variance

99 + cols = tl.arange(0, BLOCK_N)

100 + x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32)

101 + if HAS_Z and not NORM_BEFORE_GATE:

102 + z = tl.load(Z + cols, mask=cols < N).to(tl.float32)

103 + x *= z * tl.sigmoid(z)

104 + if not IS_RMS_NORM:

105 + mean = tl.sum(x, axis=0) / N

106 + tl.store(Mean + row, mean)

107 + xbar = tl.where(cols < N, x - mean, 0.)

108 + var = tl.sum(xbar * xbar, axis=0) / N

109 + else:

110 + xbar = tl.where(cols < N, x, 0.)

111 + var = tl.sum(xbar * xbar, axis=0) / N

112 + rstd = 1 / tl.sqrt(var + eps)

113 + tl.store(Rstd + row, rstd)

114 + # Normalize and apply linear transformation

115 + mask = cols < N

116 + w = tl.load(W + cols, mask=mask).to(tl.float32)

117 + if HAS_BIAS:

118 + b = tl.load(B + cols, mask=mask).to(tl.float32)

119 + x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd

120 + y = x_hat * w + b if HAS_BIAS else x_hat * w

121 + if HAS_Z and NORM_BEFORE_GATE:

122 + z = tl.load(Z + cols, mask=mask).to(tl.float32)

123 + y *= z * tl.sigmoid(z)

124 + # Write output

125 + tl.store(Y + cols, y, mask=mask)

126 +

127 +

128 +def layer_norm_fwd(

129 + x: torch.Tensor,

130 + weight: torch.Tensor,

131 + bias: torch.Tensor,

132 + eps: float,

133 + z: torch.Tensor = None,

134 + out: torch.Tensor = None,

135 + group_size: int = None,

136 + norm_before_gate: bool = True,

137 + is_rms_norm: bool = False,

138 +):

139 + M, N = x.shape

140 + if group_size is None:

141 + group_size = N

142 + assert N % group_size == 0

143 + ngroups = N // group_size

144 + assert x.stride(-1) == 1

145 + if z is not None:

146 + assert z.stride(-1) == 1

147 + assert z.shape == (M, N)

148 + assert weight.shape == (N, )

149 + assert weight.stride(-1) == 1

150 + if bias is not None:

151 + assert bias.stride(-1) == 1

152 + assert bias.shape == (N, )

153 + # allocate output

154 + if out is not None:

155 + assert out.shape == x.shape

156 + else:

157 + out = torch.empty_like(x)

158 + assert out.stride(-1) == 1

159 + mean = torch.empty((ngroups * M, ), dtype=torch.float32,

160 + device=x.device) if not is_rms_norm else None

161 + rstd = torch.empty((ngroups * M, ), dtype=torch.float32, device=x.device)

162 + # Less than 64KB per feature: enqueue fused kernel

163 + MAX_FUSED_SIZE = 65536 // x.element_size()

164 + BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size))

165 + if group_size > BLOCK_N:

166 + raise RuntimeError(

167 + "This layer norm doesn't support feature dim >= 64KB.")

168 + # heuristics for number of warps

169 + num_warps = min(max(BLOCK_N // 256, 1), 8)

170 + grid = (M, ngroups)

171 + layer_norm_fwd_kernel[grid](x,

172 + out,

173 + weight,

174 + bias,

175 + z,

176 + mean,

177 + rstd,

178 + x.stride(0),

179 + out.stride(0),

180 + z.stride(0) if z is not None else 0,

181 + M,

182 + group_size,

183 + eps,

184 + BLOCK_N=BLOCK_N,

185 + NORM_BEFORE_GATE=norm_before_gate,

186 + IS_RMS_NORM=is_rms_norm,

187 + num_warps=num_warps)

188 + return out, mean, rstd

189 +

190 +

191 +class LayerNormFn(torch.autograd.Function):

192 +

193 + @input_guard

194 + @staticmethod

195 + def forward(ctx,

196 + x,

197 + weight,

198 + bias,

199 + z=None,

200 + eps=1e-6,

201 + group_size=None,

202 + norm_before_gate=True,

203 + is_rms_norm=False):

204 + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))

205 + """

206 +

207 + x_shape_og = x.shape

208 + # reshape input data into 2D tensor

209 + x = x.reshape(-1, x.shape[-1])

210 + if x.stride(-1) != 1:

211 + x = x.contiguous()

212 + if z is not None:

213 + assert z.shape == x_shape_og

214 + z = z.reshape(-1, z.shape[-1])

215 + if z.stride(-1) != 1:

216 + z = z.contiguous()

217 + weight = weight.contiguous()

218 + if bias is not None:

219 + bias = bias.contiguous()

220 + y, mean, rstd = layer_norm_fwd(

221 + x,

222 + weight,

223 + bias,

224 + eps,

225 + z=z,

226 + group_size=group_size,

227 + norm_before_gate=norm_before_gate,

228 + is_rms_norm=is_rms_norm,

229 + )

230 + ctx.save_for_backward(x, weight, bias, mean, rstd, z)

231 + ctx.x_shape_og = x_shape_og

232 + ctx.eps = eps

233 + ctx.group_size = group_size

234 + ctx.norm_before_gate = norm_before_gate

235 + ctx.is_rms_norm = is_rms_norm

236 + return y.reshape(x_shape_og)

237 +

238 +

239 +def layernorm_fn(x,

240 + weight,

241 + bias,

242 + z=None,

243 + eps=1e-6,

244 + group_size=None,

245 + norm_before_gate=True,

246 + is_rms_norm=False):

247 + return LayerNormFn.apply(x, weight, bias, z, eps, group_size,

248 + norm_before_gate, is_rms_norm)

249 +

250 +

251 +def rmsnorm_fn(x,

252 + weight,

253 + bias,

254 + z=None,

255 + eps=1e-6,

256 + group_size=None,

257 + norm_before_gate=True):

258 + return LayerNormFn.apply(x, weight, bias, z, eps, group_size,

259 + norm_before_gate, True)

260 +

261 +

262 +class LayerNormGated(nn.Module):

263 +

264 + def __init__(

265 + self,

266 + hidden_size,

267 + eps: float = 1e-5,

268 + group_size: Optional[int] = None,

269 + norm_before_gate: bool = True,

270 + device: Optional[torch.device] = None,

271 + dtype: Optional[torch.dtype] = None,

272 + ):

273 + """If group_size is not None, we do GroupNorm with each group having group_size elements.

274 + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group).

275 + """

276 +

277 + factory_kwargs = {"device": device, "dtype": dtype}

278 + super().__init__()

279 + self.eps = eps

280 + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))

281 + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))

282 + self.group_size = group_size

283 + self.norm_before_gate = norm_before_gate

284 + self.reset_parameters()

285 +

286 + def reset_parameters(self):

287 + torch.nn.init.ones_(self.weight)

288 + torch.nn.init.zeros_(self.bias)

289 +

290 + def forward(self, x, z=None):

291 + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))

292 + """

293 + return layernorm_fn(x,

294 + self.weight,

295 + self.bias,

296 + z=z,

297 + group_size=self.group_size,

298 + eps=self.eps,

299 + norm_before_gate=self.norm_before_gate)

300 +

301 +

302 +class RMSNormGated(nn.Module):

303 +

304 + def __init__(

305 + self,

306 + hidden_size,

307 + eps: float = 1e-5,

308 + group_size: Optional[int] = None,

309 + norm_before_gate: bool = False,

310 + device: Optional[torch.device] = None,

311 + dtype: Optional[torch.dtype] = None,

312 + ):

313 + """If group_size is not None, we do GroupNorm with each group having group_size elements.

314 + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group).

315 + """

316 + factory_kwargs = {"device": device, "dtype": dtype}

317 + super().__init__()

318 + self.eps = eps

319 + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))

320 + self.register_parameter("bias", None)

321 + self.group_size = group_size

322 + self.norm_before_gate = norm_before_gate

323 + self.reset_parameters()

324 +

325 + def reset_parameters(self):

326 + torch.nn.init.ones_(self.weight)

327 +

328 + def forward(self, x, z=None):

329 + """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))

330 + """

331 + return rmsnorm_fn(x,

332 + self.weight,

333 + self.bias,

334 + z=z,

335 + eps=self.eps,

336 + group_size=self.group_size,

337 + norm_before_gate=self.norm_before_gate)

vllm/model_executor/layers/fla/ops/op.py new

+44

@@ -0,0 +1,44 @@

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:

9 +

10 +import os

11 +

12 +from vllm.triton_utils import tl, tldevice, triton

13 +

14 +if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':

15 + div = tldevice.fast_dividef

16 + exp = tldevice.fast_expf

17 + log = tldevice.fast_logf

18 + log2 = tldevice.fast_log2f

19 +else:

20 +

21 + @triton.jit

22 + def div_normal(x, y):

23 + return x / y

24 +

25 + div = div_normal

26 + exp = tl.exp

27 + log = tl.log

28 + log2 = tl.log2

29 +

30 +

31 +@triton.jit

32 +def safe_exp(x):

33 + return exp(tl.where(x <= 0, x, float('-inf')))

34 +

35 +

36 +if not hasattr(tl, 'gather'):

37 +

38 + @triton.jit

39 + def gather(src, index, axis, _builder=None):

40 + # This is a fallback implementation when tl.gather is not supported

41 + # In order to pass triton compiler, there is no actual gather operation

42 + return src

43 +else:

44 + gather = tl.gather

vllm/model_executor/layers/fla/ops/solve_tril.py new

+365

@@ -0,0 +1,365 @@

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:

9 +# ruff: noqa: E501

10 +from typing import Optional

11 +

12 +import torch

13 +

14 +from vllm.triton_utils import tl, triton

15 +

16 +from .index import prepare_chunk_indices

17 +from .utils import input_guard

18 +

19 +

20 +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None})

21 +@triton.autotune(

22 + configs=[

23 + triton.Config({}, num_warps=num_warps, num_stages=num_stages)

24 + for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4, 5]

25 + ],

26 + key=['BT'],

27 +)

28 +@triton.jit(do_not_specialize=['T'])

29 +def solve_tril_16x16_kernel(

30 + A,

31 + Ad,

32 + cu_seqlens,

33 + chunk_indices,

34 + T,

35 + H: tl.constexpr,

36 + BT: tl.constexpr,

37 + IS_VARLEN: tl.constexpr,

38 +):

39 + i_t, i_bh = tl.program_id(0), tl.program_id(1)

40 + i_b, i_h = i_bh // H, i_bh % H

41 + if IS_VARLEN:

42 + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(

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

44 + bos, eos = tl.load(cu_seqlens + i_n).to(

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

46 + T = eos - bos

47 + else:

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

49 +

50 + A = A + (bos * H + i_h) * BT

51 + Ad = Ad + (bos * H + i_h) * 16

52 +

53 + offset = (i_t * 16) % BT

54 + p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * 16, offset),

55 + (16, 16), (1, 0))

56 + p_Ai = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16),

57 + (1, 0))

58 + b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)

59 + b_A = -tl.where(

60 + tl.arange(0, 16)[:, None] > tl.arange(0, 16)[None, :], b_A, 0)

61 +

62 + o_i = tl.arange(0, 16)

63 + for i in range(1, min(16, T - i_t * 16)):

64 + b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset)

65 + b_a = b_a + tl.sum(b_a[:, None] * b_A, 0)

66 + mask = o_i == i

67 + b_A = tl.where(mask[:, None], b_a, b_A)

68 + b_A += o_i[:, None] == o_i[None, :]

69 + tl.store(p_Ai,

70 + b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"),

71 + boundary_check=(0, 1))

72 +

73 +

74 +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None})

75 +@triton.autotune(

76 + configs=[

77 + triton.Config({}, num_warps=num_warps, num_stages=num_stages)

78 + for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4, 5]

79 + ],

80 + key=['H', 'BT', 'IS_VARLEN'],

81 +)

82 +@triton.jit(do_not_specialize=['T'])

83 +def merge_16x16_to_32x32_inverse_kernel(A, Ad, Ai, cu_seqlens, chunk_indices,

84 + T, H: tl.constexpr, BT: tl.constexpr,

85 + IS_VARLEN: tl.constexpr):

86 + i_t, i_bh = tl.program_id(0), tl.program_id(1)

87 + i_b, i_h = i_bh // H, i_bh % H

88 + if IS_VARLEN:

89 + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(

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

91 + bos, eos = tl.load(cu_seqlens + i_n).to(

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

93 + T = eos - bos

94 + else:

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

96 +

97 + A += (bos * H + i_h) * 32

98 + Ad += (bos * H + i_h) * 16

99 + Ai += (bos * H + i_h) * 32

100 +

101 + p_A_21 = tl.make_block_ptr(A, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0),

102 + (16, 16), (1, 0))

103 + p_Ad_11 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 32, 0),

104 + (16, 16), (1, 0))

105 + p_Ad_22 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 32 + 16, 0),

106 + (16, 16), (1, 0))

107 + p_Ai_11 = tl.make_block_ptr(Ai, (T, 32), (H * 32, 1), (i_t * 32, 0),

108 + (16, 16), (1, 0))

109 + p_Ai_22 = tl.make_block_ptr(Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 16),

110 + (16, 16), (1, 0))

111 + p_Ai_21 = tl.make_block_ptr(Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0),

112 + (16, 16), (1, 0))

113 +

114 + A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)

115 + Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32)

116 + Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32)

117 + Ai_21 = -tl.dot(tl.dot(Ai_22, A_21, input_precision='ieee'),

118 + Ai_11,

119 + input_precision='ieee')

120 + tl.store(p_Ai_11,

121 + Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),

122 + boundary_check=(0, 1))

123 + tl.store(p_Ai_22,

124 + Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),

125 + boundary_check=(0, 1))

126 + tl.store(p_Ai_21,

127 + Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),

128 + boundary_check=(0, 1))

129 +

130 +

131 +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None})

132 +@triton.autotune(

133 + configs=[

134 + triton.Config({}, num_warps=num_warps, num_stages=num_stages)

135 + for num_warps in [2, 4, 8] for num_stages in [2, 3, 4, 5]

136 + ],

137 + key=['H', 'BT', 'IS_VARLEN'],

138 +)

139 +@triton.jit(do_not_specialize=['T'])

140 +def merge_16x16_to_64x64_inverse_kernel(A, Ad, Ai, cu_seqlens, chunk_indices,

141 + T, H: tl.constexpr, BT: tl.constexpr,

142 + IS_VARLEN: tl.constexpr):

143 + i_t, i_bh = tl.program_id(0), tl.program_id(1)

144 + i_b, i_h = i_bh // H, i_bh % H

145 + if IS_VARLEN:

146 + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(

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

148 + bos, eos = tl.load(cu_seqlens + i_n).to(

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

150 + T = eos - bos

151 + else:

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

153 +

154 + A += (bos * H + i_h) * 64

155 + Ad += (bos * H + i_h) * 16

156 + Ai += (bos * H + i_h) * 64

157 +

158 + p_A_21 = tl.make_block_ptr(A, (T, 64), (H * 64, 1), (i_t * 64 + 16, 0),

159 + (16, 16), (1, 0))

160 + p_A_32 = tl.make_block_ptr(A, (T, 64), (H * 64, 1), (i_t * 64 + 32, 16),

161 + (16, 16), (1, 0))

162 + p_A_31 = tl.make_block_ptr(A, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0),

163 + (16, 16), (1, 0))

164 + p_A_43 = tl.make_block_ptr(A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 32),

165 + (16, 16), (1, 0))

166 + p_A_42 = tl.make_block_ptr(A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 16),

167 + (16, 16), (1, 0))

168 + p_A_41 = tl.make_block_ptr(A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 0),

169 + (16, 16), (1, 0))

170 + p_Ad_11 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 64, 0),

171 + (16, 16), (1, 0))

172 + p_Ad_22 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 64 + 16, 0),

173 + (16, 16), (1, 0))

174 + p_Ad_33 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 64 + 32, 0),

175 + (16, 16), (1, 0))

176 + p_Ad_44 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 64 + 48, 0),

177 + (16, 16), (1, 0))

178 +

179 + A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)

180 + A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32)

181 + A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32)

182 + A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32)

183 + A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32)

184 + A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32)

185 +

186 + Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32)

187 + Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32)

188 + Ai_33 = tl.load(p_Ad_33, boundary_check=(0, 1)).to(tl.float32)

189 + Ai_44 = tl.load(p_Ad_44, boundary_check=(0, 1)).to(tl.float32)

190 +

191 + Ai_21 = -tl.dot(tl.dot(Ai_22, A_21, input_precision='ieee'),

192 + Ai_11,

193 + input_precision='ieee')

194 + Ai_32 = -tl.dot(tl.dot(Ai_33, A_32, input_precision='ieee'),

195 + Ai_22,

196 + input_precision='ieee')

197 + Ai_43 = -tl.dot(tl.dot(Ai_44, A_43, input_precision='ieee'),

198 + Ai_33,

199 + input_precision='ieee')

200 +

201 + Ai_31 = -tl.dot(Ai_33,

202 + tl.dot(A_31, Ai_11, input_precision='ieee') +

203 + tl.dot(A_32, Ai_21, input_precision='ieee'),

204 + input_precision='ieee')

205 + Ai_42 = -tl.dot(Ai_44,

206 + tl.dot(A_42, Ai_22, input_precision='ieee') +

207 + tl.dot(A_43, Ai_32, input_precision='ieee'),

208 + input_precision='ieee')

209 + Ai_41 = -tl.dot(Ai_44,

210 + tl.dot(A_41, Ai_11, input_precision='ieee') +

211 + tl.dot(A_42, Ai_21, input_precision='ieee') +

212 + tl.dot(A_43, Ai_31, input_precision='ieee'),

213 + input_precision='ieee')

214 +

215 + p_Ai_11 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64, 0),

216 + (16, 16), (1, 0))

217 + p_Ai_22 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 16),

218 + (16, 16), (1, 0))

219 + p_Ai_33 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 32),

220 + (16, 16), (1, 0))

221 + p_Ai_44 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 48),

222 + (16, 16), (1, 0))

223 + p_Ai_21 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 0),

224 + (16, 16), (1, 0))

225 + p_Ai_31 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0),

226 + (16, 16), (1, 0))

227 + p_Ai_32 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 16),

228 + (16, 16), (1, 0))

229 + p_Ai_41 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 0),

230 + (16, 16), (1, 0))

231 + p_Ai_42 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 16),

232 + (16, 16), (1, 0))

233 + p_Ai_43 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 32),

234 + (16, 16), (1, 0))

235 + tl.store(p_Ai_11,

236 + Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),

237 + boundary_check=(0, 1))

238 + tl.store(p_Ai_22,

239 + Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),

240 + boundary_check=(0, 1))

241 + tl.store(p_Ai_33,

242 + Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"),

243 + boundary_check=(0, 1))

244 + tl.store(p_Ai_44,

245 + Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"),

246 + boundary_check=(0, 1))

247 + tl.store(p_Ai_21,

248 + Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),

249 + boundary_check=(0, 1))

250 + tl.store(p_Ai_31,

251 + Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"),

252 + boundary_check=(0, 1))

253 + tl.store(p_Ai_32,

254 + Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"),

255 + boundary_check=(0, 1))

256 + tl.store(p_Ai_41,

257 + Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"),

258 + boundary_check=(0, 1))

259 + tl.store(p_Ai_42,

260 + Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"),

261 + boundary_check=(0, 1))

262 + tl.store(p_Ai_43,

263 + Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"),

264 + boundary_check=(0, 1))

265 +

266 + fill_zeros = tl.zeros((16, 16), dtype=tl.float32)

267 + p_Ai_12 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64, 16),

268 + (16, 16), (1, 0))

269 + p_Ai_13 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64, 32),

270 + (16, 16), (1, 0))

271 + p_Ai_14 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64, 48),

272 + (16, 16), (1, 0))

273 + p_Ai_23 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 32),

274 + (16, 16), (1, 0))

275 + p_Ai_24 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 48),

276 + (16, 16), (1, 0))

277 + p_Ai_34 = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 48),

278 + (16, 16), (1, 0))

279 + tl.store(p_Ai_12,

280 + fill_zeros.to(p_Ai_12.dtype.element_ty,

281 + fp_downcast_rounding="rtne"),

282 + boundary_check=(0, 1))

283 + tl.store(p_Ai_13,

284 + fill_zeros.to(p_Ai_13.dtype.element_ty,

285 + fp_downcast_rounding="rtne"),

286 + boundary_check=(0, 1))

287 + tl.store(p_Ai_14,

288 + fill_zeros.to(p_Ai_14.dtype.element_ty,

289 + fp_downcast_rounding="rtne"),

290 + boundary_check=(0, 1))

291 + tl.store(p_Ai_23,

292 + fill_zeros.to(p_Ai_23.dtype.element_ty,

293 + fp_downcast_rounding="rtne"),

294 + boundary_check=(0, 1))

295 + tl.store(p_Ai_24,

296 + fill_zeros.to(p_Ai_24.dtype.element_ty,

297 + fp_downcast_rounding="rtne"),

298 + boundary_check=(0, 1))

299 + tl.store(p_Ai_34,

300 + fill_zeros.to(p_Ai_34.dtype.element_ty,

301 + fp_downcast_rounding="rtne"),

302 + boundary_check=(0, 1))

303 +

304 +

305 +@input_guard

306 +def solve_tril(A: torch.Tensor,

307 + cu_seqlens: Optional[torch.Tensor] = None,

308 + output_dtype: torch.dtype = torch.float) -> torch.Tensor:

309 + """

310 + Compute the inverse of the lower triangular matrix

311 + A should be strictly lower triangular, i.e., A.triu() == 0.

312 +

313 + Args:

314 + A (torch.Tensor):

315 + [B, T, H, K]

316 + cu_seqlens (torch.Tensor):

317 + The cumulative sequence lengths of the input tensor.

318 + Default: None.

319 + output_dtype (torch.dtype):

320 + The dtype of the output tensor. Default: `torch.float`

321 +

322 + Returns:

323 + (I + A)^-1 with the same shape as A

324 + """

325 + assert A.shape[-1] in [16, 32, 64]

326 +

327 + B, T, H, BT = A.shape

328 + Ad = torch.empty(B,

329 + T,

330 + H,

331 + 16,

332 + device=A.device,

333 + dtype=torch.float if BT != 16 else output_dtype)

334 +

335 + chunk_indices = prepare_chunk_indices(

336 + cu_seqlens, 16) if cu_seqlens is not None else None

337 + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, 16)

338 + solve_tril_16x16_kernel[NT, B * H](

339 + A=A,

340 + Ad=Ad,

341 + cu_seqlens=cu_seqlens,

342 + chunk_indices=chunk_indices,

343 + T=T,

344 + H=H,

345 + BT=BT,

346 + )

347 + if BT == 16:

348 + return Ad

349 +

350 + Ai = torch.empty(B, T, H, BT, device=A.device, dtype=output_dtype)

351 + merge_fn = merge_16x16_to_32x32_inverse_kernel if BT == 32 else merge_16x16_to_64x64_inverse_kernel

352 + chunk_indices = prepare_chunk_indices(

353 + cu_seqlens, BT) if cu_seqlens is not None else None

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

355 + merge_fn[NT, B * H](

356 + A=A,

357 + Ad=Ad,

358 + Ai=Ai,

359 + cu_seqlens=cu_seqlens,

360 + chunk_indices=chunk_indices,

361 + T=T,

362 + H=H,

363 + BT=BT,

364 + )

365 + return Ai

vllm/model_executor/layers/fla/ops/utils.py new

+180

@@ -0,0 +1,180 @@

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:

9 +# ruff: noqa: E501

10 +import contextlib

11 +import functools

12 +import logging

13 +import os

14 +from enum import Enum

15 +from typing import Any, Callable, Literal, Optional

16 +

17 +import torch

18 +

19 +from vllm.triton_utils import triton

20 +

21 +logger = logging.getLogger(__name__)

22 +

23 +COMPILER_MODE = os.getenv("FLA_COMPILER_MODE") == "1"

24 +FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1"

25 +FLA_GDN_FIX_BT = os.getenv("FLA_GDN_FIX_BT", "0") == "1"

26 +

27 +SUPPRESS_LEVEL = int(os.getenv("GDN_RECOMPUTE_SUPPRESS_LEVEL", "0"))

28 +

29 +

30 +def tensor_cache(

31 + fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:

32 + """

33 + A decorator that caches the most recent results of a function with tensor inputs.

34 +

35 + This decorator will store the output of the decorated function for the most recent set of input tensors.

36 + The cache is limited to a fixed size (default is 4). When the cache is full, the oldest entry will be removed.

37 +

38 + Args:

39 + fn (Callable[..., torch.Tensor]):

40 + The function to be decorated. It should take tensor inputs and return tensor outputs.

41 +

42 + Returns:

43 + Callable[..., torch.Tensor]:

44 + A wrapped version of the input function with single-entry caching.

45 + """

46 +

47 + cache_entries: tuple[Optional[tuple], Optional[dict], Any] = []

48 + cache_size = 4

49 +

50 + @functools.wraps(fn)

51 + def wrapper(*args: Any, **kwargs: Any) -> Any:

52 + nonlocal cache_entries, cache_size

53 + for i, entry in enumerate(cache_entries):

54 + last_args, last_kwargs, last_result = entry

55 + if len(args) == len(last_args) and len(kwargs) == len(last_kwargs) \

56 + and all(a is b for a, b in zip(args, last_args)) \

57 + and all(k in last_kwargs and v is last_kwargs[k] for k, v in kwargs.items()):

58 + cache_entries = cache_entries[:i] + cache_entries[i + 1:] + [

59 + (args, kwargs, last_result)

60 + ]

61 + return last_result

62 +

63 + result = fn(*args, **kwargs)

64 +

65 + if len(cache_entries) >= cache_size:

66 + cache_entries = cache_entries[1:]

67 + cache_entries.append((args, kwargs, result))

68 + return result

69 +

70 + return wrapper

71 +

72 +

73 +def input_guard(

74 + fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:

75 + """

76 + A decorator to make sure all input tensors are contiguous and set the device based on input tensors.

77 + """

78 +

79 + @functools.wraps(fn)

80 + def wrapper(*args, **kwargs):

81 + contiguous_args = (i if not isinstance(i, torch.Tensor) else

82 + i.contiguous() for i in args)

83 + contiguous_kwargs = {

84 + k: (v if not isinstance(v, torch.Tensor) else v.contiguous())

85 + for k, v in kwargs.items()

86 + }

87 +

88 + tensor = None

89 + for arg in args:

90 + if isinstance(arg, torch.Tensor):

91 + tensor = arg

92 + break

93 + if tensor is None:

94 + for value in kwargs.values():

95 + if isinstance(value, torch.Tensor):

96 + tensor = value

97 + break

98 +

99 + if tensor is not None:

100 + ctx = torch.cuda.device(tensor.device.index)

101 + else:

102 + ctx = contextlib.nullcontext()

103 +

104 + with ctx:

105 + return fn(*contiguous_args, **contiguous_kwargs)

106 +

107 + return wrapper

108 +

109 +

110 +@functools.cache

111 +def get_available_device() -> str:

112 + try:

113 + return triton.runtime.driver.active.get_current_target().backend

114 + except BaseException:

115 + return 'cpu'

116 +

117 +

118 +@functools.cache

119 +def _check_platform() -> Literal['nvidia', 'amd', 'intel', 'musa']:

120 + device = get_available_device()

121 + mapping = {

122 + "cuda": "nvidia",

123 + "hip": "amd",

124 + "xpu": "intel",

125 + }

126 + # return the mapped value, or the original if not found

127 + return mapping.get(device, device)

128 +

129 +

130 +# For AMD GPUs, the triton backend is 'hip', while for Nvidia GPUs, the triton backend is 'cuda'.

131 +# However, the torch backend is 'cuda' for both Nvidia and AMD GPUs.

132 +# Therefore, we need to check the triton backend to determine the actual GPU vendor.

133 +device = get_available_device() if get_available_device() != 'hip' else 'cuda'

134 +device_torch_lib = getattr(torch, device)

135 +device_platform = _check_platform()

136 +

137 +is_amd = (device_platform == 'amd')

138 +is_intel = (device_platform == 'intel')

139 +is_nvidia = (device_platform == 'nvidia')

140 +is_intel_alchemist = (is_intel

141 + and 'Intel(R) Arc(TM) A' in torch.xpu.get_device_name(0))

142 +is_nvidia_hopper = (is_nvidia

143 + and ('NVIDIA H' in torch.cuda.get_device_name(0)

144 + or torch.cuda.get_device_capability()[0] >= 9))

145 +use_cuda_graph = (is_nvidia

146 + and os.environ.get('FLA_USE_CUDA_GRAPH', '0') == '1')

147 +

148 +

149 +def get_all_max_shared_mem():

150 + try:

151 + return [

152 + triton.runtime.driver.active.utils.get_device_properties(i)

153 + ['max_shared_mem'] for i in range(device_torch_lib.device_count())

154 + ]

155 + except BaseException:

156 + return [-1]

157 +

158 +

159 +class Backend(Enum):

160 + ADA = 101376 # RTX 4090

161 + AMPERE = 166912 # A100

162 + HOPPER = 232448 # H100

163 + DEFAULT = 102400 # Default

164 +

165 + @classmethod

166 + def get_shared_memory(cls, arch: str) -> int:

167 + try:

168 + return cls[arch.upper()].value

169 + except KeyError:

170 + return cls.DEFAULT.value

171 +

172 +

173 +@functools.cache

174 +def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool:

175 + try:

176 + device_shared_mem_list = get_all_max_shared_mem()

177 + max_shared_memory = device_shared_mem_list[tensor_idx]

178 + return max_shared_memory >= Backend.get_shared_memory(arch)

179 + except Exception:

180 + return False

vllm/model_executor/layers/fla/ops/wy_fast.py new

+114

@@ -0,0 +1,114 @@

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:

9 +

10 +# ruff: noqa: E501

11 +from typing import Optional

12 +

13 +import torch

14 +

15 +from vllm.triton_utils import tl, triton

16 +

17 +from .index import prepare_chunk_indices

18 +

19 +

20 +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None})

21 +@triton.autotune(

22 + configs=[

23 + triton.Config({}, num_warps=num_warps, num_stages=num_stages)

24 + for num_warps in [2, 4, 8] for num_stages in [2, 3, 4]

25 + ],

26 + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'],

27 +)

28 +@triton.jit(do_not_specialize=['T'])

29 +def recompute_w_u_fwd_kernel(k, v, beta, w, u, A, g, cu_seqlens, chunk_indices,

30 + T, H: tl.constexpr, Hg: tl.constexpr,

31 + K: tl.constexpr, V: tl.constexpr,

32 + BT: tl.constexpr, BK: tl.constexpr,

33 + BV: tl.constexpr, IS_VARLEN: tl.constexpr):

34 + i_t, i_bh = tl.program_id(0), tl.program_id(1)

35 + i_b, i_h = i_bh // H, i_bh % H

36 + if IS_VARLEN:

37 + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(

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

39 + bos, eos = tl.load(cu_seqlens + i_n).to(

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

41 + T = eos - bos

42 + else:

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

44 + p_beta = tl.make_block_ptr(beta + bos * H + i_h, (T, ), (H, ),

45 + (i_t * BT, ), (BT, ), (0, ))

46 + p_g = tl.make_block_ptr(g + (bos * H + i_h), (T, ), (H, ), (i_t * BT, ),

47 + (BT, ), (0, ))

48 + p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1),

49 + (i_t * BT, 0), (BT, BT), (1, 0))

50 + b_beta = tl.load(p_beta, boundary_check=(0, ))

51 + b_A = tl.load(p_A, boundary_check=(0, 1))

52 + b_g = tl.exp(tl.load(p_g, boundary_check=(0, )))

53 +

54 + for i_v in range(tl.cdiv(V, BV)):

55 + p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H * V, 1),

56 + (i_t * BT, i_v * BV), (BT, BV), (1, 0))

57 + p_u = tl.make_block_ptr(u + (bos * H + i_h) * V, (T, V), (H * V, 1),

58 + (i_t * BT, i_v * BV), (BT, BV), (1, 0))

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

60 + b_vb = (b_v * b_beta[:, None]).to(b_v.dtype)

61 + b_u = tl.dot(b_A, b_vb, allow_tf32=False)

62 + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))

63 +

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

65 + p_k = tl.make_block_ptr(k + (bos * Hg + i_h // (H // Hg)) * K, (T, K),

66 + (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK),

67 + (1, 0))

68 + p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (H * K, 1),

69 + (i_t * BT, i_k * BK), (BT, BK), (1, 0))

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

71 + b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype)

72 + b_w = tl.dot(b_A, b_kb)

73 + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))

74 +

75 +

76 +def recompute_w_u_fwd(

77 + k: torch.Tensor,

78 + v: torch.Tensor,

79 + beta: torch.Tensor,

80 + g_cumsum: torch.Tensor,

81 + A: torch.Tensor,

82 + cu_seqlens: Optional[torch.LongTensor],

83 +) -> tuple[torch.Tensor, torch.Tensor]:

84 + B, T, Hg, K, V = *k.shape, v.shape[-1]

85 + H = v.shape[-2]

86 + BT = A.shape[-1]

87 +

88 + chunk_indices = prepare_chunk_indices(

89 + cu_seqlens, BT) if cu_seqlens is not None else None

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

91 + BK = 64

92 + BV = 64

93 + u = torch.empty_like(v)

94 + w = k.new_empty(B, T, H, K)

95 + recompute_w_u_fwd_kernel[(NT, B * H)](

96 + k=k,

97 + v=v,

98 + beta=beta,

99 + w=w,

100 + u=u,

101 + A=A,

102 + g=g_cumsum,

103 + cu_seqlens=cu_seqlens,

104 + chunk_indices=chunk_indices,

105 + T=T,

106 + H=H,

107 + Hg=Hg,

108 + K=K,

109 + V=V,

110 + BT=BT,

111 + BK=BK,

112 + BV=BV,

113 + )

114 + return w, u

vllm/triton_utils/__init__.py

+3 -1

@@ -7,8 +7,10 @@ from vllm.triton_utils.importing import (HAS_TRITON, TritonLanguagePlaceholder,

7 if HAS_TRITON:

8 import triton

9 import triton.language as tl

10 + import triton.language.extra.libdevice as tldevice

11 else:

12 triton = TritonPlaceholder()

13 tl = TritonLanguagePlaceholder()

14 + tldevice = TritonLanguagePlaceholder()

14 -__all__ = ["HAS_TRITON", "triton", "tl"]

16 +__all__ = ["HAS_TRITON", "triton", "tl", "tldevice"]