[milesial]

Why and what changed

Scalar bytes per token cannot describe a shared Mamba cache pool. A request consumes attention pages, working states, retained checkpoints, and speculative scratch states.

This adds an observed mamba_align FPM layout and connects its physical geometry to the native vLLM replay allocator. It reserves the null page and full-state scratch copies. It preserves state ownership through restore, turnover, preemption, and release.

Fine-grained speculative lookup uses the matching attention tail and the earlier saved state. Retention now keeps both full-page boundaries for repeated and extended prompts. The pinned vLLM logits cap still applies before the speculative hash-unit drop, so some exactly aligned prompts still miss the tail.

The profile requires state_storage: full_state. ReplaySSM and RecoverSSM use different storage rules. Their profiles can use the new unmodeled layout for timing queries, while capacity sizing and replay reject that layout.

Review map

• Risk level: medium.
• Start with crates/core/src/perfmodel/fpm/resources.rs, then engine/kv_manager/vllm_backend.rs and engine/scheduler/vllm/core.rs.
• Public or serialized contract changed: additive mamba_align and unmodeled layouts, mamba_cache, and optional state_cache.num_speculative_states.
• Python exports, prediction validation, capacity materialization, and replay lowering carry the same native geometry.
• Supported scope: observed aggregated vLLM pure-TP profiles with prefix caching. Offload, disaggregated FPM replay, internal prefill checkpoints, and recommendation search remain unsupported.
• Existing profiles retain their serialized shape. Removing the new fields restores the previous contract.

Evidence

• cargo test -p aisimulate-core --lib engine::: 862 passed.
• Focused state-cache and FPM resource suites: 61 and 16 passed.
• Five Python suites covering state configuration, FPM workflow, grouped workflow, profile schema, and public imports: 562 passed.
• Rust formatting, Ruff on changed Python files, diff whitespace, and packaged legal-file comparison passed.
• Strict Clippy does not pass: 131 diagnostics remain, compared with 132 on unmodified main. The comparison found no new diagnostics by message and file.
• Negative cases cover unsupported storage policies, missing policy, fractional pages, mismatched speculation depth, manual geometry conflicts, and timing-only replay rejection.
• Expected-value derivation: boundary fixtures follow the pinned vLLM lookup and retention methods. The scheduler test covers 18 prompt/granularity combinations at each of three draft depths.
• Before/after: a 24-token prompt with 8-token pages and a 2-token hash unit reused zero tokens in the initial local implementation. It now reuses 8, matching the pinned source.
• Fast CI: passed at aa5457ce7ec3d1b106edaab512eae9273fab24f1. Full CI and performance advisory runs are pending.
• CodeRabbit reviewed commit: pending.
• Codex reviewed commit: aa5457ce7ec3d1b106edaab512eae9273fab24f1, self-review only. Independent review remains pending.

Modeling or data provenance

Allocation and retention rules derive from vLLM revision db9527a46873454610df6dbedf79a36d6bf1a7f6. The third-party notices record the source paths and modifications. Tests use synthetic geometry and timings.

At vLLM #55688, revision d9cd1fd1fcfc5a08d14ff29ba7439f76ad42d628, ReplaySSM sets separate speculative blocks to zero. Its FlashInfer replay ring still grows with draft depth. #52962 changes checkpoint precision and sharing. #62625 moves GDN recovery records outside the cache page. These policies require separate allocation models and fresh memory/timing observations.

AgentX throughput accuracy, shared-prefix junction retention, asynchronous copy lifetimes, and GPU ReplaySSM behavior remain unqualified.

Tracking

• Related AIS changes: #601 for measured MTP identity and #561 for native autoregressive CUDA timings. Their changes are not duplicated here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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