Skip to content

Reuse sparse prefill history storage after candidate selection - #247

Closed
voipmonitor wants to merge 1 commit into
astra/k31-dflash2-workspacesfrom
astra/k31-prefill-history-reuse
Closed

voipmonitor wants to merge 1 commit into
astra/k31-dflash2-workspacesfrom
astra/k31-prefill-history-reuse

Conversation

@voipmonitor

Copy link
Copy Markdown
Collaborator

Kimi K3.1 complete-chunk sparse prefill materializes a planar index history for candidate selection and a latent history for attention. At a 1,048,576-token context limit, separate storage consumes 132 MiB and 576 MiB per rank even though selection can finish before attention consumes latent records.

Complete index scoring and selected-ID exchange before producing latent history, then reuse the first 132 MiB of the latent backing for the index view. Persistent KV/index caches, peer staging, pointer tables, and synchronization signals remain independent. Update the preparation guide's memory accounting and add the complete-checkpoint-layer qualifier alongside metadata-serialization and compiled-copy tests.

This removes a 132 MiB logical allocation per rank at a 1,048,576-token context capacity without changing target arithmetic or prepared capacity. An earlier ordinary-export comparison of the combined Kimi stack had a 130 MiB net fixed-allocation reduction: another 2 MiB allocator mapping remained. That net observation is distinct from the logical tensor saving and is not a measurement of this port. No inference-speed improvement is claimed.

Validation on this branch: 68 CPU tests passed for shared storage through compiler-worker serialization, Kimi configuration/views, and serial workspaces; four CUDA-only checks were skipped. Run python -m pytest -q tests/recipes/test_kimi_prefill_histories.py tests/recipes/test_kimi31_dsa_views.py tests/recipes/test_kimi31.py tests/recipes/test_kimi_serial_workspace.py with FLARE_KIMI31_CHECKPOINT naming the target checkpoint.

The ported complete-checkpoint-layer qualifier passed all eleven cases on all sixteen Frank1 SM120 GPUs in 98.64 seconds, using this branch at 8d8582ee with preparation-only FlashInfer 6a70bdb1d. The input was target moonshotai/Kimi-K3.1-Pro revision 6d0a0202956796543364d666ab19c6ad1a494ee1, attention layer 3, 4096 input rows and 1,048,576 prepared history tokens. Cases include partial/empty input, near-limit history, shrinking active rows, and three queued replays with poisoned reusable storage. All storage, allocation, selection, and output gates passed.

Cache, queries, gates, visible history, and rank agreement are exact. Changed selected sets require exact cutoff-tie proof; an independent probability oracle checks rank-slice endpoints and active first/last queries across local heads, and the recorded compiled attention/output tail checks full-output equality in each emitted order. No tolerance is widened. This qualifies the storage transformation on the checkpoint layer, not full-model response quality or standalone Helios loading of a newly exported core. No full-model export was performed for this check.

The qualifier self-spawns sixteen ranks. On an exclusively assigned sixteen-GPU host:

result_dir=$(mktemp -d -t kimi-history.XXXXXX)
python -m qualify.kimi_history_alias \
  --checkpoint /data/kimi31-support/source-views/kimi-k31-mxfp4 \
  --config /data/kimi31-support/main-61b883b6/preparation.toml \
  --context-capacity 1048576 --output "$result_dir/result.json"

The checkpoint and configuration paths identify the measured Frank1 installation; other installations must supply the same checkpoint revision and their preparation configuration. Exact argv, source revisions, elapsed time, raw case/rank flags, and log are retained under /data/kimi31-pr-series-20261109/flare-history-tp16-1m/.

This PR is stacked on DFlash2 PR #246 at 9a3245b3, which provides the Kimi composition and deferred preparation storage. That parent depends on retained-template export PR #245 at 04c70f55; merge graph export, then DFlash2 workspaces, then history reuse. The history change does not require query-owner selection or its FlashInfer transport changes. CPU validation uses the preparation-only FlashInfer stack 6a70bdb1d.

Kimi complete-chunk prefill keeps separate global index and latent histories,
even though their consumers can run serially. At context1M these histories
occupy 708 MiB per rank.

Finish index selection and selected-ID exchange before copying latent
history. Alias index history onto the first 132 MiB of the 576 MiB latent
backing, with persistent caches and peer staging retained separately. Add
the checkpoint-layer qualifier and update documented memory accounting.

This removes a 132 MiB logical allocation per rank without changing target
arithmetic or serving capacity. No inference-speed improvement is claimed.

Sixty-eight CPU tests pass, including shared-view worker serialization,
Kimi views/configuration, and serial workspace checks; four CUDA checks are
skipped. The retained checkpoint-layer qualifier preserves exact storage,
cutoff-tie proof and recorded-order output checks. TP16 and native checks
have not yet been rerun on this port.

Design-Signed-Off-By: Codex (GPT-6); design-review=cfcf
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