GLM-5.2/GLM-5.3 PR 2: GLM-5.2 with MLA, fp8 MoE and DSA

configs/glm52_tp8.yaml

@@ -0,0 +1,30 @@
+model: "glm52"
+# GLM-5.2 on 8 GPUs of 141 GB. The fp8 checkpoint is ~750 GB; routed experts
+# stay fp8-resident, so weights are ~96 GB/rank at TP8. Resolve the
+# checkpoint through the HF cache (HF_HOME) or pass
+# model_kwargs.checkpoint_path.
+#
+# Routed experts run the fused fp8 W8A8 kernel (moe_quant_kernel: auto), which
+# is what lets decode capture. model_kwargs.moe_quant_kernel: reference serves
+# the per-expert dequant loop instead, eager, as the numerics anchor.
+#
+# Without dsa_long_context the context is capped at index_topk (2048), where
+# dense MLA equals the sparse reference exactly; the submodule refuses more.
+# See glm52_tp8_longctx.yaml to serve beyond it.
+max_seq_len: 2048
+# A 750 GB load outruns PyTorch's default process-group timeout.
+dist_timeout_s: 7200
+# A request holds at most 16 pages (17 with MTP's verify rows): 64 of them
+# fit the pool below, so a decode step never finds it empty.
+max_concurrent_requests: 64
+resources:
+  kv:
+    # MLA latent pages: 576 dims x bf16 = 1.15 KB/token/layer, ~90 KB/token
+    # across 78 layers. 1536 x 128-token pages = 196k tokens, ~17.7 GB/rank.
+    max_num_pages: 1536
+    page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0, 1, 2, 3, 4, 5, 6, 7]
+    tp_size: 8
+    graph_walks: [prefill, decode]

configs/glm52_tp8_longctx.yaml

@@ -0,0 +1,38 @@
+model: "glm52"
+# GLM-5.2 TP8 with the DSA long-context path. Same box shape as
+# glm52_tp8.yaml; the difference is the model_kwargs below.
+#
+# dsa_long_context lifts the context <= index_topk (2048) guard to
+# max_seq_len. Indexer layers then keep a per-request k-store outside the
+# paged pool: fp32 on CUDA (k_norm autocasts to fp32), ~10.75 KB/token/rank
+# across the 21 full-attention layers, ~2.1 GB/rank at 24 x 8192 tokens.
+# Decode beyond 2048 runs sparse absorbed attention over each token's
+# top-2048 latents. A request that stays within 2048 runs the dense paged
+# path bit-identically to glm52_tp8.yaml; batched with a longer one it takes
+# the sparse gather instead (fp32, outside the engine's bf16 autocast): same
+# selection, but not bitwise equal to the dense path.
+#
+# Limits, both follow-ups in mstar/engine/resources/dsa/manager.py: prefill prompts must
+# still fit index_topk, and serving is eager only, since the indexer's
+# per-request host work would not replay in a captured graph.
+max_seq_len: 8192
+dist_timeout_s: 7200
+# A full-window request holds 64 pages (8191 rows): 24 of them fit the pool
+# below, so a decode step never finds it empty.
+max_concurrent_requests: 24
+model_kwargs:
+  dsa_long_context: true
+  # Becomes config.max_seq_len: the guard bound and the KV sizing input.
+  # Must stay <= max_num_pages * page_size per rank.
+  max_seq_len: 8192
+resources:
+  kv:
+    # 24 x 64 pages + the sink page the manager holds back. At 1.15
+    # KB/token/layer, ~17.7 GB/rank.
+    max_num_pages: 1537
+    page_size: 128
+node_groups:
+  - node_names: [LLM]
+    ranks: [0, 1, 2, 3, 4, 5, 6, 7]
+    tp_size: 8
+    graph_walks: [prefill, decode]

configs/test/glm52_single_rank.yaml

