#!/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")