Skip to content

prefill: reuse embedding storage for half outputs - #1454

Closed
W1nge wants to merge 1 commit into
Niko1221:mainfrom
W1nge:codex/prefill-embedding-reuse
Closed

W1nge wants to merge 1 commit into
Niko1221:mainfrom
W1nge:codex/prefill-embedding-reuse

Conversation

@W1nge

@W1nge W1nge commented Oct 7, 2026 •

Copy link
Copy Markdown
Contributor

Reuse the prompt embedding allocation for the later half output (bo). Embedding gathers, image-row uploads, the last embedding consumer (gr_broadcast), output projections, MoE combine and HC writes are ordered on m.cs. A following chunk reloads its embeddings; a stage receiving hand_in_ does not read emb. PLE embeddings retain their separate allocation.

Remove the corresponding bo allocation from bytes_needed_impl as well as carve. This saves chunk * 2560 * sizeof(float) payload bytes (320 MiB at 32768 rows), with allocator rounding handled by the existing allocator. No new flags, default chunk changes, or numerical kernels. Independent of #1451.

Validation: reviewed all m.emb/m.bo accesses in current main d5ea713; prefill.cpp compiles with MSVC 19.44/CUDA 12.4 headers both without native/MMQ and with native/MMQ/fused enabled; git diff --check passes.

Draft: this alias was part of the locally validated 32K fork, but this isolated upstream port has not had a full-model, peer/stage or other-backend runtime comparison. The allocation saving is not an isolated throughput measurement. No broader HC scratch overlay is included.

Validation update (2026-10-08): a complete current-main integration build including #1451, #1454, #1455, #1457 and #1465 passed four real IQ3_XXS model baseline/integration pairs on RTX 2080 Ti: default, GR_UNFUSED, PREFILL_BF16X2 and legacy RING_BYTES modes. Each used INT8 KV, 194 input tokens, two prefill chunks and 8 output tokens; all runs exited successfully and output IDs matched within each pair. A 256K configured-context startup and short generation also passed (this did not fill a 256K prompt). These are combined regression checks, not an isolated performance result or exhaustive state equality. Other hardware/backends remain untested. The earlier draft-only status reflected the absence of any current-main model run; now requesting review with these limits explicit.

@W1nge
W1nge marked this pull request as ready for review October 7, 2026 23:22
Niko1221 added a commit that referenced this pull request Oct 8, 2026
…less STRATA_EMB_REUSE_ACCOUNT=1

The VRAM reuse stays, but bytes_needed_impl counts bo again by default, so the auto chunk and the borrowed cache slots
are those of 0.1.40.3 and the prompt residuals stay byte-identical (on the RTX 3060 the planner otherwise picked a
6656-token chunk instead of 6400). STRATA_EMB_REUSE_ACCOUNT=1 lets the planner use the saved bytes (+3.9% prompt there);
documented in docs/DETAILS.md.
@Niko1221

Niko1221 commented Oct 8, 2026

Copy link
Copy Markdown
Owner

Thanks, W1nge. The VRAM reuse is in 0.1.41, with your authorship kept: the half output now shares the embedding storage, about 320 MiB less at 32,768 rows, checked identical on one GPU and on a 2-GPU layer split. On a VRAM-limited card the saved bytes would also let the prompt planner pick a slightly larger chunk (an RTX 3060 went from 6,400 to 6,656 tokens, about +3.9% prompt speed), but that changes the prompt's rounding, so by default the planner keeps the old accounting and the output stays byte-identical. STRATA_EMB_REUSE_ACCOUNT=1 turns the planner gain on.

Closing as ported.

@Niko1221 Niko1221 closed this Oct 8, 2026
mobile2machan added a commit to mobile2machan/Strata that referenced this pull request Oct 8, 2026
128 commits: batch-groups auto per stage, server listen backlog 256, the
NVIDIA short-prompt CPU share, batched prompt reads (Niko1221#1323), the tokenizer
cache, the engine stall watchdog (Niko1221#1317 Niko1221#1407), conversation-cache eviction
(Niko1221#1347), prompt buffers sharing the embedding storage (Niko1221#1454), the Pascal
decode fix (Niko1221#1469), q8_1 scale clamps (Niko1221#1448), STRATA_STAGE_PIN opt-in (off:
on by default it corrupted IQ3_S decode after a long prompt), route-resident
and GDN-chunked opt-ins, gfx906 build fix, AMD docs and community benches.

main keeps 33af099 (AMD PCIe link speed/width and GPU name detection) and
a798ac0 (the k8v4 config read-back and 816 B/cell sizing).
MistVVK added a commit to MistVVK/XeStrata that referenced this pull request Oct 11, 2026
…1#1454)

埋め込みの行は最初の broadcast で読んだあとは使わず、bo への書き込みは
同じ in-order のキュー(GEMM も同じキュー)でそのあとに来るので、bo を emb と
共有する。チャンクあたり T × 2560 個の float の VRAM が空く。

チャンクの大きさを決めるときは既定でその分も数え、チャンクと借りる枠、
prefill のビットは前と同じにする。STRATA_EMB_REUSE_ACCOUNT=1 で空いた分も使う
(A/B で既定を決める)。
giridhart pushed a commit to giridhart/Strata that referenced this pull request Oct 15, 2026
…h emb unless STRATA_EMB_REUSE_ACCOUNT=1

The VRAM reuse stays, but bytes_needed_impl counts bo again by default, so the auto chunk and the borrowed cache slots
are those of 0.1.40.3 and the prompt residuals stay byte-identical (on the RTX 3060 the planner otherwise picked a
6656-token chunk instead of 6400). STRATA_EMB_REUSE_ACCOUNT=1 lets the planner use the saved bytes (+3.9% prompt there);
documented in docs/DETAILS.md.
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.

2 participants