update

tests/test_flash_attn.py

@@ -1820,7 +1820,7 @@ def test_flash_attn_output_mla(
 # @pytest.mark.parametrize("d", [32, 59, 64, 80, 96, 111, 128, 160, 192, 224, 256])
 # @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256])
 # @pytest.mark.parametrize("d", [160, 192, 224, 256])
-@pytest.mark.parametrize('d', [128])
+@pytest.mark.parametrize('d', [64, 128])
 @pytest.mark.parametrize(
     "seqlen_q,seqlen_k",
     [

tests/test_unified_attn.py

@@ -1261,7 +1261,7 @@ def cal_diff(x: torch.Tensor, y: torch.Tensor, name: str, use_fp8: bool = False,
 @pytest.mark.parametrize("use_qq_bias", [True, False])  # seqlen_q > seqlen_k 时 skip
 @pytest.mark.parametrize("use_sinks", [True, False])
 @pytest.mark.parametrize("use_mm_prefix", [True, False])
-@pytest.mark.parametrize("d", [256])
+@pytest.mark.parametrize("d", [128, 256])
 @pytest.mark.parametrize(
     "batch_size,seqlen_q,seqlen_k,block_size",
     [