DeepSeek/FlashMLA
1
0
Fork 0
mirror of https://github.com/deepseek-ai/FlashMLA synced 2025-06-26 18:15:54 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add new testcase (s_k = 16384)

Browse Source
This commit is contained in:
Shengyu Liu 2025-04-21 17:23:19 +08:00
parent 287061ec34
commit 69b64823cc
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_flash_mla.py
View File

@ -127,7 +127,7 @@ def main(torch_dtype):
causal = True causal = True
for b in [128]: for b in [128]:
for s in [4096, 8192]: for s in [4096, 8192, 16384]:
for h_q in [16, 32, 64, 128]: # TP = 8, 4, 2, 1 for h_q in [16, 32, 64, 128]: # TP = 8, 4, 2, 1
for s_q in [1, 2]: # MTP = 1, 2 for s_q in [1, 2]: # MTP = 1, 2
for varlen in [False, True]: for varlen in [False, True]: