How cytosine-code/pac26_indexer_mqa_logits works
This repository is a PAC 2026 benchmark program and optimized AArch64 implementation of paged MQA logits for compressed sparse attention. The principal workflow generates BF16 query/KV test data, invokes the SME/SVE kernel, measures throughput, and checks its output against a reference implementation. The documented kernel stages include query packing, page scheduling and K packing, SME scoring, and fused ReLU/head reduction; masking handles invalid output positions. Kernel internals are represented at a high level because the source excerpt is partial; probe programs are excluded as development tools.
Main entry points: testcase.h Generate test cases · main.cpp Run benchmark · ref_mqa_logits.h Reference logits · indexer_mqa_logits.h Pack queries
Components 17
Benchmark harness
- Run benchmark
main.cpp - Output logits
main.cpp
Test data
- Generate test cases
testcase.h - Tensor views
Tensor.h - NUMA allocator
allocator.h - Query vectors
testcase.h - Paged KV cache
testcase.h - Weights and lengths
testcase.h
Optimized kernel
- Pack queries
indexer_mqa_logits.h - Schedule pages
indexer_mqa_logits.h - Pack KV page
indexer_mqa_logits.h - SME page scores
indexer_mqa_logits.h - ReLU and reduce
indexer_mqa_logits.h - Causal output mask
indexer_mqa_logits.h
Reference and utilities
- Reference logits
ref_mqa_logits.h - Correctness checks
utils.h
Other components
- Benchmark invocation
Connections 26
- Benchmark invocation → Run benchmark: starts
- Run benchmark → Generate test cases: generates data
- Generate test cases → Tensor views: creates views
- Generate test cases → NUMA allocator: allocates
- Generate test cases → Query vectors: produces
- Generate test cases → Paged KV cache: produces
- Generate test cases → Weights and lengths: produces
- Run benchmark → Pack queries: invokes kernel
- Query vectors → Pack queries: packs
- Weights and lengths → Schedule pages: guides
- Schedule pages → Pack KV page: dispatches pages
- Paged KV cache → Pack KV page: reads pages
- Pack queries → SME page scores: supplies packed Q
- Pack KV page → SME page scores: supplies packed K
- SME page scores → ReLU and reduce: provides scores
- Weights and lengths → ReLU and reduce: supplies weights
- ReLU and reduce → Output logits: writes scores
- Weights and lengths → Causal output mask: supplies limits
- Causal output mask → Output logits: masks positions
- Run benchmark → Reference logits: computes baseline
- Query vectors → Reference logits: reads
- Paged KV cache → Reference logits: reads
- Weights and lengths → Reference logits: reads
- Reference logits → Correctness checks: provides baseline
- Output logits → Correctness checks: provides result
- Run benchmark → Correctness checks: checks accuracy