Overworld Waypoint1.5B integration

configs/waypoint.yaml

@@ -0,0 +1,28 @@
+model: "waypoint"
+
+model_kwargs:
+  variant: "waypoint-1.5-1b-720p"
+  reference_compat: true
+  compile_dit: true
+  cuda_graph: true
+  full_global_ring: false
+  # Rows per rollout step; must be <= resources.kv.num_sessions below.
+  step_batch_size: 1
+
+# The world state is a ring (not a paged KV cache); max_seq_len only satisfies the
+# conductor's config check and matches one frame's token count.
+max_seq_len: 512
+
+# Primary bound on the world pool; must be a positive int <= resources.kv.num_sessions.
+max_concurrent_requests: 1
+
+resources:
+  kv:
+    # ~816 MiB of ring per world at 720P. Raise together with
+    # max_concurrent_requests; the two are checked against each other.
+    num_sessions: 1
+
+# Single-GPU colocated: vae_encoder + dit on rank 0 (TAEHV decode is fused into the dit forward).
+node_groups:
+  - node_names: [vae_encoder, dit]
+    ranks: [0]

docs/adding_models.rst

@@ -265,19 +265,25 @@ The spec types are:
 
    * - Spec
      - What it builds
-   * - ``KVSpec(config=KVConfig(...))``
-     - A paged KV cache. ``KVConfig`` holds ``num_layers``, ``num_kv_heads``,
-       ``head_dim``, ``max_seq_len`` and ``num_qo_heads``. It also holds three fields that
-       a deployment can tune: ``max_num_pages``, ``page_size`` and ``cpu_offload_pages``
-       (the number of pinned host pages used for offload; 0 disables offload).
+   * - ``KVSpec(config=PagedKVConfig(...))``
+     - A paged KV cache. ``KVConfig`` is the abstract base for common model geometry;
+       ``PagedKVConfig`` adds ``max_seq_len`` and the deployment-tunable
+       ``max_num_pages``, ``page_size`` and ``cpu_offload_pages`` fields.
+   * - ``KVSpec(config=RingKVConfig(...))``
+     - A fixed-capacity frame ring. It adds ``tokens_per_frame``, one
+       ``RingKVLayerConfig`` per layer, and the deployment-tunable ``num_sessions``.
+       Ring storage is currently paired with FlexAttention.
    * - ``AttentionSpec(config=AttentionConfig(kv_cache=...))``
      - Self-attention planned over the named cache. ``backend`` selects
-       ``AttnBackend.FLASHINFER`` (the default) or ``AttnBackend.DENSE``.
+       ``AttnBackend.FLASHINFER`` (the default), ``AttnBackend.DENSE`` or
+       ``AttnBackend.FLEX``. FlashInfer and dense attention require a
+       ``PagedKVConfig``; FlexAttention requires a ``RingKVConfig``.
        ``flashinfer_backend`` selects a kernel generation: ``"auto"``, ``"fa2"`` or
-       ``"fa3"``.
+       ``"fa3"`` when the FlashInfer backend is selected.
    * - ``CrossAttentionSpec(config=CrossAttentionConfig(...))``
-     - Attention over a context that is written once and never extended. See
-       `Cross-attention (encoder-decoder models)`_.
+     - Attention over a paged context that is written once and never extended. Only
+       the FlashInfer backend is implemented. See `Cross-attention (encoder-decoder
+       models)`_.
    * - ``RaggedAttentionSpec(config=RaggedAttentionConfig(...))``
      - Cacheless (ragged) varlen self-attention over the segments packed into one
        forward. Nothing is paged, and nothing carries to the next step.
@@ -310,7 +316,7 @@ appears in no spec, so it receives no resources:
 
    # mstar/model/orpheus/orpheus_model.py
    def get_node_resources(self) -> list[NodeResourceSpec]:
-       kv_config = KVConfig(
+       kv_config = PagedKVConfig(
            num_layers=self.config.num_hidden_layers,
            num_kv_heads=self.config.num_key_value_heads,
            head_dim=self.config.head_dim,
@@ -925,6 +931,9 @@ Both types share the base ``CudaGraphConfig`` fields:
   engine's eager batch size for the walk. The default is ``True``, so the engine never
   batches beyond a captured size.
 - ``compile`` runs ``torch.compile`` before capture. The default is ``True``.
+- ``required`` makes every bucket in the config mandatory. If capture fails locally or
+  on another participating rank, warmup raises after rank-wide agreement instead of
+  dropping the bucket and falling back to eager execution. The default is ``False``.
 
 ``BatchedCudaGraphConfig`` also accepts ``total_tokens_multiplier``. Use it when one
 request's step commits KV across several labels that are combined into a single plan, as
@@ -1036,6 +1045,9 @@ Both types share the base ``PiecewiseCudaGraphConfig`` fields:
   default.
 - ``compile`` runs ``torch.compile`` on ``capture_fn`` before capture. The default is
   ``False``.
+- ``required`` makes every declared shape mandatory. If any participating rank cannot
+  capture one, warmup raises instead of leaving that shape on the eager path. The default
+  is ``False``.
 
 **Splitting the declaration.** When a region leases its own slot, exactly one of the two
 declarations must own each resource. The common pattern is for the outer ``declare_step``
@@ -1209,10 +1221,15 @@ that a misspelled setting is never silently ignored:
 
    * - Spec
      - Accepts
-   * - ``KVSpec``
+   * - ``KVSpec`` with ``PagedKVConfig``
      - ``max_num_pages``, ``page_size``, ``max_seq_len``, ``cpu_offload_pages``
-   * - ``AttentionSpec`` / ``CrossAttentionSpec``
-     - ``backend`` (``flashinfer`` / ``dense``), ``flashinfer_backend``
+   * - ``KVSpec`` with ``RingKVConfig``
+     - ``num_sessions``
+   * - ``AttentionSpec``
+     - ``backend`` (``flashinfer`` / ``dense`` / ``flex``),
+       ``flashinfer_backend`` (``auto`` / ``fa2`` / ``fa3``)
+   * - ``CrossAttentionSpec``
+     - ``backend`` (only ``flashinfer`` is implemented), ``flashinfer_backend``
        (``auto`` / ``fa2`` / ``fa3``)
    * - ``RaggedAttentionSpec``
      - ``flashinfer_backend`` (``auto`` / ``fa2`` / ``fa3``),
@@ -1221,8 +1238,9 @@ that a misspelled setting is never silently ignored:
        not the model.
 
 Tune the cache shape on the KV resource, not on the attention resource that reads it. For
-example, ``configs/qwen3tts.yaml`` selects FA2 under ``talker_attn``, while
-``configs/cosmos3_nano.yaml`` sets the page count under its KV key.
+example, ``configs/qwen3tts.yaml`` selects FA2 under ``talker_attn``,
+``configs/cosmos3_nano.yaml`` sets the page count under its paged KV key, and
+``configs/waypoint.yaml`` sets the resident world count under its ring KV key.
 
 .. note::
 

docs/clients.rst

@@ -34,7 +34,7 @@ NDJSON stream.
    * - ``output_modalities``
      - ``text``
      - Comma-separated desired outputs (e.g. ``text``, ``image``, ``audio``, ``video``,
-       ``action``).
+       ``video_frame``, ``action``). ``video_frame`` is streaming-only raw RGB24.
    * - ``streaming``
      - ``true``
      - ``true`` → NDJSON stream of chunks; ``false`` → one JSON document.
@@ -121,7 +121,10 @@ Result and event types live in ``mstar.client``:
 - ``AudioBuffer`` — decoded PCM with ``.sample_rate``; ``.to_wav(path)``, ``.to_numpy()``,
   ``len(...)``.
 - Stream events — ``TextChunk(text)``, ``ImageChunk(data)`` (``.save(path)``),
-  ``AudioChunk(pcm, sample_rate)``.
+  ``AudioChunk(pcm, sample_rate)``, and ``VideoFrameChunk(data, metadata)``. A
+  video-frame chunk validates its width, height, fps, pixel format and frame range;
+  ``.to_numpy()`` returns a zero-copy ``[frame_count, height, width, 3]`` uint8 view.
+  Raw ``video_frame`` requests require ``stream=True``.
 
 .. code-block:: python
 
@@ -171,8 +174,10 @@ Endpoints and model coverage:
      - ``bagel``
      - Image editing (image + prompt → image).
 
-Models without an OpenAI surface (``pi05``, ``vjepa2``, ``vjepa2_ac``) return ``404`` on
-``/v1/*``; use ``/generate`` or the SDK for them.
+Models without an OpenAI surface (``pi05``, ``vjepa2``, ``vjepa2_ac``, ``waypoint``)
+return ``404`` on ``/v1/*``; use ``/generate`` or the SDK for them. In particular,
+Waypoint emits live RGB frame chunks and is not routed through the encoded-video
+``/v1/videos/generations`` endpoint.
 
 .. code-block:: python
 

docs/installation.rst

@@ -94,28 +94,57 @@ Model families and some output formats need extra packages, exposed as pip *extr
    * - ``.[vjepa2]`` / ``.[vjepa2_ac]``
      - V-JEPA 2 runtime: ``safetensors``, ``torchcodec``, ``huggingface-hub``,
        ``mooncake-transfer-engine`` (``vjepa2_ac`` also adds ``flashinfer-python``).
+   * - ``.[waypoint]``
+     - Index-hosted Waypoint dependencies: ``huggingface-hub``, ``safetensors``,
+       and ``tensordict`` for reference validation. The pinned TAEHV implementation
+       must be installed separately as shown below; keeping its direct URL out of
+       package metadata allows ``m-star`` to be published on PyPI. **Also needs**
+       ``flash-attn-4``, which is installed separately —
+       see `flash-attn-4 (Waypoint, FA3 sm90 kernel)`_.
    * - ``.[audio]``
      - ``soundfile`` — only needed to return **non-WAV** audio containers (mp3/flac/…)
        from the OpenAI/SDK audio surfaces. WAV/PCM output works without it.
    * - ``.[dev]``
      - ``ruff`` + ``pytest`` for linting and the test suite.
    * - ``.[all]``
-     - The union of every model extra above — installs the full runtime for all model
-       families in one shot. Convenient for a machine that serves multiple models; heavier
-       and slower to install than a single family's extra. (Still excludes ``flash-attn`` —
-       see `flash-attn (Qwen3-Omni)`_.)
+     - The union of the index-hosted dependencies from every model extra above.
+       Convenient for a machine that serves multiple models; heavier and slower to
+       install than a single family's extra. It excludes the separately installed
+       TAEHV and ``flash-attn`` packages; see below and `flash-attn (Qwen3-Omni)`_.
 
 Combine extras as needed (keep ``--torch-backend=auto`` on every install):
 
 .. code-block:: bash
 
    uv pip install --torch-backend=auto -e ".[bagel,audio,dev]"
 
+Waypoint's dependencies and pinned TAEHV source are two installs. PyPI and other
+standards-conformant indices reject distributions whose metadata declares a direct-URL
+dependency, so the ``waypoint`` extra intentionally does not name TAEHV. Check the
+installer version before installing its source archive: an old pip may report success
+while producing an empty ``UNKNOWN`` wheel.
+
+.. code-block:: bash
+
+   uv --version                 # must be 0.4.0 or newer
+   uv pip install --torch-backend=auto -e ".[waypoint]"
+   uv pip install --no-deps \
+     "taehv @ https://github.com/madebyollin/taehv/archive/7dc60ec6601af2e668e31bc70acc4cb3665e4c22.zip"
+
+Or, in an existing Python 3.12 environment:
+
+.. code-block:: bash
+
+   python -m pip install --upgrade "pip>=24.3"
+   python -m pip install -e ".[waypoint]"
+   python -m pip install --no-deps \
+     "taehv @ https://github.com/madebyollin/taehv/archive/7dc60ec6601af2e668e31bc70acc4cb3665e4c22.zip"
+
 .. tip::
 
-   If you're just getting started or have the disk/time to spare, ``.[all]`` is the
-   recommended install — it pulls every model family's runtime so any model works out of
-   the box, with no need to track which extra goes with which model:
+   If you're just getting started or have the disk/time to spare, ``.[all]`` installs all
+   index-hosted model dependencies in one shot. Waypoint still needs the pinned TAEHV
+   command above, and Qwen3-Omni still needs ``flash-attn``:
 
    .. code-block:: bash
 
@@ -137,13 +166,15 @@ The GPU model families depend on:
   autoregressive backbones (every model with a ``KV_CACHE`` node runs attention through it).
 - **flash-attn** — used by Qwen3-Omni. **Not installed by any extra**; install it separately
   (see `flash-attn (Qwen3-Omni)`_).
+- **flash-attn-4** — Waypoint's flex-attention ``FLASH`` backend. **Not installed by any
+  extra**; install it separately (see `flash-attn-4 (Waypoint, FA3 sm90 kernel)`_).
 - **mooncake-transfer-engine** — RDMA tensor transport for multi-GPU, disaggregated
   deployments. Single-node deployments can use shared-memory (``SHM``) or ``TCP`` transport
   instead (see :doc:`serving`).
 
-Apart from ``flash-attn``, these are installed by the extras above. Your installed ``torch``
-must match your system CUDA toolkit — ``--torch-backend=auto`` handles that for you (next
-section).
+Apart from ``flash-attn`` and ``flash-attn-4``, these are installed by the extras above. Your
+installed ``torch`` must match your system CUDA toolkit — ``--torch-backend=auto`` handles
+that for you (next section).
 
 flash-attn (Qwen3-Omni)
 -----------------------
@@ -226,6 +257,30 @@ Three things to get right:
    FLASH_ATTN_CUDA_ARCHS="90" uv pip install flash-attn==2.8.3.post1 --no-build-isolation
    python -c "import flash_attn; print(flash_attn.__version__)"
 
+flash-attn-4 (Waypoint, FA3 sm90 kernel)
+----------------------------------------
+
+Waypoint's DiT attention runs torch flex-attention with the ``FLASH`` backend by
+default. That backend needs the **flash-attn-4** package, which provides
+``flash_attn.cute`` — the CuTe DSL rewrite of flash-attn; on H100 it runs the
+FA3-style sm90 kernel with TMA and warpgroup specialisation. It is **not on
+PyPI** as of 2026-09-17, and it is **not** pulled in by any extra.
+
+Install it from the upstream repo's ``flash_attn/cute`` subdirectory:
+
+.. code-block:: bash
+
+   git clone https://github.com/Dao-AILab/flash-attention
+   uv pip install --torch-backend=auto ./flash-attention/flash_attn/cute
+
+This is pure Python plus ``nvidia-cutlass-dsl`` — there is no CUDA extension to
+build. Its kernels are JIT-compiled on first use, which adds roughly a minute to
+the first server startup.
+
+If ``flash-attn-4`` isn't installed, set ``MSTAR_FLEX_BACKEND=TRITON`` to fall
+back to the previous Triton flex kernel. It is correct but slower — about 1.8x
+per attention call at 720p.
+
 Matching your CUDA toolkit
 --------------------------
 

examples/sdk_chat.py

@@ -11,7 +11,7 @@
 # Non-streaming
 print(client.chat("What is the capital of France?").text)
 
-# Streaming (yields TextChunk / ImageChunk / AudioChunk)
+# Streaming (yields TextChunk / ImageChunk / AudioChunk / VideoFrameChunk)
 for event in client.chat("Tell me a short story.", stream=True):
     if isinstance(event, TextChunk):
         print(event.text, end="", flush=True)

mstar/__init__.py

@@ -12,12 +12,18 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:  # for type checkers / IDEs only — no runtime import cost
-    from mstar.client import AudioBuffer, GenerateResult, MStarClient  # noqa: F401
+    from mstar.client import (  # noqa: F401
+        AudioBuffer,
+        GenerateResult,
+        MStarClient,
+        VideoFrameChunk,
+    )
 
 _LAZY: dict[str, tuple[str, str]] = {
     "MStarClient": ("mstar.client", "MStarClient"),
     "GenerateResult": ("mstar.client", "GenerateResult"),
     "AudioBuffer": ("mstar.client", "AudioBuffer"),
+    "VideoFrameChunk": ("mstar.client", "VideoFrameChunk"),
 }
 
 __all__ = list(_LAZY)