fix(vllm): classify disagg decode-engine queued requests correctly in FPM

[tedzhouhk]

Summary

InstrumentedScheduler._compute_queued only iterated self.waiting and classified everything non-PREEMPTED as queued prefill. Under disaggregated serving, decode-engine requests awaiting KV transfer from a prefill engine live in self.skipped_waiting with status WAITING_FOR_REMOTE_KVS (see vLLM's Scheduler._enqueue_waiting_request / _is_blocked_waiting_status). They were silently missed from the FPM snapshot, so the planner saw zero decode pressure on a loaded decode worker and would scale the wrong pool.

Root cause

Two-queue model in vLLM's scheduler:

• self.waiting — WAITING (new) + PREEMPTED (evicted decode).
• self.skipped_waiting — blocked-waiting states, per Scheduler._is_blocked_waiting_status:
  • WAITING_FOR_FSM — structured-output FSM compile
  • WAITING_FOR_REMOTE_KVS — disagg decode-engine KV transfer
  • WAITING_FOR_STREAMING_REQ — streaming request handshake

Previous _compute_queued only scanned self.waiting. Every WAITING_FOR_REMOTE_KVS request was invisible to the FPM even though it represents a decode-ready workload.

Fix

Iterate self.skipped_waiting alongside self.waiting:

• WAITING_FOR_REMOTE_KVS → queued decode with num_computed_tokens as the KV length. Scheduler.schedule pre-sets num_computed_tokens to the transferred KV count at the load_kv_async branch (schedule() line ~715 in vLLM 0.18.0), so that field is the correct decode-KV-context value.
• WAITING_FOR_FSM / WAITING_FOR_STREAMING_REQ → queued prefill; no KV computed yet.

Variance accumulation spans both queues (one Welford accumulator per prefill / decode bucket).

Verified against vLLM source

• RequestStatus enum (vLLM 0.18.0): WAITING, WAITING_FOR_FSM, WAITING_FOR_REMOTE_KVS, WAITING_FOR_STREAMING_REQ, RUNNING, PREEMPTED, FINISHED_*.
• Scheduler._is_blocked_waiting_status confirms the 3 blocked-waiting states route to self.skipped_waiting via _enqueue_waiting_request.
• Scheduler.schedule at load_kv_async: sets request.num_computed_tokens = num_computed_tokens before parking in skipped_waiting (comment: "When the transfer is finished, either successfully or not, request.num_computed_tokens will correctly reflect the number of computed tokens.").

Tests

Adds components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py — 10 cases covering:

• self.waiting regression paths: new prefill, preempted decode.
• self.skipped_waiting classification: WAITING_FOR_REMOTE_KVS → decode, WAITING_FOR_FSM / WAITING_FOR_STREAMING_REQ → prefill.
• Realistic mixed snapshots for both decode engine and prefill engine.
• Empty queues and cross-queue variance correctness.

Tests invoke the real _compute_queued against a stub self (via object.__new__) so any future drift in the production function body surfaces directly, without spinning up full vLLM engine state.

Pre-fix: 7/10 fail (every test exercising self.skipped_waiting).
Post-fix: 10/10 pass.

Test plan

• pytest components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py -v — 10 passed
• pre-commit run --files <changed files> — isort / black / ruff / trim-whitespace clean
• Confirmed both fix and test coverage against vLLM 0.18.0 source

Follow-ups (out of scope)

• Same bug exists on the TRT-LLM side (feat(trtllm): publish ForwardPassMetrics via FpmDirectPublisher (non-attention-DP) #8356 / TRT-LLM#13199): decode-engine queued requests in executor_request_queue are counted as queued_num_prefill. Fix will need to inspect disagg_params / generation-only markers on the request. Flagged in the PR review thread, not fixed here.

🤖 Generated with Claude Code

---

Summary by CodeRabbit

• Bug Fixes

  • Corrected the calculation of queued request metrics to accurately account for requests in different queue states and properly distinguish between prefill and decode operations.

• Tests

  • Added comprehensive unit test coverage for request classification logic, including multi-queue scenarios and edge cases to ensure metric accuracy.

[coderabbitai]

Walkthrough

The InstrumentedScheduler._compute_queued method was extended to aggregate queued metrics from both self.waiting and self.skipped_waiting queues, with new classification logic differentiating request statuses and their respective token counting methods. A comprehensive test suite was added to validate the updated behavior.

Changes

Cohort / File(s)	Summary
Core Implementation Updates components/src/dynamo/vllm/instrumented_scheduler.py	Extended _compute_queued to aggregate metrics over self.waiting and self.skipped_waiting. Added status-specific classification: PREEMPTED in waiting maps to decode_kv (using num_computed_tokens), while WAITING_FOR_REMOTE_KVS in skipped_waiting also maps to decode_kv, and WAITING_FOR_FSM / WAITING_FOR_STREAMING_REQ in skipped_waiting map to prefill (using num_tokens).
Test Coverage components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py	New test module validating _compute_queued request classification logic using minimal scheduler stubs. Includes regression tests for existing self.waiting behavior, new coverage for self.skipped_waiting classification rules, mixed queue scenarios, variance computations, and empty queue edge case.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 64.29% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the main fix: correcting classify of disaggregated decode-engine queued requests in the FPM (Fusion Planner Monitor) related to vLLM.
Description check	✅ Passed	The PR description comprehensively covers all required template sections: detailed overview of the bug, root cause analysis, the fix with technical details, verification against vLLM source, test coverage, and related follow-ups.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

components/src/dynamo/vllm/instrumented_scheduler.py (1)

486-501: LGTM — fix correctly classifies both queues.

The two-loop structure sharing a single Welford accumulator per bucket matches the PR intent: WAITING_FOR_REMOTE_KVS requests on the decode engine are now counted as queued decode using num_computed_tokens (pre-set by Scheduler.schedule at the load_kv_async branch), while WAITING_FOR_FSM / WAITING_FOR_STREAMING_REQ stay as queued prefill. Variance is correctly aggregated across both queues.

One minor defensive consideration: the else branches on lines 489 and 498 fall through to prefill for any unhandled status, so if vLLM adds a new blocked-waiting status in a future version that should count as decode (or something else), it would be silently misclassified. Given the docstring enumerates the known statuses, an explicit allowlist check with a logger.debug for unknown statuses would make future vLLM upgrades safer to audit — but the current behavior is reasonable and the docstring already notes the rationale. Optional.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@components/src/dynamo/vllm/instrumented_scheduler.py` around lines 486 - 501,
The loops iterating over self.waiting and self.skipped_waiting currently treat
any non-PREEMPTED / non-WAITING_FOR_REMOTE_KVS status as prefill, which could
silently misclassify future statuses; update the logic in
instrumented_scheduler.py (the loops that inspect RequestStatus for requests in
self.waiting and self.skipped_waiting) to use an explicit allowlist of known
statuses (e.g., RequestStatus.PREEMPTED, RequestStatus.WAITING_FOR_REMOTE_KVS,
RequestStatus.WAITING_FOR_FSM, RequestStatus.WAITING_FOR_STREAMING_REQ) and
handle them as before, and add a logger.debug (or logger.warning) branch that
logs unknown/other RequestStatus values so upgrades to vLLM that add new
statuses are visible and can be audited.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@components/src/dynamo/vllm/instrumented_scheduler.py`:
- Around line 486-501: The loops iterating over self.waiting and
self.skipped_waiting currently treat any non-PREEMPTED /
non-WAITING_FOR_REMOTE_KVS status as prefill, which could silently misclassify
future statuses; update the logic in instrumented_scheduler.py (the loops that
inspect RequestStatus for requests in self.waiting and self.skipped_waiting) to
use an explicit allowlist of known statuses (e.g., RequestStatus.PREEMPTED,
RequestStatus.WAITING_FOR_REMOTE_KVS, RequestStatus.WAITING_FOR_FSM,
RequestStatus.WAITING_FOR_STREAMING_REQ) and handle them as before, and add a
logger.debug (or logger.warning) branch that logs unknown/other RequestStatus
values so upgrades to vLLM that add new statuses are visible and can be audited.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 4e4f5c71-da0a-4060-b09a-68202c54b594

📥 Commits

Reviewing files that changed from the base of the PR and between 55a949c and 4f6287d.

📒 Files selected for processing (2)

• components/src/dynamo/vllm/instrumented_scheduler.py
• components/src/dynamo/vllm/tests/test_vllm_instrumented_scheduler.py