[Bugfix][V1] Warm attention through dummy profiles#42215
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
@ZJY0516 @qiching @tdoublep @vadiklyutiy Sorry to bother again, third one from me for the #40137 area. This time it is TurboQuant decode. I call If you have concerns about the approach or scope please let me know. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a warmup mechanism for TurboQuant decode kernels to ensure they are compiled before serving requests, thereby reducing latency on the first inference. It adds the turboquant_warmup.py module, integrates it into the kernel_warmup flow, and provides comprehensive unit tests. Feedback was provided regarding the calculation of block_table_stride, noting that the current approach might default to an incorrect value during the initial warmup phase and suggesting a more direct way to access the required constant from the model runner to avoid unnecessary re-compilation.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fb6fd58b07
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
lesj0610
force-pushed
the
lesj/tq-decode-jit-warmup-20260510
branch
from
fb6fd58 to
42aafd5
Compare
|
This PR's problem statement matches Issue #41565 exactly:
That's the workspace lock-violation we filed in #41565 ( The two PRs attack the same root cause from different angles:
These are complementary, not competing. Both landing would be belt-and-suspenders against the same regression. Either one resolves #41565. For the maintainer queue: I can validate this on 8× RTX A4000 (SM86) / Nemotron-3-Super-120B-AWQ-4bit / TurboQuant — different model class than your Qwen3-8B test (hybrid Mamba+MoE+Attention vs dense attention) and different arch generation. Worth a cross-platform second data point. Happy to run a sweep at 4K / 16K / 64K / 131K cached tokens once review opens. |
|
Following up on this — we're still pinned to a pre-PR#40941 fork on production specifically because of the workspace lock-violation this PR addresses. The bug fires reliably on:
Trace matches #41565 exactly: first-decode Today's expanded scope (full-dequant JIT warmup + V2 single-request execute path) covers more of our profile than the original — we hit JIT compile spikes on first decode under TQ regardless of the specific kernel, so the broader warmup matches the bug shape better. Happy to validate on our 8× A4000 SM_86 stack against Super-120B once you've stabilized rebases and the |
|
@MidasMining Thanks for following up and the detailed repro. Current branch covers full-dequant JIT warmup, continuation path, prefix-cache hit variants, and V2 single-request execute path. Production validation from different platform would be really helpful, let me know if you need anything. |
|
@lesj0610 Validation report from our 8× RTX A4000 SM_86 production stack. TL;DR: the patch resolves the workspace-lock crash on cold start. ✅ Setup
Results
Warmup behavior observedAll workers logged the CUDA graph memory profiling notice cleanly. No exceptions during What this validates / doesn't✅ The expanded warmup scope covers our Nemotron-H + TQ-3bit-NC + EP=8 + cyankiwi compressed-tensors INT4 combination — the previously-crashing path. ❌ Did not test: heavy concurrency stress, long-context (>32K) decode, prefix-cache eviction churn, or comparison to a pre-PR baseline reproducing the exact crash signature on this build (the prior repro was on our fork, not on upstream main at PR base SHA). Happy to add any of these if useful for review. AskIf this is the validation signal you needed, no further from us. If review surfaces config edge cases (different |
|
Additional datapoint: ran our 5-test practical-debugging quality bench ( That closes out validation from our end. Ping if you'd like a runtime check after any future rebases. |
|
@MidasMining Thanks for testing this. Your setup is quite different from mine (SM86, Nemotron-H hybrid, TP=8/EP=8) so this is really helpful. Good to see no workspace-lock crash on cold start and decode works fine. That was the main problem I was trying to fix. Concurrency / long-context / prefix-cache stress — yeah those are separate topics. Maybe worth testing later but not what this PR is about. Appreciate it. |
lesj0610
force-pushed
the
lesj/tq-decode-jit-warmup-20260510
branch
from
d8fa820 to
5a88894
Compare
|
@lesj0610 is there any chance you can simplify this PR? It is changing lots of unrelated code in the general warmup path. I would think we can resolve this issue in a more modular way with the current structures in the codebase |
Warm TurboQuant decode through the runtime decode helper so the decode Triton kernels and workspace buffers are initialized before serving requests. Co-authored-by: Codex <codex@openai.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Gemini <noreply@google.com> Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com> (cherry picked from commit 42aafd5)
Warm the TurboQuant continuation-prefill full-dequant kernel with runtime block-table constants, and stop Triton from specializing runtime stride arguments for decode/dequant kernels. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com> (cherry picked from commit c0d1a33)
Add a pre-capture old V1 attention warmup that uses the existing dummy-run path instead of a backend-specific synthetic launcher. The forced-attention decode and cached-prefill-shaped profiles let attention backends allocate workspace and compile before CUDA graph capture locks workspace and before the JIT monitor starts. Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
lesj0610
changed the title
|
@mgoin that makes sense. I reworked this to follow that direction. The PR no longer uses a TurboQuant-specific synthetic warmup. It now extends the existing old V1 dummy/profile path so attention metadata is built before CUDA graph capture and before the JIT monitor starts. This should cover TurboQuant through the normal dummy-run path instead of a backend-specific launcher. I also updated the PR description and tests accordingly. |
Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
|
This pull request has merge conflicts that must be resolved before it can be |
…rmup-20260510 # Conflicts: # vllm/model_executor/warmup/kernel_warmup.py # vllm/v1/worker/gpu/warmup.py Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
…y upstream PR #42637) After full upstream audit (vllm PRs #42637/#40798/#42215/#40914 and issues #42808/#41403/#40069/#41559), determined that 4 G4_* Genesis patches are superseded by upstream PR #42637 (Mixed-attention TurboQuant for Gemma 4, lesj0610, OPEN): - G4_30 (supports_mm_prefix gate flip) -> PR adds real mm_prefix Triton kernel mask (USE_MM_PREFIX constexpr, mm_prefix_range_tensor metadata). - G4_43 (revert forced TRITON_ATTN) -> PR makes per-layer routing the default; review after cherry-pick. - G4_44 (head_dim>256 torch SDPA fallback) -> PR's _can_use_flash_attn gate + _sdpa_causal_prefill cover this surface. - G4_45 (page-size unification diag/fix) -> PR's unify_kv_cache_spec_ page_size TQ-aware branch covers fix mode. Plus G4_50 (Genesis-native AttentionBackend) + genesis_tq/ companion (10 files) moved to sndr_private/genesis_tq_abandoned/. That direction was abandoned because upstream TurboQuantAttentionImpl + wrapper strategy (via g4_19_*) is strictly better than custom backend. G4_31 (preserve_tq_dtype - AWQ boundary) and G4_32 (validation_bypass) KEEP in main - PR #42637 does NOT touch AWQ/quant_config or validator relax. They remain Genesis-original boundary patches. Updated sndr_core/__init__.py with sndr_private fallback shim - existing env flags GENESIS_ENABLE_G4_30/43/44/45/50 still resolve via sndr_private/g4_upstream_tq_wip/ for back-compat. Production users with slim distributions get silent no-op (ImportError gracefully handled). Plan document moved to sibling Genesis_internal_docs/ per iron rule #6 (no internal docs in public repo). Production state (verified): vllm-g4-256k-packed Up 42min, smoke test correct factual answers, 5.33x compression preserved, ~102 TPS. Refs: - upstream PR vllm-project/vllm#42637 - upstream PR vllm-project/vllm#40798 - upstream PR vllm-project/vllm#42215 - my PR vllm-project/vllm#40914 - issue vllm-project/vllm#42808 - issue vllm-project/vllm#41403
…e 3)
Cherry-picks 2 upstream vllm PRs that resolve the v0.21.0 "workspace
locked at 0.00 MB" assertion family (#41565, #41726, #42544, #42808):
G4_61 (g4_61_tq_shared_workspace.py) — PR #40798 (Bot1822)
Per-layer _tq_mid_o_buf / _tq_output_buf / _tq_lse_buf allocations
replaced with shared WorkspaceManager acquisition. Layers execute
sequentially per request so one buffer set is enough.
capture_model() now pre-reserves max-shape workspace BEFORE
lock_workspace fires. _reserve_turboquant_decode_workspace iterates
ALL attn_groups (not just [0] — addresses gemini-code-assist review).
Also reserves continuation prefill cache buffer when chunked prefill
enabled and max_num_batched_tokens > _TURBOQUANT_CONTINUATION_DECODE_
THRESHOLD (128).
PR validation: Llama-3.1-70B TP=2 turboquant_3bit_nc 65536 ctx
Loading mem: 105 GiB -> 66 GiB (40 GiB recovered)
Available KV: 14.61 GiB -> 53.97 GiB (3.7x boost)
KV tokens @ 64K: 400,128 -> 1,478,384 (3.7x)
Max concurrency: 6.11x -> 22.56x
env: GENESIS_ENABLE_G4_61_TQ_SHARED_WORKSPACE
G4_62 (g4_62_tq_kernel_warmup.py) — PR #42215 (lesj0610)
Adds turboquant_decode_warmup function that walks Attention layers
with kv_cache_dtype.startswith("turboquant_"), deduplicates by
13-field _TurboQuantDecodeWarmupKey (Triton specialization criteria),
and calls impl._decode_attention with synthetic batch=max_num_decode_
tokens, seq_lens=1, block_table[:,0]=1 inputs.
Wraps kernel_warmup(worker) to call our function after deep_gemm
warmup. Reads block_size + block_table_stride from worker.model_
runner.input_batch.block_table.block_tables[0] (V1 may split KV
manager blocks into smaller attention-kernel blocks).
Complementary to G4_61: G4_61 pre-reserves max-shape, G4_62 compiles
kernels + actually-allocates. Either resolves #41565 family; both
together = belt-and-suspenders. Per MidasMining's comment.
Removes 5-25s first-request TTFT spike that comes from JIT-compiling
_tq_decode_stage1 + _tq_decode_stage2 on first real decode.
env: GENESIS_ENABLE_G4_62_TQ_KERNEL_WARMUP
Registry entries added with full credit blocks. Import-time hook in
sndr_core/__init__.py adds G4_61 and G4_62 with documented dependency
ordering (G4_61 patches launcher first, G4_62 wraps kernel_warmup after).
Smoke test (fresh container with both env flags):
G4_61 applied: True
G4_62 applied: True
Refs:
- upstream PR vllm-project/vllm#40798
- upstream PR vllm-project/vllm#42215
- issue vllm-project/vllm#41565
- issue vllm-project/vllm#42544
- issue vllm-project/vllm#42808
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
lesj0610
force-pushed
the
lesj/tq-decode-jit-warmup-20260510
branch
from
f2a738f to
8b094b3
Compare
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>
|
This pull request has merge conflicts that must be resolved before it can be |
4 participants
Problem
Old V1 startup profiling/dummy runs can skip attention metadata for request-shaped decode and cached-prefill variants. For TurboQuant, that means decode workspace and Triton decode kernels can first be touched after CUDA graph capture has locked workspace and after the JIT monitor is active.
Maintainer feedback asked to improve dummy/profile coverage instead of adding a TurboQuant-specific synthetic decode launcher. This PR now follows that direction.
Approach
_dummy_run()profile path.num_reqs_overrideto_dummy_run()so the cached-prefill shape can stay single-request without creating a large scheduler workload.do_not_specializeannotations for runtime stride/meta parameters. That is separate from warmup routing and reduces unnecessary specialization churn.This keeps warmup in the same dummy/profile mechanism used by the rest of V1 startup instead of introducing a backend-specific attention launcher.
Test Plan
.venv/bin/python -m py_compile \ vllm/v1/worker/gpu/warmup.py \ vllm/v1/worker/gpu_model_runner.py \ vllm/v1/worker/gpu_worker.py \ tests/v1/worker/test_gpu_warmup.py .venv/bin/python -m pytest tests/v1/worker/test_gpu_warmup.py -q .venv/bin/python -m pytest tests/v1/worker/test_gpu_model_runner.py -k "dummy or uniform_decode" -q git diff --checkTest Result
py_compile: passedtests/v1/worker/test_gpu_warmup.py: 2 passedtests/v1/worker/test_gpu_model_runner.py -k "dummy or uniform_decode": 1 passed, 34 deselectedgit diff --check: passedGPU runtime smoke was not rerun after this consolidation, so absence of runtime JIT warnings remains to be validated on a CUDA serving run.
Checklist
AI assistance: Codex.