Repository navigation
Conversation
The prefill indexer score ran one query row per program and padded the 4 index heads to a 16-wide dot, so every row re-read all visible keys and 3/4 of the tensor-core work was padding. qsa_mqa_paged_rows gives each program 16 rows of one request (a host-built tile map cut at the cu_seqlens request boundaries) and does one [BLOCK_N, 16] dot per head, summing the heads in the one-row kernel's order, so every score is bitwise the old one. Decode (graph-captured) keeps the one-row kernel. Microbench, one 4096-row chunk at 249K context: 18.5 -> 3.2 ms per layer, no spills (134 regs). Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Prefill QSA indexer scoring handles 16 query rows of one request per program, with no head padding. Tiles are cut at request boundaries. Decode and graph capture keep the old kernel. Stacked on #97.
Results (2 interleaved rounds, production flags)
The gain is small because late prefill chunks are bound by the expert H2D copies: the copy stream is busy about 1270 of every 1340 ms. It should grow once prefill is no longer copy-bound.
Exactness
tests/kernels/test_qsa_score_rows.py, 13 GPU cases: scores are bitwise equal as int32 bits andqsa_block_topkselects identical blocks. A negative control (a different head-sum order) fails the test, so it does detect order changes.attn_gpu_check.py: 36/36.🤖 Generated with Claude Code
https://claude.ai/code/session_01Td2gieja9Da2rPTN17Z846