Skip to content

fix(serving): prefix-cache binding, hit stats, dtype (needs #197) - #190

Merged
drunkcoding merged 3 commits into
plan/prefix-kv-reusefrom
fix/prefix-reuse-binding-and-stats
Sep 4, 2026
Merged

fix(serving): prefix-cache binding, hit stats, dtype (needs #197)#190
drunkcoding merged 3 commits into
plan/prefix-kv-reusefrom
fix/prefix-reuse-binding-and-stats

Conversation

@drunkcoding

@drunkcoding drunkcoding commented Aug 31, 2026

Copy link
Copy Markdown
Contributor

Productizing prefix-kv-reuse. Depends on #197 (fix/expert-dispatch-stream-orderingdev): without it, prefix-enabled runs are corrupted by an expert-dispatch stream race.

Fixes in this PR

  • kv_cache.set_block_store: normalize the cuda device index so cuda vs cuda:0 geometry no longer rejects the store — the prefix cache now actually binds/activates.
  • prefix_cache + engine: surface prefix_cache_hits_total / prefix_cache_matched_tokens_total.
  • attention_backend: pass q_data_type/kv_data_type to the FlashInfer plans.
  • big_modeling: eager attention for qwen3 when flash-attn is unusable (sm120).
  • prefix_cache_benchmark: subprocess-per-mode; --parity-report (tensor dumps + numerical analysis), --device-memory-ratio/--kv-cache-ratio for shared GPUs.
  • tests/integration/test_flashinfer_kernel_parity.py: full-prefill and append kernel schedules are each deterministic, within ULP-scale error of an fp32 reference, and mask partial-page tails (3/3 pass on sm120).
  • tests/serving/test_prefix_cache_cuda.py: parity test now runs one-engine-per-subprocess (kTopologyHandle-safe) with sound invariants.

HOLD resolved — what "warm not bit-exact on sm120" actually was

  1. FlashInfer kernels: exonerated. Cross-kernel bitwise equality was never an invariant (per-op delta ≈ 1 bf16 ULP; both paths equidistant from fp32 truth; partial-page masking correct). See the committed kernel-parity test.
  2. Real bug: expert-dispatch stream race (pre-existing on dev, exposed by prefix-mode timing) — found via NaN-poisoning: enabled_cold produced 100%-NaN logits from layers[0].mlp with finite inputs. Fixed in fix(dispatcher): stream-order expert exec against input producer and consumers #197.
  3. Post-fix validation (sm120, Qwen3-30B-A3B, 64-token shared prefix + 8-token suffix):
    • disabled == enabled_cold bitwise (token digest e37a9d14, logits Δ = 0.0).
    • Warm reuse active (geometry q=8, ctx=64, kv=72, hit counted). Warm step-0 argmax swaps within the shared top-2 pair (cold: 2157@22.25 / 85355@21.5; warm: 85355@22.125 / 2157@21.5) — near-tie amplification of ULP-level kernel deltas through 48 MoE layers.
    • Token-level determinism is not a property of the engine: same-mode warm reruns flip step-1 tokens because expert-output accumulation order (final_hidden_states_.add_ / index_add_) is completion-order nondeterministic (pre-existing, mode-independent). The historical warm digest 3739c3b3 reproduces exactly this way.

Merge order

dev#197 first, then #181 (rebased on dev), then this PR into #181.

…0 eager

- kv_cache.set_block_store: normalize cuda device index so 'cuda' vs 'cuda:0'
  geometry no longer rejects the store (prefix cache now binds/activates)
- prefix_cache/engine: surface prefix_cache_hits_total + matched_tokens_total
- attention_backend: pass q_data_type/kv_data_type to FlashInfer plans
- big_modeling: use eager attention for qwen3 when flash-attn unusable (sm120)
- prefix_cache_benchmark: one model per subprocess (avoids multi-model crash)
- engine: include bind exception detail in disabled reason
…ness

- benchmarks/prefix_cache_benchmark: add --parity-report (per-mode tensor
  dumps + numerical analysis instead of digest-only failure), and
  --device-memory-ratio/--kv-cache-ratio so the suite fits shared GPUs.
- tests/integration/test_flashinfer_kernel_parity: prove full-prefill and
  append kernel schedules are each deterministic, within ULP-scale error
  of an fp32 reference, and mask partial-page tails (3/3 pass on sm120).
  Cross-kernel bitwise equality is not an invariant.
- tests/serving/test_prefix_cache_cuda: run warm-vs-cold parity through
  one-engine-per-subprocess (kTopologyHandle-safe) with correct
  invariants: disabled==cold bitwise control, geometry-proven reuse,
  tolerance + near-tie-aware token comparison.
@drunkcoding drunkcoding changed the title fix(serving): prefix-cache binding, hit stats, dtype [HOLD: warm not bit-exact on sm120] fix(serving): prefix-cache binding, hit stats, dtype [HOLD: expert-fetch race] Sep 3, 2026
Warm-vs-cold token equality is unattainable: the paths use different
FlashInfer kernel schedules (ULP-level per-op deltas) and expert-output
accumulation order is nondeterministic, so near-tie argmax swaps occur
even between same-mode reruns. Keep disabled==cold bitwise as the strict
store-path regression check; at a flip require the two runs to share the
same top-2 candidate pair, which still catches real corruption.
@drunkcoding drunkcoding changed the title fix(serving): prefix-cache binding, hit stats, dtype [HOLD: expert-fetch race] fix(serving): prefix-cache binding, hit stats, dtype (needs #197) Sep 3, 2026
@drunkcoding
drunkcoding marked this pull request as ready for review September 4, 2026 08:15
@drunkcoding
drunkcoding merged commit 0d49824 into plan/prefix-kv-reuse Sep 4, 2026
3 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.

1 participant