@@ -0,0 +1,18 @@
+model: "glm52"
+# GLM-5.2 on a single rank, for dummy mode and small-scale development. The
+# real checkpoint is ~750 GB fp8, so production serving is TP8: see
+# configs/glm52_tp8.yaml.
+#
+# max_seq_len here is presence-checked only. The serving ceiling stays
+# Glm52ModelConfig.max_seq_len, and without dsa_long_context the hard limit
+# is index_topk (2048) regardless.
+max_seq_len: 2048
+model_kwargs:
+  # no checkpoint: dummy mode (no submodule is built); the registry's repo id
+  # would download all 750 GB and allocate the model on one GPU. Name a
+  # checkpoint, or config_variant: reduced with tokenizer_mode: byte, to serve.
+  model_path_hf: ""
+node_groups:
+  - node_names: [LLM]
+    ranks: [0]
+    graph_walks: [prefill, decode]

docs/adding_models.rst

@@ -477,10 +477,10 @@ three common fields:
   branch nodes share one KV pool. A node that no spec names receives no resources. Such a
   node was called "stateless" in earlier versions.
 - ``depends_on()`` returns the keys of other specs that this spec is built against.
-  ``AttentionSpec`` and ``PositionSpec`` depend on the ``kv_cache`` key they name. If a
-  spec names a dependency that the model does not declare, loading fails immediately,
-  rather than during a forward pass. ``RaggedAttentionSpec`` is the one attention spec
-  that depends on nothing: it names no cache, because it has none.
+  ``AttentionSpec``, ``PositionSpec`` and ``DSASpec`` depend on the ``kv_cache`` key they
+  name. If a spec names a dependency that the model does not declare, loading fails
+  immediately, rather than during a forward pass. ``RaggedAttentionSpec`` is the one
+  attention spec that depends on nothing: it names no cache, because it has none.
 
 The spec types are:
 
@@ -538,6 +538,19 @@ The spec types are:
        ``sliding_window=N`` lets each query attend to the last ``N`` positions, its
        own included. It needs the FlashInfer backend and causal attention, and the
        pages behind the window stay allocated.
+   * - ``DSASpec(config=DSAConfig(kv_cache=..., topk=...))``
+     - DeepSeek sparse attention (DSA) over the named ``KVLayout.MLA`` cache. A step
+       declares ``DSAStep()``, and the resource takes each request's rows from the KV plan.
+       The model's indexer layers store their index keys (``store_keys``) and pick each
+       query's top ``topk`` positions by the indexer's score (``select``); a layer that
+       reuses an earlier layer's pick reads it with ``shared_selection``. ``attend`` runs
+       absorbed MLA over just the picked latents. The keys sit in per-request buffers
+       outside the paged pool and the selection runs per request on the host, so the
+       resource is eager only, a prompt must fit in ``topk`` and past ``topk`` a request
+       decodes one row per step. Once one row of a batch passes ``topk``, every row of it
+       attends its pick. GLM-5.2 declares it under ``dsa_long_context``. Planned: the keys
+       move onto the MLA cache's pages under its page table, so the allocator, eviction and
+       prefix caching cover them, and selection gets fixed-shape buffers so decode captures.
    * - ``CrossAttentionSpec(config=CrossAttentionConfig(...))``
      - Attention over a paged context that is written once and never extended.
        ``backend`` defaults to ``AttnBackend.AUTO``. Only FlashInfer is implemented.
@@ -567,7 +580,11 @@ The spec types are:
      - Position tracking and RoPE. ``scheme`` is ``PosScheme.SEQUENTIAL`` or
        ``PosScheme.BLOCK``. The RoPE parameters are set here: ``rope_theta``,
        ``rope_scale``, and the Llama-3.1 parameters ``low_freq_factor``,
-       ``high_freq_factor`` and ``old_context_len``. A model with learned position
+       ``high_freq_factor`` and ``old_context_len``. Layers rotate with
+       ``apply_qk(q, k, label)``. ``pos_ids=`` rotates at positions the model passes
+       instead of the label's plan, and ``k=None`` rotates one tensor. ``interleave``
+       selects GPT-J pairing, ``rotary_dim`` a partial rotation, and ``rope_table_dtype``
+       a cos/sin table built once (GLM-5.2 uses float32). A model with learned position
        embeddings also declares this spec, because it needs the position counter. Such a
        model never applies RoPE.
    * - ``SamplerSpec(vocab_size=..., enable_repetion_penalty=...)``
