Skip to content

Commit 242d36a

Browse files
0x45f0x45f
authored
Fix tle error (flagos-ai#4894)
1 parent 0237c11 commit 242d36a

1 file changed

Lines changed: 127 additions & 125 deletions

File tree

‎src/flag_gems/runtime/backend/_nvidia/hopper/ops/w8a8_block_fp8_bmm.py‎

Lines changed: 127 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
try:
2828
import triton.experimental.tle.language as tle
2929

30-
HAS_TLE_W8A8_BLOCK_FP8_BMM = True
30+
HAS_TLE_W8A8_BLOCK_FP8_BMM = hasattr(tle.gpu, "alloc_barriers")
3131
except ImportError:
3232
tle = None
3333
HAS_TLE_W8A8_BLOCK_FP8_BMM = False
@@ -268,136 +268,138 @@ def _tle_w8a8_block_fp8_bmm_load_partition(
268268
)
269269

270270

271-
@libentry()
272-
@libtuner(
273-
configs=_get_tle_w8a8_block_fp8_bmm_configs(),
274-
key=["B", "M_aligned", "N", "K"],
275-
strategy=["default", "align32", "align32", "align32"],
276-
policy=_TLEW8A8BlockFP8BMMTuner,
277-
flagtune_op_name="w8a8_block_fp8_bmm",
278-
flagtune_expand_op_name="w8a8_block_fp8_bmm",
279-
)
280-
@triton.jit
281-
def w8a8_block_fp8_bmm_kernel(
282-
x_desc,
283-
y_desc,
284-
xs_ptr,
285-
z_ptr,
286-
ys_ptr,
287-
xs_sB: tl.constexpr,
288-
xs_sM: tl.constexpr,
289-
xs_sKb: tl.constexpr,
290-
z_sB: tl.constexpr,
291-
z_sM: tl.constexpr,
292-
z_sN: tl.constexpr,
293-
B: tl.constexpr,
294-
M: tl.constexpr,
295-
M_aligned: tl.constexpr,
296-
N: tl.constexpr,
297-
K: tl.constexpr,
298-
BLOCK_M: tl.constexpr,
299-
BLOCK_N: tl.constexpr,
300-
BLOCK_K: tl.constexpr,
301-
TILE_ORDER: tl.constexpr,
302-
SWAP_AB: tl.constexpr,
303-
X_ELEM_BYTES: tl.constexpr,
304-
Y_ELEM_BYTES: tl.constexpr,
305-
num_warps: tl.constexpr,
306-
num_stages: tl.constexpr,
307-
num_sms: tl.constexpr,
308-
):
309-
_ = num_warps
310-
x_smem = tle.gpu.alloc(
311-
[num_stages, 1, BLOCK_M, BLOCK_K],
312-
dtype=x_desc.dtype,
313-
layout=None,
314-
scope=tle.gpu.smem,
315-
)
316-
y_smem = tle.gpu.alloc(
317-
[num_stages, BLOCK_N, BLOCK_K],
318-
dtype=y_desc.dtype,
319-
layout=None,
320-
scope=tle.gpu.smem,
321-
)
322-
empty_x = tle.gpu.alloc_barriers(
323-
num_barriers=num_stages, arrive_count=1, init=tle.gpu.READY
324-
)
325-
empty_y = tle.gpu.alloc_barriers(
326-
num_barriers=num_stages, arrive_count=1, init=tle.gpu.READY
327-
)
328-
full_x = tle.gpu.alloc_barriers(
329-
num_barriers=num_stages,
330-
arrive_count=1,
331-
expect_bytes=BLOCK_M * BLOCK_K * X_ELEM_BYTES,
332-
)
333-
full_y = tle.gpu.alloc_barriers(
334-
num_barriers=num_stages,
335-
arrive_count=1,
336-
expect_bytes=BLOCK_N * BLOCK_K * Y_ELEM_BYTES,
271+
if HAS_TLE_W8A8_BLOCK_FP8_BMM:
272+
273+
@libentry()
274+
@libtuner(
275+
configs=_get_tle_w8a8_block_fp8_bmm_configs(),
276+
key=["B", "M_aligned", "N", "K"],
277+
strategy=["default", "align32", "align32", "align32"],
278+
policy=_TLEW8A8BlockFP8BMMTuner,
279+
flagtune_op_name="w8a8_block_fp8_bmm",
280+
flagtune_expand_op_name="w8a8_block_fp8_bmm",
337281
)
282+
@triton.jit
283+
def w8a8_block_fp8_bmm_kernel(
284+
x_desc,
285+
y_desc,
286+
xs_ptr,
287+
z_ptr,
288+
ys_ptr,
289+
xs_sB: tl.constexpr,
290+
xs_sM: tl.constexpr,
291+
xs_sKb: tl.constexpr,
292+
z_sB: tl.constexpr,
293+
z_sM: tl.constexpr,
294+
z_sN: tl.constexpr,
295+
B: tl.constexpr,
296+
M: tl.constexpr,
297+
M_aligned: tl.constexpr,
298+
N: tl.constexpr,
299+
K: tl.constexpr,
300+
BLOCK_M: tl.constexpr,
301+
BLOCK_N: tl.constexpr,
302+
BLOCK_K: tl.constexpr,
303+
TILE_ORDER: tl.constexpr,
304+
SWAP_AB: tl.constexpr,
305+
X_ELEM_BYTES: tl.constexpr,
306+
Y_ELEM_BYTES: tl.constexpr,
307+
num_warps: tl.constexpr,
308+
num_stages: tl.constexpr,
309+
num_sms: tl.constexpr,
310+
):
311+
_ = num_warps
312+
x_smem = tle.gpu.alloc(
313+
[num_stages, 1, BLOCK_M, BLOCK_K],
314+
dtype=x_desc.dtype,
315+
layout=None,
316+
scope=tle.gpu.smem,
317+
)
318+
y_smem = tle.gpu.alloc(
319+
[num_stages, BLOCK_N, BLOCK_K],
320+
dtype=y_desc.dtype,
321+
layout=None,
322+
scope=tle.gpu.smem,
323+
)
324+
empty_x = tle.gpu.alloc_barriers(
325+
num_barriers=num_stages, arrive_count=1, init=tle.gpu.READY
326+
)
327+
empty_y = tle.gpu.alloc_barriers(
328+
num_barriers=num_stages, arrive_count=1, init=tle.gpu.READY
329+
)
330+
full_x = tle.gpu.alloc_barriers(
331+
num_barriers=num_stages,
332+
arrive_count=1,
333+
expect_bytes=BLOCK_M * BLOCK_K * X_ELEM_BYTES,
334+
)
335+
full_y = tle.gpu.alloc_barriers(
336+
num_barriers=num_stages,
337+
arrive_count=1,
338+
expect_bytes=BLOCK_N * BLOCK_K * Y_ELEM_BYTES,
339+
)
338340

339-
tle.gpu.warp_specialize(
340-
[
341-
(
342-
_tle_w8a8_block_fp8_bmm_compute_partition,
341+
tle.gpu.warp_specialize(
342+
[
343343
(
344-
x_smem,
345-
y_smem,
346-
empty_x,
347-
empty_y,
348-
full_x,
349-
full_y,
350-
xs_ptr,
351-
z_ptr,
352-
ys_ptr,
353-
xs_sB,
354-
xs_sM,
355-
xs_sKb,
356-
z_sB,
357-
z_sM,
358-
z_sN,
359-
B,
360-
M,
361-
M_aligned,
362-
N,
363-
K,
364-
BLOCK_M,
365-
BLOCK_N,
366-
BLOCK_K,
367-
TILE_ORDER,
368-
SWAP_AB,
369-
num_stages,
370-
num_sms,
344+
_tle_w8a8_block_fp8_bmm_compute_partition,
345+
(
346+
x_smem,
347+
y_smem,
348+
empty_x,
349+
empty_y,
350+
full_x,
351+
full_y,
352+
xs_ptr,
353+
z_ptr,
354+
ys_ptr,
355+
xs_sB,
356+
xs_sM,
357+
xs_sKb,
358+
z_sB,
359+
z_sM,
360+
z_sN,
361+
B,
362+
M,
363+
M_aligned,
364+
N,
365+
K,
366+
BLOCK_M,
367+
BLOCK_N,
368+
BLOCK_K,
369+
TILE_ORDER,
370+
SWAP_AB,
371+
num_stages,
372+
num_sms,
373+
),
371374
),
372-
),
373-
(
374-
_tle_w8a8_block_fp8_bmm_load_partition,
375375
(
376-
x_desc,
377-
y_desc,
378-
x_smem,
379-
y_smem,
380-
empty_x,
381-
empty_y,
382-
full_x,
383-
full_y,
384-
B,
385-
M,
386-
M_aligned,
387-
N,
388-
K,
389-
BLOCK_M,
390-
BLOCK_N,
391-
BLOCK_K,
392-
TILE_ORDER,
393-
num_stages,
394-
num_sms,
376+
_tle_w8a8_block_fp8_bmm_load_partition,
377+
(
378+
x_desc,
379+
y_desc,
380+
x_smem,
381+
y_smem,
382+
empty_x,
383+
empty_y,
384+
full_x,
385+
full_y,
386+
B,
387+
M,
388+
M_aligned,
389+
N,
390+
K,
391+
BLOCK_M,
392+
BLOCK_N,
393+
BLOCK_K,
394+
TILE_ORDER,
395+
num_stages,
396+
num_sms,
397+
),
395398
),
396-
),
397-
],
398-
[1],
399-
[24],
400-
)
399+
],
400+
[1],
401+
[24],
402+
)
401403

402404

403405
@libentry()

0 commit comments

Comments
 (0)