Skip to content

Commit 45ef6f3

Browse files
valarLipvalarLip
authored
feat(gemm_a8w8_blockscale_bpreshuffle): cold-cache retune triton routing for DeepSeek MLA decode-M (#4264)
Re-measure a8w8 blockscale B-preshuffle triton routing cache-rotated (the #4202 tuning used a no-arg lambda that defeated perftest arg-rotation -> hot-cache optimistic + suboptimal configs), and validate end-to-end via the op-test. - op-test: add run_triton + triton columns; drop -m choices allowlist; fail-loud - wqkv_a (N=2048,K=7168): fix #4202 hot-cache regression (BM4/8,KSPLIT4 -> BM16, KSPLIT8,stages3); route M in {1,2,4,8,16,32,48,64,80,128}, all err=0 - wo_b (N=7168): new routing tp1 K=16384 M<=64, tp4 K=4096 M<=32 (+new JSON); tp8 K=2048 unchanged (triton loses on shallow K) - tuned CSVs: unify bw column to TB/s (asm/ck/cktile rows were GB/s) - wq_b unchanged: #4202 routing verified correct cold-cache
1 parent a773be2 commit 45ef6f3

6 files changed

Lines changed: 2243 additions & 2143 deletions

‎aiter/configs/model_configs/a8w8_blockscale_bpreshuffle_tuned_gemm_dsv3.csv‎

Lines changed: 1405 additions & 1405 deletions
Large diffs are not rendered by default.

‎aiter/configs/model_configs/dsv4_a8w8_blockscale_bpreshuffle_tuned_gemm.csv‎

Lines changed: 680 additions & 680 deletions
Large diffs are not rendered by default.

‎aiter/ops/triton/configs/gemm/gfx950-GEMM-A8W8_BLOCKSCALE_PRESHUFFLED-N=2048-K=7168.json‎

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -12,76 +12,76 @@
1212
"NUM_KSPLIT": 4
1313
},
1414
"M_LEQ_4": {
15-
"BLOCK_SIZE_M": 8,
15+
"BLOCK_SIZE_M": 16,
1616
"BLOCK_SIZE_N": 16,
1717
"BLOCK_SIZE_K": 128,
1818
"GROUP_SIZE_M": 1,
1919
"num_warps": 1,
20-
"num_stages": 2,
20+
"num_stages": 3,
2121
"waves_per_eu": 2,
2222
"matrix_instr_nonkdim": 16,
2323
"cache_modifier": ".cg",
24-
"NUM_KSPLIT": 4
24+
"NUM_KSPLIT": 8
2525
},
2626
"M_LEQ_8": {
27-
"BLOCK_SIZE_M": 4,
27+
"BLOCK_SIZE_M": 16,
2828
"BLOCK_SIZE_N": 16,
2929
"BLOCK_SIZE_K": 128,
3030
"GROUP_SIZE_M": 1,
3131
"num_warps": 1,
32-
"num_stages": 2,
33-
"waves_per_eu": 2,
32+
"num_stages": 3,
33+
"waves_per_eu": 1,
3434
"matrix_instr_nonkdim": 16,
3535
"cache_modifier": ".cg",
36-
"NUM_KSPLIT": 4
36+
"NUM_KSPLIT": 8
3737
},
3838
"M_LEQ_16": {
3939
"BLOCK_SIZE_M": 16,
4040
"BLOCK_SIZE_N": 16,
4141
"BLOCK_SIZE_K": 128,
4242
"GROUP_SIZE_M": 1,
4343
"num_warps": 1,
44-
"num_stages": 2,
44+
"num_stages": 3,
4545
"waves_per_eu": 2,
4646
"matrix_instr_nonkdim": 16,
4747
"cache_modifier": ".cg",
48-
"NUM_KSPLIT": 4
48+
"NUM_KSPLIT": 8
4949
},
5050
"M_LEQ_32": {
5151
"BLOCK_SIZE_M": 16,
5252
"BLOCK_SIZE_N": 16,
5353
"BLOCK_SIZE_K": 128,
5454
"GROUP_SIZE_M": 1,
5555
"num_warps": 1,
56-
"num_stages": 2,
56+
"num_stages": 3,
5757
"waves_per_eu": 2,
5858
"matrix_instr_nonkdim": 16,
5959
"cache_modifier": ".cg",
60-
"NUM_KSPLIT": 4
60+
"NUM_KSPLIT": 8
6161
},
6262
"M_LEQ_64": {
6363
"BLOCK_SIZE_M": 32,
6464
"BLOCK_SIZE_N": 64,
6565
"BLOCK_SIZE_K": 128,
6666
"GROUP_SIZE_M": 1,
6767
"num_warps": 4,
68-
"num_stages": 2,
68+
"num_stages": 3,
6969
"waves_per_eu": 2,
7070
"matrix_instr_nonkdim": 16,
7171
"cache_modifier": ".cg",
72-
"NUM_KSPLIT": 4
72+
"NUM_KSPLIT": 8
7373
},
7474
"M_LEQ_128": {
75-
"BLOCK_SIZE_M": 32,
75+
"BLOCK_SIZE_M": 128,
7676
"BLOCK_SIZE_N": 32,
7777
"BLOCK_SIZE_K": 128,
7878
"GROUP_SIZE_M": 1,
79-
"num_warps": 1,
80-
"num_stages": 2,
79+
"num_warps": 4,
80+
"num_stages": 3,
8181
"waves_per_eu": 2,
8282
"matrix_instr_nonkdim": 16,
8383
"cache_modifier": ".cg",
84-
"NUM_KSPLIT": 4
84+
"NUM_KSPLIT": 8
8585
},
8686
"M_LEQ_256": {
8787
"BLOCK_SIZE_M": 64,
@@ -107,4 +107,4 @@
107107
"cache_modifier": null,
108108
"NUM_KSPLIT": 1
109109
}
110-
}
110+
}

