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
[Kernel] Accelerate solve_tril with TMA (#26746)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Oct 20, 2025 at 13:39 UTC
9fce7bee745230d61c60ad467966790553b0ba48
3 files changed
+412
-301
vllm/model_executor/layers/fla/ops/op.py
+32
-11
@@ -11,29 +11,50 @@ import os
11
12
from vllm.triton_utils import tl, tldevice, triton
13
14
+from .utils import is_gather_supported
15
+
16
if os.environ.get("FLA_USE_FAST_OPS", "0") == "1":
15
- div = tldevice.fast_dividef
17
exp = tldevice.fast_expf
18
log = tldevice.fast_logf
19
log2 = tldevice.fast_log2f
20
else:
20
-
21
- @triton.jit
22
- def div_normal(x, y):
23
- return x / y
24
-
25
- div = div_normal
21
exp = tl.exp
22
log = tl.log
23
log2 = tl.log2
24
25
31
-if not hasattr(tl, "gather"):
26
+if not is_gather_supported:
27
28
@triton.jit
29
def gather(src, index, axis, _builder=None):
35
- # This is a fallback implementation when tl.gather is not supported
36
- # In order to pass triton compiler, there is no actual gather operation
37
- return src
30
+ """
31
+ Gather operation that works when tl.gather is not supported.
32
+ This is a fallback implementation that returns None.
33
+ Just to make triton compiler happy.
34
+ """
35
+ return None
36
else:
37
gather = tl.gather
38
+
39
+if hasattr(triton.language, "_experimental_make_tensor_descriptor"):
40
+ # For Triton 3.3.x
41
+ make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor
42
+elif hasattr(triton.language, "make_tensor_descriptor"):
43
+ # For Triton 3.4.x and later
44
+ make_tensor_descriptor = triton.language.make_tensor_descriptor
45
+else:
46
+ """
47
+ Fallback implementation when TMA is not supported.
48
+ Returns None to indicate TMA descriptors are unavailable.
49
+ Just make triton compiler happy.
50
+ """
51
+
52
+ @triton.jit
53
+ def make_tensor_descriptor(
54
+ base,
55
+ shape,
56
+ strides,
57
+ block_shape,
58
+ _builder=None,
59
+ ):
60
+ return None
vllm/model_executor/layers/fla/ops/solve_tril.py
+375
-290
@@ -8,12 +8,21 @@
8
# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
9
# ruff: noqa: E501
10
11
+import os
12
+
13
import torch
14
15
from vllm.triton_utils import tl, triton
16
17
from .index import prepare_chunk_indices
16
-from .utils import input_guard
18
+from .op import make_tensor_descriptor
19
+from .utils import input_guard, is_amd, is_tma_supported
20
+
21
+FLA_TRIL_PRECISION = os.environ.get("FLA_TRIL_PRECISION", "ieee")
22
+ALLOWED_TRIL_PRECISIONS = ["ieee", "tf32"] if is_amd else ["ieee", "tf32", "tf32x3"]
23
+assert FLA_TRIL_PRECISION in ALLOWED_TRIL_PRECISIONS, (
24
+ f"FLA_TRIL_PRECISION must be one of {ALLOWED_TRIL_PRECISIONS}, but got {FLA_TRIL_PRECISION}"
25
+)
26
27
28
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
@@ -28,13 +37,15 @@ from .utils import input_guard
37
@triton.jit(do_not_specialize=["T"])
38
def solve_tril_16x16_kernel(
39
A,
31
- Ad,
40
+ Ai,
41
cu_seqlens,
42
chunk_indices,
43
T,
44
H: tl.constexpr,
45
BT: tl.constexpr,
46
+ USE_TMA: tl.constexpr,
47
IS_VARLEN: tl.constexpr,
48
+ DOT_PRECISION: tl.constexpr,
49
):
50
i_t, i_bh = tl.program_id(0), tl.program_id(1)
51
i_b, i_h = i_bh // H, i_bh % H
@@ -50,30 +61,43 @@ def solve_tril_16x16_kernel(
61
T = eos - bos
62
else:
63
bos, eos = i_b * T, i_b * T + T
64
+ o_i = tl.arange(0, 16)
65
+ m_A = o_i[:, None] > o_i[None, :]
66
+ m_I = o_i[:, None] == o_i[None, :]
67
68
A = A + (bos * H + i_h) * BT
55
- Ad = Ad + (bos * H + i_h) * 16
69
+ Ai = Ai + (bos * H + i_h) * 16
70
71
offset = (i_t * 16) % BT
58
- p_A = tl.make_block_ptr(
59
- A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0)
60
- )
61
- p_Ai = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0))
62
- b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
63
- b_A = -tl.where(tl.arange(0, 16)[:, None] > tl.arange(0, 16)[None, :], b_A, 0)
72
+ if not USE_TMA:
73
+ p_A = tl.make_block_ptr(
74
+ A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0)
75
+ )
76
+ # [16, 16]
77
+ b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
78
+ else:
79
+ desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16])
80
+ desc_o = make_tensor_descriptor(Ai, [T, 16], [H * 16, 1], [16, 16])
81
+ b_A = desc.load([i_t * 16, offset]).to(tl.float32)
82
+ b_A = -tl.where(m_A, b_A, 0)
83
65
- o_i = tl.arange(0, 16)
66
- for i in range(1, min(16, T - i_t * 16)):
84
+ for i in range(2, min(16, T - i_t * 16)):
85
+ # [16]
86
b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset)
87
b_a = b_a + tl.sum(b_a[:, None] * b_A, 0)
69
- mask = o_i == i
70
- b_A = tl.where(mask[:, None], b_a, b_A)
71
- b_A += o_i[:, None] == o_i[None, :]
72
- tl.store(
73
- p_Ai,
74
- b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"),
75
- boundary_check=(0, 1),
76
- )
88
+ b_A = tl.where((o_i == i)[:, None], b_a, b_A)
89
+ b_A += m_I
90
+ if not USE_TMA:
91
+ p_Ai = tl.make_block_ptr(
92
+ Ai, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0)
93
+ )
94
+ tl.store(
95
+ p_Ai,
96
+ b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"),
97
+ boundary_check=(0, 1),
98
+ )
99
+ else:
100
+ desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne"))
101
102
103
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
@@ -88,14 +112,15 @@ def solve_tril_16x16_kernel(
112
@triton.jit(do_not_specialize=["T"])
113
def merge_16x16_to_32x32_inverse_kernel(
114
A,
91
- Ad,
115
Ai,
116
cu_seqlens,
117
chunk_indices,
118
T,
119
H: tl.constexpr,
120
BT: tl.constexpr,
121
+ USE_TMA: tl.constexpr,
122
IS_VARLEN: tl.constexpr,
123
+ DOT_PRECISION: tl.constexpr,
124
):
125
i_t, i_bh = tl.program_id(0), tl.program_id(1)
126
i_b, i_h = i_bh // H, i_bh % H
@@ -112,51 +137,93 @@ def merge_16x16_to_32x32_inverse_kernel(
137
else:
138
bos, eos = i_b * T, i_b * T + T
139
115
- A += (bos * H + i_h) * 32
116
- Ad += (bos * H + i_h) * 16
117
- Ai += (bos * H + i_h) * 32
140
+ o_i = tl.arange(0, 16)
141
+ m_A = o_i[:, None] > o_i[None, :]
142
+ m_I = o_i[:, None] == o_i[None, :]
143
+ A += (bos * H + i_h) * BT
144
+ Ai += (bos * H + i_h) * BT
145
119
- p_A_21 = tl.make_block_ptr(
120
- A, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0)
121
- )
122
- p_Ad_11 = tl.make_block_ptr(
123
- Ad, (T, 16), (H * 16, 1), (i_t * 32, 0), (16, 16), (1, 0)
124
- )
125
- p_Ad_22 = tl.make_block_ptr(
126
- Ad, (T, 16), (H * 16, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0)
127
- )
128
- p_Ai_11 = tl.make_block_ptr(
129
- Ai, (T, 32), (H * 32, 1), (i_t * 32, 0), (16, 16), (1, 0)
130
- )
131
- p_Ai_22 = tl.make_block_ptr(
132
- Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 16), (16, 16), (1, 0)
133
- )
134
- p_Ai_21 = tl.make_block_ptr(
135
- Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0)
136
- )
146
+ if not USE_TMA:
147
+ p_A_11 = tl.make_block_ptr(
148
+ A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)
149
+ )
150
+ p_A_22 = tl.make_block_ptr(
151
+ A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)
152
+ )
153
+ b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32)
154
+ b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32)
155
+ else:
156
+ desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16])
157
+ desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16])
158
+ b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32)
159
+ b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32)
160
138
- A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
139
- Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32)
140
- Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32)
141
- Ai_21 = -tl.dot(
142
- tl.dot(Ai_22, A_21, input_precision="ieee"), Ai_11, input_precision="ieee"
143
- )
144
- tl.store(
145
- p_Ai_11,
146
- Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),
147
- boundary_check=(0, 1),
148
- )
149
- tl.store(
150
- p_Ai_22,
151
- Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),
152
- boundary_check=(0, 1),
153
- )
154
- tl.store(
155
- p_Ai_21,
156
- Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),
157
- boundary_check=(0, 1),
161
+ # [16, 16]
162
+ b_Ai_11 = -tl.where(m_A, b_Ai_11, 0)
163
+ b_Ai_22 = -tl.where(m_A, b_Ai_22, 0)
164
+
165
+ for i in range(2, min(16, T - i_t * BT)):
166
+ b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i)
167
+ b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0)
168
+ b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11)
169
+ for i in range(16 + 2, min(32, T - i_t * BT)):
170
+ b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16)
171
+ b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0)
172
+ b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22)
173
+
174
+ b_Ai_11 += m_I
175
+ b_Ai_22 += m_I
176
+
177
+ if not USE_TMA:
178
+ p_A_21 = tl.make_block_ptr(
179
+ A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)
180
+ )
181
+ b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
182
+ else:
183
+ b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32)
184
+
185
+ b_Ai_21 = -tl.dot(
186
+ tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION),
187
+ b_Ai_11,
188
+ input_precision=DOT_PRECISION,
189
)
190
191
+ if not USE_TMA:
192
+ p_Ai_11 = tl.make_block_ptr(
193
+ Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)
194
+ )
195
+ p_Ai_21 = tl.make_block_ptr(
196
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)
197
+ )
198
+ p_Ai_22 = tl.make_block_ptr(
199
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)
200
+ )
201
+ tl.store(
202
+ p_Ai_11,
203
+ b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),
204
+ boundary_check=(0, 1),
205
+ )
206
+ tl.store(
207
+ p_Ai_22,
208
+ b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),
209
+ boundary_check=(0, 1),
210
+ )
211
+ tl.store(
212
+ p_Ai_21,
213
+ b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),
214
+ boundary_check=(0, 1),
215
+ )
216
+ else:
217
+ desc_o.store(
218
+ [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")
219
+ )
220
+ desc_o.store(
221
+ [i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")
222
+ )
223
+ desc_o.store(
224
+ [i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")
225
+ )
226
+
227
228
@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
229
@triton.autotune(
@@ -170,14 +237,15 @@ def merge_16x16_to_32x32_inverse_kernel(
237
@triton.jit(do_not_specialize=["T"])
238
def merge_16x16_to_64x64_inverse_kernel(
239
A,
173
- Ad,
240
Ai,
241
cu_seqlens,
242
chunk_indices,
243
T,
244
H: tl.constexpr,
245
BT: tl.constexpr,
246
+ USE_TMA: tl.constexpr,
247
IS_VARLEN: tl.constexpr,
248
+ DOT_PRECISION: tl.constexpr,
249
):
250
i_t, i_bh = tl.program_id(0), tl.program_id(1)
251
i_b, i_h = i_bh // H, i_bh % H
@@ -194,213 +262,245 @@ def merge_16x16_to_64x64_inverse_kernel(
262
else:
263
bos, eos = i_b * T, i_b * T + T
264
197
- A += (bos * H + i_h) * 64
198
- Ad += (bos * H + i_h) * 16
199
- Ai += (bos * H + i_h) * 64
265
+ o_i = tl.arange(0, 16)
266
+ m_A = o_i[:, None] > o_i[None, :]
267
+ m_I = o_i[:, None] == o_i[None, :]
268
+ A += (bos * H + i_h) * BT
269
+ Ai += (bos * H + i_h) * BT
270
201
- p_A_21 = tl.make_block_ptr(
202
- A, (T, 64), (H * 64, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0)
203
- )
204
- p_A_32 = tl.make_block_ptr(
205
- A, (T, 64), (H * 64, 1), (i_t * 64 + 32, 16), (16, 16), (1, 0)
206
- )
207
- p_A_31 = tl.make_block_ptr(
208
- A, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0)
209
- )
210
- p_A_43 = tl.make_block_ptr(
211
- A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 32), (16, 16), (1, 0)
212
- )
213
- p_A_42 = tl.make_block_ptr(
214
- A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 16), (16, 16), (1, 0)
215
- )
216
- p_A_41 = tl.make_block_ptr(
217
- A, (T, 64), (H * 64, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0)
218
- )
219
- p_Ad_11 = tl.make_block_ptr(
220
- Ad, (T, 16), (H * 16, 1), (i_t * 64, 0), (16, 16), (1, 0)
221
- )
222
- p_Ad_22 = tl.make_block_ptr(
223
- Ad, (T, 16), (H * 16, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0)
224
- )
225
- p_Ad_33 = tl.make_block_ptr(
226
- Ad, (T, 16), (H * 16, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0)
227
- )
228
- p_Ad_44 = tl.make_block_ptr(
229
- Ad, (T, 16), (H * 16, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0)
230
- )
271
+ if not USE_TMA:
272
+ p_A_11 = tl.make_block_ptr(
273
+ A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)
274
+ )
275
+ p_A_22 = tl.make_block_ptr(
276
+ A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)
277
+ )
278
+ p_A_33 = tl.make_block_ptr(
279
+ A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)
280
+ )
281
+ p_A_44 = tl.make_block_ptr(
282
+ A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)
283
+ )
284
+ b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32)
285
+ b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32)
286
+ b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32)
287
+ b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32)
288
+ else:
289
+ desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16])
290
+ desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16])
291
+ b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32)
292
+ b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32)
293
+ b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32)
294
+ b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32)
295
232
- A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
233
- A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32)
234
- A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32)
235
- A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32)
236
- A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32)
237
- A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32)
296
+ # [16, 16]
297
+ b_Ai_11 = -tl.where(m_A, b_Ai_11, 0)
298
+ b_Ai_22 = -tl.where(m_A, b_Ai_22, 0)
299
+ b_Ai_33 = -tl.where(m_A, b_Ai_33, 0)
300
+ b_Ai_44 = -tl.where(m_A, b_Ai_44, 0)
301
239
- Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32)
240
- Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32)
241
- Ai_33 = tl.load(p_Ad_33, boundary_check=(0, 1)).to(tl.float32)
242
- Ai_44 = tl.load(p_Ad_44, boundary_check=(0, 1)).to(tl.float32)
302
+ for i in range(2, min(16, T - i_t * BT)):
303
+ b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i)
304
+ b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0)
305
+ b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11)
306
+ for i in range(16 + 2, min(32, T - i_t * BT)):
307
+ b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16)
308
+ b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0)
309
+ b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22)
310
+ for i in range(32 + 2, min(48, T - i_t * BT)):
311
+ b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32)
312
+ b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0)
313
+ b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33)
314
+ for i in range(48 + 2, min(64, T - i_t * BT)):
315
+ b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48)
316
+ b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0)
317
+ b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44)
318
+ b_Ai_11 += m_I
319
+ b_Ai_22 += m_I
320
+ b_Ai_33 += m_I
321
+ b_Ai_44 += m_I
322
244
- Ai_21 = -tl.dot(
245
- tl.dot(Ai_22, A_21, input_precision="ieee"), Ai_11, input_precision="ieee"
246
- )
247
- Ai_32 = -tl.dot(
248
- tl.dot(Ai_33, A_32, input_precision="ieee"), Ai_22, input_precision="ieee"
249
- )
250
- Ai_43 = -tl.dot(
251
- tl.dot(Ai_44, A_43, input_precision="ieee"), Ai_33, input_precision="ieee"
252
- )
323
+ if not USE_TMA:
324
+ p_A_21 = tl.make_block_ptr(
325
+ A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)
326
+ )
327
+ p_A_31 = tl.make_block_ptr(
328
+ A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)
329
+ )
330
+ p_A_32 = tl.make_block_ptr(
331
+ A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)
332
+ )
333
+ p_A_41 = tl.make_block_ptr(
334
+ A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)
335
+ )
336
+ p_A_42 = tl.make_block_ptr(
337
+ A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)
338
+ )
339
+ p_A_43 = tl.make_block_ptr(
340
+ A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)
341
+ )
342
+ b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
343
+ b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32)
344
+ b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32)
345
+ b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32)
346
+ b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32)
347
+ b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32)
348
+ else:
349
+ b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32)
350
+ b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32)
351
+ b_A_32 = desc.load([i_t * BT + 32, 16]).to(tl.float32)
352
+ b_A_41 = desc.load([i_t * BT + 48, 0]).to(tl.float32)
353
+ b_A_42 = desc.load([i_t * BT + 48, 16]).to(tl.float32)
354
+ b_A_43 = desc.load([i_t * BT + 48, 32]).to(tl.float32)
355
254
- Ai_31 = -tl.dot(
255
- Ai_33,
256
- tl.dot(A_31, Ai_11, input_precision="ieee")
257
- + tl.dot(A_32, Ai_21, input_precision="ieee"),
258
- input_precision="ieee",
356
+ b_Ai_21 = -tl.dot(
357
+ tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION),
358
+ b_Ai_11,
359
+ input_precision=DOT_PRECISION,
360
)
260
- Ai_42 = -tl.dot(
261
- Ai_44,
262
- tl.dot(A_42, Ai_22, input_precision="ieee")
263
- + tl.dot(A_43, Ai_32, input_precision="ieee"),
264
- input_precision="ieee",
361
+ b_Ai_32 = -tl.dot(
362
+ tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION),
363
+ b_Ai_22,
364
+ input_precision=DOT_PRECISION,
365
)
266
- Ai_41 = -tl.dot(
267
- Ai_44,
268
- tl.dot(A_41, Ai_11, input_precision="ieee")
269
- + tl.dot(A_42, Ai_21, input_precision="ieee")
270
- + tl.dot(A_43, Ai_31, input_precision="ieee"),
271
- input_precision="ieee",
366
+ b_Ai_43 = -tl.dot(
367
+ tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION),
368
+ b_Ai_33,
369
+ input_precision=DOT_PRECISION,
370
)
371
274
- p_Ai_11 = tl.make_block_ptr(
275
- Ai, (T, 64), (H * 64, 1), (i_t * 64, 0), (16, 16), (1, 0)
276
- )
277
- p_Ai_22 = tl.make_block_ptr(
278
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 16), (16, 16), (1, 0)
279
- )
280
- p_Ai_33 = tl.make_block_ptr(
281
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 32), (16, 16), (1, 0)
282
- )
283
- p_Ai_44 = tl.make_block_ptr(
284
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 48), (16, 16), (1, 0)
285
- )
286
- p_Ai_21 = tl.make_block_ptr(
287
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 0), (16, 16), (1, 0)
288
- )
289
- p_Ai_31 = tl.make_block_ptr(
290
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 0), (16, 16), (1, 0)
291
- )
292
- p_Ai_32 = tl.make_block_ptr(
293
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 16), (16, 16), (1, 0)
294
- )
295
- p_Ai_41 = tl.make_block_ptr(
296
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 0), (16, 16), (1, 0)
297
- )
298
- p_Ai_42 = tl.make_block_ptr(
299
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 16), (16, 16), (1, 0)
300
- )
301
- p_Ai_43 = tl.make_block_ptr(
302
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 48, 32), (16, 16), (1, 0)
303
- )
304
- tl.store(
305
- p_Ai_11,
306
- Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),
307
- boundary_check=(0, 1),
308
- )
309
- tl.store(
310
- p_Ai_22,
311
- Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),
312
- boundary_check=(0, 1),
313
- )
314
- tl.store(
315
- p_Ai_33,
316
- Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"),
317
- boundary_check=(0, 1),
318
- )
319
- tl.store(
320
- p_Ai_44,
321
- Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"),
322
- boundary_check=(0, 1),
323
- )
324
- tl.store(
325
- p_Ai_21,
326
- Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),
327
- boundary_check=(0, 1),
328
- )
329
- tl.store(
330
- p_Ai_31,
331
- Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"),
332
- boundary_check=(0, 1),
333
- )
334
- tl.store(
335
- p_Ai_32,
336
- Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"),
337
- boundary_check=(0, 1),
338
- )
339
- tl.store(
340
- p_Ai_41,
341
- Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"),
342
- boundary_check=(0, 1),
343
- )
344
- tl.store(
345
- p_Ai_42,
346
- Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"),
347
- boundary_check=(0, 1),
348
- )
349
- tl.store(
350
- p_Ai_43,
351
- Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"),
352
- boundary_check=(0, 1),
372
+ b_Ai_31 = -tl.dot(
373
+ b_Ai_33,
374
+ tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION)
375
+ + tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION),
376
+ input_precision=DOT_PRECISION,
377
+ )
378
+ b_Ai_42 = -tl.dot(
379
+ b_Ai_44,
380
+ tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION)
381
+ + tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION),
382
+ input_precision=DOT_PRECISION,
383
+ )
384
+ b_Ai_41 = -tl.dot(
385
+ b_Ai_44,
386
+ tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION)
387
+ + tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION)
388
+ + tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION),
389
+ input_precision=DOT_PRECISION,
390
)
391
355
- fill_zeros = tl.zeros((16, 16), dtype=tl.float32)
356
- p_Ai_12 = tl.make_block_ptr(
357
- Ai, (T, 64), (H * 64, 1), (i_t * 64, 16), (16, 16), (1, 0)
358
- )
359
- p_Ai_13 = tl.make_block_ptr(
360
- Ai, (T, 64), (H * 64, 1), (i_t * 64, 32), (16, 16), (1, 0)
361
- )
362
- p_Ai_14 = tl.make_block_ptr(
363
- Ai, (T, 64), (H * 64, 1), (i_t * 64, 48), (16, 16), (1, 0)
364
- )
365
- p_Ai_23 = tl.make_block_ptr(
366
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 32), (16, 16), (1, 0)
367
- )
368
- p_Ai_24 = tl.make_block_ptr(
369
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 16, 48), (16, 16), (1, 0)
370
- )
371
- p_Ai_34 = tl.make_block_ptr(
372
- Ai, (T, 64), (H * 64, 1), (i_t * 64 + 32, 48), (16, 16), (1, 0)
373
- )
374
- tl.store(
375
- p_Ai_12,
376
- fill_zeros.to(p_Ai_12.dtype.element_ty, fp_downcast_rounding="rtne"),
377
- boundary_check=(0, 1),
378
- )
379
- tl.store(
380
- p_Ai_13,
381
- fill_zeros.to(p_Ai_13.dtype.element_ty, fp_downcast_rounding="rtne"),
382
- boundary_check=(0, 1),
383
- )
384
- tl.store(
385
- p_Ai_14,
386
- fill_zeros.to(p_Ai_14.dtype.element_ty, fp_downcast_rounding="rtne"),
387
- boundary_check=(0, 1),
388
- )
389
- tl.store(
390
- p_Ai_23,
391
- fill_zeros.to(p_Ai_23.dtype.element_ty, fp_downcast_rounding="rtne"),
392
- boundary_check=(0, 1),
393
- )
394
- tl.store(
395
- p_Ai_24,
396
- fill_zeros.to(p_Ai_24.dtype.element_ty, fp_downcast_rounding="rtne"),
397
- boundary_check=(0, 1),
398
- )
399
- tl.store(
400
- p_Ai_34,
401
- fill_zeros.to(p_Ai_34.dtype.element_ty, fp_downcast_rounding="rtne"),
402
- boundary_check=(0, 1),
403
- )
392
+ if not USE_TMA:
393
+ p_Ai_11 = tl.make_block_ptr(
394
+ Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0)
395
+ )
396
+ p_Ai_22 = tl.make_block_ptr(
397
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)
398
+ )
399
+ p_Ai_33 = tl.make_block_ptr(
400
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)
401
+ )
402
+ p_Ai_44 = tl.make_block_ptr(
403
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)
404
+ )
405
+ p_Ai_21 = tl.make_block_ptr(
406
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)
407
+ )
408
+ p_Ai_31 = tl.make_block_ptr(
409
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)
410
+ )
411
+ p_Ai_32 = tl.make_block_ptr(
412
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)
413
+ )
414
+ p_Ai_41 = tl.make_block_ptr(
415
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)
416
+ )
417
+ p_Ai_42 = tl.make_block_ptr(
418
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)
419
+ )
420
+ p_Ai_43 = tl.make_block_ptr(
421
+ Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)
422
+ )
423
+ tl.store(
424
+ p_Ai_11,
425
+ b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),
426
+ boundary_check=(0, 1),
427
+ )
428
+ tl.store(
429
+ p_Ai_22,
430
+ b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),
431
+ boundary_check=(0, 1),
432
+ )
433
+ tl.store(
434
+ p_Ai_33,
435
+ b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"),
436
+ boundary_check=(0, 1),
437
+ )
438
+ tl.store(
439
+ p_Ai_44,
440
+ b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"),
441
+ boundary_check=(0, 1),
442
+ )
443
+ tl.store(
444
+ p_Ai_21,
445
+ b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),
446
+ boundary_check=(0, 1),
447
+ )
448
+ tl.store(
449
+ p_Ai_31,
450
+ b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"),
451
+ boundary_check=(0, 1),
452
+ )
453
+ tl.store(
454
+ p_Ai_32,
455
+ b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"),
456
+ boundary_check=(0, 1),
457
+ )
458
+ tl.store(
459
+ p_Ai_41,
460
+ b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"),
461
+ boundary_check=(0, 1),
462
+ )
463
+ tl.store(
464
+ p_Ai_42,
465
+ b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"),
466
+ boundary_check=(0, 1),
467
+ )
468
+ tl.store(
469
+ p_Ai_43,
470
+ b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"),
471
+ boundary_check=(0, 1),
472
+ )
473
+ else:
474
+ desc_o.store(
475
+ [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")
476
+ )
477
+ desc_o.store(
478
+ [i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")
479
+ )
480
+ desc_o.store(
481
+ [i_t * BT + 32, 32], b_Ai_33.to(desc_o.dtype, fp_downcast_rounding="rtne")
482
+ )
483
+ desc_o.store(
484
+ [i_t * BT + 48, 48], b_Ai_44.to(desc_o.dtype, fp_downcast_rounding="rtne")
485
+ )
486
+ desc_o.store(
487
+ [i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")
488
+ )
489
+ desc_o.store(
490
+ [i_t * BT + 32, 0], b_Ai_31.to(desc_o.dtype, fp_downcast_rounding="rtne")
491
+ )
492
+ desc_o.store(
493
+ [i_t * BT + 32, 16], b_Ai_32.to(desc_o.dtype, fp_downcast_rounding="rtne")
494
+ )
495
+ desc_o.store(
496
+ [i_t * BT + 48, 0], b_Ai_41.to(desc_o.dtype, fp_downcast_rounding="rtne")
497
+ )
498
+ desc_o.store(
499
+ [i_t * BT + 48, 16], b_Ai_42.to(desc_o.dtype, fp_downcast_rounding="rtne")
500
+ )
501
+ desc_o.store(
502
+ [i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne")
503
+ )
504
505
506
@input_guard
@@ -410,62 +510,47 @@ def solve_tril(
510
output_dtype: torch.dtype = torch.float,
511
) -> torch.Tensor:
512
"""
413
- Compute the inverse of the lower triangular matrix
513
+ Compute the inverse of the matrix I + A
514
A should be strictly lower triangular, i.e., A.triu() == 0.
515
516
Args:
517
A (torch.Tensor):
418
- [B, T, H, K]
518
+ [B, T, H, BT], where BT should only be 16, 32, or 64.
519
cu_seqlens (torch.Tensor):
420
- The cumulative sequence lengths of the input tensor.
421
- Default: None.
520
+ The cumulative sequence lengths of the input tensor. Default: `None`.
521
output_dtype (torch.dtype):
423
- The dtype of the output tensor. Default: `torch.float`
522
+ The dtype of the output tensor. Default: `torch.float`.
523
+ If `None`, the output dtype will be the same as the input dtype.
524
525
Returns:
526
(I + A)^-1 with the same shape as A
527
"""
528
assert A.shape[-1] in [16, 32, 64]
529
+ output_dtype = A.dtype if output_dtype is None else output_dtype
530
531
B, T, H, BT = A.shape
431
- Ad = torch.empty(
432
- B, T, H, 16, device=A.device, dtype=torch.float if BT != 16 else output_dtype
433
- )
434
-
435
- chunk_indices = (
436
- prepare_chunk_indices(cu_seqlens, 16) if cu_seqlens is not None else None
437
- )
438
- NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, 16)
439
- solve_tril_16x16_kernel[NT, B * H](
440
- A=A,
441
- Ad=Ad,
442
- cu_seqlens=cu_seqlens,
443
- chunk_indices=chunk_indices,
444
- T=T,
445
- H=H,
446
- BT=BT,
447
- )
448
- if BT == 16:
449
- return Ad
450
-
451
- Ai = torch.empty(B, T, H, BT, device=A.device, dtype=output_dtype)
452
- merge_fn = (
453
- merge_16x16_to_32x32_inverse_kernel
454
- if BT == 32
455
- else merge_16x16_to_64x64_inverse_kernel
456
- )
532
chunk_indices = (
533
prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
534
)
535
NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
536
+
537
+ Ai = torch.zeros_like(A, dtype=output_dtype)
538
+ if BT == 16:
539
+ merge_fn = solve_tril_16x16_kernel
540
+ elif BT == 32:
541
+ merge_fn = merge_16x16_to_32x32_inverse_kernel
542
+ elif BT == 64:
543
+ merge_fn = merge_16x16_to_64x64_inverse_kernel
544
+
545
merge_fn[NT, B * H](
546
A=A,
463
- Ad=Ad,
547
Ai=Ai,
548
cu_seqlens=cu_seqlens,
549
chunk_indices=chunk_indices,
550
T=T,
551
H=H,
552
BT=BT,
553
+ USE_TMA=is_tma_supported,
554
+ DOT_PRECISION=FLA_TRIL_PRECISION,
555
)
556
return Ai
vllm/model_executor/layers/fla/ops/utils.py
+5
@@ -150,6 +150,11 @@ is_nvidia_hopper = is_nvidia and (
150
or torch.cuda.get_device_capability()[0] >= 9
151
)
152
use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1"
153
+is_gather_supported = hasattr(triton.language, "gather")
154
+is_tma_supported = (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 9) and (
155
+ hasattr(triton.language, "_experimental_make_tensor_descriptor")
156
+ or hasattr(triton.language, "make_tensor_descriptor")
157
+)
158
159
160
def get_all_max_shared_mem():