[yuchenwang3]

Packed QSA currently scores every query against every complete block in the pack, then masks blocks from other documents. For packs containing several documents, most of those scores are discarded.

This scores each document against its own complete blocks while retaining the existing query chunk budget, causal visibility, stable tie order, pack-space indices, and padded output width. It leaves key pooling, RoPE, and partial-block tails in the existing path. The document loop adds a host read of cu_seqlens.

The regression tests compare selected IDs with an independent globally masked reference, including empty documents, short documents, tails, ties, and different chunk sizes. A four-document test records 1,024 query/block pairs instead of 4,096. This is a work-count check, not a throughput measurement.

Validation:

• 17 CPU helper cases passed on macOS, PyTorch 2.13.0. The local runner loads the actual helper definitions without CUDA-only Transformer Engine imports; it does not validate the full module.
• pre-commit run --all-files passed.
• 18 tests passed with no skips on NVIDIA L20C, Python 3.12.3, PyTorch 2.11.0+cu130 / CUDA 13.0: the new selection cases and the existing real TELinear packed-indexer unit. Document-local allowed token sets matched the per-document reference.
• A BF16 full-selector benchmark with real TELinear, matching seeded inputs/weights, five warmups and 20 synchronized CUDA-event samples gave the following results. Selected IDs matched exactly in both packs. These are 8K synthetic selector measurements; they do not establish full-model or multi-node speedups. CP/TP and full-model training have not been validated here.

Pack	Base	Patched	Base / patched
16 documents × 512 tokens	18.783 ms	6.954 ms	2.70
4 documents × 2048 tokens	3.203 ms	3.108 ms	1.03

Implementation and tests were generated with OpenAI Codex.