docs: scope the OpenXLA SSM / hybrid / recurrent track (#502)#565
Merged
Conversation
Add spike/openxla/SSM_HYBRID_DESIGN.md, the Window D scoping design for epic #493. It grounds the recurrent/state-space track in the current OpenXLA/IREE backend: the recurrent/state-space mixing primitives (selective scan, linear recurrence, delta rule) that do not fit the shared attention core, the per-slot recurrent state allocation and carry through the continuous-batching engine (distinct from the KV cache, fixed-size, threaded as parallel device buffers through the same shim mechanism the KV cache uses), and the coexistence with KV-cache attention layers in hybrid models. Records an explicit decision to defer the implementation to the follow-up epic #564, per the second option of the issue acceptance criteria, with a staged plan (scan spike, primitive, state carry, Mamba2 reference, Nemotron-H hybrid, then breadth).
14 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Scoping design for Window D of epic #493: the SSM / hybrid / recurrent architecture track on the OpenXLA/IREE backend (Mamba/Mamba2, Jamba, Falcon-H1, RWKV7, Qwen3-Next, Plamo2, Nemotron-H, Kimi-Linear, RecurrentGemma). Design-doc only; no code changes.
What changed
spike/openxla/SSM_HYBRID_DESIGN.md, matching theSTAGE2_DESIGN.mdconvention. It grounds the track in the current codebase and covers the three deliverables the issue asks for:AttnLayout.kcache_b/vcache_b,xla_ensure_batch_kv).emit_transformer_layer, with a slot carrying both a KV region and a state region.Design decision
Deferred to a dedicated follow-up epic, #564, per the second option of the issue acceptance criteria. A first reference (Mamba2) requires a new scan graph primitive (whose prefill form is an unsolved lowering question), a new C-ABI per-slot state resource, a new per-layer mixer dispatch, and a new state-aware validation harness, comparable in size to the whole Stage 2 continuous-batching effort.
Test plan
AttnLayout,emit_ffn_body,moe_block,XlaBatchEngine,IreeRaggedLlama,xla_ensure_batch_kv,hybrid_ssm.rs,Mamba2Cache,kv_arch.rs) verified against the current tree.Closes #502