feat(dflash): converge DFlash spec-decode stack → dev (supersedes #140)

README.md

@@ -63,6 +63,7 @@ MoE-Infinity supports HuggingFace MoE checkpoints registered in [`moe_infinity/c
 | [Mixtral](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) | `mistralai/Mixtral-8x7B-Instruct-v0.1`, `Mixtral-8x22B` |
 | [Qwen3-MoE](https://huggingface.co/Qwen/Qwen3-30B-A3B) | `Qwen/Qwen3-30B-A3B` |
 | [Qwen3.5-MoE](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) | `Qwen/Qwen3.5-35B-A3B` (text-only; see note) |
+| [GLM-5.2](https://huggingface.co/zai-org/GLM-5.2-FP8) | `zai-org/GLM-5.2-FP8` (requires `transformers` >= 5.12) |
 | [GPT-OSS](https://huggingface.co/models?search=gpt-oss) | `openai/gpt-oss-*` |
 | [DBRX](https://huggingface.co/models?search=dbrx) | `databricks/dbrx-instruct` |
 | [Jamba](https://huggingface.co/models?search=jamba) | `ai21labs/Jamba-*` |
@@ -73,6 +74,8 @@ MoE-Infinity supports HuggingFace MoE checkpoints registered in [`moe_infinity/c
 
 > Qwen3.5-MoE (`Qwen3_5MoeForConditionalGeneration`, requires `transformers` >= 5.12) is a vision-language checkpoint served **text-only**: its 256 routed experts are offloaded while the small text backbone — token embeddings, the hybrid linear (GatedDeltaNet) / full attention layers, shared expert, and `lm_head` — stays resident on GPU. The v5 packed expert tensors are expanded to per-expert on load. Vision and MTP weights are present but unused for text generation.
 
+> GLM-5.2 (`GlmMoeDsaForCausalLM`, `model_type="glm_moe_dsa"`) requires `transformers` >= 5.12 and is registered only when that class is importable (otherwise skipped automatically). Its 256 routed FP8 experts (block-scale e4m3, dequantized to BF16 on load) are offloaded; the 3 dense layers, shared expert, MLA attention, DSA indexer, and MTP layer stay resident. Sparse attention uses `attn_implementation="eager"`.
+
 ## Installation
 
 We recommend installing MoE-Infinity in a virtual environment. To install MoE-Infinity, you can either install it from PyPI or build it from source.
@@ -248,6 +251,21 @@ torchrun --nproc-per-node 4 examples/deepseek_v4_flash_example.py \
 
 **Suggested hardware / environment (Path B):** **4x GPUs** (tensor-parallel mp4; mp1 exceeds the sparse-attention kernel's shared-memory limit), **>= ~140 GB pinned host RAM**, and the `v4flash` docker image (tilelang `fp4_gemm`; on Blackwell/SM120 the native `moe_infinity._v4_fp4` CUDA path is auto-selected and is 1.5–3.2x faster). The checkpoint must first be converted to the official mp-sharded format. See [`moe_infinity/models/deepseek_v4/README.md`](./moe_infinity/models/deepseek_v4/README.md) for checkpoint conversion, kernel selection, and validation details.
 
+### GLM-5.2 (FP8 Expert Offloading)
+
+GLM-5.2 (`zai-org/GLM-5.2-FP8`) runs through the drop-in `MoE` class:
+
+```python
+from moe_infinity import MoE
+
+model = MoE("zai-org/GLM-5.2-FP8", {
+    "offload_path": "/ssd/moe-infinity/glm-5.2",
+    "device_memory_ratio": 0.5,
+})
+```
+
+> **Memory note:** the FP8 block-scaled routed experts are kept FP8 in the host store (~753 GB) and dequantized on-device by the expert dispatcher, which requires the native `moe_infinity._v4_fp4` extension. Weights that run in PyTorch rather than the dispatcher — MLA attention, the DSA indexer, the dense-layer MLPs, and the shared expert — are dequantized to BF16 on load. Requires `transformers` >= 5.12.
+
 ### Benchmarking
 
 For correct throughput and latency measurement, it is critical to separate **prefill time (TTFT)** from **decode throughput**. Including prefill in your throughput calculation will produce misleadingly low numbers.

benchmarks/performance_model/__init__.py

@@ -0,0 +1,8 @@
+from benchmarks.performance_model.roofline import predict_decode
+from benchmarks.performance_model.types import (
+    DemandResult,
+    ModelParams,
+    WorkloadPoint,
+)
+
+__all__ = ["ModelParams", "WorkloadPoint", "DemandResult", "predict_decode"]

benchmarks/performance_model/bench_glm.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import os
+
+import torch
+
+
+def measure_tiny_glm(tmp_dir: str, gen_len: int = 16) -> dict:
+    from benchmarks.performance_model.model_config import extract_model_params
+    from benchmarks.performance_model.roofline import predict_decode
+    from benchmarks.performance_model.types import WorkloadPoint
+    from moe_infinity import MoE
+    from moe_infinity.spec_decode.glm_mtp import GlmMtpSpeculator
+    from tests.python.integration._glm_tiny import build_tiny_glm
+
+    ckpt_dir = os.path.join(tmp_dir, "tiny_glm_ckpt")
+    off_dir = os.path.join(tmp_dir, "tiny_glm_off")
+    os.makedirs(ckpt_dir, exist_ok=True)
+    os.makedirs(off_dir, exist_ok=True)
+
+    build_tiny_glm(ckpt_dir)
+
+    model = MoE(ckpt_dir, {"offload_path": off_dir, "device_memory_ratio": 0.8})
+
+    input_ids = torch.tensor([[1, 2, 3, 4]], device="cuda")
+    batch = input_ids.shape[0]
+    seq_len = input_ids.shape[1]
+
+    torch.cuda.reset_peak_memory_stats()
+    for _ in range(2):
+        with torch.no_grad():
+            model.generate(input_ids, max_new_tokens=gen_len)
+    torch.cuda.synchronize()
+
+    torch.cuda.reset_peak_memory_stats()
+    start_evt = torch.cuda.Event(enable_timing=True)
+    end_evt = torch.cuda.Event(enable_timing=True)
+
+    start_evt.record()
+    with torch.no_grad():
+        model.generate(input_ids, max_new_tokens=gen_len)
+    end_evt.record()
+    torch.cuda.synchronize()
+
+    elapsed_ms = start_evt.elapsed_time(end_evt)
+    decode_tok_s = gen_len / (elapsed_ms / 1000.0)
+    peak_mem_bytes = torch.cuda.max_memory_allocated()
+
+    spec = GlmMtpSpeculator(model)
+    mtp_input = input_ids.clone()
+
+    mtp_start = torch.cuda.Event(enable_timing=True)
+    mtp_end = torch.cuda.Event(enable_timing=True)
+
+    mtp_start.record()
+    spec.generate(mtp_input, max_new_tokens=gen_len, temperature=0.0)
+    mtp_end.record()
+    torch.cuda.synchronize()
+
+    mtp_elapsed_ms = mtp_start.elapsed_time(mtp_end)
+    mtp_tok_s = gen_len / (mtp_elapsed_ms / 1000.0)
+    mean_accept_len = spec.last_stats.get("mean_accept_len", 1.0)
+
+    params = extract_model_params(ckpt_dir)
+    wp = WorkloadPoint(batch=batch, seq_len=seq_len, gen_len=gen_len)
+    demand = predict_decode(params, wp)
+
+    return {
+        "model": "tiny_glm",
+        "batch": batch,
+        "seq_len": seq_len,
+        "gen_len": gen_len,
+        "decode_tok_s": decode_tok_s,
+        "mtp_tok_s": mtp_tok_s,
+        "mean_accept_len": mean_accept_len,
+        "peak_mem_bytes": peak_mem_bytes,
+        "pred_flops_per_token": demand.flops_per_token,
+        "pred_hbm_bytes_per_token": demand.hbm_bytes_per_token,
+        "pred_bound": demand.bound,
+    }
+
+
+def run(out_csv: str, quick: bool = True, gen_len: int = 16) -> None:
+    import tempfile
+
+    os.makedirs(
+        os.path.dirname(out_csv) if os.path.dirname(out_csv) else ".",
+        exist_ok=True,
+    )
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        row = measure_tiny_glm(tmp_dir, gen_len=gen_len)
+
+    fieldnames = [
+        "model",
+        "batch",
+        "seq_len",
+        "gen_len",
+        "decode_tok_s",
+        "mtp_tok_s",
+        "mean_accept_len",
+        "peak_mem_bytes",
+        "pred_flops_per_token",
+        "pred_hbm_bytes_per_token",
+        "pred_bound",
+    ]
+
+    with open(out_csv, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerow(row)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--out", default="results/perf_model/glm_bench.csv")
+    parser.add_argument("--quick", action="store_true", default=True)
+    parser.add_argument("--gen", type=int, default=16)
+    args = parser.parse_args()
+    run(args.out, quick=args.quick, gen_len=args.gen)

benchmarks/performance_model/model_config.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from benchmarks.performance_model.types import ModelParams
+
+
+def _expert_dtype_from_config(config) -> str:
+    qc = getattr(config, "quantization_config", None)
+    if qc is None:
+        return "bf16"
+    if isinstance(qc, dict):
+        method = qc.get("quant_method", "")
+    else:
+        method = getattr(qc, "quant_method", "")
+    return "fp8" if "fp8" in str(method).lower() else "bf16"
+
+
+def _extract_glm(config) -> ModelParams:
+    return ModelParams(
+        name=getattr(config, "_name_or_path", "glm_moe_dsa"),
+        num_layers=config.num_hidden_layers,
+        num_attn_heads=config.num_attention_heads,
+        num_kv_heads=getattr(
+            config, "num_key_value_heads", config.num_attention_heads
+        ),
+        head_dim=getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        ),
+        hidden_size=config.hidden_size,
+        vocab_size=config.vocab_size,
+        num_experts=config.n_routed_experts,
+        top_k=config.num_experts_per_tok,
+        shared_experts=getattr(config, "n_shared_experts", 1),
+        expert_intermediate_size=config.moe_intermediate_size,
+        first_k_dense=getattr(config, "first_k_dense_replace", 0),
+        expert_dtype=_expert_dtype_from_config(config),
+        attn_dtype="bf16",
+        kv_lora_rank=getattr(config, "kv_lora_rank", None),
+        q_lora_rank=getattr(config, "q_lora_rank", None),
+    )
+
+
+def _extract_generic(config) -> ModelParams:
+    num_experts = (
+        getattr(config, "num_local_experts", None)
+        or getattr(config, "n_routed_experts", None)
+        or getattr(config, "num_experts", 1)
+    )
+    top_k = getattr(config, "num_experts_per_tok", None) or getattr(
+        config, "top_k", 1
+    )
+    return ModelParams(
+        name=getattr(config, "_name_or_path", "unknown"),
+        num_layers=config.num_hidden_layers,
+        num_attn_heads=config.num_attention_heads,
+        num_kv_heads=getattr(
+            config, "num_key_value_heads", config.num_attention_heads
+        ),
+        head_dim=getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        ),
+        hidden_size=config.hidden_size,
+        vocab_size=config.vocab_size,
+        num_experts=num_experts,
+        top_k=top_k,
+        shared_experts=getattr(config, "n_shared_experts", 0),
+        expert_intermediate_size=getattr(
+            config,
+            "moe_intermediate_size",
+            getattr(config, "intermediate_size", config.hidden_size * 4),
+        ),
+        first_k_dense=getattr(config, "first_k_dense_replace", 0),
+        expert_dtype=_expert_dtype_from_config(config),
+        attn_dtype="bf16",
+        kv_lora_rank=getattr(config, "kv_lora_rank", None),
+        q_lora_rank=getattr(config, "q_lora_rank", None),
+    )
+
+
+def extract_model_params(model_name_or_path: str) -> ModelParams:
+    from transformers import AutoConfig
+
+    config = AutoConfig.from_pretrained(
+        model_name_or_path, trust_remote_code=True
+    )
+    arch = (getattr(config, "architectures", None) or [""])[0].lower()
+
+    if "glmmoedsa" in arch:
+        return _extract_glm(config)
+    return _extract_generic(config)