Unify module perf-table head-axis keying: [native][local] everywhere except GQA

[tianhaox]

Unify head-axis keying across all module perf tables: [native][local] everywhere except GQA

Background

#1429/#1431 pinned the DSV4 module convention: rows persist rank-local num_heads + mandatory tp_size; the model-native count derives as num_heads * tp_size; loaded tables key [native][local]. That fix was scoped to DSV4, but the audit and two follow-up findings show the underlying problem is family-wide:

1. Latent collision in every local-only module table. MLA/DSA/MiniMax loaders key the head axis by local count only. Today that is unambiguous by luck: e.g. every model in the MLA table is 128-native, so no two natives can shard to the same local. The moment a second native geometry lands in any of these tables (a 64- or 96-head MLA model), it reproduces DSV4 module perf files carry two num_heads semantics (sglang=native, vllm=local); unify the collector convention and re-collect sglang 0.5.10 #1429's mixed-bucket problem — local 32 could be native-64@tp2 or native-128@tp4, which are different modules. The 128-native assumption is even hardcoded: load_context_mla_data (operations/mla.py:1567) backfills legacy rows with num_heads = 128 // tp_size — the MLA twin of the Flash-only NATIVE_HEADS = 64 fallback that fix: unify DSV4 num_heads to the rank-local convention — data migration, strict loaders, collector guard (#1429) #1431 killed in the DSV4 collector, and it silently mislabels the first non-128-native MLA model. Note the loader also ignores tp_size whenever num_heads is present, so native is currently not derived at all — the migration reads both columns and keys [native][local].
2. Model-name alias duplication (separate issue, but it lands on the same key path): vllm 0.22.0 MLA carries the same 5888-row grid under 3 model names, DSA under 4. The model string is doing double duty as both geometry proxy and provenance. With [native][local] as the geometry key (+ existing dtype columns), the model column can be demoted to provenance and aliases dedupe naturally — which is what finally makes the 0.22.0 duplicates deletable.
3. Coverage accounting (separate issue): the MLA DSv3 128-native×tp deployment grid only exists in old versions (sglang 0.5.10, trtllm rc10). Keying [native][local] makes the per-native coverage gap auditable directly from the table structure instead of ad-hoc row archaeology.

The rule (proposed pin)

A module table keys its head axis [native][local] iff the local key is not computation-complete — i.e. the module's work depends on model variables that are not otherwise in the key (lora ranks, hidden size, indexer geometry). If the local key fully determines the physical computation, local-only stays.

Applied per family:

Table family	Head-axis key	Rationale
DSV4 modules (csa/hca × ctx/gen)	[native][local] — done in #1431	Flash-64 vs Pro-128 coexist; module bundles indexer/hidden geometry
MLA (context_mla / generation_mla)	migrate to [native][local]	q/kv lora ranks are model variables outside the key; today single-native by luck
DSA	migrate to [native][local]	same as MLA; table already carries 4 alias names, i.e. multiple identities per geometry
MiniMax (linear attn modules)	migrate to [native][local]	model-geometry-dependent module, same argument
GQA attention (context/generation attention)	keep local-only	local q-heads, kv-heads, head_dim (and window) are all in the key already; same local key ⇒ same physical computation; native adds no information. No collision exists or can exist.
Sparse-kernel tables (paged_mqa_logits / hca_attn / csa_topk_calib)	unchanged (native-keyed by contract, #1431)	—

Non-goals

• No data-file rewrites: every family already persists local + tp_size (the DSV4 module perf files carry two num_heads semantics (sglang=native, vllm=local); unify the collector convention and re-collect sglang 0.5.10 #1429 audit confirmed this); this is a consumer-side loader/keying change only.
• Alias normalization (canonical geometry identity, model column → provenance) is step 2, tracked separately — it depends on this issue but changes lookup semantics, not just key structure.
• Backfilling the MLA 128-native×tp coverage gap is a collection decision, tracked separately.

Work items

1. Migrate MLA loaders (Python perf_database + Rust perf_database/mla.rs) to [native][local], byte-equal predicates, real-data oracle parity tests (same discipline as fix: unify DSV4 num_heads to the rank-local convention — data migration, strict loaders, collector guard (#1429) #1431).
2. Same for DSA and MiniMax module loaders.
3. Port the fix: unify DSV4 num_heads to the rank-local convention — data migration, strict loaders, collector guard (#1429) #1431 stale-fingerprint guard shape to each family where it is expressible (heads constant across tp within one (model, version) group ⇒ reject) — cheap insurance even though these families were never native-keyed.
4. Guardrail test scanning all shipped module parquet per family (mirror of test_shipped_dsv4_module_tables_are_rank_local).
5. Document the pin (extend the convention section written by fix: unify DSV4 num_heads to the rank-local convention — data migration, strict loaders, collector guard (#1429) #1431) including the GQA exception and its computation-completeness rationale.

Sequencing

After #1431 merges (it defines the convention and the reference implementation). Families are independent; MLA first (it unblocks the alias-dedup and coverage issues, which both cite the MLA table).