sgl-project/sglang · error · NotImplementedError

DSPARK aux hidden capture requires PP=1.

Error message

DSPARK aux hidden capture requires PP=1.

What it means

KimiK3.set_dspark_layers_to_capture enables DSPARK aux-hidden-state capture on specific layers by setting a flag that only takes effect on the last pipeline-parallel rank. With PP>1, requested layers on earlier ranks would be silently skipped, so the code raises NotImplementedError to force PP=1 rather than producing a broken capture.

Source

Thrown at python/sglang/srt/models/kimi_k3.py:2888

                config.hidden_size,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
                use_attn_tp_group=get_parallel().enable_dp_lm_head,
            )
        else:
            self.lm_head = PPMissingLayer()
        logit_scale = getattr(config, "logit_scale", 1.0)
        self.logits_processor = LogitsProcessor(config=config, logit_scale=logit_scale)
        self.capture_aux_hidden_states = False

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None:
        if self.pp_group.world_size > 1:
            # Capture layers living on non-last PP ranks would be silently
            # skipped (the flag is only set on the last rank).
            raise NotImplementedError("DSPARK aux hidden capture requires PP=1.")
        if not self.pp_group.is_last_rank:
            return
        if layer_ids is None:
            raise ValueError(
                "DSPARK requires explicit layer_ids for aux hidden capture."
            )
        self.capture_aux_hidden_states = True
        self.model.dspark_layers_to_capture = list(layer_ids)

    @torch.no_grad()
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        forward_batch: ForwardBatch,
        input_embeds: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,

View on GitHub (pinned to 0132848349)

Solutions

  1. Relaunch with --pipeline-parallel-size 1 (tensor parallel / expert parallel are fine).
  2. If PP is required, disable the DSPARK aux-hidden capture feature.
  3. Track upstream support for PP>1 capture (flag propagation across ranks) before retrying.

Example fix

# before
python -m sglang.launch_server --model kimi-k3 --pp 2 --speculative-algorithm DSPARK ...

# after
python -m sglang.launch_server --model kimi-k3 --pp 1 --tp 8 --speculative-algorithm DSPARK ...
Defensive patterns

Strategy: validation

Validate before calling

from sglang.srt.distributed import get_pp_group
if get_pp_group().world_size > 1:
    raise RuntimeError("disable DSPARK aux capture or set --pp 1")

Prevention

When it happens

Trigger: Calling set_dspark_layers_to_capture(layer_ids) (done by the DSPARK/speculative runtime during draft setup) while pp_group.world_size > 1 — i.e. launching the server with --pipeline-parallel-size > 1 together with DSPARK aux hidden capture.

Common situations: Enabling DSPARK speculative training/serving on a multi-node or PP-sharded Kimi K3 deployment; reusing PP flags from a non-DSPARK launch script when turning on DSPARK.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/ce31b31dd9a787ba. Report an issue: GitHub.