添加测试用例

tests/test_flash_mla_kvfp8.py

@@ -210,13 +210,13 @@ def main(torch_dtype, is_prof=False):
     # '''
     for b in [3, 6, 9, 12, 15, 18, 21, 24]:
         for s in [111, 112, 123, 1234, 432, 4325, 4000, 8192, 11111]:
-            for h_q in [16]:
+            for h_q in [16, 128]:
                 for s_q in [1, 2, 3]:  # MTP = 1, 2
                     for varlen in [False, True]:
                         test_flash_mla_fp8_e5m2(b, s_q, s, h_q, h_kv, d, dv, causal, varlen,True)
     for b in [3, 6, 9, 12, 15, 18, 21, 24, 32, 64, 128, 256]:
         for s in [4000]:
-            for h_q in [16]:
+            for h_q in [16, 128]:
                 for s_q in [1]:  # MTP = 1, 2
                     for varlen in [False]:
                         test_flash_mla_fp8_e5m2(b, s_q, s, h_q, h_kv, d, dv, causal, varlen)