fix(deepseek-v4): let DSpark survive mixed-context batches (tool-eval-bench)

The production-3.8 dspark38 serve booted and generated, but died under
tool-eval-bench --parallel 4: the fork's DSpark proposer hard-requires every
request in a batch to have the SAME flattened/effective context length, and an
unhandled ValueError in the worker killed the EngineCore (500 EngineDeadError
on all subsequent requests).

patch_dspark_uniform.py restructures dspark_proposer._trim_rejected_target_context
to ALWAYS split per-request chunks via query_start_loc, trim rejected tokens
(when given), and pad every shorter request to the batch-max effective length by
repeating its last valid token (hidden+position). This satisfies both uniform
checks (_trim_rejected/effective and _view_by_request/flattened) and keeps the
draft continuation point [:, -1] = real last token.

Verified on ds4: 4 concurrent requests with widely divergent context lengths
([1..2052]-class) all return HTTP 200, engine stays healthy, follow-up request
200. Server left running for the benchmark.

docker-images/deepseek-v4-flash-dspark38/Dockerfile

@@ -30,3 +30,9 @@ RUN sed -i \
 # scale to (n_groups, o_lora_rank, -1) before the bmm (see patch script).
 COPY patches/patch_dspark_woa.py /tmp/patch_dspark_woa.py
 RUN /opt/env/bin/python3 /tmp/patch_dspark_woa.py
+
+# Patch 3: DSpark mixed-context batch — pad shorter requests to the batch-max
+# effective length so a parallel benchmark (tool-eval-bench --parallel 4) with
+# divergent request contexts doesn't kill the engine (see patch script).
+COPY patches/patch_dspark_uniform.py /tmp/patch_dspark_uniform.py
+RUN /opt/env/bin/python3 /tmp/patch_dspark_uniform.py

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

@@ -0,0 +1,141 @@
+#!/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")