A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Misc] parametrize 'dtype' in test_flash_mla (#22641)
Signed-off-by: RUTHLESS-BOT <wujiafeng@cmbchina.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
RUTHLESS-BOT committed
Aug 13, 2025 at 04:31 UTC
53c730286c5ad86a6d78d4a4d8a2cd7042725d24
1 file changed
+3
-4
tests/kernels/attention/test_flashmla.py
+3
-4
@@ -35,11 +35,10 @@ FLASH_MLA_UNSUPPORTED_REASON = is_flashmla_supported()[1] \
35
@pytest.mark.parametrize("block_size", [64])
36
@pytest.mark.parametrize("causal", [True])
37
@pytest.mark.parametrize("varlen", [False, True])
38
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
39
@torch.inference_mode()
40
def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,
40
- varlen):
41
- # TODO: parametrize using pytest
42
- dtype = torch.bfloat16
41
+ varlen, dtype):
42
device = torch.device("cuda:0")
43
torch.set_default_dtype(dtype)
44
torch.set_default_device(device)
@@ -48,7 +47,7 @@ def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,
47
random.seed(0)
48
49
print(f"{b=}, {s_q=}, {mean_sk=}, {h_q=}, {h_kv=}, "
51
- f"{d=}, {dv=}, {causal=}, {varlen=}")
50
+ f"{d=}, {dv=}, {causal=}, {varlen=}, {dtype=}")
51
52
cache_seqlens = torch.full((b, ), mean_sk, dtype=torch.int32)
53
if varlen: