[emerardd]

Description

Closes #1954.

Enable native KV-group destinations in the existing direct-path patching functions, following the approved scope. Q destinations and all source indices remain query/output-head indices. K/V sweeps use the uniform destination KV count, returning [n_layers, n_kv_heads] for a single source and [n_layers, n_heads, n_layers, n_kv_heads] for all sources. MHA retains its existing shapes and results.

Destination counts honor per_layer_num_key_value_heads, falling back to the scalar KV count or query-head count. Heterogeneous destination counts are refused before any forward; earlier source layers do not determine the destination axis. The underlying input forks, raw-norm routing and existing source-architecture restrictions are unchanged.

Tests replace the former GQA refusal cases with positive coverage and check raw/compatibility MQA, GQA and MHA against an independent projection-level oracle, including out-of-group invariance, zero-delta interventions, all-sources consistency, flag restoration, metadata handling and heterogeneous-count refusal. The analysis-tools guide now explains the result axes and includes a pinned SmolLM2 example outside the unit suite.

Type of change

• New feature (non-breaking change which adds functionality)

Validation

• Related four-file regression surface: 305 passed. Existing SWIG deprecations and two expected gated-hook warnings remain.
• Full mypy: 390 source files passed, including the final incremental check.
• Pinned pycln/isort/Black checks and git diff --check passed.
• Additional one-layer MQA/GQA/MHA diagnostics: nine cases checked both entry points, zero placeholders, no forwards/metric calls and restored flags.
• Executed the exact pinned SmolLM2-135M documentation snippet offline, then two additional prompt contrasts for K/V: five sweeps total, each with shape [30, 3] and restored flags.
• The modified Sphinx page built successfully in isolation. Missing cross-document targets are warnings from that isolated build; the full docs build was not run.

The complete unit/test-pr suite and GPU/MPS validation were not run locally.

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

The full-suite item is unchecked because only the affected regression surface was run. The GQA refusal tests were updated as requested for the newly supported behavior; existing MHA oracle coverage remains.