Skip to content

Commit a25429f

Browse files
authored
[UT][CHUNK_PREFILL]add ut mini scope for fmha (vllm-project#46)
Signed-off-by: Yizhou Wang <yizhou.wang@intel.com>
1 parent b26c1bf commit a25429f

1 file changed

Lines changed: 10 additions & 0 deletions

File tree

tests/flash_attn/test_flash_attn_varlen_func.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,16 @@ def ref_paged_attn(query: torch.Tensor,
7575
return torch.cat(outputs, dim=0)
7676

7777

78+
#override pytest parameters when enable mini pytest
79+
MINI_PYTEST_PARAMS = {
80+
"default": {
81+
"seq_lens": [[(1, 1328), (5, 18), (129, 463)]],
82+
"head_size": [64, 128],
83+
"num_blocks": [64],
84+
}
85+
}
86+
87+
7888
@pytest.mark.parametrize("seq_lens",
7989
[[(1, 1328), (5, 18),
8090
(129, 463)], [(1, 523), (1, 37), (1, 2011)]])

0 commit comments

Comments
 (0)