[Bugfix][QSA] Reuse bounded prefill logits workspace - #56500
jacklin78911-collab wants to merge 1 commit into
Conversation
Reserve logits and top-k scratch during QSA profiling and reuse compact views across prefill chunks without changing the selection geometry. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Liqian Lin <jacklin78911@gmail.com> Signed-off-by: jacklin78911-collab <jacklin78911@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
Ran this on GB10 / DGX Spark, sm_121, aarch64, single node — the hardware class #56457 was Stock and patched in one run, toggling the source between arms and reading the marker
The patch imports and runs on sm_121, and a 170k real-weights prefill completes with no functional What this does NOT show. I could not reproduce #56457 on a single node at all: 250,010 tokens at On your stated risk — reserving the budget during profiling increasing live memory. At startup: Two caveats a reviewer should weigh. The serving venv is not stock upstream — four deviations at this Happy to run further cells on this box if a specific configuration would help. Testing and this write-up were done with AI assistance; every number is from the run logs above and I reviewed them. |
Note from a single-GB10 (DGX Spark) validation attempt — negative result + a packaging caveatWe tried to validate this PR on a DGX Spark (GB10, sm_121, ARM64, 128 GB unified) and ended up not 1. The main hunk targets a file that neither shippable base has. The chunk-loop rewrite lands in
Both bases still carry the older shape, where the per-chunk budget is a module constant rather than _LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024 # v0.29.0 ops/qsa.py:14
...
rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1))
topk_workspace = torch.empty((_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device)
2. On the 128 MB-capped shape, the symptom does not reproduce. One fresh server, growing chunked prefills at 8K → 16K → 24K → 32K → 48K → 64K prompt tokens, each
Flat plateau: the before-step readings sit within 15 MB of each other across the whole sweep, and Caveat on the measurement: on GB10 the memory is unified, so So this is a negative result for this shape, not a claim about the Raw output and tooling: https://github.com/k3net/docai-evals/tree/b1f14a36c5bcc02cf2cd65705031e676fa9cb73f/experiments/2026-09-12-qwen38-flash-next-prefix-cache-cross-request-gb10 |
Thanks for testing both bases and documenting the version mismatch. This PR targets the current Additional validation at head
An allocation-only probe with contexts growing from 3,200 to 256,000 tokens measured:
Scoring and top-k kernels were skipped in this probe. The patch retained 528.65 MiB of live allocations versus 14.65 MiB after baseline calls, so the persistent-workspace cost remains relevant. These measurements support the allocator-reuse mechanism. They do not establish serving performance, model quality, or resolution of the original two-GB10 TP=2 failure, which remains unverified. |
|
@jschmied Thanks for testing the exact patch on sm_121 with real weights and documenting the overlays and measurement limits. This adds useful coverage beyond the dummy-model checks. I agree that single-node completion does not confirm resolution of #56457 or establish a speedup. Additional H20 validation now passes all 103 QSA/config/workspace tests, with baseline/patched outputs matching 24/24 tokens for each of V1 and V2 in the small dummy-model checks. The repeated stock/512 MiB, patched/512 MiB, and stock/64 MiB startup measurements you mentioned would be particularly useful for assessing the persistent workspace’s impact on KV capacity. Please share those results when available. |
3 participants
Purpose
Chunked QSA prefills allocate a progressively larger logits tensor as
max_seq_lengrows. The caching allocator can retain the previous sizes, so a per-tensor logits budget does not bound the memory retained across prefill steps. Addresses #56457.Use the existing WorkspaceManager for a bounded logits buffer and disjoint top-k scratch. Reserve capacity in the QSA owner's profiling path, which returns before calling the indexer, and use compact tensor views for each inference chunk. This preserves the logical logits width, chunk sizes, kernel launches and top-k dispatch introduced by #54915. The existing manager supplies separate ubatch/workspace-lane storage; no new global cache or synchronization mechanism is added.
The reservation includes at least one logits row when the configured budget is smaller than a row. Reserving the configured budget during profiling moves that memory cost up front; it can increase live memory for short-only workloads when another operator has not already reserved a larger shared workspace.
Validation
Base:
dc07f1638f73814b95776832b85df1cc92850416. Local environment: RTX 4060 Laptop GPU (SM89, 8 GiB), PyTorch 2.11.0+cu130, existing prebuilt native extensions. This is not the repository's current PyTorch 2.13 CI environment.Commands for the full and final sweeps:
.venv/bin/python -m pytest \ tests/models/qwen4_exp/test_qsa_reference.py \ tests/models/qwen4_exp/test_qsa_pre_indexer.py \ tests/models/qwen4_exp/test_config.py \ tests/v1/worker/test_workspace.py -q .venv/bin/python -m pytest \ tests/models/qwen4_exp/test_qsa_reference.py \ tests/models/qwen4_exp/test_qsa_pre_indexer.py \ tests/models/qwen4_exp/test_config.py \ tests/v1/worker/test_workspace.py -q \ -k 'not (prefill_selection and dtype1) and not tp1_r2048 and not (tiled and indexer_dtype1)'Allocation and timing probes
Allocation-only probe: real wrapper and CUDA allocations, with scoring/top-k skipped to isolate allocation behavior; 3,200 query rows, 12 layer calls per step, 16 growing contexts from 3,200 to 51,200 tokens, default 512 MiB budget and native allocator.
The larger live reservation is intentional. These numbers establish bounded storage reuse in this probe; they are not a GB10 hang reproduction or a universal memory-saving claim.
Separate real-kernel probe: BF16, four query heads, head dimension 128, identical inputs and native top-k. Valid logits and selected score multisets were bit-exact across the two implementations. For the 3,200-token cold-prefill case, selected indices can differ at ties; repeated baseline runs also differ, consistent with the existing top-k determinism work in #55122.
Median GPU times from three alternating rounds, FlashInfer CUPTI, CUDA graphs and cold L2:
A separate warmed eager-wrapper wall-time check, including Python/allocation overhead, was approximately unchanged: median differences -0.2%, +0.0%, +1.0% for these shapes. These are local microbenchmarks, not serving-latency or model-throughput claims.
Remaining validation and related work
RuntimeError: UVA is not available, before reaching QSA.Duplicate checks covered the issue discussion, open PRs referencing #56457, and QSA workspace/allocation/allocator searches. No open PR implementing this storage-reuse fix was found as of 2026-09-12. #54915 is the existing compact-width optimization being preserved; #55122 concerns top-k determinism; #56240 adds a different attention backend.
AI assistance: OpenAI Codex assisted with investigation, implementation, tests, benchmarks and a second pass through the complete diff and actual call chain. No claim is made that a human has already reviewed every changed line or run these commands.