‎aiter/ops/triton/configs/gemm/gfx950-GEMM-A8W8_BLOCKSCALE_PRESHUFFLED-N=7168-K=16384.json‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"M_LEQ_8": {
3-
"BLOCK_SIZE_M": 8,
3+
"BLOCK_SIZE_M": 16,
44
"BLOCK_SIZE_N": 32,
55
"BLOCK_SIZE_K": 128,
66
"GROUP_SIZE_M": 1,
@@ -13,12 +13,12 @@
1313
},
1414
"M_LEQ_16": {
1515
"BLOCK_SIZE_M": 16,
16-
"BLOCK_SIZE_N": 16,
16+
"BLOCK_SIZE_N": 32,
1717
"BLOCK_SIZE_K": 128,
1818
"GROUP_SIZE_M": 1,
19-
"num_warps": 1,
20-
"num_stages": 3,
21-
"waves_per_eu": 8,
19+
"num_warps": 2,
20+
"num_stages": 2,
21+
"waves_per_eu": 4,
2222
"matrix_instr_nonkdim": 16,
2323
"cache_modifier": ".cg",
2424
"NUM_KSPLIT": 8
@@ -29,20 +29,20 @@
2929
"BLOCK_SIZE_K": 128,
3030
"GROUP_SIZE_M": 1,
3131
"num_warps": 1,
32-
"num_stages": 3,
33-
"waves_per_eu": 4,
32+
"num_stages": 2,
33+
"waves_per_eu": 2,
3434
"matrix_instr_nonkdim": 16,
3535
"cache_modifier": ".cg",
3636
"NUM_KSPLIT": 8
3737
},
3838
"M_LEQ_64": {
3939
"BLOCK_SIZE_M": 64,
40-
"BLOCK_SIZE_N": 64,
40+
"BLOCK_SIZE_N": 128,
4141
"BLOCK_SIZE_K": 128,
4242
"GROUP_SIZE_M": 1,
43-
"num_warps": 1,
44-
"num_stages": 2,
45-
"waves_per_eu": 1,
43+
"num_warps": 4,
44+
"num_stages": 3,
45+
"waves_per_eu": 2,
4646
"matrix_instr_nonkdim": 16,
4747
"cache_modifier": ".cg",
4848
"NUM_KSPLIT": 8

‎aiter/ops/triton/configs/gemm/gfx950-GEMM-A8W8_BLOCKSCALE_PRESHUFFLED-N=7168-K=4096.json‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
{
2+
"M_LEQ_4": {
3+
"BLOCK_SIZE_M": 16,
4+
"BLOCK_SIZE_N": 32,
5+
"BLOCK_SIZE_K": 128,
6+
"GROUP_SIZE_M": 1,
7+
"num_warps": 2,
8+
"num_stages": 3,
9+
"waves_per_eu": 4,
10+
"matrix_instr_nonkdim": 16,
11+
"cache_modifier": ".cg",
12+
"NUM_KSPLIT": 8
13+
},
14+
"M_LEQ_8": {
15+
"BLOCK_SIZE_M": 16,
16+
"BLOCK_SIZE_N": 32,
17+
"BLOCK_SIZE_K": 128,
18+
"GROUP_SIZE_M": 1,
19+
"num_warps": 2,
20+
"num_stages": 3,
21+
"waves_per_eu": 4,
22+
"matrix_instr_nonkdim": 16,
23+
"cache_modifier": ".cg",
24+
"NUM_KSPLIT": 8
25+
},
26+
"M_LEQ_16": {
27+
"BLOCK_SIZE_M": 16,
28+
"BLOCK_SIZE_N": 32,
29+
"BLOCK_SIZE_K": 128,
30+
"GROUP_SIZE_M": 1,
31+
"num_warps": 2,
32+
"num_stages": 3,
33+
"waves_per_eu": 4,
34+
"matrix_instr_nonkdim": 16,
35+
"cache_modifier": ".cg",
36+
"NUM_KSPLIT": 8
37+
},
38+
"M_LEQ_32": {
39+
"BLOCK_SIZE_M": 32,
40+
"BLOCK_SIZE_N": 32,
41+
"BLOCK_SIZE_K": 128,
42+
"GROUP_SIZE_M": 1,
43+
"num_warps": 2,
44+
"num_stages": 3,
45+
"waves_per_eu": 2,
46+
"matrix_instr_nonkdim": 16,
47+
"cache_modifier": ".cg",
48+
"NUM_KSPLIT": 4
49+
},
50+
"M_LEQ_64": {
51+
"BLOCK_SIZE_M": 32,
52+
"BLOCK_SIZE_N": 128,
53+
"BLOCK_SIZE_K": 128,
54+
"GROUP_SIZE_M": 1,
55+
"num_warps": 4,
56+
"num_stages": 2,
57+
"waves_per_eu": 4,
58+
"matrix_instr_nonkdim": 16,
59+
"cache_modifier": ".cg",
60+
"NUM_KSPLIT": 4
61+
},
62+
"M_LEQ_128": {
63+
"BLOCK_SIZE_M": 64,
64+
"BLOCK_SIZE_N": 64,
65+
"BLOCK_SIZE_K": 128,
66+
"GROUP_SIZE_M": 1,
67+
"num_warps": 4,
68+
"num_stages": 3,
69+
"waves_per_eu": 4,
70+
"matrix_instr_nonkdim": 16,
71+
"cache_modifier": ".cg",
72+
"NUM_KSPLIT": 2
73+
},
74+
"M_LEQ_256": {
75+
"BLOCK_SIZE_M": 128,
76+
"BLOCK_SIZE_N": 128,
77+
"BLOCK_SIZE_K": 128,
78+
"GROUP_SIZE_M": 1,
79+
"num_warps": 8,
80+
"num_stages": 2,
81+
"waves_per_eu": 2,
82+
"matrix_instr_nonkdim": 16,
83+
"cache_modifier": ".cg",
84+
"NUM_KSPLIT": 2
85+
},
86+
"any": {
87+
"BLOCK_SIZE_M": 128,
88+
"BLOCK_SIZE_N": 128,
89+
"BLOCK_SIZE_K": 128,
90+
"GROUP_SIZE_M": 1,
91+
"num_warps": 8,
92+
"num_stages": 2,
93+
"waves_per_eu": 4,
94+
"matrix_instr_nonkdim": 16,
95+
"cache_modifier": ".cg",
96+
"NUM_KSPLIT": 1
97+
}
98+
}

‎op_tests/test_gemm_a8w8_blockscale.py‎

Lines changed: 31 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,26 @@ def run_gemm_bpreshuffle(x, weightshuffle, x_scale, w_scale, dtype=dtypes.bf16):
6060
)
6161

