Skip to content

feat(qwen35): add joint prefix cache - #836

Open
Ke-Wng wants to merge 11 commits into
pegainfer-project:mainfrom
Ke-Wng:feat/qwen35-prefix-cache
Open

Ke-Wng wants to merge 11 commits into
pegainfer-project:mainfrom
Ke-Wng:feat/qwen35-prefix-cache

Conversation

@Ke-Wng

@Ke-Wng Ke-Wng commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Description

This PR is a follow-up of #257 and #423. It adds opt-in prefix caching for Qwen3.5 on both single-GPU and tensor-parallel serving.

Qwen3.5 is a hybrid model, so full-attention KV alone is not a valid reusable prefix. A cache hit is reported only when full-attention KV and the complete recurrent/conv state are available at the same token boundary.

Prefix caching remains disabled by default and can be enabled with:

--qwen35-prefix-cache-mib <MiB>

Type of Change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Documentation update

What changed

  • Added opt-in Qwen3.5 prefix caching for both single-GPU and tensor-parallel serving.
  • A hit requires content-matched full-attention KV and a complete recurrent/conv snapshot at the same 256-token boundary.
  • Warm requests restore that joint boundary and prefill only the remaining suffix.
  • Added a bounded GPU snapshot cache with pinning, LRU replacement, and cache statistics.
  • Snapshot residency is independent of KV residency: released KV blocks return to kvbm's reusable inactive pool and remain reclaimable for cold admission.
  • Added TokenEvent::Scheduled.cached_tokens reporting for accepted prefix hits.
  • Corrected the shared decode KV lifecycle so only KV completed by the current forward is registered; the newly sampled token remains unregistered until its next forward.

Design

Qwen35PrefixCache coordinates two independently managed resources:

  • RequestKv and KvCacheManager manage content-addressed full-attention KV pages.
  • SnapshotCache and RecurrentStateStore manage recurrent/conv snapshots.

A snapshot entry identifies a token lineage, boundary, and recurrent-state slot; it does not retain KV guards. During lookup, the cache selects the longest boundary for which both matching KV and a snapshot are currently available. A KV-only or snapshot-only candidate is treated as a miss.

This separation preserves joint-hit correctness without allowing cached snapshots to pin the KV pool and block new admissions.

Request flow

  1. begin_request() creates request-local KV state and probes the longest valid joint boundary.
  2. On a hit, the executor restores the referenced recurrent-state slot and verifies that KV and recurrent positions match before reporting cached_tokens.
  3. Prefill and decode reserve KV pages transactionally; failed forwards revert their reservations.
  4. At aligned prefill boundaries, the executor reserves a snapshot slot, copies recurrent state, and publishes the snapshot only after all required copies succeed.
  5. Request completion, rejection, and disconnect release request-owned KV and recurrent state.

Non-aligned prompt tails are processed normally but do not create snapshots.

Tensor-parallel behavior

  • The controller owns logical KV state, snapshot metadata, request admission, and the shared page-ID namespace.
  • Each rank owns its physical KV buffers and recurrent-state storage.
  • The controller sends identical logical page tables and snapshot slot IDs to every rank.
  • A restore or publication succeeds only when every rank confirms the same boundary.
  • The usable cache capacity is limited by the smallest rank-local capacity.
  • Public prefill plans are fully validated before any request is admitted. If a failure occurs after distributed mutation has begun, the executor is poisoned rather than exposing partially committed controller/worker state.

Consistency guarantees

  • A reported hit always contains matching full-attention KV and recurrent/conv state at the same token boundary and lineage.
  • Snapshot publication occurs only after the physical recurrent-state copy succeeds on every required rank.
  • KV-only and snapshot-only candidates are never reported as hits.
  • Reclaiming inactive KV does not require evicting the corresponding snapshot; if its KV is no longer available, that snapshot simply cannot produce a hit.
  • Failed forwards revert scheduled KV pages, and failed snapshot copies abort their reservations.
  • Request lifecycle exits release all request-owned state.

KvState migration defense

This PR replaces KvState with the RequestKv/KvView lifecycle coordinated by Qwen35PrefixCache; the table maps each protection in the old path to its successor in the new design.

