Skip to content

Commit e6bd49b

Browse files
Tiannn967Tiannn967
authored
Optimize MM kernels and autotuning (#5407)
* add gemv path for mm * add split-k mm path * add general mm stride and boundary fast paths * specialize contiguous mm layout * fix split-k autotune output reset * separate general mm from split-k * remove contiguous mm kernel * add mm flagtune expand config * optimize MetaX mm dispatch and tuning * support specialized mm tuning configs * refine dense mm autotune pruning * add specialized mm tuning spaces * add bf16 nt mm kernel * add k-parallel gemv path * add two-step split-k path * dispatch specialized mm kernels * format MetaX mm implementation * fix(metax): remove global Triton swizzle override * clean up mm tuning and shape logging * isolate sparse attention compiler switches * trim specialized mm default configs * generalize NN and NT mm dtype support Support FP16, BF16, FP32, and mixed input dtypes in the specialized kernels, and trim their default tuning candidates.
1 parent e2cb329 commit e6bd49b

7 files changed

Lines changed: 1737 additions & 116 deletions

File tree

‎src/flag_gems/runtime/backend/_metax/fused/sparse_attention.py‎

Lines changed: 41 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -29,18 +29,34 @@
2929
"""
3030

3131
import os
32-
33-
os.environ.setdefault("TRITON_DISABLE_SWIZZLE", "1")
34-
# MetaX (mcTriton) compiler-pass enable flags
35-
os.environ.setdefault("TRITON_ENABLE_MACA_OPT_MOVE_DOT_OPERANDS_OUT_LOOP", "1")
36-
os.environ.setdefault("TRITON_ENABLE_MACA_MERGE_CONVERT_LAYOUT", "1")
37-
os.environ.setdefault("TRITON_ENABLE_SMEM_OFFSET_CACHE", "1")
38-
os.environ.setdefault("TRITON_ENABLE_BSM_INDEX_OPT", "1")
32+
from contextlib import contextmanager
3933

4034
import torch # noqa: E402
4135
import triton # noqa: E402
4236
import triton.language as tl # noqa: E402
4337

38+
_SPARSE_ATTN_COMPILER_ENV = {
39+
"TRITON_DISABLE_SWIZZLE": "1",
40+
"TRITON_ENABLE_MACA_OPT_MOVE_DOT_OPERANDS_OUT_LOOP": "1",
41+
"TRITON_ENABLE_MACA_MERGE_CONVERT_LAYOUT": "1",
42+
"TRITON_ENABLE_SMEM_OFFSET_CACHE": "1",
43+
"TRITON_ENABLE_BSM_INDEX_OPT": "1",
44+
}
45+
46+
47+
@contextmanager
48+
def _sparse_attn_compiler_env():
49+
previous = {name: os.environ.get(name) for name in _SPARSE_ATTN_COMPILER_ENV}
50+
os.environ.update(_SPARSE_ATTN_COMPILER_ENV)
51+
try:
52+
yield
53+
finally:
54+
for name, value in previous.items():
55+
if value is None:
56+
os.environ.pop(name, None)
57+
else:
58+
os.environ[name] = value
59+
4460

4561
# ===========================================================================
4662
# Kernel 1: SIMPLE bf16 kernel (fallback)
@@ -416,7 +432,7 @@ def _sparse_attn_kernel_v_chunked(
416432
# ===========================================================================
417433
# Python wrapper — multi-tier dispatch
418434
# ===========================================================================
419-
def sparse_attn_triton(
435+
def _sparse_attn_triton_impl(
420436
q: torch.Tensor,
421437
kv: torch.Tensor,
422438
attn_sink: torch.Tensor,
@@ -634,3 +650,20 @@ def sparse_attn_triton(
634650
num_warps=2,
635651
)
636652
return o
653+
654+
655+
def sparse_attn_triton(
656+
q: torch.Tensor,
657+
kv: torch.Tensor,
658+
attn_sink: torch.Tensor,
659+
topk_idxs: torch.Tensor,
660+
softmax_scale: float,
661+
) -> torch.Tensor:
662+
with _sparse_attn_compiler_env():
663+
return _sparse_attn_triton_impl(
664+
q,
665+
kv,
666+
attn_sink,
667+
topk_idxs,
668+
softmax_scale,
669+
)

‎src/flag_gems/runtime/backend/_metax/heuristics_config_utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,7 @@ def index_select_heur_block_n(args):
206206

207207

208208
def mm_heur_even_k(args):
209-
return args["K"] % (args["BLOCK_K"] * args["SPLIT_K"]) == 0
209+
return args["K"] % args["BLOCK_K"] == 0
210210

211211

212212
def ones_heur_block_size(args):

‎src/flag_gems/runtime/backend/_metax/mm_metax_expand.yaml‎

Lines changed: 276 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,276 @@
1+
mm:
2+
- config:
3+
param_map:
4+
META:
5+
BLOCK_M: block_m
6+
BLOCK_N: block_n
7+
BLOCK_K: block_k
8+
pipeline: pipeline
9+
scenario: scenario
10+
num_stages: stages
11+
num_warps: warps
12+
block_m:
13+
- 16
14+
- 32
15+
- 64
16+
- 128
17+
- 256
18+
block_n:
19+
- 32
20+
- 64
21+
- 128
22+
- 256
23+
block_k:
24+
- 16
25+
- 32
26+
- 64
27+
- 128
28+
stages:
29+
- 2
30+
- 3
31+
- 4
32+
warps:
33+
- 2
34+
- 4
35+
- 8
36+
pipeline:
37+
- "null"
38+
- "basic"
39+
- "cpasync"
40+
scenario:
41+
- ""
42+
- strategy:
43+
M: default
44+
N: default
45+
K: default
46+
stride_am: default
47+
stride_bk: default
48+
49+
mm_nt:
50+
- config:
51+
param_map:
52+
META:
53+
BLOCK_M: block_m
54+
BLOCK_N: block_n
55+
BLOCK_K: block_k
56+
pipeline: pipeline
57+
scenario: scenario
58+
num_stages: stages
59+
num_warps: warps
60+
block_m:
61+
- 16
62+
- 32
63+
- 64
64+
- 128
65+
- 256
66+
block_n:
67+
- 32
68+
- 64
69+
- 128
70+
- 256
71+
block_k:
72+
- 32
73+
- 64
74+
- 128
75+
pipeline:
76+
- "basic"
77+
- "cpasync"
78+
warps:
79+
- 4
80+
- 8
81+
stages:
82+
- 2
83+
- 3
84+
- 4
85+
scenario:
86+
- ""
87+
- strategy:
88+
M: align32
89+
N: default
90+
K: default
91+
92+
mm_splitk:
93+
- config:
94+
param_map:
95+
META:
96+
BLOCK_M: block_m
97+
BLOCK_N: block_n
98+
BLOCK_K: block_k
99+
SPLIT_K: split_k
100+
pipeline: pipeline
101+
num_stages: stages
102+
num_warps: warps
103+
block_m:
104+
- 16
105+
block_n:
106+
- 16
107+
- 64
108+
block_k:
109+
- 64
110+
- 128
111+
- 256
112+
split_k:
113+
- 8
114+
- 16
115+
stages:
116+
- 2
117+
- 3
118+
- 6
119+
- 8
120+
warps:
121+
- 2
122+
- 4
123+
pipeline:
124+
- "null"
125+
- "basic"
126+
- strategy:
127+
M: align32
128+
N: default
129+
K: default
130+
stride_am: default
131+
stride_bk: default
132+
133+
mm_splitk_two_step:
134+
- config:
135+
param_map:
136+
META:
137+
BLOCK_M: block_m
138+
BLOCK_N: block_n
139+
BLOCK_K: block_k
140+
pipeline: pipeline
141+
num_stages: stages
142+
num_warps: warps
143+
block_m:
144+
- 16
145+
- 32
146+
block_n:
147+
- 16
148+
- 32
149+
- 64
150+
- 128
151+
block_k:
152+
- 64
153+
stages:
154+
- 2
155+
- 4
156+
warps:
157+
- 2
158+
- 4
159+
pipeline:
160+
- "null"
161+
- "basic"
162+
- "cpasync"
163+
- strategy:
164+
M: default
165+
N: default
166+
K: default
167+
stride_am: default
168+
stride_bk: default
169+
170+
gemv:
171+
- config:
172+
param_map:
173+
META:
174+
BLOCK_M: block_m
175+
BLOCK_K: block_k
176+
num_stages: stages
177+
num_warps: warps
178+
block_m:
179+
- 8
180+
- 16
181+
- 32
182+
block_k:
183+
- 256
184+
stages:
185+
- 2
186+
- 3
187+
- 4
188+
- 5
189+
- 6
190+
- 7
191+
- 8
192+
warps:
193+
- 1
194+
- 2
195+
- 4
196+
- 8
197+
- strategy:
198+
M: align32
199+
K: default
200+
stride_am: default
201+
stride_bk: default
202+
203+
gemv_k_parallel:
204+
- config:
205+
param_map:
206+
META:
207+
BLOCK_M: block_m
208+
BLOCK_K: block_k
209+
num_stages: stages
210+
num_warps: warps
211+
block_m:
212+
- 4
213+
- 8
214+
- 16
215+
- 32
216+
block_k:
217+
- 128
218+
- 256
219+
stages:
220+
- 2
221+
- 4
222+
warps:
223+
- 1
224+
- 2
225+
- 4
226+
- strategy:
227+
M: default
228+
K: default
229+
stride_am: default
230+
stride_bk: default
231+
232+
mm_nn:
233+
- config:
234+
param_map:
235+
META:
236+
BLOCK_M: block_m
237+
BLOCK_N: block_n
238+
BLOCK_K: block_k
239+
pipeline: pipeline
240+
scenario: scenario
241+
num_stages: stages
242+
num_warps: warps
243+
block_m:
244+
- 16
245+
- 32
246+
- 64
247+
- 128
248+
- 256
249+
block_n:
250+
- 32
251+
- 64
252+
- 128
253+
- 256
254+
block_k:
255+
- 16
256+
- 32
257+
- 64
258+
- 128
259+
stages:
260+
- 2
261+
- 3
262+
- 4
263+
warps:
264+
- 2
265+
- 4
266+
- 8
267+
pipeline:
268+
- "null"
269+
- "basic"
270+
- "cpasync"
271+
scenario:
272+
- ""
273+
- strategy:
274+
M: default
275+
N: default
276+
K: default

0 commit comments

Comments
 (0)