6262

63+
@perftest(num_iters=TEST_NUM_ITERS)
64+
def run_triton(x, weightshuffle, x_scale, w_scale, dtype=dtypes.bf16, backend=None):
65+
# Direct call into the triton preshuffle kernel, mirroring the dispatch in
66+
# gemm_a8w8_blockscale_bpreshuffle: reshape the (n, k) preshuffled weight to
67+
# (n // 16, k * 16) and pass the transposed x_scale.
68+
from aiter.ops.triton.gemm.basic.gemm_a8w8_blockscale import (
69+
gemm_a8w8_blockscale_preshuffle,
70+
)
71+
72+
n, k = weightshuffle.shape
73+
return gemm_a8w8_blockscale_preshuffle(
74+
x,
75+
weightshuffle.reshape(n // 16, k * 16),
76+
x_scale,
77+
w_scale,
78+
dtype=dtype,
79+
backend=backend,
80+
)
81+
82+
6383
@benchmark()
6484
def test_gemm(dtype, m, n, k, ck_preshuffle=True):
6585
ret = {}
@@ -97,6 +117,17 @@ def test_gemm(dtype, m, n, k, ck_preshuffle=True):
97117
ret[f"{tag} err"] = err_asm
98118
ret["asm/ck"] = avg_c / avg_b
99119

120+
# Triton path requires a preshuffled weight. When not preshuffled we simply omit
121+
# these columns; pd.DataFrame NaN-fills them for those rows in the summary.
122+
if ck_preshuffle:
123+
d, avg_d = run_triton(x, gemm_weight, x_scale_t, w_scale, dtype)
124+
err_triton = checkAllclose(a, d, msg="triton", catastrophic_check=True)
125+
ret["triton us"] = avg_d
126+
ret["triton TFLOPS"] = m * n * k * 2 / avg_d / 1e6
127+
ret["triton TB/s"] = (x.nbytes + weight.nbytes) / avg_d / 1e6
128+
ret["triton err"] = err_triton
129+
ret["triton/ck"] = avg_d / avg_b
130+
100131
return ret
101132

102133

@@ -201,35 +232,6 @@ def test_splitk_correctness(m=4, n=2112, k=7168, dtype=dtypes.bf16, splitK=1):
201232
"-m",
202233
type=int,
203234
nargs="*",
204-
choices=[
205-
1,
206-
2,
207-
4,
208-
8,
209-
16,
210-
32,
211-
64,
212-
96,
213-
128,
214-
160,
215-
192,
216-
224,
217-
256,
218-
288,
219-
320,
220-
352,
221-
384,
222-
416,
223-
448,
224-
480,
225-
512,
226-
1024,
227-
2048,
228-
4096,
229-
6144,
230-
8192,
231-
10240,
232-
],
233235
default=[
234236
1,
235237
2,

0 commit comments

Comments
 (0)