MiniMax-M3 support — GQA + MSA block-sparse attention, converter (follow-up to #418)

.github/workflows/ci.yml

@@ -17,6 +17,27 @@ jobs:
       - name: C test suite
         run: cd c && make test-c
 
+  # MiniMax-M3 rides the colibri binary (ARCH_M3 auto-detect), so a GLM-side
+  # refactor of the shared hot paths (rmsnorm/act_glu/expert_ffn/MSA) can break
+  # M3 while every GLM test stays green. This gate is the only thing that
+  # notices: tiny random M3 checkpoint -> converter -> numpy oracle -> the
+  # engine's teacher-forcing gate, token-exact on every position (#601).
+  m3-tiny:
+    name: MiniMax-M3 (generated tiny oracle)
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.13'
+          cache: pip
+          cache-dependency-path: c/tools/requirements-m3-tiny.txt
+      - name: Install pinned tiny-fixture dependencies
+        run: python -m pip install -r c/tools/requirements-m3-tiny.txt
+      - name: Generate fixture and run the M3 oracle gate
+        run: make -C c m3-tiny-check
+
   # Every engine, on every platform a release archive is built for. Before this,
   # CI built colibri + inkling on Linux only and kimi_k3 on nothing -- so the fact
   # that inkling and kimi_k3 did not compile under msys2/UCRT64 at all was invisible

c/Makefile

@@ -602,7 +602,7 @@ ifeq ($(VK),1)
 CFLAGS   += -DCOLI_VULKAN
 LDFLAGS  += -lvulkan
 VK_OBJ    = backend_vulkan.o
-VK_SPV    = shaders/qmatmul.spv shaders/qmatmul_gate_up.spv shaders/attention_absorb.spv shaders/rmsnorm.spv
+VK_SPV    = shaders/qmatmul.spv shaders/qmatmul_gate_up.spv shaders/attention_absorb.spv shaders/attention_gqa.spv shaders/rmsnorm.spv
 endif
 
 all: colibri$(EXE)
@@ -668,6 +668,22 @@ deepseek-v4-tiny-generate:
 	@echo "SKIP deepseek-v4-tiny-generate: V4 runtime requires x86-64/aarch64 Linux or Windows/MSYS2"
 endif
 
+# MiniMax-M3 CI gate (#601): tiny random checkpoint -> int8 conversion -> numpy
+# oracle -> the engine's REF/TF gate, token-exact on every prefill and decode
+# position (the driver enforces the counts; the engine itself always exits 0).
+# Deps: tools/requirements-m3-tiny.txt (torch generates, numpy+safetensors
+# convert and run the oracle). M3 rides the main colibri binary, so this needs
+# no engine platform gate beyond colibri's own.
+.PHONY: m3-tiny-generate m3-tiny-check
+m3-tiny-generate:
+	rm -rf m3tiny m3tiny_i8
+	$(PYTHON) tools/make_m3tiny.py m3tiny
+	$(PYTHON) tools/convert_fp8_to_int4.py --arch m3 --indir m3tiny --outdir m3tiny_i8 --ebits 8 --io-bits 8
+	$(PYTHON) tools/oracle_m3.py m3tiny_i8
+
+m3-tiny-check: m3-tiny-generate colibri$(EXE)
+	$(PYTHON) tests/test_m3_tiny.py --binary ./colibri$(EXE) --snap m3tiny_i8 --ref ref_m3.json
+
 # Convenience aliases — 'glm' kept for backward compatibility.
 # On POSIX EXE is empty, so `colibri: colibri$(EXE)` becomes a self-dependency
 # and marking it phony also forces the real binary to rebuild on every invocation.

c/backend_vulkan.h

@@ -29,6 +29,16 @@ void coli_vk_mem_info(size_t *used_bytes, size_t *tensor_count);
  * mem_budget reports device-local usage/budget in GB (VK_EXT_memory_budget);
  * returns 0 if unavailable. */
 void coli_vk_alloc_priority(float p);
+/* Fused gate+up activation (model-global): 0 = silu(gate)*up (GLM), 1 = swigluoai (M3). */
+void coli_vk_set_activation(int act, float alpha, float limit);
+/* Decode GQA attention core (MiniMax-M3): K/V rows mirrored via coli_vk_kv_row
+ * (K in the L buffer, V in the R buffer, both NK*hd wide). Returns 0 -> CPU fallback. */
+int coli_vk_gqa_attn(float *ctx, const float *q, int layer, int S, int H, int NK, int hd,
+                     int st0, int T, float scale);
+/* Fused GQA core + o-projection: ctx stays on-device, only [S,Dout] returns. */
+int coli_vk_gqa_attn_project(float *out, const float *q, ColiVkTensor **ot, const void *ow,
+                             const float *osc, int ofmt, int ogrp, int layer, int S, int H,
+                             int NK, int hd, int st0, int T, float scale, int Dout);
 int  coli_vk_mem_budget(double *used_gb, double *budget_gb);
 
 /* y[S,O] = (x[S,I] @ dequant(W[O,I])^T) * scale[O].

c/coli

@@ -68,8 +68,9 @@ except ModuleNotFoundError:
     except ModuleNotFoundError:
         _version = "unknown"
 
-from family_registry import (FamilyConfigError, UnknownFamilyError, all_families,
-                             family_by_id, resolve_model, tuning_replay_prompt)
+from family_registry import (COLIBRI_CORE_GROUP, FamilyConfigError, UnknownFamilyError,
+                             all_families, family_by_id, resolve_model,
+                             tuning_replay_prompt)
 
 _TUNE_ROTATION_PROMPT = (
     "A service has high median throughput but poor tail latency. Explain how "
@@ -313,10 +314,22 @@ def redraw_prompt_box(message, width):
 def model_arch(model):
     return resolve_model(model).descriptor.id
 
+def colibri_core(family):
+    """Is this family served by the colibri binary itself, rather than by a
+    sister engine?
+
+    engine_group is the registry's statement that two families share an engine
+    implementation, so they also share that engine's env contract, its CAP
+    channel and its byte-protocol REPL. Keyed on the group rather than on
+    id == "glm" so a future third family on this binary inherits the same
+    handling by declaring the group, not by being added to a tuple here."""
+    if isinstance(family, str): family = family_by_id(family)
+    return family.engine_group == COLIBRI_CORE_GROUP
+
 def engine_for(model):
     family = resolve_model(model).descriptor
     if os.environ.get("COLI_ENGINE"):
-        if (family.id != "glm" and
+        if (not colibri_core(family) and
                 os.environ.get("COLI_DOCKER_GLM_ONLY") == "1"):
             raise UnknownFamilyError(
                 f"this image contains only the GLM engine; {family.display_name} "
@@ -396,7 +409,7 @@ def cap_for_launch(explicit_cap, env, fallback):
 def operator_cap(a, arch):
     """Return a real operator cap, including GLM's documented CAP channel."""
     cap=getattr(a,"cap",None)
-    if cap is not None or arch!="glm":
+    if cap is not None or not colibri_core(arch):
         return cap
     try:
         cap=int(os.environ.get("CAP","0"))
@@ -405,7 +418,7 @@ def operator_cap(a, arch):
     return cap if cap else None
 
 def env_for_engine(a, arch, plan=None):
-    if arch == "glm":
+    if colibri_core(arch):
         return env_for(a)
     explicit_env = set(os.environ)
     env = os.environ.copy()
@@ -1220,11 +1233,19 @@ def cmd_run(a):
             [engine, str(cap_for_launch(a.cap,e,16)), "8"],
             input=prompt+"\n", text=True, env=e, check=False)
         sys.exit(result.returncode)
-    # template ufficiale GLM-5.2: niente \n dopo i ruoli; <think></think> = risposta diretta (nothink).
-    # THINK=1 lascia <think> aperto, stessa convenzione del serve mode (glm.c). EN: THINK=1 leaves
-    # <think> open so the engine emits its reasoning block; the default stays nothink.
-    tk="<think>" if os.environ.get("THINK","0")=="1" else "<think></think>"
-    e=env_for(a); e["PROMPT"]=f"[gMASK]<sop><|user|>{prompt}<|assistant|>{tk}"
+    # Official GLM-5.2 template: no \n after the role markers; <think></think> means a
+    # direct answer (nothink). THINK=1 leaves <think> open so the engine emits its
+    # reasoning block -- same convention as serve mode (glm.c); the default stays nothink.
+    think=os.environ.get("THINK","0")=="1"
+    if arch=="minimax_m3":
+        # MiniMax-M3 chat_template.jinja: ]~!b[ once, ]~b]<role>\n<content>[e~[\n blocks,
+        # ai turn opens bare (adaptive thinking); </mm:think> prefix = nothink, like history.
+        tk="" if think else "</mm:think>"
+        p=f"]~!b[]~b]user\n{prompt}[e~[\n]~b]ai\n{tk}"
+    else:
+        tk="<think>" if think else "<think></think>"
+        p=f"[gMASK]<sop><|user|>{prompt}<|assistant|>{tk}"
+    e=env_for(a); e["PROMPT"]=p
     sys.exit(subprocess.call([GLM, str(cap_for_launch(a.cap,e,0))], env=e))
 
 def server_probe(base, api_key=None, timeout=1.5):
@@ -1396,7 +1417,7 @@ def cmd_chat(a):
     arch=family.id
     if not family.has_gateway_adapter:
         sys.exit(f"{family.display_name}: gateway adapter is not wired")
-    if arch!="glm":
+    if not colibri_core(family):
         engine=engine_for(a.model)
         need_model(a.model,engine)
         model_id=family.default_model_id