@@ -652,6 +669,9 @@ appears in no spec, so it receives no resources:
    * - Decoder of an encoder-decoder model
      - The three or four specs above, plus a second ``KVSpec`` for the encoder context
        and a ``CrossAttentionSpec`` over that second cache.
+   * - An LLM with DeepSeek sparse attention (GLM-5.2)
+     - The specs above over an MLA cache (``KVLayout.MLA``, ``AttnBackend.FLASHINFER_MLA``),
+       plus a ``DSASpec`` over the same cache.
    * - A node that carries per-request state between steps of a loop (an RNN-T
        decoder, recurrent attention)
      - A ``RecurrentStateSpec``, read and written with ``gather`` / ``scatter_``.

docs/installation.rst

@@ -149,6 +149,10 @@ 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``).
+   * - ``.[glm]``
+     - GLM runtime: ``transformers``, ``flashinfer-python``, ``safetensors``,
+       ``triton``, ``huggingface-hub``, ``mooncake-transfer-engine``, and ``ninja``
+       (speeds up the JIT build of the vendored MoE align kernel).
    * - ``.[waypoint]``
      - Index-hosted Waypoint dependencies: ``huggingface-hub``, ``safetensors``,
        and ``tensordict`` for reference validation. The pinned TAEHV implementation

docs/models.rst

@@ -61,6 +61,9 @@ Registry keys live in ``mstar/model/registry.py`` (``MODEL_REGISTRY`` / ``HF_MOD
      - Full-duplex speech-to-speech (NemotronLabs VoiceChat-11B): Fast-Conformer
        encoder + Nemotron-H hybrid Mamba-2/attention backbone + EarTTS talker +
        RVQ codec; user speech in, agent text and 22.05 kHz speech out.
+   * - ``glm52``
+     - ``zai-org/GLM-5.2-FP8``
+     - GLM-5.2 (753B/40B): MLA + DSA MoE causal LM, fp8 block-scale weights, TP8.
    * - ``orpheus``
      - ``canopylabs/orpheus-3b-0.1-ft``
      - TTS: Llama 3.2 3B LLM emitting audio tokens + SNAC 24 kHz decoder.

mstar/api_server/openai/adapters.py

@@ -1338,6 +1338,56 @@ def transcription_to_request(self, req: TranscriptionRequest, audio_path: str) -
         )
 
 
