[NPU] Support GLM-5.2 DSpark loading

python/sglang/srt/configs/speculators.py

@@ -0,0 +1,154 @@
+"""Translate the Speculators dense DSpark export into SGLang's config layout."""
+
+from copy import deepcopy
+from typing import Any, Optional
+
+
+def _require_int(config: dict, name: str, minimum: int = 1) -> int:
+    value = config.get(name)
+    if type(value) is not int or value < minimum:
+        raise ValueError(f"Speculators DSpark {name} must be an integer >= {minimum}.")
+    return value
+
+
+def _check_alias(config: dict, name: str, expected: Any, source: str) -> None:
+    value = config.get(name)
+    if value is not None and value != expected:
+        raise ValueError(f"Speculators DSpark {name} conflicts with {source}.")
+
+
+def normalize_speculators_dspark_config(config: dict) -> Optional[dict]:
+    """Return a non-mutating translation, or None for an existing HF layout.
+
+    Speculators stores the decoder under ``transformer_layer_config`` and
+    numbers auxiliary hidden states one above SGLang's DFlash capture IDs.
+    The result uses the native HF layout, so reloading it never subtracts
+    twice. This only handles the Qwen3-style dense draft, not the verifier.
+    """
+    if "transformer_layer_config" not in config:
+        return None
+    if config.get("architectures") != ["DSparkDraftModel"]:
+        return None
+    if config.get("speculators_model_type", "dspark") != "dspark":
+        raise ValueError("Speculators DSpark speculators_model_type must be 'dspark'.")
+
+    speculators = config.get("speculators_config", {})
+    if not isinstance(speculators, dict):
+        raise ValueError("Speculators DSpark speculators_config must be an object.")
+    if speculators.get("algorithm", "dspark") != "dspark":
+        raise ValueError(
+            "Speculators DSpark speculators_config.algorithm must be 'dspark'."
+        )
+
+    decoder = config["transformer_layer_config"]
+    if not isinstance(decoder, dict) or decoder.get("model_type") != "qwen3":
+        raise ValueError(
+            "Speculators DSpark requires a qwen3 transformer_layer_config."
+        )
+    if config.get("model_type") not in (None, "dspark", "qwen3"):
+        raise ValueError("Speculators DSpark model_type conflicts with the decoder.")
+    if config.get("text_config") is not None:
+        raise ValueError(
+            "Speculators DSpark cannot use both text_config "
+            "and transformer_layer_config."
+        )
+
+    normalized = deepcopy(config)
+    for key, value in decoder.items():
+        # The outer architecture selects DSpark; the inner type selects its
+        # decoder configuration. Neither should select the target network.
+        if key in ("architectures", "model_type", "auto_map", "_name_or_path"):
+            continue
+        _check_alias(normalized, key, value, f"transformer_layer_config.{key}")
+        normalized[key] = deepcopy(value)
+    normalized["model_type"] = "qwen3"
+
+    for name in (
+        "hidden_size",
+        "intermediate_size",
+        "num_hidden_layers",
+        "num_attention_heads",
+        "num_key_value_heads",
+        "vocab_size",
+        "block_size",
+        "markov_rank",
+    ):
+        _require_int(normalized, name)
+    if "head_dim" in normalized:
+        # head_dim need not equal hidden_size // num_attention_heads.
+        _require_int(normalized, "head_dim")
+    # The input-only mask can use a padded row of the shared target embedding;
+    # the runtime checks its upper bound against the actual target vocabulary.
+    _require_int(normalized, "mask_token_id", minimum=0)
+    _check_alias(normalized, "draft_vocab_size", normalized["vocab_size"], "vocab_size")
+    _check_alias(
+        normalized, "target_hidden_size", normalized["hidden_size"], "hidden_size"
+    )
+    if normalized.get("markov_head_type") not in ("vanilla", "gated", "rnn"):
+        raise ValueError("Speculators DSpark has an unsupported markov_head_type.")
+    for name in (
+        "enable_confidence_head",
+        "confidence_head_with_markov",
+        "sample_from_anchor",
+    ):
+        if name in normalized and type(normalized[name]) is not bool:
+            raise ValueError(f"Speculators DSpark {name} must be a boolean.")
+
+    aux_ids = normalized.get("aux_hidden_state_layer_ids")
+    if (
+        not isinstance(aux_ids, list)
+        or not aux_ids
+        or any(type(layer) is not int or layer < 1 for layer in aux_ids)
+    ):
+        raise ValueError(
+            "Speculators DSpark aux_hidden_state_layer_ids must be a non-empty "
+            "list of positive integers."
+        )
+    # Target hooks append features in model execution order. Sorting a trained
+    # list here would silently change the FC input, and duplicates cannot be
+    # captured by their membership-based hooks.
+    if any(left >= right for left, right in zip(aux_ids, aux_ids[1:])):
+        raise ValueError(
+            "Speculators DSpark aux_hidden_state_layer_ids must be strictly increasing."
+        )
+    # Capture count need not equal draft depth.
+    target_ids = [layer - 1 for layer in aux_ids]
+    _check_alias(
+        normalized, "target_layer_ids", target_ids, "aux_hidden_state_layer_ids"
+    )
+    normalized["target_layer_ids"] = target_ids
+
+    # Existing readers give these legacy aliases priority over the fields above.
+    # Accept matching aliases, but never silently override the export contract.
+    for alias, source in (
+        ("dspark_target_layer_ids", "target_layer_ids"),
+        ("dspark_block_size", "block_size"),
+        ("dspark_noise_token_id", "mask_token_id"),
+        ("dspark_markov_rank", "markov_rank"),
+        ("dspark_markov_head_type", "markov_head_type"),
+    ):
+        _check_alias(normalized, alias, normalized[source], source)
+    for section in ("dflash_config", "dspark_config"):
+        aliases = normalized.get(section, {})
+        if not isinstance(aliases, dict):
+            raise ValueError(f"Speculators DSpark {section} must be an object.")
+        for name in (
+            "target_layer_ids",
+            "block_size",
+            "mask_token_id",
+            "markov_rank",
+            "markov_head_type",
+        ):
+            # Native readers use dict.get(key, canonical), so an explicit null
+            # would hide the canonical value rather than use that fallback.
+            # In particular, null target IDs would silently enable auto-picking.
+            if name in aliases and aliases[name] is None:
+                del aliases[name]
+            _check_alias(aliases, name, normalized[name], f"top-level {name}")
+
+    # Keep the canonical IDs as provenance, but emit just one decoder layout.
+    # HF expands defaults (e.g. RoPE parameters) during construction, so keeping
+    # a second raw decoder would make a saved config conflict with itself.
+    normalized.pop("transformer_layer_config")
+    normalized.pop("auto_map", None)
+    return normalized

