[BugFix][Speculative Decoding] MTP logprobs path drops enable_pd_reorder and raises TypeError

[Anai-Guo]

Motivation

fastdeploy/spec_decode/mtp_cuda.py:303 calls recover_batch_index_for_sampler_output with two arguments, but the function requires three:

# fastdeploy/worker/input_batch.py:1291
def recover_batch_index_for_sampler_output(sampler_output, index_to_batch_id, enable_pd_reorder):

TypeError: recover_batch_index_for_sampler_output() missing 1 required positional argument: 'enable_pd_reorder'

The guard above it is tensor_parallel_rank == 0 and substep == 0 and sampler_output.logprobs_tensors is not None, so this fires as soon as MTP speculative decoding runs with logprobs requested.

The value is already at hand — the very next statement in the same block passes it to the sibling helper:

recover_batch_index_for_sampler_output(
    sampler_output,
    self.model_inputs.index_to_batch_id,
)                                              # <-- missing enable_pd_reorder
recover_model_output_map = recover_batch_index_for_output(
    self.model_inputs,
    self.model_inputs.index_to_batch_id,
    self.model_inputs.enable_pd_reorder,       # <-- passed here
    [...],
)

Every other call site in the repo passes it: pre_and_post_process.py:389, 420, 574, 624 (as model_output.enable_pd_reorder) and all five call sites in tests/worker/test_recover_batch_index_sampling_mask.py.

self.model_inputs is a ProposerInputBatch (input_batch.py:774), which sets self.enable_pd_reorder in __init__ (L785) and refreshes it from the target batch at L791, so the attribute is always present.

Modifications

Pass self.model_inputs.enable_pd_reorder as the third argument. One line.

Verification

No GPU here, so I verified at the signature level: parsed input_batch.py with ast, rebuilt recover_batch_index_for_sampler_output's inspect.Signature, and replayed every call site through Signature.bind.

Before:

def recover_batch_index_for_sampler_output(sampler_output, index_to_batch_id, enable_pd_reorder)
  spec_decode/mtp_cuda.py:303                          TypeError: missing a required argument: 'enable_pd_reorder'
  model_executor/pre_and_post_process.py:389           OK
  model_executor/pre_and_post_process.py:420           OK
  model_executor/pre_and_post_process.py:574           OK
  model_executor/pre_and_post_process.py:624           OK
  tests/worker/test_recover_batch_index_sampling_mask.py:38/54/74/88/102   OK

After: all ten bind cleanly.

black (25.1.0 config, line-length 119) and ruff check --config pyproject.toml both report no changes for the touched file.

🤖 Generated with Claude Code

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}