Skip to content

Commit e463220

Browse files
authored
Fix mm tune key error (#2199)
1 parent 9bb5f0a commit e463220

1 file changed

Lines changed: 2 additions & 2 deletions

File tree

  • src/flag_gems/runtime/backend/_nvidia/hopper/ops

src/flag_gems/runtime/backend/_nvidia/hopper/ops/mm.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -179,10 +179,10 @@ def matmul_get_configs(pre_hook=matmul_tma_set_block_size_hook):
179179
configs=matmul_get_configs(pre_hook=None)
180180
if os.environ.get("USE_FLAGTUNE") == "1" and get_expand_config("matmul") != -1
181181
else runtime.get_tuned_config("mm"),
182-
key=["M", "N", "K", "stride_am", "stride_bk", "dtype"],
182+
key=["M", "N", "K", "stride_am", "stride_bk"],
183183
strategy=get_expand_config("matmul")["strategy"]
184184
if os.environ.get("USE_FLAGTUNE") == "1" and get_expand_config("matmul") != -1
185-
else ["default", "default", "default", "default", "default", "default"],
185+
else ["default", "default", "default", "default", "default"],
186186
warmup=5,
187187
rep=10,
188188
)

0 commit comments

Comments
 (0)