fix(router): account for all output blocks in response chunks

[bzsuni]

Overview:

With --router-track-output-blocks enabled, the router undercounts output blocks when a response chunk crosses more than one block boundary.
This makes its local decode-load estimate depend on how the backend groups tokens into responses.

OutputBlockTracker advances its total to include all newly crossed blocks, but RequestGuard sends only one AddOutputBlock update. Later chunks cannot recover the missing blocks because the tracker has already advanced.

The fix carries the number of new blocks in OutputBlockUpdate and enqueues an update for each through the existing booking-checked scheduler path.
It leaves the decay formula, cleanup, and local-only load updates unchanged. Output-block tracking remains disabled by default; no backend or wire-format changes are needed.

Details:

The regression uses the actual request guard and scheduler, without GPU inference.
For a 16-token prompt, 48 output tokens, and block size 16, single-token chunks count 4 blocks before and after the fix. Grouping the same tokens as [1, 32, 15] exposes the bug.
SGLang's real output accumulator also produces this layout with stream_interval=32 in the CPU producer probe.

Recorded output for the grouped case:

Before — the regression fails:

track_output_blocks=true prompt_tokens=16 output_tokens=48 block_size=16 chunks=[1, 32, 15] initial_blocks=1 observed_blocks=3 expected_blocks=4

After — the same regression passes:

track_output_blocks=true prompt_tokens=16 output_tokens=48 block_size=16 chunks=[1, 32, 15] initial_blocks=1 observed_blocks=4 expected_blocks=4

Where should the reviewer start?

Related Issues

🚫 This PR is NOT linked to an issue:

• Confirmed — no related issue

Summary by CodeRabbit

• Bug Fixes

  • Improved scheduler block accounting across multiple block boundaries.
  • Preserved decay behavior and stopped processing after the first failed scheduler update.
  • Ensured output accounting remains consistent regardless of response chunking or prompt length.
  • Corrected cleanup behavior so active requests and tracked blocks reset properly.

• Tests

  • Added coverage for boundary crossings, partial prompt blocks, decay behavior, and disabled tracking.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[github-actions]

👋 Hi bzsuni! Thank you for contributing to ai-dynamo/dynamo.

Just a reminder: The NVIDIA Test Github Validation CI runs an essential subset of the testing framework to quickly catch errors.Your PR reviewers may elect to test the changes comprehensively before approving your changes.

🚀

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ef8dfd4c-2968-4456-bbec-3f96a2d17f85

📥 Commits

Reviewing files that changed from the base of the PR and between 287a8e7 and 6607d6c.

📒 Files selected for processing (2)

• lib/llm/src/kv_router/routing_host/request_guard.rs
• lib/llm/src/kv_router/routing_host/tests.rs

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

Output block tracking now counts every crossed scheduler block. Request handling enqueues one update per crossed block. Tests cover boundary crossings, decay, disabled tracking, chunk invariance, scheduler load, cleanup, and configurable router setup.

Changes

Output block accounting

Layer / File(s)	Summary
Crossed block update processing lib/llm/src/kv_router/routing_host/request_guard.rs	OutputBlockTracker reports newly crossed blocks. Request processing enqueues one scheduler update for each crossed block and stops after the first enqueue failure. Tests cover boundaries, partial blocks, decay, zero output, and disabled tracking.
Chunk-invariant accounting validation lib/llm/src/kv_router/routing_host/tests.rs	Tests verify consistent accounting across prompt sizes, response chunk layouts, and tracking settings. Router test construction now accepts a caller-provided KvRouterConfig and removes duplicated configuration setup.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 6607d

The change corrects output-block accounting across scheduler boundaries without changing default behavior or wire formats. Regression coverage verifies boundary crossings, cleanup, decay, and disabled tracking, so it is merge-ready.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 12 functions across 2 files.	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 identifies the router fix for accounting for all output blocks in response chunks. It accurately summarizes the main change.
Description check	✅ Passed	The description explains the bug, implementation, regression coverage, expected behavior, and confirms that no issue is linked. The "Where should the reviewer start?" section is still empty, but this …
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.

• Fix all pre-merge checks with AI

---

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

[jthomson04]

Reviewed c5fc0c5. No new findings from the source review. Both earlier review comments are addressed. Tests were not run.

[jthomson04]

/ok to test c5fc0c5

[jthomson04]

/ok to test 9664347

[jthomson04]

Non-blocking follow-up: with output tracking and agent_hints.osl enabled, this loop scans all output blocks once per new block to apply the same decay value. Now that #14666 has merged, we can batch the direct updates and apply decay once per chunk. I’d keep that optimization separate from this accounting fix.

Tracked in #15446.