Fix arguments of plan for split QK/VO head dims

[abmfy]

#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

[yzh119]

Thanks for the update, yes we should update the sparse attention plan functions.