Add Kimi K3

[kernelpool]

Model: https://huggingface.co/moonshotai/Kimi-K3
Tested: https://huggingface.co/kernelpool/Kimi-K3-2bit-UVMAX

mlx.launch --verbose --backend jaccl --hostfile hosts-jaccl.json --env MLX_METAL_FAST_SYNCH=1 -- /Users/optimus/repo/mlx-lm/.venv/bin/python /Users/optimus/repo/mlx-lm/mlx_lm/examples/sharded_generate.py --model /Users/optimus/models/catalyst/Kimi-K3-2bit-UVMAX --prompt "Who is Albert Einstein?" -m 1024
[INFO] Running /Users/optimus/repo/mlx-lm/.venv/bin/python /Users/optimus/repo/mlx-lm/mlx_lm/examples/sharded_generate.py --model /Users/optimus/models/catalyst/Kimi-K3-2bit-UVMAX --prompt 'Who is Albert Einstein?' -m 1024
The user is asking a simple factual question about Albert Einstein. This is a straightforward request for information about a well-known historical figure. Let me think about what I know about Einstein and how to present it well.

Albert Einstein (1879-1955) was one of the most influential physicists in history. Key facts:

- Born March 14, 1879, in Ulm, Germany
- Died April 18, 1955, in Princeton, New Jersey, USA
- Theoretical physicist
- Developed the theory of relativity (special relativity in 1905, general relativity in 1915)
- Famous equation E=mc² (mass-energy equivalence)
- Won the Nobel Prize in Physics in 1921 (for his work on the photoelectric effect, not relativity)
- His "miracle year" was 1905 when he published four groundbreaking papers
- Worked at the patent office in Bern, Switzerland when he made many of his early discoveries
- Emigrated to the United States in 1933 due to Nazi persecution (he was Jewish)
- Worked at the Institute for Advanced Study in Princeton
- Contributed to quantum mechanics (though he had famous reservations, "God does not play dice")
- Wrote a letter to Roosevelt about the possibility of atomic bombs (leading to the Manhattan Project), though he didn't work on the bomb himself
- Was offered the presidency of Israel in 1952 but declined
- Known for his iconic appearance (wild hair, mustache)
- Famous quotes and cultural icon status
- His brain was famously (and controversially) preserved/studied after his death

The question is simple and casual - "Who is Albert Einstein?" This suggests the user wants a clear, informative answer but not necessarily an exhaustive biography. I should provide a good overview covering the essentials:

1. Who he was (theoretical physicist, one of the most famous scientists)
2. His major contributions (relativity, E=mc²)
3. Key biographical facts (dates, Nobel Prize)
4. His cultural significance

Format considerations: This is a conversational question, so I should respond in a natural, informative way. A moderate-length response with perhaps some structure would work well. I could use a few paragraphs or a mix of prose and maybe some key points. Given the casual nature of the question, I shouldn't write an overwhelming essay, but I should give enough information to be genuinely useful. A couple of paragraphs covering the essentials would be appropriate. Maybe 150-250 words is a good target - enough to be informative but not overwhelming.

Let me draft a response:

Albert Einstein (1879-1955) was a German-born theoretical physicist widely regarded as one of the greatest scientists of all time. He's best known for developing the theory of relativity, which revolutionized our understanding of space, time, and gravity.

His famous equation E=mc² showed that mass and energy are equivalent, laying the conceptual foundation for nuclear physics. In 1905, often called his "miracle year," he published four groundbreaking papers while working as a patent clerk in Switzerland.

He won the Nobel Prize in Physics in 1921, though interestingly it was awarded for his explanation of the photoelectric effect rather than relativity.

Einstein was Jewish and fled Nazi Germany in 1933, settling in the United States where he worked at Princeton's Institute for Advanced Study. Beyond his scientific work, he was known for his humanitarian views, his advocacy for civil rights, and his iconic appearance with wild hair.

That's a good structure. Let me refine it to be natural and complete.<|close|>think<|sep|><|open|>response<|sep|>Albert Einstein (1879–1955) was a German-born theoretical physicist widely considered one of the greatest scientists in history. He's best known for developing the **theory of relativity**, which fundamentally transformed our understanding of space, time, and gravity.

Some key facts about him:

- **E=mc²**: His famous equation showing that mass and energy are equivalent — a concept that became foundational to modern physics.
- **"Miracle Year" (1905)**: While working as a patent clerk in Switzerland, he published four groundbreaking papers that revolutionized physics.
- **Nobel Prize (1921)**: He won the Nobel Prize in Physics, though interestingly it was awarded for his explanation of the photoelectric effect rather than relativity.
- **Life in America**: As a Jewish scientist, he fled Nazi Germany in 1933 and settled in the United States, working at Princeton's Institute for Advanced Study.

Beyond science, Einstein was known for his humanitarian and pacifist views, his sharp wit, and his iconic appearance — the wild hair and mustache that made him a pop culture symbol of genius. His name has become so synonymous with intelligence that people often say "Einstein" as shorthand for a brilliant person.

Is there a particular aspect of his life or work you'd like to know more about?<|close|>response<|sep|>
==========
Prompt: 93 tokens, 5.044 tokens-per-sec
Generation: 1024 tokens, 11.738 tokens-per-sec
Peak memory: 420.215 GB

[guruswami-ai]

Independent validation on 4× M3 Ultra in TP4, plus one note for anyone applying this to the current release.

We ran this PR's model file on a 4-node Apple silicon cluster and it works. Some
data in case it is useful, and one compatibility observation.

What was actually tested — please read before weighing this

We did not run the complete PR on main. The deployment is mlx-lm
v0.31.3 with this PR's model, tokenizer and tool-parser changes backported
from 7d505c285b801108a52c23353c7fb6af07204717. Because v0.31.3 lacks
make_text_state_machine, it also carries a local streaming-aware
response-layer marker filter in place of your generate.py change.

So these fixtures validate our backport plus this PR's model and tool-parser
pieces. They do not validate your generate.py marker handling directly.

Setup

• 4× Mac Studio M3 Ultra 512 GB, macOS 26.6, Thunderbolt 5 full mesh, JACCL/RDMA
• Kimi-K3 2.78 T MXFP4, tensor-parallel across 4 ranks, 420.8 GB/rank
• prefill_step_size=512, prompt cache size 1

What worked

A four-fixture suite over the HTTP API, all passing:

fixture	result
non-streaming	content exactly 391, no XTML control tokens
streaming	same, clean across SSE chunk boundaries
reasoning separation	answer in content, reasoning in its own channel
tool call	structured tool_calls → get_weather(city="Tokyo")

Also: load ~5 min at 420.8 GB/rank; prefill throughput flat at ~108–114 tok/s
from 2.5 k to 38.5 k tokens, which is the expected shape when 69 of 93 layers
are KDA linear attention.

Durability boundary: this validates model loading, TP4 sharding and
wire-level protocol correctness — not sustained serving reliability. Longer
JACCL soaks hit a separate MLX RDMA failure, reported on
mlx#3933. No model-output
correctness failure was observed before transport failure.

A_log — independent confirmation that your handling is right

We initially diverged from this PR here and were wrong, so it seems worth
recording the evidence.

We reasoned from shapes: released shards store A_log as [128], num_heads
is 96, head_dim is 128 — so we treated it as per-channel and replicated it
across ranks. Then we read the values instead of reasoning about them, across
all 69 KDA tensors:

shapes: {(128,)}
nonzero-count distribution: [96]
tensors whose [96:] is EXACTLY zero: 69/69
layer 0 A_log[92:100] = [-0.5055, -0.0233, -0.4065, -0.1461, 0.0, 0.0, 0.0, 0.0]

96 live values padded to head_dim. Your [:num_heads] truncation,
(num_heads, 1) reshape and per-head sharding are correct; our reading was not.
Worth noting the failure is silent either way — 128 divides evenly by a
world size of 4, so every shape still broadcasts and nothing raises.

