[lkh-bug]

Motivation

BlockedAllocator.free() validates every block for validity before mutating the free list, but the "already free" check cannot detect duplicates within a single call: marking happens in the second (mutation) loop, so both passes see the stale -1 marker for a repeated id.

Repro (pure CPU, no model or GPU needed):

allocator = BlockedAllocator(4)
a = allocator.allocate(3)      # blocks [0, 1, 2]
allocator.free([0, 0, 1])      # passes validation
# free_blocks is now 4 (should be 2): block 0 is in the free list twice
b = allocator.allocate(2)      # returns [1, 0]
# block 0 is handed out while sequence 1 still owns it -> two sequences
# share one KV block and silently overwrite each other's cache state

The same call path is reachable from production code: DSStateManager.flush_sequence() passes seq.all_block_ids() (a Tensor) into BlockedKVCache.free() -> BlockedAllocator.free(). Iterating that Tensor yields 0-dim tensors, which hash by identity -- so even a set-based duplicate check would be defeated without normalizing to int.

Changes

1. In the validation pass, reject ids that appear more than once in the same free() call (Block X provided multiple times to free).
2. Normalize each element with int(block) before checking, so Tensor inputs go through the same code path.

No behavior change for valid inputs: the free-list mutation loop is untouched.

Tests

Added test_free_duplicate_block_rejected in tests/unit/inference/v2/ragged/test_blocked_allocator.py:

• duplicate free() raises ValueError
• allocator state is unchanged by the failed call (free_blocks stays consistent)
• the blocks can still be freed exactly once afterwards, and a full allocate/free cycle yields no duplicate ids

Test commands / results

• Pre-fix repro on CPU (pin_memory shim): free([0, 0, 1]) succeeded and a subsequent allocate(2) returned [1, 0] -- block 0 allocated twice.
• Post-fix: the duplicate call raises ValueError: Block 0 provided multiple times to free; free_blocks stays consistent.
• flake8 on changed files: clean.
• No GPU was required for this change: the allocator is host-side pure-Python logic; the added test runs under the existing inference_v2 marker and only needs BlockedAllocator construction (CPU with pinned-memory shim locally; will exercise natively on the GPU CI runners).

Duplicate check

Searched issues and PRs for BlockedAllocator, "already free", "allocated twice", "duplicate block" (both deepspeedai/DeepSpeed and legacy microsoft/DeepSpeed) -- no existing report. PR #8503 (fixing #5407) touches only engine_v2.py and does not include this fix.

AI assistance

This fix was prepared with AI coding assistance; the submitting human has reviewed every changed line and the verification above.

Signed-off-by: likehong likehong@kingsoft.com

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-28T03:10:04.400657Z	09fc930	PR opened

ℹ️ 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 09fc93030c

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Use normalized block IDs in the mutation pass

When block 1 is allocated and a caller passes free(True) (or a 0-D bool tensor), the conversion validates ID 1, but the mutation loop still uses the original True. PyTorch treats that as a boolean mask, so self._blocks[True] = ... overwrites the whole free-list tensor while only incrementing _free_blocks once, corrupting allocator state; materialize the normalized IDs and use them in both passes (or reject booleans).

Useful? React with 👍 / 👎.