Skip to content

Commit af5291a

Browse files
authored
replace Gluon with TLE and add general and split-K paths for w8a8_block_fp8_bmm (flagos-ai#4794)
1 parent 5a95a79 commit af5291a

5 files changed

Lines changed: 790 additions & 409 deletions

File tree

src/flag_gems/runtime/backend/_nvidia/hopper/general_ops_hopper_expand.yaml

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,3 +315,80 @@ mul:
315315
n_elements: align32
316316
n_cols: default
317317
dtype: default
318+
w8a8_block_fp8_bmm_general:
319+
- config:
320+
param_map:
321+
META:
322+
BLOCK_M: block_m
323+
BLOCK_N: block_n
324+
BLOCK_K: block_k
325+
GROUP_M: group_m
326+
num_stages: stages
327+
num_warps: warps
328+
block_m:
329+
- 16
330+
- 64
331+
block_n:
332+
- 16
333+
- 32
334+
- 64
335+
block_k:
336+
- 128
337+
group_m:
338+
- 1
339+
- 2
340+
- 4
341+
- 8
342+
stages:
343+
- 3
344+
- 4
345+
- 5
346+
- 6
347+
warps:
348+
- 4
349+
- strategy:
350+
B: default
351+
M: align32
352+
N: align32
353+
K: align32
354+
stride_xm: align32
355+
stride_yk: align32
356+
w8a8_block_fp8_bmm_splitk:
357+
- config:
358+
param_map:
359+
META:
360+
BLOCK_M: block_m
361+
BLOCK_N: block_n
362+
BLOCK_K: block_k
363+
SPLIT_K: split_k
364+
num_stages: stages
365+
num_warps: warps
366+
block_m:
367+
- 8
368+
- 16
369+
- 64
370+
block_n:
371+
- 32
372+
- 64
373+
block_k:
374+
- 128
375+
split_k:
376+
- 2
377+
- 4
378+
- 8
379+
- 16
380+
- 32
381+
stages:
382+
- 2
383+
- 3
384+
- 4
385+
warps:
386+
- 4
387+
- 8
388+
- strategy:
389+
B: default
390+
M: align32
391+
N: align32
392+
K: align32
393+
stride_xm: align32
394+
stride_yk: align32

0 commit comments

Comments
 (0)