[Spec Decode] Defer clearing KV connector metadata for EAGLE3 speculative decode + prefill / decode disagg setup

[zixi-qi]

…tive decoding

Purpose

• When speculative decoding is enabled with disaggregated prefill/decode, clear_connector_metadata() was called after the target model forward but before the draft model forward. This caused has_connector_metadata() to return False during the draft model's attention layers, skipping save_kv_layer calls and preventing draft KV cache from being transferred.
• The fix defers clearing connector metadata until after the draft model runs by adding a clear_metadata parameter to post_forward and _get_kv_connector_output, and explicitly clearing after draft token proposal completes.
• This fix is only applicable for the case where draft model prefill happens on the prefill host, the original issue does not occur if draft prefill runs on decode host

Test Plan

• Verified disaggregated prefill/decode with EAGLE3 speculative decoding works end-to-end
• Confirmed draft model KV caches are saved by the KV connector after the fix

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request addresses an issue with speculative decoding in a disaggregated prefill/decode setup, where KV connector metadata was being cleared prematurely. The fix defers this clearing until after the draft model has run. This is achieved by introducing a clear_metadata parameter to the post_forward and _get_kv_connector_output methods, allowing for conditional clearing. A new method is also added to explicitly clear the metadata when needed. The changes are consistently applied across different model runner implementations and appear to be correct and robust. The use of a default value for the new parameter ensures backward compatibility. Overall, this is a well-executed fix.

[mergify]

Hi @zixi-qi, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[houseroad]

Looks good. Better to add some unittest.

[dosubot]

Related Documentation

Checked 0 published document(s) in 1 knowledge base(s). No updates required.

^{How did I do? Any feedback?}  

[NickLucche]

will follow up with a small refactor on top of this with @ZhanqiuHu here #34926

Better to add some unittest

I guess we'll pick this up..