Muse Glimmer 30B on a shared VLM scaffold (2/3)

README.md

@@ -117,6 +117,7 @@ _Note_: The **first request(s) on a fresh environment can be slow** — often te
 | [BAGEL](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT) | Unified multimodal | text, image → text, image | `/v1/chat/completions`, `/v1/images/generations` |
 | [Qwen3-Omni](https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct) | Omni | text, image, audio, video → text, speech | `/v1/chat/completions` |
 | [Qwen3.5 0.8B–27B](https://huggingface.co/Qwen/Qwen3.5-4B) / [Qwen3.8-27B](https://huggingface.co/Qwen/Qwen3.8-27B) | Vision-language model | text, image → text (hybrid gated-DeltaNet attention) | `/v1/chat/completions`, `/generate` |
+| [Muse Glimmer 30B](https://huggingface.co/meta-models/Muse-Glimmer-30B) *(Beta)* | Vision-language model | text, image → text (sliding-window and NoPE attention layers 3:1) | `/v1/chat/completions`, `/generate` |
 | [Nemotron VoiceChat](https://huggingface.co/nvidia/NVIDIA-NemotronLabs-VoiceChat-11B) *(Beta)* | Full-duplex speech LM | speech → text, speech (streamed) | `/generate` |
 | [Orpheus](https://huggingface.co/canopylabs/orpheus-3b-0.1-ft) | Speech LM | text → speech | `/v1/audio/speech` |
 | [Fish Audio S2-Pro](https://huggingface.co/fishaudio/s2-pro) | Speech LM | text, reference audio → speech (streamed; weights under the Fish Audio Research License, non-commercial use only) | `/v1/audio/speech` |

benchmark/base.py

@@ -315,7 +315,7 @@ def get_supported_modalities(self):
         return {RequestType.T2T, RequestType.I2T, RequestType.V2T}
 
     def get_model_kwargs(self, request_type: RequestType):
-        # Greedy everywhere (mstar's sampler defaults to 0.6), and the same prompt
+        # Greedy everywhere (mstar's default sampling isn't greedy), and the same prompt
         # render: mstar opens a <think> block by default, while the 0.8B/2B templates
         # only do when asked. mstar reads the flat enable_thinking and vLLM/SGLang
         # read chat_template_kwargs.
@@ -328,6 +328,35 @@ class Qwen3_8_Dense(Qwen3_5_Dense):
     DEFAULT_MODEL_ID = "Qwen/Qwen3.8-27B"
 
 
+class MuseGlimmer(Model):
+    """Muse Glimmer 30B (text + image in, text out) benchmark metadata."""
+
+    HF_URL = "meta-models/Muse-Glimmer-30B"
+
+    def get_hf_url(self):
+        return self.HF_URL
+
+    def get_supported_modalities(self):
+        return {RequestType.T2T, RequestType.I2T}
+
+    def get_model_kwargs(self, request_type: RequestType):
+        # the checkpoint samples at temperature 1 by default; every engine
+        # decodes greedily here, so they generate the same tokens
+        return {"temperature": 0.0}
+
+    def get_tokenizer(self):
+        # tokenizer_config.json is in a format older transformers cannot read
+        if self._tokenizer is None:
+            from huggingface_hub import snapshot_download
+
+            from mstar.model.muse.processing import load_tokenizer
+
+            self._tokenizer = load_tokenizer(snapshot_download(
+                self.HF_URL, allow_patterns=["tokenizer.json", "tokenizer_config.json", "chat_template.jinja"],
+            ))
+        return self._tokenizer
+
+
 class Qwen3TTS1p7B(Qwen3TTS):
     HF_URL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
 
@@ -495,6 +524,7 @@ class ModelType(Enum):
     NEMOTRON_ASR = "nemotron_asr"
     QWEN3_5 = "qwen3.5"
     QWEN3_8 = "qwen3.8"
+    MUSE_GLIMMER = "muse_glimmer"
 
     def inst(self, **kwargs) -> Model:
         if self == ModelType.BAGEL:
@@ -539,4 +569,6 @@ def inst(self, **kwargs) -> Model:
             return Qwen3_5_Dense(**kwargs)
         if self == ModelType.QWEN3_8:
             return Qwen3_8_Dense(**kwargs)
+        if self == ModelType.MUSE_GLIMMER:
+            return MuseGlimmer(**kwargs)
         raise NotImplementedError(f"Unknown model type {self}")

benchmark/kv_capacity.py

@@ -0,0 +1,179 @@
+"""How many requests of L tokens a paged KV pool that keeps a window holds, with and without it.
+
+Builds the KV manager as a server would: the model's own KV spec with the serving
+config's ``resources:`` block applied (``apply_yaml_overrides``), sharded to one rank of
+its node's tp * sp heads, on the GPU. It then admits requests one after another as a
+server does and grows each one: a prefill of L - 128 tokens in one step, then 128 one-token
+decode steps. It stops at the first request admission makes wait, or whose step is refused.
+At one rank a request waits until it can be allotted its pages (`ask` / `reserve`); above
+one rank no pool allots, so only a refused step stops it. Every request stays resident.
+A one-step prefill holds all of a prompt's sliding K/V until it commits, so above one rank
+a server that chunks a prompt's prefill can hold more requests of long prompts than this
+counts (at one rank a request is allotted its whole prompt's pages, however it is chunked).
+"without" is the same pool keeping every layer's tokens, with the same max_num_pages, so
+the same bytes.
+
+Only the checkpoint's metadata (configs, tokenizer) is read, from the registry's repo.
+``--model-path`` reads another repo or a local directory instead:
+
+    python benchmark/kv_capacity.py --config configs/muse_glimmer.yaml \\
+        --model-path /path/to/checkpoint --out glimmer.json
+"""
+import argparse
+import dataclasses
+import json
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import torch
+import yaml
+
+# run as a file, import the tree this file is in
+REPO = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(REPO))
+
+from mstar.communication.tensors import LocalTransferEngine  # noqa: E402
+from mstar.engine.resources.kv.config import (  # noqa: E402
+    KVReqConfig,
+    KVSpec,
+    KVStep,
+    PagedKVConfig,
+    PrefixSpan,
+    StreamLayout,
+)
+from mstar.engine.resources.kv.manager import KVManager  # noqa: E402
+from mstar.engine.resources.kv.transfer import TransferEngineInfo  # noqa: E402
+from mstar.engine.resources.spec import apply_yaml_overrides  # noqa: E402
+from mstar.engine.resources.step import Segment, StepContext  # noqa: E402
+
+DECODE = 128
+
+
+def kv_configs(config, model_path=None):
+    """The serving config's sliding KV pool for one rank, with and without its sliding layers."""
+    from mstar.model.registry import get_model_class, model_init_kwargs
+
+    serving = yaml.safe_load(Path(config).read_text())
+    kwargs = {**model_init_kwargs(serving["model"]), **(serving.get("model_kwargs") or {})}
+    if model_path:
+        kwargs["model_path_hf"] = model_path
+    specs = get_model_class(serving["model"])(**kwargs).get_node_resources()
+    apply_yaml_overrides(specs, serving)
+    sliding = [s for s in specs if isinstance(s, KVSpec) and isinstance(s.config, PagedKVConfig)
+               and len(s.config.groups) > 1]
+    if len(sliding) != 1:
+        raise ValueError(f"{config} declares {len(sliding)} paged KV pools with full and sliding layers, not one")
+    (spec,) = sliding
+    cfg = spec.config
+    # A copy declares what the pool did and keeps every token, since the pool derives its window
+    # from its attentions; max_num_pages counts the same bytes either way.
+    without = dataclasses.replace(cfg)
+    # one rank's share of the heads, as the KV manager shards it
+    groups = [g for g in serving["node_groups"] if spec.nodes & set(g["node_names"])]
+    if not groups:
+        raise ValueError(f"{config}: no node group holds {sorted(spec.nodes)}")
+    shards = {g.get("tp_size", 1) * g.get("sp_size", 1) for g in groups}
+    if len(shards) != 1:
+        raise ValueError(f"{config}: the node groups holding {sorted(spec.nodes)} "
+                         f"({[g['node_names'] for g in groups]}) shard its heads {sorted(shards)} ways, not one way")
+    (shards,) = shards
+    cfg.shard(shards)
+    without.shard(shards)
+    return cfg, without, int(serving.get("max_seq_len", cfg.max_seq_len)), shards
+
+
+def manager(cfg, shards=1, device="cuda"):
+    # Built at one rank, so above one it is told not to lead: like a pool there, it then
+    # doesn't decide admission (`KVManager._decides`).
+    return KVManager(
+        cfg=cfg, name="kv", joint_comm_group=None,
+        transfer_engine_info=TransferEngineInfo("h", "h", LocalTransferEngine("h")),
+        device=torch.device(device), dtype=torch.bfloat16, leads=shards == 1,
+    )
+
+
+def step(m, rid, span):
+    kv = KVStep(segments=(Segment(rid, "main", span),))
+    ctx = StepContext(request_ids=(rid,), graph_walk="walk", slot=0, capture=False)
+    if not m.admit(kv, ctx).ok:
+        return False
+    m.plan(kv, ctx)
+    m.commit(kv, ctx)
+    return True
+
+
+def admitted(m, rid, prompt, decode):
+    """Whether admission lets ``rid`` in: a pool that decides allots it its prompt's pages,
+    and one that doesn't asks nothing."""
+    layout = StreamLayout(spans=(PrefixSpan(prompt, prompt, "walk"),), growth=(("decode", 1),))
+    m.ingest_request(rid, KVReqConfig(kv_layout={"main": layout}, loop_caps={"decode": decode}))
+    asked = m.ask(rid, "node", "walk")
+    if asked is None:
+        return True
+    # a request that can never fit fails its ask
+    return asked.ok and m.reserve(rid, "node", "walk", None) is None
+
+
+def fill(cfg, length, shards=1, decode=DECODE, device="cuda"):
+    m = manager(cfg, shards, device)
+    n = 0
+    while True:
+        rid = f"r{n}"
+        if not (
+            admitted(m, rid, length - decode, decode)
+            and step(m, rid, length - decode) and all(step(m, rid, 1) for _ in range(decode))
+        ):
+            m.remove_request(rid)
+            break
+        n += 1
+    counts = m.page_counts()
+    return {"requests": n, "used_pages": counts.used, "total_pages": counts.total,
+            "kv_bytes": m.kv_cache.tensor.nbytes, "page_size": cfg.tokens_per_page,
+            "groups": [(len(g.layers), g.page_size, None if g.retention is None else g.retention.context_budget)
+                       for g in cfg.groups]}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--config", required=True, help="serving config, e.g. configs/muse_glimmer.yaml")
+    ap.add_argument("--out", help="JSON results, rewritten after each cell (required unless --one)")
+    ap.add_argument("--lengths", default="8192,32768,131072",
+                    help=f"request lengths; those past the config's max_seq_len, or not past {DECODE} "
+                         "decode tokens, are skipped")
+    ap.add_argument("--model-path", help="read the checkpoint's configs from this repo or directory")
+    ap.add_argument("--one", help="length,arm: run one cell in this process and print its JSON")
+    a = ap.parse_args()
+    if a.one:
+        length, arm = a.one.split(",")
+        with_cfg, without_cfg, _, shards = kv_configs(a.config, a.model_path)
+        cfg = with_cfg if arm == "with" else without_cfg
+        print("RESULT " + json.dumps(fill(cfg, int(length), shards)), flush=True)
+        return
+    if not a.out:
+        ap.error("--out is required")
+    _, _, max_seq_len, _ = kv_configs(a.config, a.model_path)
+    lengths = [n for n in map(int, a.lengths.split(",")) if DECODE < n <= max_seq_len]
+    # a process per cell: a built KV manager's pool isn't returned to the GPU when it's dropped
+    results = {}
+    for length in lengths:
+        for arm in ("without", "with"):
+            t0 = time.time()
+            cmd = [sys.executable, str(Path(__file__).resolve()), "--config", a.config, "--one", f"{length},{arm}"]
+            if a.model_path:
+                cmd += ["--model-path", a.model_path]
+            proc = subprocess.run(cmd, capture_output=True, text=True, check=False)
+            if proc.returncode:
+                sys.stderr.write(proc.stderr)
+                raise SystemExit(f"cell {length} {arm} exited {proc.returncode}")
+            (line,) = [x for x in proc.stdout.splitlines() if x.startswith("RESULT ")]
+            r = json.loads(line[len("RESULT "):])
+            r["seconds"] = round(time.time() - t0, 1)
+            results[f"{length} {arm}"] = r
+            print(length, arm, r, flush=True)
+            Path(a.out).write_text(json.dumps(results, indent=1))
+
+
+if __name__ == "__main__":
+    main()

benchmark/muse_glimmer/serve_checks.py

@@ -0,0 +1,134 @@
+"""Live checks of a running Muse Glimmer server through its OpenAI API.
+
+Against a server started with ``mstar serve muse_glimmer``; exits nonzero on any
+failed check:
+
+  stream      a greedy answer streamed equals the same answer non-streamed
+  repeat      two identical greedy requests return identical text
+  image       an image's color is named; a two-image prompt answers (the checkpoint
+              itself calls a red and a blue square "both red": HF does too)
+  bad input   an unknown reasoning effort and an unsupported video part are 400s,
+              and the server answers normally afterwards
+  abort       a stream dropped after a few chunks leaves the server serving
+
+  python benchmark/muse_glimmer/serve_checks.py --port 8000
+"""
+import argparse
+import base64
+import io
+import json
+import sys
+import urllib.error
+import urllib.request
+
+MODEL = "muse_glimmer"
+GREEDY = {"temperature": 0.0, "reasoning_effort": "low"}
+
+
+def _png(color: str) -> str:
+    from PIL import Image
+
+    buf = io.BytesIO()
+    Image.new("RGB", (96, 64), color).save(buf, format="PNG")
+    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
+
+
+def _mp4() -> str:
+    """A real 8-frame clip, so the request reaches the model's own refusal."""
+    import av
+    import numpy as np
+
+    buf = io.BytesIO()
+    with av.open(buf, "w", format="mp4") as container:
+        stream = container.add_stream("libx264", rate=8)
+        stream.width = stream.height = 64
+        stream.pix_fmt = "yuv420p"
+        for i in range(8):
+            frame = np.full((64, 64, 3), 30 * i, dtype=np.uint8)
+            for packet in stream.encode(av.VideoFrame.from_ndarray(frame, format="rgb24")):
+                container.mux(packet)
+        for packet in stream.encode():
+            container.mux(packet)
+    return "data:video/mp4;base64," + base64.b64encode(buf.getvalue()).decode()
+
+
+class Client:
+    def __init__(self, port: int):
+        self.url = f"http://localhost:{port}/v1/chat/completions"
+
+    def _request(self, body: dict):
+        req = urllib.request.Request(
+            self.url, json.dumps({"model": MODEL, **body}).encode(),
+            {"Content-Type": "application/json"},
+        )
+        return urllib.request.urlopen(req, timeout=600)
+
+    def chat(self, content, **body) -> str:
+        with self._request({"messages": [{"role": "user", "content": content}], **body}) as r:
+            return json.load(r)["choices"][0]["message"]["content"] or ""
+
+    def stream(self, content, max_chunks: int | None = None, **body) -> list[str]:
+        chunks = []
+        with self._request({"messages": [{"role": "user", "content": content}], "stream": True, **body}) as r:
+            for raw in r:
+                line = raw.decode().strip()
+                if not line.startswith("data:") or line == "data: [DONE]":
+                    continue
+                delta = json.loads(line[5:])["choices"][0].get("delta", {}).get("content")
+                if delta:
+                    chunks.append(delta)
+                if max_chunks is not None and len(chunks) >= max_chunks:
+                    break   # leaving the `with` drops the connection mid-stream
+        return chunks
+
+    def status_of(self, content, **body) -> int:
+        try:
+            with self._request({"messages": [{"role": "user", "content": content}], **body}) as r:
+                return r.status
+        except urllib.error.HTTPError as e:
+            return e.code
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--port", type=int, default=8000)
+    args = ap.parse_args()
+    c = Client(args.port)
+    failures = []
+
+    def check(name: str, ok: bool, detail: str) -> None:
+        print(f"{'PASS' if ok else 'FAIL'} {name}: {detail}", flush=True)
+        if not ok:
+            failures.append(name)
+
+    q = "Name three primary colors."
+    whole = c.chat(q, max_tokens=64, **GREEDY)
+    streamed = "".join(c.stream(q, max_tokens=64, **GREEDY))
+    check("stream", whole == streamed and bool(whole), f"{whole[:60]!r} vs {streamed[:60]!r}")
+    check("repeat", c.chat(q, max_tokens=64, **GREEDY) == whole, "same request, same text")
+
+    for color in ("red", "blue"):
+        one = c.chat([{"type": "image_url", "image_url": {"url": _png(color)}},
+                      {"type": "text", "text": "What color is this image? One word."}], max_tokens=64, **GREEDY)
+        check(f"image {color}", f"to=user<|message|>{color}" in one.lower(), repr(one[-60:]))
+    two = c.chat([{"type": "image_url", "image_url": {"url": _png("red")}},
+                  {"type": "image_url", "image_url": {"url": _png("blue")}},
+                  {"type": "text", "text": "Name the colors of the two images."}], max_tokens=64, **GREEDY)
+    check("two images", bool(two.strip()), repr(two[-60:]))
+
+    check("bad reasoning effort", c.status_of(q, max_tokens=4, temperature=0.0, reasoning_effort="ultra") == 400,
+          "reasoning_effort='ultra'")
+    video = [{"type": "video_url", "video_url": {"url": _mp4()}},
+             {"type": "text", "text": "What happens?"}]
+    check("video refused", c.status_of(video, max_tokens=4, **GREEDY) == 400, "video part")
+    check("serves after bad input", c.chat(q, max_tokens=64, **GREEDY) == whole, "the stream check's answer")
+
+    c.stream("Count from 1 to 200, one number per line.", max_chunks=3, max_tokens=512, **GREEDY)
+    check("serves after abort", c.chat(q, max_tokens=64, **GREEDY) == whole, "the stream check's answer")
+
+    print("SERVE CHECKS", "FAIL " + ", ".join(failures) if failures else "PASS")
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())