lablup / lablup/mlxcel

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

Open
#564 0 comments 0 reactions 1 assignee Claimed by @inureyes View on GitHub
area:architecture area:inference priority:low status:ready type:enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.