改动 / #3306

#3306 [Engine] each GPU can choose its own prompt chunk size

open · @gopinath87607 · 0 评论 · 去 GitHub 看

Server & APIMulti-GPUAMD / HIPNVIDIA / CUDALinux

说明

> **Why this PR was recreated.** The original (#2829) was opened from the `gopinath87607/Strata` fork. Deleting that fork deleted its branches, and GitHub closes every pull request whose branch lived in a deleted fork - a closed PR whose source branch is gone cannot be reopened. This is the same work from a new fork, rebased onto current `main` (`e88ee39d`); the conflict with upstream's `chunk_vram` prompt planner in `generate.cpp` was resolved. It replaces the closed #2829, which was closed by the fork's deletion, not by a maintainer.

## Problem and result

Strata's layer split is one prompt path: every stage reads prompt chunks of the chain's size. A card whose expert cache cannot fund that chunk therefore pulls every stage down to its number (the `#448` warning), even when the other cards could read more. After this change, `STRATA_PREFILL_MAX` gives each prompt-path participant its own opt-in ceiling: a capped stage lays out at `min(chain, its cap)` and reads each larger hand-off in its own steps, while the other stages keep the chain's chunk. A ceiling only lowers a stage and never raises one. Unset, no number changes.

## Related discussion

The original pull request and its discussion: #2829 (closed when its fork was deleted; this PR replaces it). The work comes from the Maya checkout (`project-maya`), local commit `ef3b6e7` (2026-10-15).

## Approach

`STRATA_PREFILL_MAX="8192,4096,8192,8192"` - an opt-in ceiling on the prompt chunk of each prompt-path participant, in the split's order. Entry 0 is the first card's path (it reads the chunk's embeddings and hands its rows on), entry *i* the *i*-th stage; `0` or an absent entry is no ceiling, and values are multiples of 256. `STRATA_GLM_PREFILL_MAX` is accepted as the same knob.

The ceiling is the other choice to one grid: a capped stage lays out at `min(chain, its cap)` and `Prefill::run_impl`'s `c0 += m.T` loop reads each larger hand-off in steps of its own, forwarding each sub-step, while the other stages keep the chain's chunk.

Every chunk-pricing site prices a part at its own chunk: the auto scan (`fits_one`, `ring_room`, the old candidate list), the serve loan (`lend`), the resident-RAM lend regions, stage init (at the stage's own maximum - `relayout` requires `chunk <= T_max`), the init retry, split-small's extra slots, the own-buffers headroom test and the head-alias candidates. The start prints one line per capped entry: `strata serve: prompt chunk: CUDA%d reads N tokens a step (STRATA_PREFILL_MAX N)`, with `, inside a larger chain chunk` when the cap is below the chain's.

Unset, not one number changes: the list is empty, `part_chunk` / `part_ceiling` return the chain chunk, `max_head` is 0, and the new print loop runs zero times. This is by construction, not by a measurement.

## Validation

Tested revision: `7a4175d5`, one commit on current `main` (`e88ee39d`). Environment: Linux, CUDA 13.2 (V13.2.86), `sm_86` + `sm_120`; the build machine is the Q8_0 4-way rig (2x RTX 3060 12 GB, 2x RTX 5060 8 GB).

**Passed:**
- `git rebase e88ee39d` - the conflict in `generate.cpp` with upstream's `chunk_vram` prompt planner was resolved; the branch's diff is this feature only.
- `ninja -C build strata` at `7a4175d5` - exit 0 (CUDA 13.2, `sm_86` + `sm_120`).
- `-fsyntax-only` on `generate.cpp` with the build's own flags - clean; one pre-existing unused-variable warning in unrelated code. This was run before the rebase; the full build above covers the rebased head.
**Failed:**
- none observed.

**Not run:**
- The binding-cap A/B on our rig (`STRATA_PREFILL_MAX="8192,8192,8192,2048"`, expecting the chain to rise toward CUDA0's own 6912 while CUDA3 reads its hand-offs in 2048-token steps). Not run; no speedup is claimed.
- Serving the rebased head: only the build was re-run after the rebase; no GPU run was made at `7a4175d5`.
- HIP and SYCL builds: no toolchain on this machine. `generate.cpp` carries no per-backend code, so both builds compile the same change.

Measured elsewhere: the Maya checkout's local commit `ef3b6e7` (not this engine, not run by us here) measured moving every part as a loss - 94,840 ms for the capped 16k/16k/8k/7k shape against 91,718.9 ms for 13 equal chunks of 7,744, on its byte-identical 100,535-token read (an earlier nesting took 98,568.1 ms). That negative result is why the ceiling only lowers a stage and why it is opt-in.

On the start log of the four-way Q8_0 rig this was written for: the chain settles at 4096 because CUDA3's 690 slots cannot fund 6912 (`WARNING: prompt chunk 4096 tokens, not 6912 ...`, `#448`). That is the case the knob is for; the A/B that would price it has not been run.

## Limits

- The knob only helps when one participant's cache is the binding limit. The one measurement we have says capping every part is a loss on its rig.
- No number is claimed for this engine; the A/B above is how the gain or loss on a given rig is measured.
- HIP and SYCL are unverified beyond the fact that the file has no per-backend code.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

本站相关内容

相关页面的快捷入口。