bugfix: Fix arguments of plan for split QK/VO head dims (#795)

#765 introduced changes to the API of `plan`, including renaming
`head_dim` to `head_dim_qk` and adding `head_dim_vo`. However, some
calling sites were not updated to reflect these changes, resulting in
failing unit tests.

This PR addresses the issue by updating the relevant calls, which should
resolve the following unit test failures after merging:

- `tests/test_block_sparse.py::test_block_sparse_attention`
-
`tests/test_non_contiguous_prefill.py::test_batch_paged_prefill_packed_input`

---------

Signed-off-by: abmfy <[email protected]>

Author: abmfy

flashinfer/sparse.py

@@ -355,6 +355,7 @@ def plan(
                 q_data_type,
                 indptr.dtype,
                 head_dim,
+                head_dim,
                 PosEncodingMode[pos_encoding_mode].value,
                 False,  # use_sliding_window
                 logits_soft_cap > 0,  # use_logits_soft_cap
@@ -374,6 +375,7 @@ def plan(
                     -1,  # window_left
                     logits_soft_cap,  # logits_soft_cap
                     head_dim,
+                    head_dim,
                     torch.empty(0, dtype=q_data_type),
                     torch.empty(0, dtype=kv_data_type),
                     get_cuda_stream(device),
@@ -442,6 +444,7 @@ def plan(
                     self.C,  # page_size
                     False,  # is_cuda_graph_enabled,
                     head_dim,
+                    head_dim,
                     causal,
                     get_cuda_stream(device),
                 )

tests/test_non_contiguous_prefill.py

@@ -177,7 +177,7 @@ def test_batch_paged_prefill_packed_input(
         paged_kv_last_page_len=paged_kv_last_page_len,
         num_qo_heads=num_qo_heads,
         num_kv_heads=num_kv_heads,
-        head_dim=head_dim,
+        head_dim_qk=head_dim,
         page_size=page_size,
         causal=causal,
     )