epic: OpenXLA SSM / hybrid / recurrent architecture support (Mamba2, Nemotron-H, Jamba, and more)

[inureyes]

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: epic: OpenXLA backend architecture-coverage parity with the MLX engine #493. Scoping design: design: scope the SSM / hybrid / recurrent track (Mamba2, Jamba, RWKV7, Qwen3-Next, and more) #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.