feat(server): configure prefill-probe routing - #186
Conversation
Signed-off-by: asurla1998 <asurla1998@users.noreply.huggingface.co>
Signed-off-by: asurla1998 <asurla1998@users.noreply.huggingface.co>
Signed-off-by: asurla1998 <asurla1998@users.noreply.huggingface.co>
WalkthroughChangesLearned prefill probe
Estimated code review effort: 5 (Critical) | ~120 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (4)
crates/libsy-llm-client/src/vllm_hidden_state_probe.rs (2)
441-442: 🔒 Security & Privacy | 🔵 Trivial | 💤 Low valueValidate the artifact path before opening it.
open_synchronized_artifactopens the reported path (and its.lockcompanion) beforevalidate_hidden_states_pathconfirms the canonical target is insiderootand is a regular file. Reordering keeps the containment check strictly ahead of any file handle on an attacker-influenced path.♻️ Proposed reordering
- let (mut artifact, synchronization_lock) = open_synchronized_artifact(path, timeout)?; - let artifact_path = validate_hidden_states_path(root, path)?; + let (mut artifact, synchronization_lock) = open_synchronized_artifact(path, timeout)?; + // Re-validate after the writer released the lock: the file must still + // canonicalize inside `root` and be a regular file before it is read. + let artifact_path = validate_hidden_states_path(root, path)?;Note the current order is already safe for reads (validation precedes
read_to_end), so this is hardening only.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/libsy-llm-client/src/vllm_hidden_state_probe.rs` around lines 441 - 442, Reorder the hidden-state artifact flow so validate_hidden_states_path(root, path) completes before open_synchronized_artifact(path, timeout) is called. Preserve the existing artifact and synchronization_lock bindings and subsequent read behavior after validation.
168-169: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winStale sweep runs on every extraction and its failure aborts the probe.
reap_stale_artifactsdoes a fullread_dirplus per-entrymetadata/lock attempts before each HTTP call, and any single I/O error (e.g. one unreadable entry another writer owns) propagates and fails the whole extraction. Consider decoupling the sweep from the request path (interval-gated or background) and logging sweep errors instead of failing the probe.♻️ Sketch: don't fail extraction on sweep errors
- self.reap_stale_artifacts().await?; + // A best-effort sweep must not fail routing; the artifact read below is + // what the caller actually depends on. + if let Err(error) = self.reap_stale_artifacts().await { + tracing::debug!(%error, "hidden-state stale sweep failed"); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/libsy-llm-client/src/vllm_hidden_state_probe.rs` around lines 168 - 169, Update VllmHiddenStateProbe::extract so stale-artifact cleanup is not performed on every extraction and cleanup I/O failures do not abort the probe request. Gate reap_stale_artifacts by an appropriate interval or move it to background maintenance, and log any sweep error while allowing extraction to continue to the HTTP call.crates/libsy/src/algorithms/prefill_probe.rs (1)
210-226: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffNo request coalescing: concurrent identical tasks each issue a probe call.
The cache is only consulted before and written after the await, so N in-flight requests for the same task produce N vLLM prefill calls plus N inference jobs. If duplicate concurrent tasks are expected (retries, fan-out agents), a per-key in-flight map or a single-flight helper would cut probe load without changing behavior.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/libsy/src/algorithms/prefill_probe.rs` around lines 210 - 226, Update select_for_task to coalesce concurrent requests sharing the same cache_key by introducing or reusing a per-key in-flight future/task map. Ensure only the first caller runs probe.extract and routing.select, while concurrent callers await the same result; preserve existing cache reads, result caching, and error behavior, and remove the in-flight entry when processing completes.crates/libsy/src/algorithms/prefill_probe/artifact.rs (1)
19-24: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffHard-pinned architecture/output constants make every retrained checkpoint a code change.
PCA_DIM,TRUNK_HIDDEN,ENSEMBLE_SIZE, and especially the literalOUTPUT_NAMESmodel ids are enforced with strict equality, so a checkpoint trained with a different ensemble size or a renamed head fails to load even though the loader could derive those values fromrouter.json(pca_dim,trunk_hidden,ensemble_size,output_namesare all present in metadata). Consider validating self-consistency (shapes vs. metadata) plusformat_version, and keeping only the structural invariants as constants.Also applies to: 312-334
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/libsy/src/algorithms/prefill_probe/artifact.rs` around lines 19 - 24, Replace the hard-pinned PCA_DIM, TRUNK_HIDDEN, ENSEMBLE_SIZE, and OUTPUT_NAMES checks in the artifact loader with values parsed from router.json metadata (pca_dim, trunk_hidden, ensemble_size, and output_names). Validate checkpoint tensor shapes and metadata self-consistency, while retaining only genuine structural invariants as constants and validating format_version.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/switchyard-server/CONFIGURATION.md`:
- Around line 24-26: Update the checklist step for constructing the algorithm to
reference the appropriate target resolver, either resolve_target or
resolve_targets, instead of naming only resolve_targets. Keep the guidance
applicable to RouteConfig::PrefillProbe and other implementations.
In `@docs/operations/vllm_hidden_state_probe.md`:
- Around line 27-31: Update the artifact-directory creation command in the probe
setup to apply explicit restrictive permissions rather than relying on the
process umask. Use mode 0700 for same-user processes, or 0770 with an
appropriate shared group when the processes run under different users.
- Around line 33-59: Update the vllm serve example command to include the
--no-enable-chunked-prefill option, preserving the existing hidden-state
extraction and KV connector configuration.
---
Nitpick comments:
In `@crates/libsy-llm-client/src/vllm_hidden_state_probe.rs`:
- Around line 441-442: Reorder the hidden-state artifact flow so
validate_hidden_states_path(root, path) completes before
open_synchronized_artifact(path, timeout) is called. Preserve the existing
artifact and synchronization_lock bindings and subsequent read behavior after
validation.
- Around line 168-169: Update VllmHiddenStateProbe::extract so stale-artifact
cleanup is not performed on every extraction and cleanup I/O failures do not
abort the probe request. Gate reap_stale_artifacts by an appropriate interval or
move it to background maintenance, and log any sweep error while allowing
extraction to continue to the HTTP call.
In `@crates/libsy/src/algorithms/prefill_probe.rs`:
- Around line 210-226: Update select_for_task to coalesce concurrent requests
sharing the same cache_key by introducing or reusing a per-key in-flight
future/task map. Ensure only the first caller runs probe.extract and
routing.select, while concurrent callers await the same result; preserve
existing cache reads, result caching, and error behavior, and remove the
in-flight entry when processing completes.
In `@crates/libsy/src/algorithms/prefill_probe/artifact.rs`:
- Around line 19-24: Replace the hard-pinned PCA_DIM, TRUNK_HIDDEN,
ENSEMBLE_SIZE, and OUTPUT_NAMES checks in the artifact loader with values parsed
from router.json metadata (pca_dim, trunk_hidden, ensemble_size, and
output_names). Validate checkpoint tensor shapes and metadata self-consistency,
while retaining only genuine structural invariants as constants and validating
format_version.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 949d5b25-2c93-4bed-9621-2e707916fb5e
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock,!Cargo.lock
📒 Files selected for processing (15)
crates/libsy-llm-client/Cargo.tomlcrates/libsy-llm-client/src/lib.rscrates/libsy-llm-client/src/vllm_hidden_state_probe.rscrates/libsy/Cargo.tomlcrates/libsy/src/algorithms.rscrates/libsy/src/algorithms/prefill_probe.rscrates/libsy/src/algorithms/prefill_probe/artifact.rscrates/libsy/src/algorithms/prefill_probe/policy.rscrates/libsy/src/lib.rscrates/switchyard-server/CONFIGURATION.mdcrates/switchyard-server/Cargo.tomlcrates/switchyard-server/README.mdcrates/switchyard-server/src/config.rsdocs/operations/vllm_hidden_state_probe.mdmkdocs.yml
| 2. Add its TOML fields as a `RouteConfig` variant in `src/config.rs`. | ||
| 3. Construct it in the `build_algorithm` match, resolving target names with `resolve_targets`. | ||
| 4. Add a parsing test and an end-to-end server test when the algorithm makes LLM calls. |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Make target-resolution guidance match the implementation.
The checklist names resolve_targets, but RouteConfig::PrefillProbe resolves strong_target and weak_target individually with resolve_target in crates/switchyard-server/src/config.rs (Lines 295–399). Say “the appropriate target resolver (resolve_target/resolve_targets)” to avoid misleading future implementations.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/switchyard-server/CONFIGURATION.md` around lines 24 - 26, Update the
checklist step for constructing the algorithm to reference the appropriate
target resolver, either resolve_target or resolve_targets, instead of naming
only resolve_targets. Keep the guidance applicable to RouteConfig::PrefillProbe
and other implementations.
| Create a directory used only for probe artifacts: | ||
|
|
||
| ```bash | ||
| mkdir -p /dev/shm/switchyard-prefill | ||
| ``` |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟠 Major | ⚡ Quick win
Create the artifact directory with explicit restrictive permissions.
mkdir -p relies on the process umask and can leave the /dev/shm directory readable by other local users, while the guide correctly identifies hidden-state artifacts as sensitive. Use 0700 when both processes share a user, or a dedicated group with 0770 when they do not.
Suggested setup
- mkdir -p /dev/shm/switchyard-prefill
+ install -d -m 0700 /dev/shm/switchyard-prefill📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| Create a directory used only for probe artifacts: | |
| ```bash | |
| mkdir -p /dev/shm/switchyard-prefill | |
| ``` | |
| Create a directory used only for probe artifacts: | |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/operations/vllm_hidden_state_probe.md` around lines 27 - 31, Update the
artifact-directory creation command in the probe setup to apply explicit
restrictive permissions rather than relying on the process umask. Use mode 0700
for same-user processes, or 0770 with an appropriate shared group when the
processes run under different users.
| Start vLLM with the hidden-state extraction method and disk connector. Replace the model and layer | ||
| IDs with the values used to train the checkpoint: | ||
|
|
||
| ```bash | ||
| vllm serve Qwen/Qwen3-8B \ | ||
| --speculative_config '{ | ||
| "method": "extract_hidden_states", | ||
| "num_speculative_tokens": 1, | ||
| "draft_model_config": { | ||
| "hf_config": { | ||
| "eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4] | ||
| } | ||
| } | ||
| }' \ | ||
| --kv_transfer_config '{ | ||
| "kv_connector": "ExampleHiddenStatesConnector", | ||
| "kv_role": "kv_producer", | ||
| "kv_connector_extra_config": { | ||
| "shared_storage_path": "/dev/shm/switchyard-prefill", | ||
| "use_synchronization_lock": true | ||
| } | ||
| }' | ||
| ``` | ||
|
|
||
| Keep `use_synchronization_lock` enabled. Switchyard acquires the companion `.lock` file before | ||
| reading, so it cannot parse a partially written safetensors artifact. Chunked prefill is | ||
| incompatible with vLLM hidden-state extraction and must be disabled. |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the referenced documentation section with line numbers.
sed -n '1,120p' docs/operations/vllm_hidden_state_probe.md | cat -nRepository: NVIDIA-NeMo/Switchyard
Length of output: 4682
🌐 Web query:
vLLM extract_hidden_states chunked prefill disabled no-enable-chunked-prefill hidden-state extraction documentation
💡 Result:
In vLLM, the "hidden state extraction" feature allows you to capture intermediate layer activations during inference [1][2]. Historically, this feature was incompatible with chunked prefill, requiring users to explicitly disable chunked prefill using the --no-enable-chunked-prefill flag to ensure correct operation and prevent errors [1][3][4]. However, recent improvements (notably in PR #43805, dated May 2026) have addressed this limitation by modifying how hidden states are saved [5]. By moving the save operation to the get_finished stage—which occurs only once per request after generation completes—vLLM is now able to support hidden state extraction without strictly requiring that chunked prefill be disabled [5]. Regarding the --no-enable-chunked-prefill flag: - In older versions of vLLM or specific configurations, this flag was the standard way to disable chunked prefill [6][7][4]. - In the newer V1 engine architecture, chunked prefill is enabled by default [8]. While some versions allowed the --no-enable-chunked-prefill argument to override this, vLLM may issue warnings that disabling this feature is not officially supported and could lead to performance issues or crashes [8][9]. - If you are using a version of vLLM that predates the improvements mentioned in PR #43805, you must continue to use --no-enable-chunked-prefill (or enable_chunked_prefill=False in the Python API) when using extract_hidden_states [1][4]. If you are on a current version, this manual override may no longer be necessary [5]. It is recommended to consult the documentation for your specific vLLM version, as the requirement for this flag depends on the underlying engine implementation and version [8][10].
Citations:
- 1: https://docs.vllm.ai/en/stable/features/speculative_decoding/extract_hidden_states/
- 2: https://docs.vllm.ai/en/latest/features/speculative_decoding/extract_hidden_states/
- 3: [Perf] Optimize hidden state extraction logic vllm-project/vllm#37374
- 4: https://docs.vllm.ai/en/v0.23.0/features/speculative_decoding/extract_hidden_states/
- 5: Hidden states extraction improvements vllm-project/vllm#43805
- 6: https://docs.vllm.ai/en/stable/configuration/engine_args/
- 7: https://docs.vllm.ai/en/v0.8.5/serving/engine_args.html
- 8: https://discuss.vllm.ai/t/vllm-error-unrecognized-arguments-no-enable-chunked-prefill/1183
- 9: [Usage]: vllm v1 not use
--no-enable-chunked-prefillto disable chunked prefill. vllm-project/vllm#42360 - 10: https://docs.vllm.ai/en/latest/api/vllm/model%5Fexecutor/models/extract_hidden_states/
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,120p' docs/operations/vllm_hidden_state_probe.md | cat -nRepository: NVIDIA-NeMo/Switchyard
Length of output: 4682
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search the repository for chunked prefill references around hidden-state extraction docs.
rg -n --hidden --glob '!**/.git/**' -C 2 'chunked prefill|no-enable-chunked-prefill|extract_hidden_states|hidden-state extraction' docs .Repository: NVIDIA-NeMo/Switchyard
Length of output: 12703
Disable chunked prefill in the example launch command. Hidden-state extraction is incompatible with chunked prefill here, so add --no-enable-chunked-prefill to the vllm serve example.
Suggested command change
vllm serve Qwen/Qwen3-8B \
+ --no-enable-chunked-prefill \
--speculative_config '{📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| Start vLLM with the hidden-state extraction method and disk connector. Replace the model and layer | |
| IDs with the values used to train the checkpoint: | |
| ```bash | |
| vllm serve Qwen/Qwen3-8B \ | |
| --speculative_config '{ | |
| "method": "extract_hidden_states", | |
| "num_speculative_tokens": 1, | |
| "draft_model_config": { | |
| "hf_config": { | |
| "eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4] | |
| } | |
| } | |
| }' \ | |
| --kv_transfer_config '{ | |
| "kv_connector": "ExampleHiddenStatesConnector", | |
| "kv_role": "kv_producer", | |
| "kv_connector_extra_config": { | |
| "shared_storage_path": "/dev/shm/switchyard-prefill", | |
| "use_synchronization_lock": true | |
| } | |
| }' | |
| ``` | |
| Keep `use_synchronization_lock` enabled. Switchyard acquires the companion `.lock` file before | |
| reading, so it cannot parse a partially written safetensors artifact. Chunked prefill is | |
| incompatible with vLLM hidden-state extraction and must be disabled. | |
| Start vLLM with the hidden-state extraction method and disk connector. Replace the model and layer | |
| IDs with the values used to train the checkpoint: | |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/operations/vllm_hidden_state_probe.md` around lines 33 - 59, Update the
vllm serve example command to include the --no-enable-chunked-prefill option,
preserving the existing hidden-state extraction and KV connector configuration.
1 participant
What
prefill_proberoute type to the Rust server TOML configuration.Why
Completes the server integration needed to configure and run learned prefill routing end to end.
Depends on #174 and #185. Until those PRs merge, GitHub will also show their commits in this stacked PR.
How tested
cargo fmt --all --checkcargo clippy --workspace --all-targets -- -D warningscargo test --workspaceuv run ruff check .uv run mypy switchyarduv run pytest tests/ -v -m "not integration"— 1718 passed, 12 skipped, 28 deselectedcd docs && make publishChecklist
Notes for reviewers
The intended delta is the single
feat(server)commit on top of #185. Token-aware pricing remains a separate follow-up.Summary by CodeRabbit
New Features
Documentation