Repository navigation
Reuse sparse prefill history storage after candidate selection - #247
Closed
voipmonitor wants to merge 1 commit into
Closed
voipmonitor wants to merge 1 commit into
voipmonitor wants to merge 1 commit into
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.pywithFLARE_KIMI31_CHECKPOINTnaming 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
8d8582eewith preparation-only FlashInfer6a70bdb1d. The input was targetmoonshotai/Kimi-K3.1-Prorevision6d0a0202956796543364d666ab19c6ad1a494ee1, 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:
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 at04c70f55; 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 stack6a70bdb1d.