Previous defense Failure mode covered Successor
KvState owned an OwnedPagePermit, releasing request ownership on drop. Pages remaining pinned after request completion, cancellation, or failure. Replaced: RequestKv owns kvbm block assignments. Single-GPU scheduler cleanup drops the owning backend state and relies on RequestKv RAII; the direct executor and TP controller call Qwen35PrefixCache::release_request explicitly. Unapplied reservations are returned by revert_schedule.
KvPool reserved a dedicated padding_page_id() and excluded it from request capacity. CUDA Graph padding rows overwriting a live request page. Replaced: BlockPool reserves padding_block_id(), max_request_blocks() excludes it from admission, and the controller passes the same ID to single-GPU and TP decode buffers.
Prefill and decode called KvState::ensure_capacity before advancing state. Kernels running without enough KV pages or allocation failure leaving partially advanced logical state. Replaced and strengthened: RequestKv::schedule_prefill and schedule_decode reserve pages before forward and produce immutable KvViews. apply_* commits only after success; failures call revert_schedule.
KvState::seq_len, last_page_len, and advance defined KV positions and page extents. Attention writing to the wrong position or KV and recurrent state progressing independently. Replaced: RequestKv::kv_position owns committed progress and KvView carries the scheduled extent. Prefill, decode, restore, and snapshot publication validate recurrent positions against them.
checked_prefill_end_pos and RoPE coverage were validated before entering the prefill chunk loop. Overflow or context-limit failure after partially mutating a request. Inherited: full-range validation remains in prefill.rs; public single-GPU and TP entry points also reject prompts that leave no output position before creating request state.
Decode validated batch shape and used bounded page metadata with a dedicated padding page. Misaligned request/KV/recurrent rows, metadata buffer overflow, or invalid padding access. Inherited and strengthened: the checks now validate KvView and recurrent rows; metadata capacity is sized for per-row references, checked before upload, and padded rows use the reserved block.
Scheduler admission excluded padding capacity and reserved each request’s remaining KV lifetime. Admitting a request that later fails when decode needs another page. Replaced: admission uses max_request_blocks() and request_lifetime_pages(), including kvbm’s transient dangling decode page. Jointly reusable prefix pages are credited without double-charging shared active blocks.
TP validated equal page geometry, limited capacity to the smallest rank, and guarded worker request identity and phase. Page IDs exceeding a rank-local buffer or controller/worker state diverging after partial mutation. Replaced and strengthened: the controller owns one logical KvCacheManager sized to the smallest rank and sends identical page IDs to every worker. Plans are validated before admission; rank responses and restore positions are checked, and post-mutation failures take the fail-closed path.

Validation

  • Release Qwen3.5 library tests passed.
  • Release Clippy passed with -D warnings.
  • Prefix-cache coverage includes cold and warm requests, longest-boundary restore, suffix-only prefill, extension, logprob parity, LRU behavior, KV-only fallback, and reclaiming cached KV under full-pool pressure.
  • Single-GPU and TP2 restore paths passed scheduler, chunked-prefill, HF golden, and serving coverage.
  • A real TP2 regression verifies that an invalid multi-request prefill plan creates no controller or worker state and that the rejected request ID can be retried successfully.
  • Audited the normal decode callers in Qwen3, Qwen3.5, Kimi-K2, and GLM5.2/KV-store. kvbm-logical, Qwen3, KV-cache, and KV-store tests passed, including the pad-and-seal roundtrip.
  • apply_speculative remains unchanged because Qwen3.5 does not use it and its terminal/truncated contracts require separate cross-model validation.

Performance

Measured on 2026-08-06, using local Qwen3.5-4B weights and idle RTX 4090s. TP1 used GPU 1 with CUDA Graphs enabled; TP2 used GPUs 1/2 with --cuda-graph=false. The prefix A/B matrix used a 1,024 MiB cache, 2 warmups, and 10 measured iterations.

For single-token generation, the table shows TTFT p50 in milliseconds (cache off -> cache on):

