feat(models): add Kyutai Helium (helium) text model support

docs/supported-models.md

@@ -53,6 +53,7 @@ Implemented model families include:
 - GPT-BigCode (`gpt_bigcode`, BigCode: the StarCoder / SantaCoder code models, a GPT-2 derivative whose defining change is Multi-Query Attention. With `multi_query: true` a single KV head is shared by every query head, so the fused `c_attn` projection produces `n_embd + 2 * kv_dim` features rather than `3 * n_embd` and is split at `[n_embd, n_embd + kv_dim]`, which is not an even three-way split; on `bigcode/gpt_bigcode-santacoder` that is 2048 + 2 * 128 = 2304 for 16 query heads of width 128 over one KV head. MQA maps onto the existing GQA attention path with `n_kv_heads = 1`, so no new attention code is involved. The rest of the block follows GPT-2: learned absolute position embeddings added to the token embeddings at the input boundary with no RoPE anywhere, `LayerNorm` with bias, and a tied output head. Unlike GPT-2 there is no `Conv1D` weight layout. HuggingFace `GPTBigCode` builds its projections with `nn.Linear`, so `c_attn` / `c_proj` / `c_fc` are already stored `[out, in]` (`transformer.h.0.attn.c_attn.weight` is `[2304, 2048]`, where the GPT-2 equivalent is `[768, 2304]`) and the loader must not transpose them; a weight that arrives in the transposed orientation is rejected by name rather than accepted. The MLP width comes from `n_inner`, not a hardcoded `4 * n_embd`, and there is no `h.N.attn.bias` causal-mask buffer to strip because HuggingFace registers this family's mask non-persistently. Loading rejects a `wpe` or `wte` table with fewer rows than `config.json`'s `n_positions` / `vocab_size` claims, a projection or norm whose shape disagrees with the config in any layer, and an untied config with no `lm_head` tensor, since the position clamp and the token lookup both trust config-derived bounds and an embedding gather does not range-check a positive index. A quantized `c_attn` is packed on the input axis only, so its output row count is checked against the config-derived width on the same path as an unquantized checkpoint; without that check, a 4-bit or 8-bit conversion whose `config.json` omits `multi_query` (which then defaults to `true`) against a packed weight sized for full multi-head attention would leave every slice in bounds while `Attention::forward`'s config-derived offsets read K and V from the wrong channels. Validated checkpoint: `bigcode/gpt_bigcode-santacoder` (f16, 1.1B). Code checkpoints ship no chat template, so the CLI renders the prompt verbatim.)
 - GPT-NeoX (`gpt_neox`, EleutherAI: the decoder behind the Pythia suite and the NeoX-derived checkpoints. Three things define it. First, the fused `query_key_value` projection is **interleaved per head**: it is `[3 * hidden_size, hidden_size]` in plain `nn.Linear` orientation, but its output is not three contiguous `hidden_size` blocks. The projection is reshaped to `(..., num_heads, 3 * head_dim)` and split on *that* axis, so the layout is head-major, `[q_i | k_i | v_i]` per head; on `EleutherAI/pythia-1b` that is 8 heads of width 256 contributing 768 contiguous channels each. A flat three-way split of the last axis (the GPT-2 and GPT-BigCode pattern) yields Q, K and V of exactly the right shape assembled from the wrong channels, so nothing throws and the model still emits fluent English out of a scrambled attention; the unit tests therefore pin the layout by channel value, not by shape. Second, **partial RoPE**: only `int(head_dim * rotary_pct)` channels of each head rotate and the rest pass through untouched (64 of 256 for Pythia's `rotary_pct` 0.25), expressed through the `dims` argument of `fast_rope` with `traditional = false`. Third, the **optional parallel residual**: with `use_parallel_residual: true` (every Pythia checkpoint) the attention and MLP sub-layers both read the same pre-norm input and their outputs are summed into the residual, `x + attn(input_layernorm(x)) + mlp(post_attention_layernorm(x))`; with `false` the block is the ordinary chained form where the MLP norm reads the post-attention residual. Both layouts are implemented, and running the wrong one is a silent quality regression rather than a crash. The rest of the block is `LayerNorm` with bias (not RMSNorm), no learned position embeddings at all, an MLP of `dense_4h_to_h(gelu(dense_h_to_4h(x)))` with no gate/up pattern, and an untied `embed_out` head. HuggingFace registers three per-layer PyTorch buffers that reach the checkpoint and are dropped at load: `attention.bias` (a `[1, 1, n, n]` causal mask whose name collides with the projection-bias namespace, exactly as GPT-2's `h.N.attn.bias` does), `attention.masked_bias`, and `attention.rotary_emb.inv_freq` (a precomputed RoPE table that `fast_rope` recomputes from `rotary_emb_base`). Both the raw HuggingFace key layout (`gpt_neox.layers.N`, top-level `embed_out`) and the layout produced by upstream mlx-lm's `sanitize` (`model.h.N`, `model.embed_out`) load. Loading rejects an `embed_in` table with fewer rows than `config.json`'s `vocab_size` claims, since an embedding gather wraps a negative index but does not range-check a positive one; a projection, bias or norm whose shape disagrees with the config in any layer, including on the quantized path, where packing compresses the input axis only and the output row count is still the width the fused reshape depends on; a `rotary_pct` whose rotary dimension count is zero, negative, odd or larger than `head_dim`, because MLX enforces exactly that contract on `rope`'s `dims` argument by throwing, and an MLX C++ exception crossing the cxx bridge is an uncatchable `std::terminate` at the first forward pass rather than a load error; a non-positive or non-finite `rotary_emb_base`, which would make every rotated channel NaN; a non-finite, negative or zero `layer_norm_eps`, since `fast::layer_norm` never looks at `eps` and computes `x * rsqrt(mean(x^2) + eps)`, so a bad value turns every hidden state into NaN without anything throwing and the checkpoint would otherwise load cleanly and generate uniform garbage; and a `quantization` block whose `bits` falls outside `1..=32` or whose `group_size` is non-positive, since both reach `quantized_matmul` unreconciled and MLX's own check on the derived unpacked width divides by zero or collapses to zero before throwing, the same uncatchable-`std::terminate` shape the rope guard exists to prevent. Validated checkpoint: `EleutherAI/pythia-1b` (f16, 244 tensors). Pythia ships no chat template, so the CLI renders the prompt verbatim and `--no-chat-template` is optional.)
 - GPT-OSS
+- Kyutai Helium (`helium`, Kyutai: a dense Llama-shaped decoder. RMSNorm before attention and before the MLP, grouped-query attention, a SwiGLU MLP over `gate_proj` / `up_proj` / `down_proj`, no QK-norm, no MoE, no sliding window, and Llama's own weight key names, so the decoder block, attention and MLP are the existing dense path unchanged. **The one architectural difference is the RoPE convention**: upstream builds `nn.RoPE(head_dim, traditional=True, base=rope_theta)`, so Helium rotates interleaved channel pairs `(2i, 2i+1)` where every other Llama-family model in this tree rotates the split-half pairs `(i, i + dims/2)`. The two produce identically shaped tensors from identical weights, so running the wrong one is a silent quality regression rather than a crash, and no shape assertion can catch it. The flag is threaded from the loader through to every RoPE call the family can reach, including the batched decode path; it also disables the two fused quantized RoPE fast paths (`MLXCEL_ENABLE_FUSED_QKV_SPLIT_ROPE` and `MLXCEL_ENABLE_FUSED_CAUSAL_PREFILL_ATTENTION`), whose C++ launchers hardcode the split-half rotation and take no flag, so a quantized Helium would otherwise be mis-rotated under either of those opt-in gates. Tensor parallelism is deliberately refused: the TP runtime builds its per-rank model by parsing `config.json` straight into the shared Llama config, which cannot carry the flag. `head_dim` comes from `hidden_size / num_attention_heads` as upstream does, and a config whose declared `head_dim` disagrees is rejected rather than one being silently preferred. Loading also rejects an embedding table or output head with fewer rows than `config.json`'s `vocab_size` claims, since an embedding gather wraps a negative index but does not range-check a positive one; a projection, bias or norm whose shape disagrees with the config in any layer, on both axes and including on the quantized path, where packing compresses the input axis only so the output row count is still the width the fused QKV concatenation and the attention reshape depend on and the input width is recovered from the scales the way MLX recovers it, `scales.shape(-1) * group_size`; quantization `biases` whose shape disagrees with the `scales` they are the zero points for; an attention block whose `q_proj`, `k_proj` and `v_proj` disagree on whether they are quantized, on whether they carry affine `biases`, or on whether they carry a dense `bias`, since the fused loader decides from `q_proj` alone and silently drops a whole bias set when one is missing; an odd or non-positive head width, because MLX enforces exactly that contract on `rope`'s `dims` argument by throwing and an MLX C++ exception crossing the cxx bridge is an uncatchable `std::terminate` at the first forward pass rather than a load error; a non-positive or non-finite `rope_theta`, which would make every rotated channel NaN; a non-finite, negative or zero `rms_norm_eps` (Helium's own `1e-08` is unusually small for this family and stays accepted); and a `quantization` block whose `bits` falls outside `1..=32` or whose `group_size` is non-positive. Stop tokens come from `config.json`'s `eos_token_id`, since Helium's `tokenizer_config.json` declares neither an EOS token nor a chat template. Validated checkpoint: `kyutai/helium-1-preview-2b` (mlx-community 4-bit, 559 tensors). Helium 1 preview is an English/French base model and ships no chat template, so the CLI renders the prompt verbatim.)
 
 Many of these families have checkpoint-specific config or weight-layout
 requirements. If a checkpoint fails detection or loading, inspect its

src/distributed/tensor_parallel/inference.rs

@@ -240,6 +240,18 @@ fn fallback_architecture(model_type: ModelType) -> &'static str {
         // V rather than between whole heads. This arm keeps the dispatch table
         // total.
         ModelType::GptNeoX => "gpt_neox",
+        // Helium is structurally a shardable dense Llama, but TP is refused for
+        // it on purpose and this arm only keeps the dispatch table total. The
+        // tensor-parallel runtime builds its per-rank model by parsing
+        // `config.json` straight into `llama3::ModelArgs`, and
+        // `llama3::ModelArgs::rope_traditional` is deliberately not
+        // deserializable, so a sharded Helium would silently rotate split-half
+        // pairs while the single-process path rotates interleaved pairs. The
+        // refusal comes from `validate_supported_runtime`, whose
+        // `runtime_kind_for` match has no arm for `ModelType::Helium` and so
+        // returns `None`. Enabling TP here means threading the flag through the
+        // rank-local config first.
+        ModelType::Helium => "helium",
         ModelType::StarCoder2 => "starcoder2",
         ModelType::Mellum => "mellum",
         ModelType::MiniCPM | ModelType::MiniCPMOVLM => "minicpm",

src/distributed/tensor_parallel/llama_runtime_tests.rs

@@ -52,6 +52,9 @@ fn make_test_model_args() -> LlamaModelArgs {
         rope_scaling: None,
         quantization: None,
         tie_word_embeddings: false,
+        // Every Llama-family checkpoint rotates split-half; only the Helium
+        // loader sets this.
+        rope_traditional: false,
     }
 }
 
@@ -913,6 +916,7 @@ fn local_llama_args_preserves_computed_head_dim_when_config_omits_it() {
         rope_scaling: None,
         quantization: None,
         tie_word_embeddings: true,
+        rope_traditional: false,
     };
     let plan = generate_shard_plan("llama", 1, &ShardConfig::with_tp_size(2)).unwrap();
 

src/loaded_model.rs

@@ -164,6 +164,7 @@ pub enum LoadedModel {
     Gpt2(models::Gpt2Model),
     GptBigCode(models::GptBigCodeModel),
     GptNeoX(models::GptNeoxModel),
+    Helium(models::HeliumModel),
     StarCoder2(models::StarCoder2Model),
     // Mixed full/sliding caches use a wrapper
     Mellum(models::MellumWrapper),
@@ -313,6 +314,7 @@ macro_rules! delegate_language_model {
             LoadedModel::Gpt2(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::GptBigCode(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::GptNeoX(inner) => LanguageModel::$method(inner, $($arg),*),
+            LoadedModel::Helium(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::StarCoder2(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::Mellum(inner) => LanguageModel::$method(inner, $($arg),*),
             LoadedModel::MiniCPM(inner) => LanguageModel::$method(inner, $($arg),*),

src/model_metadata.rs

@@ -203,6 +203,7 @@ macro_rules! for_each_model_registration {
             Gpt2 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::Gpt2Model::load, args: models::gpt2::ModelArgs, weight_builder: models::Gpt2Model::from_weights, wrap: LoadedModel::Gpt2 } };
             GptBigCode => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::GptBigCodeModel::load, args: models::gpt_bigcode::ModelArgs, weight_builder: models::GptBigCodeModel::from_weights, wrap: LoadedModel::GptBigCode } };
             GptNeoX => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::GptNeoxModel::load, args: models::gpt_neox::ModelArgs, weight_builder: models::GptNeoxModel::from_weights, wrap: LoadedModel::GptNeoX } };
+            Helium => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::HeliumModel::load, args: models::helium::ModelArgs, weight_builder: models::HeliumModel::from_weights, wrap: LoadedModel::Helium } };
             StarCoder2 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::StarCoder2Model::load, args: models::starcoder2::StarCoder2Config, weight_builder: models::StarCoder2Model::from_weights, wrap: LoadedModel::StarCoder2 } };
             Mellum => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MellumModel::load, args: models::mellum::ModelArgs, weight_builder: models::MellumModel::from_weights, wrap: |m| LoadedModel::Mellum(models::MellumWrapper::new(m)) } };
             MiniCPM => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MiniCPMModel::load, args: models::minicpm::ModelArgs, weight_builder: models::MiniCPMModel::from_weights, wrap: LoadedModel::MiniCPM } };

src/models/detection.rs

@@ -200,6 +200,7 @@ pub fn get_model_type(model_path: &Path) -> Result<ModelType> {
         "gpt2" => Ok(ModelType::Gpt2),
         "gpt_bigcode" => Ok(ModelType::GptBigCode),
         "gpt_neox" => Ok(ModelType::GptNeoX),
+        "helium" => Ok(ModelType::Helium),
         "starcoder2" => Ok(ModelType::StarCoder2),
         "mellum" => Ok(ModelType::Mellum),
         "minicpm" => Ok(ModelType::MiniCPM),

src/models/detection_tests.rs

@@ -182,6 +182,42 @@ fn gpt_neox_model_type_is_detected() {
     fs::remove_dir_all(model_dir).unwrap();
 }
 
+#[test]
+fn helium_model_type_is_detected() {
+    // Helium's config is field-for-field a Llama config apart from `model_type`
+    // and the smaller `rms_norm_eps`, so detection must key off `model_type`
+    // alone and must not fall through to the Llama arm. The two decode
+    // differently: Helium rotates interleaved RoPE pairs, Llama split-half ones.
+    let model_dir = temp_path("helium_text");
+    fs::create_dir_all(&model_dir).unwrap();
+    fs::write(
+        model_dir.join("config.json"),
+        r#"{
+            "model_type": "helium",
+            "architectures": ["HeliumForCausalLM"],
+            "hidden_size": 2560,
+            "num_attention_heads": 20,
+            "num_key_value_heads": 20,
+            "num_hidden_layers": 24,
+            "intermediate_size": 7040,
+            "head_dim": 128,
+            "max_position_embeddings": 4096,
+            "rms_norm_eps": 1e-08,
+            "rope_theta": 100000.0,
+            "attention_bias": false,
+            "mlp_bias": false,
+            "tie_word_embeddings": false,
+            "vocab_size": 48000
+        }"#,
+    )
+    .unwrap();
+
+    let detected = super::detection::get_model_type(&model_dir).unwrap();
+    assert_eq!(detected, ModelType::Helium);
+
+    fs::remove_dir_all(model_dir).unwrap();
+}
+
 #[test]
 fn fastvlm_model_type_is_detected() {
     for model_type in ["fastvlm", "llava_qwen2"] {