Skip to content

fix(vllm): bind V1 hybrid Mamba pages on 0.28 - #500

Merged
RixinLiu merged 4 commits into
ovg-project:mainfrom
shipiyouniao:feature/vllm-028-v1-hybrid
Sep 25, 2026
Merged

RixinLiu merged 4 commits into
ovg-project:mainfrom
shipiyouniao:feature/vllm-028-v1-hybrid

Conversation

@shipiyouniao

@shipiyouniao shipiyouniao commented Sep 17, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Original implementation by @CMQY. The commit retains CMQY <919517564@qq.com> as its Git author.

Complete the vLLM 0.28 Model Runner V1 hybrid-cache contract: provide raw Mamba pages, use native partial-prefix metadata, and copy logical blocks without interpreting the entire reserved FTensor address range as live KV data. This does not add a scheduler policy, environment variable, or change to the configured hash granularity.

Dependencies and Scope

Part of #490 and the release checklist in #509. #492, #493, #511 and #512 are merged. This branch is now rebased onto main at 0ed269fe9057508699ceb95fab071f20ab3e26e5, so it no longer includes the profiling and packed-storage parent changes in its diff. The four remaining commits retain CMQY's original implementation and all three version-boundary corrections by Rixin.

This is Model Runner V1 only. The rebase preserves the merged cross-layer-sharing work in #472 and reuses the native block-pool mixin from #491. Following Rixin's source audit, raw Mamba binding starts at 0.27; logical-block copy and native prefix metadata start at 0.26. The native-pool compatibility range is 0.26 through 0.29. It does not extend the V2 adapter to 0.28 or claim V2 Mamba support.

Root Cause

There are three distinct contracts on the startup and prefix-hit paths.

  1. Mamba binding changed. vLLM 0.28 MambaBase.bind_kv_cache expects one raw [blocks, 1, 1, page_bytes] tensor and derives the typed recurrent-state views. Returning the older list of state tensors is not sufficient.

  2. Partial-prefix entries have multiple hash aliases and retained copy references. The native BlockPool methods own insertion, promotion, alias removal and movement between copy endpoints. Keeping only one custom key per block misses this contract. The existing native mixin was previously enabled only for 0.29, leaving 0.28 on the older path.

  3. Native V1 copy assumes compact physical backing. copy_kv_cache_blocks_inplace obtains the whole storage and divides it into scheduler blocks:

    blocks.set_(tensor.untyped_storage())
    assert blocks.numel() % num_blocks == 0
    blocks = blocks.view(num_blocks, -1)

    An FTensor reserves more virtual address space than the exposed logical KV view. Split K/V storage also has a reserved gap, and interleaved layer views share a backing allocation. The native calculation can fail the divisibility assertion or derive the wrong byte range; copying once per layer can copy shared backing repeatedly.

These are correctness problems, not hash-size tuning. In particular, a hash granularity smaller than the allocation block must still use vLLM's native hash-boundary conversion and partial-to-full promotion.

The coordinator must also preserve the hash granularity of its original native pool. A missing field is not evidence that it equals the allocation block size: substituting 16 for an actual 4-token hash size can associate a physical block with the wrong prefix. vLLM 0.12's BlockPool explicitly stores the independent hash size. Earlier 0.9-0.11 pools lack that field and use the common allocation size for the supported prefix-cache path; see the 0.11.2 manager.

Changes

  • Return a strided raw Mamba page view from 0.27, preserving the logical block stride, layer offset and backing-storage offset. Earlier V1 releases keep the old typed-state-list format, with copy metadata attached to their views.
  • Reuse NativeBlockPoolMixin from 0.26 through 0.29. Native metadata methods manage hash aliases and copy lifetimes; the existing elastic pool still owns physical allocation, cache caps and eviction policy. No second prefix-index implementation is introduced.
  • Read hash granularity from the original native pool before initializing KVCached. Missing or invalid metadata on modern/unknown versions raises RuntimeError, which the coordinator patch propagates rather than swallowing. Only the explicit 0.9-0.11 compatibility branch may use the allocation block size when the field is absent. A coordinator field cannot mask missing native-pool metadata. The pre-0.28 caching path also uses native hash-boundary conversion when hash and allocation sizes differ.
  • Attach byte-page metadata when constructing packed, MLA, split K/V and Mamba views, including reconstructed heterogeneous views. Split K/V metadata excludes the unused gap. Shared backing is copied once.
  • Patch only the V1 runner's imported copy helper. The native runner call site still determines when the copy is submitted. Disabled KVCached delegates to the original helper. Invalid indices, duplicate destinations or missing metadata fail before any pool is modified.

Previous Path

sequenceDiagram
    participant S as Scheduler
    participant P as Elastic block pool
    participant W as V1 worker
    participant F as FTensor storage
    alt Partial-prefix metadata path
        S->>P: Cache or reuse a partial prefix
        P-->>S: Missing native alias/lifetime contract
    else A block copy reaches the worker
        S->>W: Submit logical block copies
        W->>F: Divide whole reserved storage by num_blocks
        F-->>W: Wrong logical extent or divisibility assertion
    end
Loading

Updated Path

sequenceDiagram
    participant S as Scheduler
    participant P as Native metadata + elastic allocation
    participant W as V1 worker
    participant F as Logical byte-page views
    S->>P: Cache, promote or move prefix hashes
    P-->>S: Native aliases and retained copy endpoints
    S->>W: Submit block copies in normal worker order
    W->>F: Validate metadata and copy each shared backing once
    W-->>S: Complete batch using the native lifecycle
    S->>P: Release references and evict all owned aliases