+class TextChatAdapter(OpenAIAdapter):
+    """Text-only LLMs: chat messages in, text out. The model's ``process_prompt``
+    applies its own chat template; sampling fields reach it as model_kwargs."""
+
+    supports_chat = True
+
+    def chat_to_request(self, req: ChatCompletionRequest, upload_dir: Path) -> SubmitArgs:
+        _refuse_tool_calls(req)
+        text, _, modalities, parts = flatten_messages(req.messages, upload_dir, allow_remote=False)
+        # an image or audio part was decoded to disk and then dropped: the model
+        # answered as if it had seen it
+        other = sorted(set(modalities) - {"text"})
+        if other:
+            raise ValueError(f"this model takes text only; the messages carry {other}")
+        if not text:
+            raise ValueError("messages carry no text content")
+        for name in ("max_tokens", "max_completion_tokens"):
+            value = getattr(req, name, None)
+            if value is not None and value < 1:
+                raise ValueError(f"{name} must be at least 1, got {value}")
+        # what the model cannot serve is refused, not dropped: off, a field passes
+        if req.n not in (None, 1):
+            raise ValueError(f"n={req.n}: this model returns one choice")
+        if req.stop:
+            raise ValueError("stop sequences are not supported; the reply ends at a stop token or max_tokens")
+        mk = _passthrough(req)
+        for name in ("logprobs", "top_logprobs", "logit_bias", "presence_penalty", "frequency_penalty",
+                     "stop_token_ids"):
+            if mk.get(name):
+                raise ValueError(f"{name} is not supported; got {mk[name]!r}")
+            mk.pop(name, None)
+        if mk.pop("best_of", None) not in (None, 1):
+            raise ValueError("best_of is not supported; this model returns one choice")
+        if mk.pop("response_format", None) not in (None, {"type": "text"}):
+            raise ValueError("response_format needs constrained decoding; only {'type': 'text'} is served")
+        _apply_sampling(req, mk)
+        # extra_body can override max_tokens with anything; a string reached the
+        # conductor's min() and took its whole message batch down
+        _check_numeric_kwargs(mk)
+        if mk.get("max_output_tokens") is not None and mk["max_output_tokens"] < 1:
+            raise ValueError(f"max_output_tokens must be at least 1, got {mk['max_output_tokens']}")
+        return SubmitArgs(
+            text=text,
+            input_modalities=["text"],
+            output_modalities=["text"],
+            model_kwargs=mk,
+            prompt_parts=parts,
+        )
+
+
 class OmnilingualASRAdapter(OpenAIAdapter):
     """Omnilingual ASR CTC: speech-to-text in 1,600+ languages, no language code.
 
@@ -1419,6 +1469,7 @@ def transcription_to_request(self, req: TranscriptionRequest, audio_path: str) -
     "flux2_klein": DiffusionImageAdapter(),
     "flux2_klein_9b": DiffusionImageAdapter(),
     "z_image_turbo": DiffusionImageAdapter(),
+    "glm52": TextChatAdapter(),
 }
 
 # One key per size; every size takes the same adapter.

mstar/cli/main.py

@@ -73,6 +73,7 @@
     "ltx2_5": "ltx2_5.yaml",
     "flux2_klein_9b": "flux2_klein_9b.yaml",
     "z_image_turbo": "z_image_turbo.yaml",
+    "glm52": "glm52_tp8.yaml",
 }
 
 # qwen 3.5 series (dense)

mstar/engine/resources/__init__.py

@@ -29,6 +29,7 @@
     cross_label,
 )
 from mstar.engine.resources.base import CGSlotSpec, PublishedInfo, Resource
+from mstar.engine.resources.dsa.config import DSAConfig, DSASpec, DSAStep
 from mstar.engine.resources.kv.config import (
     KVConfig,
     KVLayout,
@@ -90,6 +91,9 @@
     "CGSlotSpec",
     "CrossAttentionConfig",
     "CrossAttentionSpec",
+    "DSAConfig",
+    "DSASpec",
+    "DSAStep",
     "KVConfig",
     "KVLayout",
     "KVReqConfig",

mstar/engine/resources/dsa/config.py

@@ -0,0 +1,40 @@
+"""What a model declares for DeepSeek sparse attention (DSA): its config, spec and step.
+
+Kept free of the manager so a submodule can declare a step without importing it.
+"""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from mstar.engine.resources.spec import NodeResourceSpec
+from mstar.engine.resources.step import ResourceStep
+
+if TYPE_CHECKING:
+    from mstar.engine.resources.base import Resource
+
+
+@dataclass
+class DSAConfig:
+    # the paged MLA cache: a step's rows come from its plan, the attended latents from its pages
+    kv_cache: str
+    # the keys a query attends once it sees more: its top ``topk`` by the indexer's score
+    topk: int
+
+
+@dataclass
+class DSASpec(NodeResourceSpec):
+    config: DSAConfig
+
+    def depends_on(self) -> set[str]:
+        return {self.config.kv_cache}
+
+    @property
+    def resource_class(self) -> "type[Resource]":
+        from mstar.engine.resources.dsa.manager import DSAManager
+
+        return DSAManager
+
+
+@dataclass(frozen=True)
+class DSAStep(ResourceStep):
+    pass