model: add Kimi-K3.1-Pro text model (#31779)
* model: add Kimi-K3.1-Pro text model
K3 with the full-attention MLA layers replaced by DSA: on this checkpoint 24
of the 93 layers carry a lightning indexer (32 heads x 128, top_k 2048, no
rope). The KDA (linear-attention) layers and the rest of the K3 skeleton -
cross-layer residuals, latent MoE, situ, MLA output gate, full-rank KDA gate -
are unchanged, so this reuses the K3 graph and converter and adds the indexer.
Indexer:
- the selection is an additive mask over the whole cache, not a gather of
the top-k rows, so it reuses the existing sparse flash attention path;
off that path the mask stays correct but still costs a kq over the cache
- the indexer key LayerNorm takes its eps from the model's rms_norm_eps,
because dsa_config declares no eps of its own
Graph and cache:
- new arch with four indexer tensors per DSA layer, and a hybrid memory
that keeps a K-only MLA cache for those layers, recurrent state for the
KDA layers, and the indexer key cache
- the graph node budget shares the K3 entry
Conversion:
- KimiK31Model(KimiK3Model), picked by KimiK31ForConditionalGeneration
- dsa_config becomes attention.indexer.{head_count,key_length,top_k,types}
plus the indexer LayerNorm eps; the layer list is 1-indexed like
linear_attn_config's, and the converter asserts that it covers exactly
the full-attention layers
- the routed experts stay MXFP4 and are repacked as in K3, not dequantized
- the vision tower and mm_projector are skipped: text only for now
- the arch key follows the checkpoint's own model_type ("kimi_k31") and the
loader registers KimiK31ForConditionalGeneration; the release name carries
a -Pro suffix that nothing in the config or the tensor names uses
Quantisation: the indexer's q_b, k and proj are kept out of the quant
overrides, since the top-k selection is sensitive to their precision.
Verified on an 8-layer slice against the HF checkpoint: 38 tensors checked
exactly, the indexer metadata against the config, and the selection against
full attention - identical PPL at top_k beyond the context length, divergent
once top_k is forced below it. Arch test fixture and the full arch suite pass.
* converter: avoid clutter in the k3 base
* narrow comment to match model
* update metal fusion expectations
* restore
* sync with PR 31733 refactoring
