Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
github-actions
Bot
added
the
documentation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fc23081700
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| def _draft_load_format_fields() -> dict: | ||
| draft_load_format = get_spec().speculative_draft_load_format | ||
| if draft_load_format is None: | ||
| return {} | ||
| logger.info(f"Using draft model load_format: '{draft_load_format}'") | ||
| return dict(load_format=draft_load_format) |
There was a problem hiding this comment.
When a plugin-registered speculative algorithm is selected with --speculative-draft-load-format, the scheduler now passes the original target ServerArgs directly to the plugin worker, while this new override is only invoked by private helpers used by the built-in workers. CustomSpecAlgo.create_worker does not apply either helper, so existing plugin workers that read server_args.load_format will silently receive the target format instead of the configured draft format; previously the scheduler applied it before every worker factory. Preserve the universal handoff behavior or provide the custom-worker path with an equivalently overridden copy.
Useful? React with 👍 / 👎.
| @@ -1010,10 +1005,8 @@ def __init__( | |||
| server_args.speculative_algorithm | |||
There was a problem hiding this comment.
[suggestion] All four workers rebind the local parameter server_args = draft_server_args_copy(...) after already storing the target instance on self.server_args. Later statements in the same __init__ that still read server_args (e.g. EAGLE's adaptive gate / frozen-KV's adaptive assert) now silently read the draft copy while self.server_args remains the target. Values currently coincide for those fields, so behavior is fine today, but the dual meaning of the name is a footgun for the next edit that assumes one or the other.
Suggestion: Use a distinct name, e.g. draft_server_args = draft_server_args_copy(...), pass that into the draft worker, and keep server_args / self.server_args unambiguously as the target instance. Apply the same pattern in standalone_worker_v2.py, multi_layer_eagle_worker_v2.py, and frozen_kv_mtp_worker_v2.py.
| self.assertIsNone(target.context_length) | ||
| self.assertEqual(target.load_format, "auto") | ||
|
|
||
| def test_the_draft_load_format_applies_only_when_configured(self): |
There was a problem hiding this comment.
[suggestion] The regression that motivated moving load_format off the scheduler was "draft load_format lands on the shared ServerArgs." test_the_draft_load_format_applies_only_when_configured asserts the draft receives "dummy" but never asserts the target still has "auto". test_the_target_instance_is_left_alone only seeds load_format="auto" without speculative_draft_load_format, so it would still pass if the helper mutated the target when a draft format is configured. The worker isolation test also only pins context_length.
Suggestion: In the configured branch of test_the_draft_load_format_applies_only_when_configured (or in test_the_target_instance_is_left_alone with both fields seeded), add self.assertEqual(target.load_format, "auto") after the copy. Optionally seed speculative_draft_load_format in the four-worker isolation test and assert the caller's instance keeps its load_format as well.
| `build_draft_tp_worker()` get private bags (a preserved publish of the rewritten copy); | ||
| drafts constructed directly with `is_draft_worker=True` skip publish and **share the | ||
| target's bags** — a draft-side write there poisons the target. | ||
| target's bags** — a draft-side write there poisons the target. Their `ServerArgs` |
There was a problem hiding this comment.
[nit] The added sentence says v2 draft workers get a private ServerArgs via draft_server_args_copy(). That is true for the four direct-construction workers this PR changes, but dflash/dspark still get their private instance from build_draft_tp_worker's own deepcopy + draft_server_args_overrides, not from draft_server_args_copy(). "Either way" is right; naming only one helper over-generalizes.
Suggestion: Phrase as: private ServerArgs via draft_server_args_copy() (eagle/standalone/multi-layer/frozen-kv) or via build_draft_tp_worker()'s deepcopy (dflash/dspark).
fc23081 to
e544954
Compare
e544954 to
c4fbeb9
Compare
|
Round 2 — the finding is real and is fixed in @chatgpt-codex-connector is right that handing the draft its own instance is not enough: model-level weight loading reads the bags, not the instance it was constructed with (
with get_context().preserve_config():
get_context().set_server_args(draft_server_args)
self.draft_worker = DraftWorkerClass(**draft_worker_kwargs)so every algorithm — built-in, dflash/dspark (which nest their own publish) and plugin — loads its draft under the draft config, and the target's config is back in the slot afterwards. A test asserts |
Summary
Testing
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c4fbeb91a6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| draft_server_args = draft_server_args_copy( | ||
| self.server_args, self.tp_worker.model_runner.model_config | ||
| ) |
There was a problem hiding this comment.
When a registered plugin factory inspects server_args.load_format—for example, to select a loader-specific worker—the factory still receives self.server_args in the subsequent create_worker call, even though the constructed worker now receives this draft copy. Previously the scheduler applied --speculative-draft-load-format before invoking the factory, so a target format of auto and draft format of dummy now makes the factory select behavior for auto while its worker is initialized for dummy. Fresh evidence in the current tree is that the worker handoff has been corrected to use draft_server_args, but the factory handoff remains on the target object; pass the same copy to create_worker as well.
Useful? React with 👍 / 👎.
Re-review (Round 2 —
|
c4fbeb9 to
32be9c0
Compare
|
Round 3 — one genuinely new finding, fixed in Factory handoff (@chatgpt-codex-connector): right, and it is the same regression one call earlier.
The other comments in this batch are re-posts of the round-2 items after the force-push; those fixes are on the current tips — |
Summary
Testing
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 32be9c01d6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| draft_server_args.override( | ||
| "draft_worker.copy", | ||
| context_length=target_model_config.context_len, | ||
| **_draft_load_format_fields(), |
There was a problem hiding this comment.
When the target runner dynamically sets disable_chunked_prefix_cache=True because its MLA attention backend does not support the cached-prefix layout (model_runner_components/misc_utils.py:28-37), this copy republishes the pristine ServerArgs value instead of that resolved bag value. Direct EAGLE drafts such as MistralLarge3ForCausalLMEagle construct DeepSeek attention while this stale copy is published, and DeepseekMHAForwardMixin.init_mha_forward permanently captures the false value (forward_mha.py:147-149), so long-prefix requests can later select the unsupported MHA chunked-KV path. draft_server_args_overrides already preserves this resolved gate for DFLASH/DSPARK; apply the same override to this direct-draft copy.
Useful? React with 👍 / 👎.
32be9c0 to
3ab0285
Compare
|
Round 4 — the finding is real and the root cause is wider than the one field; fixed in @chatgpt-codex-connector is right that the copy carried the pristine So rather than adding one more field to the list, Two tests: a load-time |
Summary
Testing
|
3ab0285 to
1385e39
Compare
1385e39 to
a0176dd
Compare
|
Commit message and the SKILL bullet corrected (
Also fixed the double "Using draft model load_format" log the review spotted: it is emitted once, in |
EAGLEWorkerV2, StandaloneWorkerV2, MultiLayerEagleWorkerV2 and FrozenKVMTPWorkerV2 wrote the draft's context_length onto the ServerArgs instance they share with the target worker, and the scheduler wrote the draft's load_format onto that same object just before creating them. The target's config carried draft values from then on, and anything constructed later in the process inherited them. Scheduler.maybe_init_draft_worker now makes one draft copy through draft_server_args_copy() and hands it to both the worker factory and the worker, so every algorithm gets it — the four built-ins, dflash/dspark (which deepcopy it again inside build_draft_tp_worker), and anything registered through SpeculativeAlgorithm.register. The copy starts from the config the process resolved, not from the pristine seed, so load-time overrides made before this point (the chunked-prefix gate, the SM100 GDN prefill default) are part of what the draft sees; context_length and load_format are applied on top. The construction runs under a preserved publish of that copy, the shape build_draft_tp_worker already used. Weight loading reads the bags rather than the instance it was handed — Inkling's ModelOpt scale normalization keys on load_format — so the draft has to be built with its own config published, and the target's is back in the slot when construction returns. The EAGLE hot-token-map write is deleted, not moved. init_token_map runs from alloc_memory_pool, long after the draft's TpModelWorker built its ModelConfig, and hot_vocab_size is only ever read off model_config.hf_config, which json_model_override_args reaches at ModelConfig construction. The write could not affect the draft model; only the shared instance saw it. hot_token_id is unchanged, so a draft checkpoint that declares hot_vocab_size behaves as before. Tests: draft_server_args_copy carries the target context_length, a configured draft load_format and any load-time override while leaving the target's instance alone; and the scheduler handoff pins that the factory and the worker both receive the copy, that the copy is the published config during construction, and that the target's is restored afterwards. Writer ratchet 31 -> 26.
a0176dd to
1c7bb56
Compare
ch-wan
added
run-ci
ready-to-merge
|
Closing unmerged and reopening against GitHub classifies a chained-base series as a stack, and in that mode it refuses base retargeting ( The replacement PR carries the identical commit; the review history, the six rounds of comment triage and the validation notes stay here for reference. Link posted below. |
|
Reopened as #33335 (base |
1 participant
The problem
EAGLEWorkerV2,StandaloneWorkerV2,MultiLayerEagleWorkerV2andFrozenKVMTPWorkerV2wrote the draft'scontext_lengthonto theServerArgsinstance they share with the target worker, and the scheduler wrote the draft's
load_formatonto that same object just before constructing them:Neither the worker nor
EagleDraftWorkercopies, so the draft'sTpModelWorkerreads the values from the shared object — which then carries draft values for the
rest of the process.
dflashanddsparknever had this:build_draft_tp_workerdeepcopies first.
The change
draft_server_args_copy(server_args, target_model_config)gives the four workersthe same treatment: deepcopy, then apply the per-draft values through the audited
mutation point.
ModelConfig.from_server_argsandbuild_load_configboth read theinstance they are handed, so the values land exactly where they did before, while
the target keeps what the launcher resolved.
load_formatmoves into both draft-copy paths.build_draft_tp_worker's overrideset needs it too: the copy those workers make used to inherit the scheduler's write,
so removing that write without this would silently drop
--speculative-draft-load-formatfor dflash and dspark.The EAGLE hot-token-map write is deleted, not moved
init_token_mapsetjson_model_override_argsto{"hot_vocab_size": N}, but itruns from
alloc_memory_pool— long afterEagleDraftWorker.__init__built thedraft's
TpModelWorkerand with it theModelConfig.hot_vocab_sizeis only everread off
model_config.hf_config(llama_eagle.py, the two draft-extend graphrunners), and
json_model_override_argsreacheshf_configonly atModelConfigconstruction. The write therefore could not affect the draft model; only the shared
instance saw it.
hot_token_idis unchanged, so a draft checkpoint that declareshot_vocab_sizein its own config behaves exactly as before.If the intent was to trim the draft vocabulary at load time, that needs the token map
resolved before the draft model is built — a behavioural change that deserves its
own PR and an end-to-end run with a token-map checkpoint.
Validation
test_draft_server_args_copy.py: the copy carriescontext_length/load_format, the target instance is untouched, and thebuild_draft_tp_workeroverride set carries the draft load format.
test_spec_worker_draft_isolation.py: each of the four workers hands its draft adistinct
ServerArgsand leaves the target's alone.test/registered/unit/{spec,model_executor}and the config ratchets pass; fullregistered CPU battery shows no new failures against the base commit.
machine this was written on, so no EAGLE/MTP end-to-end run was made. CI's
speculative suites are the gate — please look at them before merging.
Writer ratchet 31 → 26.
CI States
Latest PR Test (Base): ❌ Run #30771343684
Latest PR Test (Extra): 🚫 Run #30784011979