A high-throughput and memory-efficient inference and serving engine for LLMs
[Bug] Fix test flashmla for DSv4 (#45052)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 11, 2026 at 16:22 UTC
5a6c7b7ab569f49491b5428a7983be5b17b85378
1 file changed
+5
-3
tests/kernels/attention/test_flashmla_sparse.py
+5
-3
@@ -29,8 +29,10 @@ def test_sparse_flashmla_metadata_smoke():
29
topk=topk,
30
is_fp8_kvcache=True,
31
)
32
- assert tile_md.dtype == torch.int32
33
- assert num_splits.dtype == torch.int32
32
+ assert isinstance(tile_md, fm.FlashMLASchedMeta)
33
+ assert tile_md.tile_scheduler_metadata is None
34
+ assert tile_md.num_splits is None
35
+ assert num_splits is None
36
37
38
def test_sparse_flashmla_decode_smoke():
@@ -116,7 +118,7 @@ def test_sparse_flashmla_prefill_smoke():
118
kv = torch.zeros((s_kv, h_kv, d_qk), dtype=torch.bfloat16, device=device)
119
indices = torch.zeros((s_q, h_kv, topk), dtype=torch.int32, device=device)
120
119
- out, max_logits, lse = fm.flash_mla_sparse_prefill(q, kv, indices, 1.0, d_v)
121
+ out, max_logits, lse = fm.flash_mla_sparse_fwd(q, kv, indices, 1.0, d_v)
122
assert out.shape == (s_q, h_q, d_v)
123
assert max_logits.shape == (s_q, h_q)
124
assert lse.shape == (s_q, h_q)