The one thing we would suggest: assert the tail is zero before truncating,
rather than truncating whenever shape[0] > num_heads. A future checkpoint with
a non-zero tail would otherwise be silently discarded.

Compatibility note for the current release

The marker-stripping relies on structural_markers in
make_text_state_machine. That function exists in main, so this PR is
self-consistent — but it is not in v0.31.3, which is still the latest
release. Backporting #1626 to v0.31.3 without its generate.py change
therefore leaves the response markers in content:

'<|open|>response<|sep|>391<|close|>response<|sep|><|close|>message<|sep|>'

Reasoning/response separation still works there; only the response body is
affected. We worked around it with a streaming-aware exact-marker filter at the
response layer — it has to buffer partial marker prefixes, since a marker can
span SSE chunks and a per-chunk str.replace misses "<|op" + "en|>".

Not a defect in this PR, but it may be worth a line in the description, since
anyone trying K3 on the current release will hit it and the output looks
plausible enough to miss.

Offer

Happy to share the fixture suite, or to re-run any of this on the 4-node TP4
setup if that would help review.

[guruswami-ai]

Thanks again for the Kimi-K3 implementation. We have a small tests-led follow-up
prepared from an independent TP4 validation, if it would be useful. It is based
on 43f5e87 and is three commits, the first of which carries the only proposed
production change: a guard against silently discarding a non-zero A_log
padding tail.

The coverage is deliberately independent of the full checkpoint — no downloads,
and only the distributed file needs more than one process:

• A_log padding — the released checkpoint's zero-padded tail is accepted
  and trimmed to the live prefix; a non-zero tail raises instead of being
  truncated away, naming the tensor key, observed length and expected head
  count. sanitize() currently truncates unconditionally.
• Kimi tensor-parallel equivalence on a randomized tiny model at TP2 and
  TP4, with direct assertions on the KDA A_log and dt_bias slices, the
  depthwise convolution groups and weight rows, projection_dim, and that
  head-dimension-only parameters (o_norm, f_a_proj) keep their unsharded
  values. Full rank tensors are compared rather than a checksum, and cached
  prefill plus a decode step are compared against an unsharded reference so the
  recurrent KDA state and MLA KV cache are covered.
• A separate MoE fixture for the latent-MoE reduction. This one needs
  explaining: in the main fixture the control for "reduce after normalize"
  cannot fire at all. RMSNorm is w·y/√(mean(y²)+eps), and at a conventional
  ±0.02 init the routed activations sit far below √eps, so the norm is nearly
  linear — and a linear map commutes with summation. The separate fixture keeps
  the real epsilon and drives the routed tensor to ~158× √eps, where correct
  sharding gives max_abs 2.98e-07 and moving the norm before the all-reduce
  gives 7.096e-01. The ratio is asserted so the fixture cannot quietly lose its
  power.
• XTML streaming state machine — every structural marker fed split at every
  internal offset across two step() calls, plus preservation of legitimate
  partial-prefix text ("a<|" then "b" → "a<|b"). The markers are inferred
  from the vocabulary through TokenizerWrapper rather than hardcoded.

We ran negative controls for wrong A_log slicing, stale convolution groups and
the MoE reduction ordering, so the tests demonstrably detect the mistakes they
are meant to guard against. The wrong-slice control is the interesting one: it
produces max_abs 2.3e-10 and passes a numerical comparison, and is caught only
by the structural assertions.

Tolerances are rtol=1e-5, atol=1e-6; the whole-model comparison measured 0.0
at TP2 and 1.164e-10 at TP4, and the isolated MoE comparison 2.980e-07 at
both. TP8 was not run — we only have four nodes. Black 25.1.0 and isort 6.0.0
(--profile=black) pass on every touched file.

The three commits are independent and touch disjoint files, so they can be taken
separately if only some are of interest.

The PR already has good parser-level Kimi tool-call cases, so we have not
duplicated those. Would you prefer cherry-pickable commits based on the current
add-kimi-k3 head, or a follow-up PR after #1626 merges?