recipes/docker-images/deepseek-v4-flash-dspark38/patches/patch_dspark_uniform.py

#!/usr/bin/env python3
"""Patch dspark_proposer.py so a DSpark batch with mixed per-request context
lengths does not kill the engine (observed under `tool-eval-bench --parallel 4`).

Stock behaviour: `_trim_rejected_target_context` + `_view_by_request` require
every request in a DSpark batch to have the SAME flattened length. When
requests have divergent context lengths (e.g. [1, 2052]) this raises ValueError
in the worker -> EngineCore dies -> EngineDeadError on every later request.

There are two separate uniform checks that both crash:
  1. `_trim_rejected_target_context` requires uniform *effective* lengths after
     rejection-trimming (raises only when any token was rejected).
  2. `_view_by_request` requires the flattened input to be a multiple of
     batch_size (crashes when all drafts are ACCEPTED and contexts differ,
     because the raw per-request lengths are returned unpadded).

Fix: restructure `_trim_rejected_target_context` to ALWAYS split the flattened
per-request chunks via `query_start_loc`, trim rejected tokens (when provided),
and pad every shorter request to the batch-max effective length by repeating its
last valid token (hidden state + position). This keeps the batch rectangular /
uniform for `_view_by_request` and `prefill_main` while the draft continuation
point `[:, -1]` stays the real last token. Only falls back to returning the raw
tensor if per-request boundaries are unavailable.
"""
import pathlib

P = pathlib.Path(
    "/opt/env/lib/python3.12/site-packages/vllm/v1/spec_decode/dspark_proposer.py"
)
src = P.read_text()

OLD = """        if num_rejected_tokens_gpu is None or common_attn_metadata is None:
            return target_hidden_states, target_positions

        rejected = num_rejected_tokens_gpu.detach().cpu().tolist()
        if not any(int(value) for value in rejected):
            return target_hidden_states, target_positions

        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
        if query_start_loc_cpu is None:
            query_start_loc_cpu = common_attn_metadata.query_start_loc.detach().cpu()
        query_starts = query_start_loc_cpu.tolist()
        if len(query_starts) != len(rejected) + 1:
            raise ValueError(
                "DSpark rejected-context trimming requires query_start_loc "
                "to have batch_size + 1 entries; got "
                f"{len(query_starts)} starts for {len(rejected)} requests."
            )

        hidden_chunks: list[torch.Tensor] = []
        position_chunks: list[torch.Tensor] = []
        effective_lengths: list[int] = []
        flat_positions = target_positions.reshape(-1)
        for req_index, num_rejected in enumerate(rejected):
            start = int(query_starts[req_index])
            end = int(query_starts[req_index + 1]) - int(num_rejected)
            if end <= start:
                raise ValueError(
                    "DSpark rejected-context trimming removed every token for "
                    f"request {req_index}: start={start}, end={end}."
                )
            hidden_chunks.append(target_hidden_states[start:end])
            position_chunks.append(flat_positions[start:end])
            effective_lengths.append(end - start)

        if len(set(effective_lengths)) != 1:
            raise ValueError(
                "DSpark currently requires uniform effective per-request "
                "target context lengths after rejection trimming; got "
                f"{effective_lengths}."
            )

        return torch.cat(hidden_chunks, dim=0), torch.cat(position_chunks, dim=0)"""

NEW = """        # sm121-dspark38: if we cannot determine per-request boundaries, fall
        # back to returning the tensors untouched.
        qsl_cpu = getattr(common_attn_metadata, "query_start_loc_cpu", None)
        if common_attn_metadata is None or (
            qsl_cpu is None
            and getattr(common_attn_metadata, "query_start_loc", None) is None
        ):
            return target_hidden_states, target_positions
        if qsl_cpu is None:
            qsl_cpu = common_attn_metadata.query_start_loc.detach().cpu()
        query_starts = qsl_cpu.tolist()
        nreq = len(query_starts) - 1
        if nreq < 1:
            return target_hidden_states, target_positions

        rejected = (
            num_rejected_tokens_gpu.detach().cpu().tolist()
            if num_rejected_tokens_gpu is not None
            else [0] * nreq
        )
        if len(query_starts) != len(rejected) + 1:
            raise ValueError(
                "DSpark rejected-context trimming requires query_start_loc "
                "to have batch_size + 1 entries; got "
                f"{len(query_starts)} starts for {len(rejected)} requests."
            )

        hidden_chunks: list[torch.Tensor] = []
        position_chunks: list[torch.Tensor] = []
        effective_lengths: list[int] = []
        flat_positions = target_positions.reshape(-1)
        for req_index, num_rejected in enumerate(rejected):
            start = int(query_starts[req_index])
            end = int(query_starts[req_index + 1]) - int(num_rejected)
            if end <= start:
                raise ValueError(
                    "DSpark rejected-context trimming removed every token for "
                    f"request {req_index}: start={start}, end={end}."
                )
            hidden_chunks.append(target_hidden_states[start:end])
            position_chunks.append(flat_positions[start:end])
            effective_lengths.append(end - start)

        # sm121-dspark38: pad every shorter request to the batch-max effective
        # length by repeating its last valid token (hidden state + position),
        # so the batch stays rectangular/uniform for _view_by_request and
        # prefill_main while the draft continuation point ([:, -1]) remains the
        # real last token. The stock fork hard-rejects non-uniform lengths.
        max_eff = max(effective_lengths)
        for i, (eff, chunk, pos_chunk) in enumerate(
            zip(effective_lengths, hidden_chunks, position_chunks)
        ):
            if eff < max_eff:
                pad = max_eff - eff
                hidden_chunks[i] = torch.cat(
                    [chunk, chunk[-1:].expand(pad, -1).contiguous()], dim=0
                )
                position_chunks[i] = torch.cat(
                    [pos_chunk, pos_chunk[-1:].expand(pad).contiguous()], dim=0
                )
        effective_lengths[:] = [max_eff] * len(effective_lengths)

        return torch.cat(hidden_chunks, dim=0), torch.cat(position_chunks, dim=0)"""

assert src.count(OLD) == 1, f"expected 1 block, got {src.count(OLD)}"
P.write_text(src.replace(OLD, NEW))
print("patched dspark_proposer.py: always-pad to uniform context")