Prompt tokens Cached tokens TP1 TTFT p50 TP1 reduction TP2 TTFT p50 TP2 reduction
320 256 29.22 -> 15.56 46.7% 46.18 -> 18.75 59.4%
576 512 46.85 -> 16.92 63.9% 72.97 -> 18.82 74.2%
1,088 1,024 88.75 -> 15.93 82.1% 126.14 -> 19.36 84.7%
2,112 2,048 161.00 -> 16.40 89.8% 230.07 -> 19.82 91.4%
4,160 4,096 308.64 -> 17.71 94.3% 439.37 -> 21.68 95.1%

For 128 generated tokens, cache reuse primarily reduces prefill/TTFT; decode remains the dominant part of E2E latency:

Prompt tokens TP1 TTFT p50 off -> on TP1 E2E p50 off -> on TP2 TTFT p50 off -> on TP2 E2E p50 off -> on
1,088 94.58 -> 24.13 ms 1,542.02 -> 1,452.78 ms 127.05 -> 20.38 ms 1,445.78 -> 1,324.97 ms
2,112 166.25 -> 24.23 ms 1,700.09 -> 1,535.84 ms 231.56 -> 21.89 ms 1,640.17 -> 1,421.95 ms
4,160 313.51 -> 24.30 ms 2,006.08 -> 1,696.04 ms 442.01 -> 24.88 ms 2,036.57 -> 1,606.08 ms

TP1 serving-load checks also completed:

  • At a 2,112-token prompt and 128 output tokens, concurrency 1/4/8 produced TTFT p50 of 24.67/94.41/152.25 ms, request throughput of 83.11/75.87/69.04 tok/s, and a 2,048-token cache hit for every request.
  • Under mixed load, baseline ITL p50/p99 was 12.03/12.25 ms; mixed-load ITL was 12.03/36.12 ms. After initial insertion, 4,096-token warm-prefix injections restored 3,840 tokens with 41.42–51.00 ms prefill and no warnings.
  • TP1 HTTP warm TTFT p50 for prompt lengths 320/576/1088/2112/4160 improved from 31.84/48.71/93.34/146.05/282.10 ms to 18.51/21.42/21.34/29.16/25.51 ms.

@Ke-Wng
Ke-Wng marked this pull request as draft August 3, 2026 09:20
@Ke-Wng
Ke-Wng force-pushed the feat/qwen35-prefix-cache branch 2 times, most recently from 7d46f71 to 40d56ff Compare August 6, 2026 11:56
@Ke-Wng
Ke-Wng marked this pull request as ready for review August 6, 2026 11:58

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 40d56ffdff

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread pegainfer-qwen35/src/prefix_cache.rs Outdated
if reservation.was_replacement() {
self.stats.evictions += 1;
}
self.state.publish(reservation, kv_guards);

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Reclaim cached KV before admission stalls

When prefix caching is enabled and a near-capacity prompt publishes a snapshot, these guards can retain most of the KV pool after the request finishes. scheduler_loop checks available_blocks() before calling begin_request, so a repeated request that could reuse those blocks is deferred as though it needed a cold allocation; because the only eviction path runs from reserve_prefix after a request has already been admitted and prefilled, an idle engine can defer that request forever. Admission must account for reusable blocks or evict unpinned prefix entries under KV pressure.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed in 0e65237 by decoupling snapshot residency from KV residency. Snapshot entries no longer retain KV guards; released KV blocks return to kvbm’s inactive pool, remain reusable for prefix matches, and count toward available_blocks(). Cold admission can therefore reclaim them without first evicting the snapshot. A regression test covers both warm reuse and full-pool cold reclamation.

