sgl-project/sglang · error · NotImplementedError

DSPARK aux hidden capture requires PP=1.

Error message

DSPARK aux hidden capture requires PP=1.

What it means

KimiLinearModel.set_dspark_layers_to_capture (kimi_linear.py:793) enables DSPARK aux-hidden-state capture for training/distillation hooks. Aux states are gathered on the last pipeline stage only, so running with pipeline parallelism (pp_group.world_size > 1) is unsupported and raises NotImplementedError.

Source

Thrown at python/sglang/srt/models/kimi_linear.py:793

        if self.pp_group.is_last_rank:
            self.lm_head = ParallelLMHead(
                self.config.vocab_size,
                self.config.hidden_size,
                quant_config=quant_config,
                prefix=maybe_prefix(prefix, "lm_head"),
            )
        else:
            self.lm_head = PPMissingLayer()
        logit_scale = getattr(self.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:
            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,
        inputs_embeds: Optional[torch.Tensor] = None,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
    ) -> torch.Tensor:

View on GitHub (pinned to 0132848349)

Solutions

1. Run with PP=1 (use TP/EP for sharding instead) when DSPARK capture is needed
2. Disable the DSPARK capture feature / do not call the setter in PP deployments
3. Request/patch upstream support for gathering aux states across PP stages

Example fix

# before
python -m sglang.launch_server --pp-size 2 ...  (+ dspark capture)
# after
python -m sglang.launch_server --pp-size 1 --tp-size 4 ...

assert model.pp_group.world_size == 1  # before set_dspark_layers_to_capture

try:
    model.set_dspark_layers_to_capture(ids)
except NotImplementedError:
    ids = None  # disable DSPARK under PP

When it happens

Trigger: Calling set_dspark_layers_to_capture while --pp-size (pipeline parallel) > 1, e.g. in a multi-node or memory-constrained PP deployment that also enables DSPARK capture.

Common situations: Large Kimi-Linear deployments that shard with PP for memory, then trying to attach DSPARK hidden-capture hooks.

Related errors

• Currently DSpark speculative decoding only supports pp_size
• --enable-linear-replayssm-spec with DSPARK/DFLASH requires a
• DSPARK aux hidden capture requires PP=1.
• DSPARK requires explicit layer_ids for aux hidden capture.
• /v1/models ${response.status}

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