feat(kimi_linear): context parallelism and packed-sequence support (#3243)

* Add Kimi Linear model support

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* Address Claude review feedback for Kimi Linear

Remove unused gated RMSNorm fallback and document the tensor layouts for Kimi
Linear padding helpers.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* Fix Kimi Linear review findings

Respect explicit buffer_device on CPU-only hosts, document Kimi MoE tensor
layouts, and cover checkpoint-free initialization plus the HF-order inference
MoE path in a CPU unit test.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* Guard unsupported MoE routing weights backends

Fail fast when route_weight_after_down_proj is used with DeepEP or TE
expert backends, and cover those unsupported combinations with unit tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* Avoid mutating Kimi backend config

Copy the caller-provided backend before defaulting Kimi gate precision, and
extend checkpoint-free initialization coverage to assert KDA fp32 parameters
are finite after initialization.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* Preserve MoE topk ordering defaults

Keep the shared MoE topk ordering default aligned with the previous unsorted
behavior, and add numerical coverage that Kimi's eval-only HF-order MoE path
matches the standard grouped experts path.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* Document unweighted SwiGLU tensor layout

Add the tensor contract for the fused gate/up helper used by the shared MoE
expert paths.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* Document Kimi expert split tensor layout

Add the tensor contract for the Kimi state-dict adapter expert split override.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* Cover Kimi DTensor expert split adapter path

Add focused tests for Kimi's DTensor expert split override, covering
replicated weights, uneven Shard(0) expert distribution, and ep-submesh rank
selection.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* feat(kimi_linear): context parallelism and packed-sequence support

Kimi Linear interleaves KDA linear-attention layers with MLA full-attention
layers, so context parallelism has to satisfy both at once.

KDA carries a sequential recurrent state, so FLA's context-parallel kernels
need every rank to own one contiguous slice of the token stream and take
document boundaries through cu_seqlens. PyTorch's load-balanced context_parallel
layout does not satisfy that, so the model shards the batch itself through
_cp_make_batch_fn and drives both layer types from a single [batch, sequence]
document-id map.

MLA attends globally: each rank all-gathers the compressed KV latent
(kv_lora_rank + qk_rope_head_dim per token, ~18x smaller than the expanded
per-head K/V on the 48B checkpoint) and runs FlexAttention with a causal,
per-document block mask against the full-sequence keys.

The same document map fixes packed sequences without CP: MLA now gets a
document-blocked causal mask instead of a plain causal one, and KDA derives
per-document cu_seqlens from the indexed packing mask instead of per-row ones.

Also wires up the pieces this needs: MoE layers named block_sparse_moe are now
found by the MoE parallelizer (CP mesh, and expert parallelism), and apply_cp
no longer hard-imports Transformer Engine for models that own their CP.

Signed-off-by: khazic <khazzz1c@gmail.com>

* refactor(kimi_linear): move CP hook to the ContextParallelSharder contract

The CP dispatch no longer reads the private `_cp_make_batch_fn` batch key that
this branch was written against, so after merging main the hook returned a key
nothing consumes: the batch went through unsharded, with no error and no
`kimi_packed_context`.

`prepare_model_inputs_for_cp` now takes the batch positionally, matching how the
dispatch calls it, and returns a `ContextParallelSharder` under `cp_sharder`.
`shard_batch_for_kimi_cp` reports a `ShardLayout` alongside the sharded batch, so
the sharder's token verbs can shard and gather side tensors on the same
contiguous layout, and it declares `contiguous_local_indices` as its index map
(KDA's recurrent state needs contiguous per-rank slices, not the framework's
default head/tail round-robin).

A dispatch-level test covers the wiring end to end: it drives
`ContextParallelSharder(model, mesh, batch)` and asserts rank 1 owns the
contiguous second half, which is what silently regressed above.

Signed-off-by: khazic <khazzz1c@gmail.com>

* docs(kimi_linear): add Kimi Linear model coverage page

The registry gained KimiLinearForCausalLM without a matching card under
docs/model-coverage/, which failed test_doc_coverage. Adds the model card,
the LLM index row, and the nightly nav entry.

Signed-off-by: khazic <khazzz1c@gmail.com>

* docs(kimi_linear): apply review wording and drop dead init_weights branch

Address review feedback on PR #3243: adopt the docs wording suggestions
for the Kimi Linear model-coverage page and the LLM index, and collapse
the identical if/else branches in KimiLinearDecoderLayer.init_weights
into a single unconditional call.

Signed-off-by: khazic <khazzz1c@gmail.com>

* refactor(kimi_linear): give Kimi Linear 48B its own checkpoint identity

Moonshot publishes both this model and the Kimi K3 text backbone under
model_type kimi_linear with architectures [KimiLinearForCausalLM], so neither
field identified the model on its own. The previous approach disambiguated on
the architecture name, which does not actually separate the two families.

Name the 48B A3B implementation distinctly instead: model_type
kimi_linear_48b_a3b, architecture KimiLinear48BForCausalLM, config class
KimiLinear48BConfig. K3 keeps kimi_linear -> KimiK3TextConfig unchanged, and
_CUSTOM_CONFIG_ARCH_REGISTRATIONS plus the architecture-aware config resolution
in model_init.py are gone.

The config class stamps its own identity onto the instance so a published
Moonshot checkpoint resolves to the right model class, and the example recipes
name the config class explicitly, matching the deepseek_v4 recipe pattern.

Signed-off-by: khazic <khazzz1c@gmail.com>

* test(kimi_linear): cover the context-parallel path with distributed tests

The CP forward/backward paths (KimiMLAAttention._forward_with_cp,
KimiDeltaAttention._forward_with_cp, the _AllGatherSequence autograd Function
and document_causal_flex_attention) had only single-process coverage: every
test in test_cp.py drives a fake mesh with no process group, so a
collective-ordering hang, rank mismatch or gradient-domain error could not be
caught. Add coverage at two levels.

tests/unit_tests/models/kimi_linear/test_cp_distributed.py runs two gloo ranks
on CPU and pins the collective Kimi Linear owns end to end: all_gather_sequence
rebuilds the global sequence in rank order, its backward sums every rank's
contribution before narrowing back to the local shard, and shard_batch_for_kimi_cp
lays the batch out to match on a real device mesh.

tests/functional_tests/context_parallel/run_kimi_linear_cp.py follows the
existing run_qwen3_5_moe_linear_attn_cp.py precedent and asserts full-model
parity on 2 GPUs: it drives the real sharder for both the CP=1 reference and
the CP=2 run and compares logits and per-parameter gradients over three
document layouts (single document, packed documents straddling the shard
boundary, and a padding tail).

Also document why KimiDecoderLayer._moe_infer_hf_order sits next to the
canonical MoE path, and what pins the two together.

Signed-off-by: khazic <khazzz1c@gmail.com>

* chore(examples): hand the CP recipe's ci owner to huiyingl

Signed-off-by: khazic <khazzz1c@gmail.com>

* refactor(moe): drop the router_topk_sorted flag

Requested in review. The shared Gate keeps PyTorch's default sorted top-k in
every branch that had it, and sigmoid_with_bias keeps its explicit sorted=False.
Kimi Linear needs no replacement: weights and indices stay paired, and expert
dispatch groups by expert id before weighting and reducing, so the slot order
does not reach the output.

Signed-off-by: khazic <khazzz1c@gmail.com>

* docs(model-coverage): add the Kimi Linear page slug

The new page had no frontmatter slug, so its Fern route did not resolve and
the model-coverage link check failed alongside the organization slug check.

Signed-off-by: khazic <khazzz1c@gmail.com>

* ci: run the Kimi Linear CP test in its own L2 job

The Kimi Linear CP test compiles a large set of FLA Triton kernels before
it can run its three document-layout cases, which takes longer than the
headroom left in the shared L2_Context_Parallel job (20 minute budget,
already ~14 minutes without it). All three attempts of that job timed out
mid-test.

Move the test into tests/functional_tests/context_parallel_kimi_linear/
and add a matching L2_CP_KimiLinear matrix entry, so the shared job keeps
its previous runtime and the Kimi Linear test gets a budget of its own.

Signed-off-by: khazic <khazzz1c@gmail.com>

* ci(kimi_linear): drop the L2 CP functional job

Kimi Linear is a single-model addition, so its context-parallel path
belongs in the unit tests rather than the L2 functional suite. The
distributed CP behavior is already covered by
tests/unit_tests/models/kimi_linear/test_cp_distributed.py.

Signed-off-by: khazic <khazzz1c@gmail.com>

---------

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
Signed-off-by: khazic <khazzz1c@gmail.com>
Co-authored-by: HuiyingLi <willwin.lee@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>