.filter(|&v| v > 0)
.unwrap_or(pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS);
let handle = pegainfer_qwen35::start_engine_with_capacity_and_policy(
let handle = pegainfer_qwen35::launch_with_options_and_policy(

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Preserve the requested benchmark seed

When a Qwen3.5 benchmark uses non-greedy sampling with --seed other than 42, this new call routes through launch_with_options_and_policy, whose EngineLoadOptions still hardcodes seed: 42; the previous path passed command_seed(&cli). Consequently the advertised seed no longer controls engine sampling for Qwen3.5, breaking reproducibility and making different seed runs use the same sampling stream. Pass the command seed through the launch options or retain the explicit EngineLoadOptions path.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

This is no longer applicable after the frontend consolidation: pegainfer-server/src/bin/bench_serving/main.rs
and its per-command benchmark seed path no longer exist. Qwen3.5 now launches through ModelLine with an
intentionally fixed engine seed of 42, so no user-provided benchmark seed is silently discarded.

@xiaguan xiaguan self-assigned this Aug 16, 2026
@github-actions

Copy link
Copy Markdown

This pull request has been inactive for 14 days. It will be closed after another 30 days unless there is new activity.

@github-actions github-actions Bot added the stale Automatically marked after inactivity label Aug 30, 2026
@Ke-Wng
Ke-Wng force-pushed the feat/qwen35-prefix-cache branch from 40d56ff to ad9d43e Compare September 6, 2026 06:47
@github-actions github-actions Bot removed the stale Automatically marked after inactivity label Sep 6, 2026
@Ke-Wng
Ke-Wng force-pushed the feat/qwen35-prefix-cache branch from ad9d43e to f0a1936 Compare September 8, 2026 08:38
@FeathBow

FeathBow commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator

@codex review please

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f0a19368ef

ℹ️ 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".

Comment thread pegainfer-qwen35/src/tp_executor.rs Outdated
Comment on lines +990 to +999
for request in plan.requests {
self.begin_request(
request.request_id,
&request.prompt_tokens,
self.max_position_embeddings
.saturating_sub(request.prompt_tokens.len())
.max(1),
None,
false,
)?;

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Validate all prefill items before beginning requests

When the public TP executor receives an invalid plan such as an empty prompt or duplicate request IDs, this loop can successfully broadcast RestoreRequest and insert state for earlier items before a later iteration fails. Because validate_prefill_chunks runs only afterward and this error path does not drop already-begun requests, the failed call leaves controller and worker state consuming capacity, and retrying the same ID then fails as a duplicate. Validate the complete plan before this loop or roll back every request begun by the call.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed in 0e65237 by validating the complete prefill layout—including empty prompts, duplicate or existing request IDs, capacity, and context limits—before broadcasting any RestoreRequest. If a failure occurs after an earlier request has already committed, the executor is poisoned rather than exposing partially admitted state. CPU validation and a real TP2 retry regression cover this path.

@Ke-Wng
Ke-Wng force-pushed the feat/qwen35-prefix-cache branch from d3f79b9 to 0e65237 Compare September 9, 2026 06:29
@Ke-Wng

Ke-Wng commented Sep 9, 2026

Copy link
Copy Markdown
Contributor Author

@codex review please

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0e65237eec

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread pegainfer-qwen35/src/tp_executor.rs Outdated
let graph_state = model.create_batch_decode_graph_state_with_capacity(
slots,
model.kv_buffer().num_blocks(),
(model.kv_buffer().num_blocks() - 1) as i32,

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Use the controller-reserved TP padding page

When TP CUDA Graph decode pads a non-bucket batch after the logical pool has allocated its highest page ID, this selects num_blocks() - 1 as the padding page even though BlockPool::build reserves page 0 and leaves the highest page allocatable. On at least the minimum-capacity rank, padding rows then write K/V into a page owned by a real or cached request, silently corrupting its attention state. Pass the common logical pool's actual padding_block_id() to every worker instead of deriving the physical last page.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed in bea4e1e. The controller now obtains the reserved padding block ID from the common logical KV pool and passes it to every TP worker for CUDA Graph decode.

/// Reserve KV capacity for the next one-token decode forward.
pub(crate) fn schedule_decode(&self, request: &mut RequestKv) -> Result<()> {
request
.schedule_decode(self.kv.pool())

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Reserve the extra dangling decode page at admission

When a multi-token request's prompt_len + max_tokens - 1 lands exactly on a page boundary under a full KV budget, the unchanged admission calculation reserves only that many pages, but the newly adopted RequestKv::schedule_decode can transiently draw one additional generation page. Such a request is admitted as servable and later fails allocation; in production decode_step responds by terminating the entire active batch. Size lifetime reservations using prompt_len + max_tokens for multi-token requests, as required by the RequestKv peak draw.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed in bea4e1e. Admission now reserves prompt_len + max_tokens for multi-token requests, covering RequestKv’s transient peak allocation. One-token requests remain prompt-only because they finish directly after prefill.

@Ke-Wng
Ke-Wng force-pushed the feat/qwen35-prefix-cache branch from 0e65237 to bea4e1e Compare September 10, 2026 13:24
@Ke-Wng

Ke-Wng commented Sep 10, 2026

Copy link
Copy Markdown
Contributor Author

@codex review please

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bea4e1e2dc

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment on lines +170 to 172
for view in views {
all_page_indices.extend_from_slice(view.page_indices());
indptr.push(all_page_indices.len() as i32);

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Size decode metadata for shared prefix references

When several concurrent requests restore the same cached prefix, concatenating every KvView can produce more page-index entries than the fixed page_indices_d capacity of max_total_pages + max_batch_size: shared physical pages count once toward the pool capacity but once per request in this vector. The scheduler can legitimately admit this state because subsequent requests reuse the already-active prefix and need only distinct suffix pages; for example, three requests can each reference a 200-page sequence while using only 400 unique pages. Once the combined views exceed the buffer, memcpy_htod fails and the scheduler terminates the decode batch, so this metadata buffer must be sized for per-row references rather than unique physical pages.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed in 309eca7. Decode page-index metadata is now sized by per-request references (max_batch × max pages per context) rather than unique physical KV pages, covering shared-prefix reuse across concurrent requests.

@Ke-Wng

Ke-Wng commented Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

@codex review please

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 309eca7b70

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

}

let request_pages = pages_needed(max_kv_tokens(prompt_len, max_tokens), page_size);
let request_pages = request_lifetime_pages(prompt_len, max_tokens, page_size);

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Credit shared prefix pages during admission

When concurrent requests share a long cached prefix and the pool cannot cover each request's full logical lifetime independently, this debits request_pages for every request before begin_request attaches the reusable prefix. Unlike the resolved idle-cache case, available_blocks() counts only reset and inactive pages, so once the first warm request promotes the shared blocks to active, later requests are deferred even though attaching those same active blocks would require only suffix/output pages. This unnecessarily serializes shared-prefix traffic under KV pressure; admission should probe and credit the jointly reusable pages.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed in 17b38c8. Admission now subtracts reusable active prefix pages from each request’s KV requirement. After the first request activates a cached prefix, waiting requests sharing it are reconsidered immediately instead of being serialized.

@Ke-Wng

Ke-Wng commented Sep 12, 2026

Copy link
Copy Markdown
Contributor Author

@codex review please

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 17b38c84f9

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread pegainfer-qwen35/src/prefix_cache.rs Outdated
Comment on lines +454 to +458
// Retain registered KV only when prefix caching is enabled.
if !self.enabled() {
request.mark_blocks_reset_on_release();
}
applied?;

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Mark newly registered blocks after applying the step

When prefix caching is disabled (the default), this marks only blocks that were assigned before apply_prefill; the apply then registers the newly completed blocks, leaving them unmarked. The single-GPU scheduler retires requests by dropping their state because drop_prefill_state and drop_active_state are no-ops, so the final block registered by a request enters the inactive prefix cache rather than the reset pool. With distinct traffic these useless entries eventually fill the pool and force LRU/registry eviction work despite caching being disabled; mark blocks after a successful apply or explicitly release single-GPU requests through release_request (the same ordering affects apply_decode).

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed in d4e2e1b. Newly registered KV blocks are now marked after successful prefill and decode applies, so disabling prefix caching returns them to the reset pool on release.

@Ke-Wng

Ke-Wng commented Sep 14, 2026

Copy link
Copy Markdown
Contributor Author

@codex review please

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. More of your lovely PRs please.

Reviewed commit: d4e2e1bf22

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

@Ke-Wng

Ke-Wng commented Sep 14, 2026

Copy link
Copy Markdown
Contributor Author

Hi @FeathBow, this PR is ready for review. Could you please take a look when you have time? Thanks!

@FeathBow FeathBow left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Btw for some tests, with one slot and no subsequent eviction, removing the duplicate-path timestamp update would leave every assertion passing, so the claimed LRU-refresh behavior is currently untested. Besides, Please remove the leftover new_requests capacity calculation now that prefill no longer creates worker requests? It's a redundant scan. And please remove the unused things, e.g. active/prefilling arguments from both available_pages methods and their callers?

Some tests we could fold this into the existing CLI validation cases. This only parses arguments and validates launch options; it does not exercise prefix caching on either TP configuration.

pub(crate) fn apply_decode(&self, request: &mut RequestKv, token: u32) -> Result<()> {
request.apply_decode(token, self.kv.pool())?;
if !self.enabled() {
request.mark_blocks_reset_on_release();

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Could we ensure that decode registers only blocks whose KV has actually been computed? The delegated SchedulableSequence::apply_decode appends the newly sampled token before registering completed token blocks. With a 254-token prompt and two output tokens, the request ends at KV position 255 but registers the 256-token prefix. If a snapshot for that prefix survives eviction of its original KV, a later request can pair the retained snapshot with this newly registered, incomplete page and restore to 256, reading an unwritten final KV entry. Both single-GPU and TP serving use this path. Please defer registration of the dangling-token block until its KV is complete.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed by reversing the operations in SchedulableSequence::apply_decode: it now registers KV completed by the current forward first, then appends the newly sampled token. Therefore, a newly crossed block remains unregistered until the next forward computes that token’s KV. Tests were updated to cover this deferred registration.

Comment thread pegainfer-qwen35/src/prefix_cache.rs Outdated
}

#[test]
fn released_prefix_kv_is_reclaimable_without_snapshot_eviction() {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Could this regression exercise Qwen35PrefixCache itself? The test currently creates a BlockPool and a separate SnapshotCache, so it still passes if the joint coordinator starts retaining KV guards again. The allocation assertions mainly verify the dependency's reclaim behavior, while the final lookup checks an unrelated map. Please cover reclamation through the production coordinator, or remove this test's claim to cover the joint-cache regression.

@Ke-Wng Ke-Wng Sep 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I removed the misleading reclamation test, added real LRU eviction coverage, removed redundant capacity scans and arguments, and folded the CLI case into existing validation.

Signed-off-by: wangke <364517893@qq.com>
Signed-off-by: wangke <364517893@qq.com>
Signed-off-by: wangke <364517893@qq.com>
Signed-off-by: wangke <364517893@qq.com>
Signed-off-by: wangke <364517893@qq.com>
Signed-off-by: wangke <364517893@qq.com>
Signed-off-by: wangke <364517893@qq.com>
@Ke-Wng
Ke-Wng requested a review from FeathBow September 15, 2026 06:59
@FeathBow

Copy link
Copy Markdown
Collaborator

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9119897069

ℹ️ 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".

Comment thread pegainfer-qwen35/src/executor.rs Outdated
Comment on lines +176 to +179
self.model
.config()
.max_position_embeddings
.saturating_sub(req.prompt_tokens.len()),

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Reject prompts that fill the entire context

When the direct single-GPU Qwen35Executor receives a prompt whose length equals max_position_embeddings, this computes max_output_tokens as zero. The prefill forward itself accepts that end position, but apply_prefill(Some(first_token)) subsequently calls RequestSequence::append_token, which asserts that an output token remains and panics. The TP executor already rejects this exact case; add the equivalent validation here so this public API returns an error instead of aborting.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed. The single-GPU executor now rejects prompts that leave no output-token capacity before mutating state, matching the TP validation path.

Signed-off-by: wangke <364517893@qq.com>
@Ke-Wng
Ke-Wng force-pushed the feat/qwen35-prefix-cache branch from 9119897 to bd17753 Compare September 15, 2026 09:17
@Ke-Wng

Ke-Wng commented Sep 15, 2026

Copy link
Copy Markdown
Contributor Author

@codex review, thanks

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bd17753d90

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

}

let request_pages = pages_needed(max_kv_tokens(prompt_len, max_tokens), page_size);
let request_pages = request_lifetime_pages(prompt_len, max_tokens, page_size);

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Reject zero-token requests before prefill

When a GenerateRequest has max_tokens == 0, this admission path treats it like a one-token request and admits it. Both single-GPU and TP final-prefill paths then pass a sampled token to RequestKv::apply_prefill, whose underlying RequestSequence::append_token asserts because its output-token limit is already zero, panicking the Qwen3.5 scheduler thread. Handle zero-token requests as an immediate completion or reject them before materializing KV state.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Fixed. Qwen3.5 admission now rejects max_tokens == 0 before materializing KV state.

Signed-off-by: wangke <364517893@qq.com>
@Ke-Wng

Ke-Wng commented Sep 15, 2026

Copy link
Copy Markdown
Contributor Author

@codex review please

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Delightful!

Reviewed commit: fee3726e17

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

}
};

// The forward computed KV for the existing dangling token. Register

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

SchedulableSequence::apply_decode is reached through RequestKv::apply_decode, whose production callers include qwen3 , glm52, kimi-k2, and pegainfer-kv-store. Deferring registration by one step is a semantic change for all of them, and the only assertions updated are the three inside kvbm.

Two things I'd like recorded before this merges.

First, which consumers did you check? qwen3 has its own prefix cache and most likely benefits from the same fix, but that is a silent behavior change and the description doesn't mention it. The kv-store seal path is the one I'd look at specifically: pad_tail_block is a no-op when kv_position sits on a block boundary, which is now exactly the state where the block that just crossed is still unregistered.

Second, is apply_speculative deliberately left on the old ordering? At scheduled.rs it still appends every accepted token, then registers, then advances kv_position. Its callers are qwen3's dspark path and glm52's MTP path. If the last accepted token's KV is also uncomputed in that step, it has the same asymmetry this commit just fixed for decode; if the accepted contract guarantees computed KV for every token, then it is correct as is. Either answer works, but the two paths in one file now disagree about what "registered" means, so I recommend saying which in a comment.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Thanks for highlighting the shared impact. I traced the normal decode paths in Qwen3, Qwen3.5, Kimi-K2, and GLM5.2/KV-store. They use the same contract: the forward computes KV for the existing dangling token and returns a newly sampled token whose KV is not yet computed. Registering before appending therefore preserves the intended invariant.

For KV-store sealing, if decode advances kv_position to a page boundary, that page is registered after its final KV row is computed and before the new dangling token is appended to the next page. Therefore, the boundary no-op in pad_tail_block does not skip a completed page.

I also audited apply_speculative. Its result contains computed accepted drafts followed by an uncomputed correction/bonus token, so it has the same underlying issue. Qwen3.5 does not use that path, and changing it would require separate validation speculative decoding, so I am leaving it unchanged in this PR.

I’ve updated the PR description to record the shared apply_decode behavior change, the consumers reviewed, the tests actually run, and why apply_speculative remains out of scope for this PR.


pub(crate) fn log_prefix_cache_stats(&self) {
let stats = self.kv_cache.stats();
log::info!(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

log records should to drop the module or model prefix, naming Qwen3.5 specifically, because the text layout already prints each record's target. Only anyhow! and bail! messages keep theirs, since those surface to callers without a target.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Done.

Signed-off-by: wangke <364517893@qq.com>

@FeathBow FeathBow left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks!

The migration defense table is the bigger one. KvState appears in nine qwen35 files on the base commit and in none here, so docs/conventions/migration-defense.md applies, and what it asks for is an audit rather than a paragraph: each old guard, the failure mode it covered, and its successor. Eight fix-up commits on this branch put guards back one at a time; the table's value now is the ones nobody has hit yet.

For the TP pass, half-fixing it also made the two summaries disagree — the TP path now logs prefix cache summary: ranks=… while the single-GPU path logs Qwen3.5 prefix cache summary: …, so a log rule written against one matches only one of the two paths.

To be clear on scope: the crate has plenty of pre-existing prefixed lines (tp_executor.rs:, scheduler/mod.rs, weights.rs), and I am not suggesting this PR touch those. Just the three it adds.

Signed-off-by: wangke <364517893@qq.com>
@Ke-Wng
Ke-Wng requested a review from FeathBow September 16, 2026 01:24
@FeathBow

Copy link
Copy Markdown
Collaborator

@codex review, temp ignore conflicts

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. 👍

Reviewed commit: abc071ab69

ℹ️ 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".

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.

3 participants