Loading

Validation

Tested revision: 086118324935d40842679d0c3ed0bb744b2143f0, rebased onto main at 0ed269fe9057508699ceb95fab071f20ab3e26e5. Rixin's three code changes are preserved; the CPU-manifest conflict retains both the main-branch tests and Mamba-binding tests. All results below were rerun on this exact candidate.

  • Full pre-commit and GPU-free C++ checks passed. The CPU suite passed 579 tests on each of Python 3.9-3.13, and all five MyPy checks passed. Regression coverage includes native 4-token hashing with 16-token allocation blocks, absent/invalid metadata, legacy compatibility, propagation before initialization, and the merged FTensor capacity guard.
  • On a Tesla T4 with PyTorch 2.13.0+cu129 and vLLM 0.28.0, 276 focused runtime tests and 25 independent compatibility-contract tests passed. The native extension was reused only after verifying identical native sources and build configuration.
  • Real VMM copy tests covered NHD/HND, separate/interleaved layers and logical/kernel block ratios 1/2: 8/8 passed. Each mapped 600 blocks, copied shared attention/Mamba backing exactly once, forced a real physical-capacity rejection, released the pressure, recovered 600 blocks and verified their contents again.

The initial pressure probe assumed immediate recovery after freeing external Torch memory. Main now caches physical capacity for 100 ms, so that assertion failed. The retained failure record is not counted as a pass. The corrected probe checks fresh native capacity and bounded retry recovery; all eight cases recovered after about 100 ms. No production cache policy was changed for this test.

The serving probe used a small random-weight Qwen3.5 hybrid attention/GDN model, FP16, TP=1, V1 explicitly selected, eager execution and Triton attention. Multimodal inputs were disabled. The allocation block was 784 tokens; prompts of 48, 64, 783, 784, 785 and 1,087 tokens covered repeated-prefix hits and allocation boundaries. Hash-size 16 was a functional test case, not a recommended production setting.

Case Complete requests Output tokens Match to native token IDs Injected allocation failures
Native control, hash size 16, synchronous 24/24 768/768 Reference None
Native control, hash size 16, asynchronous 24/24 768/768 Reference None
KVCached, hash size 16, synchronous 24/24 768/768 Exact None
KVCached, hash size 16, asynchronous 24/24 768/768 Exact None
Partial-prefix CoW admission failure, synchronous 24/24 768/768 Exact 2
Partial-prefix CoW admission failure, asynchronous 24/24 768/768 Exact 2

Both fault cases armed the injection only after a real partial-prefix hit, rejected two consecutive allocator attempts, exercised the scheduler miss/retry path, then completed subsequent worker copies and requests (16 synchronous and 17 asynchronous copy submissions). Total, including native controls: 144/144 requests and 4,608/4,608 output tokens, with token-by-token equality to the corresponding native run. There was no EngineCore failure or CUDA execution error during generation, and no test-container restart.

Limitations: tokenizer-free token-ID tests produce a native chat-template warmup warning. Earlier probes recorded FTensor cleanup errors after completed generation; this PR does not claim to fix teardown. The 0.26/0.27 boundaries have source and unit-test coverage, not full serving tests on those runtimes. The real-device results here are for 0.28 only, not model-quality measurements, full-size FP8 acceptance, TP/MPS validation, or a throughput comparison.

Follow-up

Our goal is for this to be the last manually prepared vLLM 0.28 compatibility PR. Once the release-triggered workflow in #423, #495 and #496 is deployed, subsequent compatibility work should go through that automation, retaining GPU validation and human review. This is the intended handoff, not a claim that the workflow is already deployed or every model/backend is covered.

Copilot AI lite review requested due to automatic review settings September 17, 2026 02:44

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@RixinLiu

Copy link
Copy Markdown
Collaborator

LGTM, i pushed three commits below, they are only about applying the fix before 0.28, after you took a look at them, and finished the rebase, i will merge it:

  • 3a8cab0: Enable raw Mamba page binding from 0.27, when vLLM started expecting raw pages instead of state tensor lists.
  • 100ee15: Enable logical-block copying from 0.26, and attach copy metadata to its legacy state tensor views so copies use the actual block boundaries.
  • 684818a: Enable native prefix-cache metadata and copy-lifetime handling from 0.26, covering partial-prefix entries, hash migration, and cleanup on eviction.

Validation: I only ran CPU tests, no GPU tests

@shipiyouniao
shipiyouniao force-pushed the feature/vllm-028-v1-hybrid branch from 684818a to 0861183 Compare September 25, 2026 11:41
@shipiyouniao
shipiyouniao marked this pull request as ready for review September 25, 2026 11:48
@shipiyouniao

Copy link
Copy Markdown
Collaborator Author

Rebased onto current main, with your three version-boundary fixes preserved. The PR now contains just the four hybrid/prefix commits; #492/#493 are no longer carried in its diff. CMQY's authorship is unchanged.

On the rebased head, all five Python CPU/type-check jobs and pre-commit passed privately, and all 11 upstream checks are green. The T4 rerun also passed: eight real VMM copy/pressure-recovery cases, plus synchronous/asynchronous hybrid serving and partial-prefix admission-failure recovery, totaling 144 requests and 4,608 output tokens with exact native token parity.

The description now reflects the 0.26/0.27 boundaries and current results. Those older boundaries have source/unit coverage; the fresh serving run is specifically 0.28, not TP/MPS qualification. Marked ready for your final look; I haven't merged it.

@RixinLiu
RixinLiu merged commit 9564931 into ovg-project:main Sep 25, 2026
11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants