feat: add GLM-5.3-Flash MTP support

atom/config.py

@@ -639,6 +639,18 @@ def glm5_kpool_block_size(index_kpool: int) -> int:
     return index_kpool * _MQA_LOGITS_PRESHUFFLE_ROWS
 
 
+def _glm5_next_unsupported_features(config: "Config") -> list[str]:
+    """Return parallel modes that do not yet preserve GLM-5.3 k-pool state."""
+    unsupported = []
+    if config.prefill_context_parallel_size > 1:
+        unsupported.append("PCP")
+    if config.decode_context_parallel_size > 1:
+        unsupported.append("DCP")
+    if config.enable_tbo or config.enable_tbo_decode:
+        unsupported.append("TBO")
+    return unsupported
+
+
 _CONFIG_REGISTRY: dict[str, str] = {
     "deepseek_v32": "deepseek_v3",
     "deepseek_v4": "deepseek_v3",  # V4 reuses V3 schema; V4-specific fields
@@ -1134,6 +1146,8 @@ class SpeculativeConfig:
         "qwen3_5_moe_text": "qwen3_5_mtp",
         "mimo_v2": "mimo_v2_mtp",
         "mimo_v2_flash": "mimo_v2_mtp",
+        "glm5_next": "glm5_next_mtp",
+        "glm5_next_text": "glm5_next_mtp",
     }
 
     # mtp_model_type → (n_predict_attr, architecture)
@@ -1142,6 +1156,7 @@ class SpeculativeConfig:
         "deepseek_v4_mtp": ("num_nextn_predict_layers", "DeepseekV4MTPModel"),
         "qwen3_next_mtp": ("num_nextn_predict_layers", "Qwen3NextMTPModel"),
         "qwen3_5_mtp": ("mtp_num_hidden_layers", "Qwen3_5MTPModel"),
+        "glm5_next_mtp": ("num_nextn_predict_layers", "Glm5NextMTPModel"),
     }
 
     def use_dspark(self) -> bool:
@@ -2079,15 +2094,7 @@ def __post_init__(self):
         # set here and the attention builder sizes the index cache from it.
         is_glm5_next = any("Glm5Next" in str(a) for a in arches)
         if is_glm5_next:
-            unsupported_features = []
-            if self.prefill_context_parallel_size > 1:
-                unsupported_features.append("PCP")
-            if self.decode_context_parallel_size > 1:
-                unsupported_features.append("DCP")
-            if self.speculative_config is not None:
-                unsupported_features.append("speculative decoding")
-            if self.enable_tbo or self.enable_tbo_decode:
-                unsupported_features.append("TBO")
+            unsupported_features = _glm5_next_unsupported_features(self)
             if unsupported_features:
                 raise ValueError(
                     "GLM-5.3-Flash text serving does not yet support "

atom/model_ops/attentions/kimi_mla_gdn_attn.py

@@ -197,14 +197,26 @@ def _index_rows_per_block(self) -> int:
             )
         return rows
 
+    def _kpool_history_size(self) -> int:
+        """Return ring rows needed to survive speculative rejection."""
+        pool = self._kpool_size()
+        spec = self.model_runner.config.speculative_config
+        if spec is None:
+            return pool
+        # Two speculative windows plus the incomplete pool survive rejection.
+        window = spec.num_speculative_tokens + 1
+        return 1 << (pool + 2 * window - 1).bit_length()
+
     def _kpool_tail_bytes(self) -> int:
         """Per-request tail bytes across every indexer-owning layer."""
         kpool = self._kpool_size()
         if kpool <= 1 or not getattr(self.model_runner, "has_mla_indexer", False):
             return 0
         hf = self.model_runner.config.hf_config
         index_cache_layer_ids, _ = self._index_cache_layout()
-        per_layer = 2 * kpool * hf.index_head_dim * torch.bfloat16.itemsize
+        per_layer = (
+            2 * self._kpool_history_size() * hf.index_head_dim * torch.bfloat16.itemsize
+        )
         return len(index_cache_layer_ids) * per_layer
 
     def _kpool_tail_plane_shape(self) -> tuple[int, int] | None:
@@ -263,7 +275,7 @@ def allocate_per_req_cache(self, entries: dict[str, int]) -> dict:
                 len(index_cache_layer_ids),
                 entries.get(STATE_SLOT_CLASS, 0),
                 2,  # 0 = K, 1 = gate score
-                self._kpool_size(),
+                self._kpool_history_size(),
                 hf.index_head_dim,
             ),
             dtype=torch.bfloat16,

atom/model_ops/glm5_next/indexer.py

@@ -141,8 +141,36 @@ def _sparse_attn_indexer_kpool(
         raise NotImplementedError(
             "GLM-5.3 kpool does not support DCP/PCP; use dcp=pcp=1"
         )
-    if not context.is_prefill and attn_metadata.max_seqlen_q > 1:
-        raise NotImplementedError("GLM-5.3 kpool does not support speculative decode")
+    # Prefer explicit scheduler metadata. The ragged-query condition covers
+    # metadata implementations that do not expose ``num_spec_decodes``.
+    is_speculative_verify = getattr(attn_metadata, "num_spec_decodes", 0) > 0 or (
+        not context.is_prefill and attn_metadata.max_seqlen_q > 1
+    )
+    if is_speculative_verify:
+        from .speculative import run_speculative_kpool_indexer
+
+        run_speculative_kpool_indexer(
+            attn_metadata,
+            kv_cache,
+            q_fp8,
+            k,
+            gate_score,
+            weights,
+            compress_ape,
+            tail_cache,
+            state_slot_idx_in,
+            state_slot_idx,
+            positions,
+            sparse_kv_indices_buffer,
+            index_kpool,
+            topk_tokens,
+            topk_out_width,
+            get_current_atom_config().kv_cache_block_size,
+            max_model_len,
+            scale_fmt,
+            stable_topk,
+        )
+        return result
 
     device = hidden_states.device
     block_size = get_current_atom_config().kv_cache_block_size