[shiaho777]

The sparse kernel used the query count as the key bound (offs_k < T) and built the block bitmap with ceil(T / block_size). Context parallelism holds a 1/cp_size query shard and the full key sequence. The CP path therefore copied the local queries into a buffer as long as the keys and filled the rest with -1. Those rows launch tiles and allocate a bitmap for the whole sequence, then the kernel skips them.

The key bound and the bitmap width now follow k.shape[0]. The query grid still follows the query count. The CP path passes the local queries and the gathered keys directly. When the two lengths are equal, the bounds match the previous kernel.

Checked with python3 -m pytest tests/test_qsa_key_length.py. Two queries and eight keys, block size 4, flag key index 6 in block 1. flake8 is clean. The Triton kernels were not executed here; triton is not installed.

[hjh0119]

NameError: name 'T' is not defined

[shiaho777]

The index path still read T after the bitmap switched to the key length. qsa_sparse_attention_from_indices now sizes those single-sequence bases from q.shape[0].

tests/test_qsa_key_length.py calls that path with 2 queries and 8 keys and checks the bases are one zero per query. python3 -m pytest tests/test_qsa_key_length.py is 2 passed. flake8 is clean on the two files. Head is ed5a867.

[shiaho777]

yapf rejected the single-element shape checks. They are now (2, ). python3 -m pytest tests/test_qsa_key_length.py is still 2 passed, and yapf -d is clean on the two files. Head is f0827a3.