bugfix: fix batch prefill attention kernel unittests (#781)

yzh119 · web-flow · commit 74a40541f934 · 2025-02-04T14:14:15.000-05:00
This pull request includes a small change to the
`tests/test_batch_prefill_kernels.py` file. The change modifies the
`custom_mask` parameter in the
`test_batch_prefill_with_paged_kv_cache_custom_mask` function to have a
default value.

*
[`tests/test_batch_prefill_kernels.py`](diffhunk://#diff-f04bd675bf627f5739bd24c7e2eed84ae66aba380b1855a9b13b69c4a000de6fL550-R550):
Changed the `custom_mask` parameter to `custom_mask=custom_mask` to set
a default value.
diff --git a/tests/test_batch_prefill_kernels.py b/tests/test_batch_prefill_kernels.py
@@ -547,7 +547,7 @@ def test_batch_prefill_with_paged_kv_cache_custom_mask(
         num_kv_heads,
         head_dim,
         page_size,
-        custom_mask,
+        custom_mask=custom_mask,
         pos_encoding_mode=pos_encoding_mode,
         logits_soft_cap=logits_soft_cap,
     )

Original file line number	Diff line number	Diff line change
`@@ -547,7 +547,7 @@ def test_batch_prefill_with_paged_kv_cache_custom_mask(`
`547`	`547`	`num_kv_heads,`
`548`	`548`	`head_dim,`
`549`	`549`	`page_size,`
`550`		`- custom_mask,`
	`550`	`+ custom_mask=custom_mask,`
`551`	`551`	`pos_encoding_mode=pos_encoding_mode,`
`552`	`552`	`logits_soft_cap=logits_soft_cap,`
`553`	`553`	`)`