Skip to content

Proposal: support Kyutai's REKI (Qwen3.5-2B + DiT expert with register memory) #45

Open

Description

@ravediamond

Hi, after ACT I'd like to add REKI (paper, code). It's a Qwen3.5-2B VLM with a 0.2B flow-matching DiT expert, and it has a memory: 32 register tokens per frame in the expert, cached for the last 3 frames (1 s apart). The VLM keeps no state, only the registers persist.

Checkpoint: kyutai/reki-2b, plus LIBERO, RoboTwin and RoboCasa fine-tunes (BF16 safetensors). Weights are CC-BY-NC-SA 4.0 and the code is MIT, so the converter only reads weights users download.

Why: it's one of the few open VLAs with memory, and the released stack only runs on linux-64 with CUDA 12, so not on Jetson, Mac or CPU-only boards today.

Plan:

  1. Qwen3.5 text stack (Gated DeltaNet + full attention, via ggml_gated_delta_net / ggml_ssm_conv, both in b11223) and its vision tower, as a module next to qwen3_lm.
  2. The expert (AdaRMS with gated residuals, self- and cross-attention, SwiGLU) and the register path.
  3. Converter from the released safetensors, vla_predict_check against the reference on reki-2b-libero over 3 consecutive frames, then a benchmark row.

One question before I start: the memory needs per-handle state. Would you accept vla_reset() in the C API (the register cache stored in the model handle, appended after each predict) and a reset flag on PredictRequest? Or would you prefer the caller to own the state (pass the cached register K/V in and get them back out)?

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions