[KV Offload] Fix num_tokens_after_batch for different termination types

[Alex-ai-future]

Summary

Fix incorrect num_tokens_after_batch computation in offloading scheduler by using num_computed_tokens for aborted requests (where partial KV data may exist) instead of always using num_tokens.

Problem

The original code used num_tokens for all finished requests:

if req.is_finished():
    num_tokens_after_batch = req.num_tokens

For aborted requests, num_computed_tokens < num_tokens (partial KV data). Using num_tokens would attempt to store chunks without actual KV data on GPU.

Solution

Distinguish aborted requests from other finished requests:

if req.status is RequestStatus.FINISHED_ABORTED:
    num_tokens_after_batch = req.num_computed_tokens
elif req.is_finished():
    num_tokens_after_batch = req.num_tokens

Rationale:

• Aborted requests: Use num_computed_tokens (actual KV data), never store chunks without KV data
• Other finished requests (EOS, length cap, repetition, error): Use num_tokens to handle async scheduling delay correctly

The clamp in storable_chunks (min(num_chunks, num_allocated_chunks)) remains as defense-in-depth, handling all block_ids edge cases (empty, partial capacity, etc.).

Test Plan

All existing tests pass (112 passed):

• test_abort_queued_request_does_not_build_store_job ✅ (queued abort)
• test_request_preemption ✅ (preempt handling)
• test_last_block_offloaded_at_request_finish ✅ (normal EOS)
• test_two_groups_full_and_sliding_window ✅ (multi-group)
• All other offloading scheduler tests ✅

Impact

• Semantic correctness: Aborted requests only store chunks with actual KV data
• Simplicity: Minimal code change, handles the root cause
• Safety: Clamp provides defense-in-depth for block_ids edge cases

AI Assistance

Used Claude for code analysis, root cause identification, and test validation. All changes reviewed and approved by the human author.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Alex-ai-future]

@orozery PTAL

[orozery]

I think this is the only change that is relevant.
All other changes can be reverted.

[Alex-ai-future]

@orozery
Thank you for your advice. I finally have a clear understanding and have modified it according to your requirements.

[orozery]

Leftover unrelated change? Let's revert

[orozery]

This is more long and fragile than simply:

Suggested change

	if req.status in (
	RequestStatus.FINISHED_STOPPED,
	RequestStatus.FINISHED_LENGTH_CAPPED,
	RequestStatus.FINISHED_REPETITION,
	):
	# Normal termination: use num_tokens to handle scheduling
	# delay with async scheduling.
	num_tokens_after_batch = req.num_tokens
	elif req.is_finished():
	# Abnormal termination (abort/error): store all computed KV data.
	# num_computed_tokens represents actual computed tokens.
	num_tokens_after_batch = req.num_computed_tokens
	if req.status is RequestStatus.FINISHED_ABORTED:
	num_tokens_after_batch = req.num_computed_tokens
	elif req.is_finished():
	num_tokens_after_batch = req.num_tokens