Skip to content

Add perf model for vllm::rocm_aiter_sparse_attn_indexer (GLM5) - #937

Open
spandoesai wants to merge 2 commits into
mainfrom
perf-model/dsa-sparse-attn-indexer
Open

spandoesai wants to merge 2 commits into
mainfrom
perf-model/dsa-sparse-attn-indexer

Conversation

@spandoesai

Copy link
Copy Markdown
Collaborator

The DeepSeek sparse-attention "lightning indexer" kernel had no perf model, so it reported 0 FLOPs/bytes and was dropped from roofline scoring despite being a meaningful chunk of GLM5 decode time. This adds rocm_aiter_sparse_attn_indexer (subclassing InferenceAttention to reuse the vLLM annotation parsing for exact QK-pair aggregates).

  • FLOPs: fp8 MQA index-score QK only (no PV — the kernel scores + selects top-k, no value aggregation).
  • Bytes: fused MQA traffic — keys read once across heads (cached fp8 rows + current bf16 key), fp32 weights read, int32 top-k write.
  • Registers the op and adds it to the inference report's annotation name-filter so it gets the seqlen aggregates its model needs.
  • Verified end-to-end: the indexer now reports non-zero GFLOPS/Data Moved on GLM5 with no regressions.

The DeepSeek sparse-attention "lightning indexer" kernel previously had no
perf model, so it reported 0 FLOPs / 0 bytes and was excluded from roofline
scoring despite being a sizeable fraction of decode time.

Add rocm_aiter_sparse_attn_indexer (subclassing InferenceAttention to reuse
the vLLM iteration-annotation parsing for exact QK-pair aggregates). FLOPs
count the fp8 MQA index-score QK only (no value aggregation). Bytes reflect
the fused MQA traffic: keys read once across the H index heads, split into
cached fp8 rows (key + per-block scales) and the current-chunk bf16 key, plus
the fp32 weights read and the top-k index write. Register the op in the
pseudo-op mapping, and add it to the inference report's annotation name-filter
so the kernel receives the seqlen aggregates its perf model needs.

Co-authored-by: Cursor <cursoragent@cursor.com>
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 13.20755% with 46 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
...odel/extensions/attention_perf_model_extensions.py 13.20% 46 Missing ⚠️

📢 Thoughts on this report? Let us know!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants