[ROCm] Add tuned selective_state_update config for AMD MI350 (#48159)
Signed-off-by: Giuseppe Grossi <ggrossi@amd.com>
Giuseppe Grossi committed
Jul 15, 2026 at 03:18 UTC
d119beb1b90804f55e5c1d70f1c1961b8a852bd1
2 files changed
+102
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI350_OAM,cache_dtype=float16.json
new
+51
@@ -0,0 +1,51 @@
1
+{
2
+ "triton_version": "3.6.0",
3
+ "128": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "num_warps": 4
6
+ },
7
+ "256": {
8
+ "BLOCK_SIZE_M": 32,
9
+ "num_warps": 4
10
+ },
11
+ "1024": {
12
+ "BLOCK_SIZE_M": 16,
13
+ "num_warps": 2
14
+ },
15
+ "2048": {
16
+ "BLOCK_SIZE_M": 8,
17
+ "num_warps": 1
18
+ },
19
+ "4096": {
20
+ "BLOCK_SIZE_M": 8,
21
+ "num_warps": 1
22
+ },
23
+ "8192": {
24
+ "BLOCK_SIZE_M": 16,
25
+ "num_warps": 2
26
+ },
27
+ "16384": {
28
+ "BLOCK_SIZE_M": 8,
29
+ "num_warps": 1
30
+ },
31
+ "32768": {
32
+ "BLOCK_SIZE_M": 8,
33
+ "num_warps": 1
34
+ },
35
+ "65536": {
36
+ "BLOCK_SIZE_M": 8,
37
+ "num_warps": 1
38
+ },
39
+ "131072": {
40
+ "BLOCK_SIZE_M": 16,
41
+ "num_warps": 1
42
+ },
43
+ "196608": {
44
+ "BLOCK_SIZE_M": 16,
45
+ "num_warps": 1
46
+ },
47
+ "262144": {
48
+ "BLOCK_SIZE_M": 8,
49
+ "num_warps": 1
50
+ }
51
+}
\ No newline at end of file
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI350_OAM,cache_dtype=float32.json
new
+51
@@ -0,0 +1,51 @@
1
+{
2
+ "triton_version": "3.6.0",
3
+ "128": {
4
+ "BLOCK_SIZE_M": 16,
5
+ "num_warps": 4
6
+ },
7
+ "256": {
8
+ "BLOCK_SIZE_M": 16,
9
+ "num_warps": 4
10
+ },
11
+ "1024": {
12
+ "BLOCK_SIZE_M": 8,
13
+ "num_warps": 1
14
+ },
15
+ "2048": {
16
+ "BLOCK_SIZE_M": 8,
17
+ "num_warps": 1
18
+ },
19
+ "4096": {
20
+ "BLOCK_SIZE_M": 4,
21
+ "num_warps": 1
22
+ },
23
+ "8192": {
24
+ "BLOCK_SIZE_M": 32,
25
+ "num_warps": 2
26
+ },
27
+ "16384": {
28
+ "BLOCK_SIZE_M": 4,
29
+ "num_warps": 1
30
+ },
31
+ "32768": {
32
+ "BLOCK_SIZE_M": 4,
33
+ "num_warps": 1
34
+ },
35
+ "65536": {
36
+ "BLOCK_SIZE_M": 64,
37
+ "num_warps": 8
38
+ },
39
+ "131072": {
40
+ "BLOCK_SIZE_M": 64,
41
+ "num_warps": 8
42
+ },
43
+ "196608": {
44
+ "BLOCK_SIZE_M": 64,
45
+ "num_warps": 2
46
+ },
47
+ "262144": {
48
+ "BLOCK_SIZE_M": 64,
49
+ "num_warps": 8
50
+ }
51
+}
\ No newline at end of file