[Discussion] Tile indexed-attention scoring independently of the global prefill chunk
- Ngôn ngữ chính
- C
- Star
- 22.3k
- Fork
- 2.1k
- Merge trung bình
- 1 ngày 3 giờ
- Pull request đã merge (30 ngày)
- 4
Mô tả
> **Disclosure:** This discussion was prepared and filed by **Codex (OpenAI)** on behalf of `@apetersson` while attempting to build a 128 GB-optimized, selectively native-MXFP4 DeepSeek V4 Flash quant on the experimental `ds4f-mxfp4` branch. The audited revision is `4893e0c40fba03dbc85555faeb035799aa04e0b6`. The memory figures are source-derived estimates and should be independently measured.
## Summary
At very long DeepSeek V4 Flash contexts, indexed-attention score scratch scales with:
```text
compressed_cache_rows × global_prefill_chunk
```
The current way to bound that scratch is to lower `--prefill-chunk`, which also lowers the batch used by the MoE and other dense prefill stages. Would it be practical to tile only indexer score/top-k work inside a larger global prefill chunk?
This is related to the allocation/lifecycle discussion in #643, but is separated because it is a throughput-oriented execution change rather than a simple allocation audit.
## 128 GB / 1M-context motivation
For DeepSeek V4 Flash at `ctx=1,048,576`:
- Ratio-4 compressed capacity is approximately 262,146 rows.
- With global prefill chunk `4096`, `indexer_scores` alone is approximately 4 GiB.
- Source-audited full chunk-scaled workspace is approximately:
- `11.959 GiB` at chunk 4096
- `5.980 GiB` at chunk 2048
- `2.991 GiB` at chunk 1024
For the candidate 128 GB hybrid package, estimated total residency changes from approximately `123.24 GiB` at chunk4096 to `117.10 GiB` at chunk2048 and `114.02 GiB` at chunk1024.
Chunk1024 is therefore a workable capacity setting, but a one-million-token prompt requires four times as many chunks as chunk4096 and may give up useful GPU efficiency in routed MoE, dense projections and other stages that do not need the full indexer score matrix.
## Possible design
Keep a global prefill chunk such as 4096 for the normal layer pipeline, but process indexed-attention score/top-k work in bounded row tiles, for example 512 or 1024 query rows at a time:
1. Produce or view the tile's indexer queries and weights.
2. Score that tile against the current compressed index cache.
3. Run top-k into the compact `comp_selected` representation.
4. Consume the selected rows for indexed mixed attention.
5. Reuse the same score scratch for the next tile.
The exact scheduling may need ping-pong buffers or careful overlap so tiling does not serialize the rest of the layer.
## Correctness questions
- Can score/top-k tiles consume the existing batch indexer-query tensors without copying?
- Are compressor and indexer frontiers invariant across subtiles inside one logical prefill chunk?
- Does any zero-prefix or resumed-prefill path require all selected indices to exist simultaneously?
- Can absolute positions, raw-window boundaries and ratio-4 emission ordering remain identical?
- Would tiled execution change floating-point results beyond the differences already seen between global chunk sizes?
## Suggested acceptance and benchmark matrix
Correctness:
- Compare full logits or top-logprobs against the current chunk4096 path.
- Cover zero-prefix prefill, resumed prefill and a chunk that crosses compressor/indexer emission boundaries.
- Run `--decode-consistency` after a long tiled prefill.
- Test contexts where compressed rows are below and above indexer top-k.
Performance/memory on the same M1 Ultra:
| Global chunk | Indexer tile | Purpose |
| ---: | ---: | --- |
| 4096 | 4096 | Current high-throughput/high-memory baseline |
| 2048 | 2048 | Current compromise |
| 1024 | 1024 | Current capacity-safe baseline |
| 4096 | 2048 | Proposed intermediate tile |
| 4096 | 1024 | Proposed capacity-oriented tile |
Record:
- Peak runtime tensor bytes and process physical footprint.
- Indexer score/top-k time.
- Whole-prompt prefill tokens/second.
- Routed MoE prefill time.
- Logit/retrieval consistency.
The useful outcome would be memory close to the global chunk1024 configuration while retaining materially more of chunk4096's overall prefill throughput.
Hướng dẫn đóng góp
Đánh giá
Issue này chưa được đánh giá.