python/sglang/srt/environ.py

@@ -921,6 +921,11 @@ class Envs:
     # ===================================================================
     # Ascend NPU
     # ===================================================================
+    # "original" opts a GLM DSA QuaRot target + dense DSpark draft into
+    # checkpoint-local vocab and load-time FC rotation. The draft checkpoint
+    # must contain original (not already converted) weights. Empty keeps the
+    # existing loading behavior.
+    SGLANG_NPU_GLM_DSPARK_QUAROT = EnvStr("")
     SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False)
     SGLANG_NPU_USE_MULTI_STREAM = EnvBool(False)
     SGLANG_NPU_USE_MLAPO = EnvBool(False)

python/sglang/srt/hardware_backend/npu/dspark_quarot.py

@@ -0,0 +1,171 @@
+"""Explicit original-coordinate GLM DSpark loading on NPU.
+
+The configuration lives only around one draft construction. FC conversion is a
+load-time candidate, not an exact inverse for a non-orthogonal stored Q.
+"""
+
+import logging
+import time
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import torch
+from safetensors import safe_open
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class GlmDSparkQuaRotConfig:
+    rotation_path: str
+    hidden_size: int
+    target_model_path: str
+
+
+_config: ContextVar[Optional[GlmDSparkQuaRotConfig]] = ContextVar(
+    "glm_dspark_quarot_config", default=None
+)
+
+
+def get_glm_dspark_quarot_config() -> Optional[GlmDSparkQuaRotConfig]:
+    return _config.get()
+
+
+@contextmanager
+def glm_dspark_quarot_scope(config: Optional[GlmDSparkQuaRotConfig]):
+    token = _config.set(config)
+    try:
+        yield
+    finally:
+        _config.reset(token)
+
+
+def build_glm_dspark_quarot_config(
+    *, device, mode, target_model_config, target_model
+) -> Optional[GlmDSparkQuaRotConfig]:
+    """Read metadata only after matching the explicit target-specific path.
+
+    ``original`` is the caller's declaration about the draft checkpoint. Target
+    QuaRot metadata alone cannot establish the draft's coordinate convention.
+    """
+    if not mode or str(device).split(":", 1)[0] != "npu":
+        return None
+    hf_config = getattr(target_model_config, "hf_text_config", None)
+    architectures = getattr(hf_config, "architectures", None)
+    if (
+        not isinstance(architectures, (list, tuple))
+        or "GlmMoeDsaForCausalLM" not in architectures
+    ):
+        return None
+    quant_config = getattr(target_model, "quant_config", None)
+    get_name = getattr(quant_config, "get_name", None)
+    if not callable(get_name) or get_name() != "modelslim":
+        return None
+    if mode != "original":
+        raise ValueError("SGLANG_NPU_GLM_DSPARK_QUAROT must be 'original' when enabled")
+
+    description = getattr(quant_config, "quant_description", None)
+    if not isinstance(description, dict) or description.get("is_rot_used") is not True:
+        raise ValueError("GLM DSpark original mode requires a ModelSlim QuaRot target")
+    try:
+        relative_path = description["optional"]["quarot"]["rotation_map"][
+            "global_rotation"
+        ]
+    except (KeyError, TypeError) as exc:
+        raise ValueError(
+            "GLM DSpark target has no QuaRot global_rotation path"
+        ) from exc
+    if (
+        not isinstance(relative_path, str)
+        or not relative_path
+        or Path(relative_path).is_absolute()
+    ):
+        raise ValueError("QuaRot global_rotation must be a relative target file path")
+    width = getattr(hf_config, "hidden_size", None)
+    if type(width) is not int or width <= 0:
+        raise ValueError("GLM DSpark target hidden_size must be a positive integer")
+    target_path = Path(target_model_config.model_path).resolve(strict=True)
+    rotation_path = (target_path / relative_path).resolve(strict=True)
+    # get_slice reads the header, not the full matrix during draft construction.
+    with safe_open(str(rotation_path), framework="pt", device="cpu") as reader:
+        if "global_rotation" not in reader.keys():
+            raise ValueError(f"Missing global_rotation tensor in {rotation_path}")
+        if reader.get_slice("global_rotation").get_shape() != [width, width]:
+            raise ValueError(
+                f"QuaRot global_rotation must have shape [{width}, {width}]: {rotation_path}"
+            )
+    return GlmDSparkQuaRotConfig(str(rotation_path), width, str(target_path))
+
+
+@torch.no_grad()
+def fold_glm_dspark_fc(
+    weight: torch.Tensor, config: GlmDSparkQuaRotConfig
+) -> torch.Tensor:
+    """Return fresh CPU FC blocks F_i @ Q in the original weight dtype.
+
+    This never transforms an existing model Parameter in place, rescales Q,
+    applies R, or changes CPU thread settings. Call it on newly read checkpoint
+    weights so reloading does not rotate already converted parameters again.
+    """
+    width = config.hidden_size
+    if (
+        type(width) is not int
+        or width <= 0
+        or weight.ndim != 2
+        or weight.shape[0] != width
+        or weight.shape[1] == 0
+        or weight.shape[1] % width
+        or not weight.is_floating_point()
+    ):
+        raise ValueError(
+            "GLM DSpark FC must be floating [hidden_size, K * hidden_size]"
+        )
+
+    started = time.monotonic()
+    logger.info(
+        "GLM DSpark original mode: folding FC %s with Q=%s from target=%s "
+        "on CPU in FP32, storing %s; no inverse or scale correction",
+        tuple(weight.shape),
+        config.rotation_path,
+        config.target_model_path,
+        weight.dtype,
+    )
+    with safe_open(config.rotation_path, framework="pt", device="cpu") as reader:
+        if "global_rotation" not in reader.keys():
+            raise ValueError(
+                f"Missing global_rotation tensor in {config.rotation_path}"
+            )
+        q = reader.get_tensor("global_rotation")
+    if q.shape != (width, width) or not q.is_floating_point():
+        raise ValueError(
+            "QuaRot global_rotation must be a floating hidden_size square matrix"
+        )
+    q = q.to(dtype=torch.float32, device="cpu")
+    if not torch.isfinite(q).all():
+        raise ValueError("QuaRot global_rotation contains non-finite FP32 values")
+
+    folded = torch.empty(weight.shape, dtype=weight.dtype, device="cpu")
+    # Keep only Q, the output, and small FP32 tiles; do not allocate a full FP32
+    # copy of the 6144 x 30720 checkpoint or a block-diagonal rotation matrix.
+    for start in range(0, weight.shape[1], width):
+        for row in range(0, width, 128):
+            source = (
+                weight[row : row + 128, start : start + width]
+                .detach()
+                .to(device="cpu", dtype=torch.float32)
+            )
+            if not torch.isfinite(source).all():
+                raise ValueError("GLM DSpark FC contains non-finite FP32 values")
+            converted = (source @ q).to(dtype=weight.dtype)
+            if not torch.isfinite(converted).all():
+                raise ValueError("GLM DSpark converted FC contains non-finite values")
+            folded[row : row + 128, start : start + width].copy_(converted)
+    del source, converted, q
+    logger.info(
+        "GLM DSpark FC load-time Q folding finished in %.3f seconds",
+        time.monotonic() - started,
+    )
+    return folded