Skip to content

GitLab

Commit 3357c074 authored by zhuwenwen's avatar zhuwenwen
Browse files

update test_flash_mla.py

parent 1de31a24
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
tests/kernels/test_flashmla.py
View file @ 3357c074
...@@ -33,7 +33,8 @@ FLASH_MLA_UNSUPPORTED_REASON = is_flashmla_supported()[1] \ ...@@ -33,7 +33,8 @@ FLASH_MLA_UNSUPPORTED_REASON = is_flashmla_supported()[1] \
@pytest.mark.parametrize("dv", [512]) @pytest.mark.parametrize("dv", [512])
@pytest.mark.parametrize("block_size", [64]) @pytest.mark.parametrize("block_size", [64])
@pytest.mark.parametrize("causal", [True]) @pytest.mark.parametrize("causal", [True])
@pytest.mark.parametrize("varlen", [False, True]) # @pytest.mark.parametrize("varlen", [False, True])
@pytest.mark.parametrize("varlen", [True])
@torch.inference_mode() @torch.inference_mode()
def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,
varlen): varlen):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment