epic: OpenXLA SSM / hybrid / recurrent architecture support (Mamba2, Nemotron-H, Jamba, and more)
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem / Background
Parent epic: #493 (OpenXLA backend architecture-coverage parity with the MLX engine). Scoping design: #502, recorded in `spike/openxla/SSM_HYBRID_DESIGN.md`.
Window D of #493 covers state-space (Mamba/Mamba2), gated-linear-attention, and recurrent architectures, and the hybrid models that interleave them with attention. The OpenXLA/IREE backend serves only attention-based dense and MoE decoder families today; `Config::from_json` rejects every recurrent `model_type` at the unsupported-architecture error (`src/lib/mlxcel-xla/src/emitter/config.rs`). The MLX engine already implements all of these families (see `src/server/prompt_cache/hybrid_ssm.rs`, `HYBRID_SSM_MODEL_TYPES`), so this epic brings the OpenXLA path to parity.
Target set: Mamba / Mamba2, Jamba, Falcon-H1, RWKV7, Qwen3-Next, Plamo2, Nemotron-H, Kimi-Linear, RecurrentGemma.
This work was deferred out of the #502 design issue by an explicit, recorded decision: a first reference architecture needs a new scan graph primitive, a new per-slot state resource in the C shim, a new per-layer mixer dispatch, and a new validation harness, which is comparable in size to the entire Stage 2 continuous-batching effort rather than to the cheap per-family dense work.
## Goal
Land recurrent / state-space / hybrid architectures on the OpenXLA/IREE backend, reusing the config-driven emit, the continuous-batching engine, the IREE runtime, and the serve worker unchanged, by adding the two genuinely new pieces:
1. A recurrent/state-space mixing graph primitive (selective scan / linear recurrence) that does not fit the shared attention core.
2. Explicit per-slot recurrent STATE allocation and carry through the continuous-batching engine, distinct from the KV cache (fixed-size, does not grow with position), coexisting with the KV cache in hybrid models.
## New graph primitives and state handling
The design (`spike/openxla/SSM_HYBRID_DESIGN.md`) reduces the nine target families to three primitive families (selective SSM/SSD, gated linear attention / delta rule, real-gated linear recurrent unit) that all share one abstraction: a per-token linear recurrence over a fixed-size state, with a chunked/parallel prefill form and a one-step decode form. The recurrent state threads through the same functional-graph + shim mechanism the KV cache already uses (`kcache_b` / `vcache_b`, `xla_ensure_batch_kv` in `csrc/xla_iree.c`), as parallel per-slot buffers with their own fixed shape and lifecycle. Hybrid models carry both a KV region (attention layers) and a state region (recurrent layers) side by side, advanced together in one `pump` step.
## Sub-issues
### Phase 1: Foundation
- [ ] Spike: the selective-scan graph (decode step + chunked-SSD prefill) and the causal `conv1d`, plus their IREE-CUDA (GB10 / sm_121) and IREE-Metal lowering, with a token-exact single-layer probe vs an HF reference (mirrors Stage 2a spike-first).
- [ ] `emitter/ssm.rs`: the `conv1d` + `selective_scan` graph primitives, an `SsmConfig`, and the per-layer `layer_mixers` schedule so `emit_transformer_layer` dispatches attention vs recurrent per layer; `weight_specs` for the SSM tensors.
- [ ] State carry: shim `sscache` / `sscache_b` + `xla_ensure_batch_state`, state-threaded prefill-slot and ragged-decode graphs, and the device-side per-slot state seed on admit.
### Phase 2: Reference architecture
- [ ] Mamba2 end to end: loads and generates via `MLXCEL_BACKEND=xla`, serves via `mlxcel-server`, token-exact single-seq vs the HF fp32 oracle, and reference-exact serve (the Stage 2a reference-equivalence gate extended to prove the state carry).
### Phase 3: Hybrid coexistence
- [ ] Nemotron-H (Mamba2 + attention + MLP/MoE): the per-layer mixer dispatch and the dual KV + state carry, with no new primitive.
- [ ] Jamba, Falcon-H1, Plamo2 on the SSM primitive + hybrid dispatch.
### Phase 4: Linear-attention / delta and recurrent families
- [ ] Qwen3-Next (GatedDeltaNet), Kimi-Linear, RWKV7: the linear-attention / delta-rule recurrence kernels on the same state-carry machinery.
- [ ] RecurrentGemma (RG-LRU) over the existing sliding-window mask.
## Acceptance Criteria
- [ ] The selective-scan primitive and the per-slot state carry are implemented and covered by the reference-equivalence gate.
- [ ] Mamba2 loads and generates via the CLI `MLXCEL_BACKEND=xla` path and serves via `mlxcel-server`, token-exact single-seq vs the HF fp32 oracle and reference-exact serve.
- [ ] At least one hybrid family (Nemotron-H) validates the dual KV + state carry end to end.
- [ ] The OpenXLA recurrent path never participates in prefix-hash sharing (the recurrent state cannot be reconstructed from a token-prefix hash), matching the MLX carve-out in `src/server/prompt_cache/hybrid_ssm.rs`.
## Technical Considerations
- Parent epic: #493. Scoping design: #502 (`spike/openxla/SSM_HYBRID_DESIGN.md`).
- The prefill scan lowering (chunked-SSD or associative) on IREE-CUDA / Metal is the one genuinely new graph problem and must be spike-validated first, exactly as Stage 2a validated the ragged KV write.
- The recurrence may need f32 scan accumulation even on the f16 GPU precision path for numerical stability; the epic's sub-0.01-logit near-tie tolerance applies.
- Sequencing: Mamba2 then Nemotron-H are the foundation; Qwen3-Next, Kimi-Linear, and RWKV7 are recent and each add their own kernel, so they follow the foundation rather than lead it.
Contributor guide
Assessment
This issue has not been assessed yet.