perf(qwen): reduce decode input preparation overhead - #97
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughDecode hidden states are now optional for device-embedding executors. Qwen3-14B decode preparation writes directly to persistent compiled buffers with block-table caching, while batching tests cover the updated buffers and ChangesDecode pipeline
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant WorkerProcess
participant Qwen314BModelRunner
participant CompiledDecodeBuffers
participant DecodeResult
WorkerProcess->>Qwen314BModelRunner: run_decode with optional hidden states
Qwen314BModelRunner->>CompiledDecodeBuffers: populate persistent decode buffers
Qwen314BModelRunner->>DecodeResult: return logits and hidden_states=None
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![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 optimizes the decode path for device-embedding executors (such as Qwen) by omitting decode hidden states (passing None instead of zero placeholder tensors) and writing active decode metadata directly into persistent kernel buffers. It also introduces caching for block-table rows to avoid rewriting them unless page allocations change. The feedback suggests a performance optimization in _write_cached_decode_block_table_row to move the cache check to the top of the function, thereby avoiding unnecessary tensor slicing and .numel() calls on cache hits.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
b037b7b to
ad8d401
Compare
| # Host-embedding executors (for example DeepSeek) consume hidden_states. | ||
| # Device-embedding executors (for example Qwen) gather embeddings from | ||
| # token_ids inside the decode kernel and leave this as None. | ||
| hidden_states: torch.Tensor | None |
There was a problem hiding this comment.
现在两个模型的输入差距很大吧，共用这个基类还有意义吗？可用后续整改，本次先识别下deepseek路径下面有没有未判空的调用hidden_states
04eb9cc to
7e89943
Compare
Remove unused hidden-state placeholders from device-embedding decode while preserving the real hidden-state inputs used by DeepSeek. Write Qwen token IDs, sequence lengths, block tables, and slot mappings directly into persistent compiled buffers. Cache block-table rows until page allocations change, check cache hits before creating tensor row views, retain row-0 replication for fixed-batch padding, and honor the padded token ABI by updating only column zero. Enforce DeepSeek host-hidden-state requirements across standard and MTP decode consumers so the shared nullable DecodeBatch contract fails clearly instead of dereferencing None. Extend batching and DeepSeek coverage for optional hidden states, direct buffer reuse, block-table cache invalidation, SAMPLED_IDS_PAD lanes, and missing DeepSeek hidden-state validation. The optimized trace reduces average PyptoExecutor.run_decode time from 41.48 ms to 38.14 ms. Pre/post host overhead falls from 3.88 ms to 0.50 ms while decode kernel time remains approximately 37.6 ms. The old trace included verbose input dumping, so this trace delta is an upper bound. An isolated B16 no-hidden preparation benchmark improves from 469 us to 175 us after direct-buffer caching. Short-circuiting cache hits before tensor slicing further improves the cache-hit median from 175.6 us to 127.8 us.
2 participants
Summary
[B, SAMPLED_IDS_PAD]ABI by updating only token column 0Performance
The optimized trace contains 133
PyptoExecutor.run_decodecalls:PyptoExecutor.run_decodekernel.decode_fwdThe unchanged kernel time isolates the improvement to host-side decode handling. The older trace included verbose per-step input dumping, so the full trace delta should be treated as an upper bound. An isolated B16, sequence-length-3338 preparation benchmark without hidden-state work improved from 469 us to 175 us median (about 63%).
Correctness
Production Qwen uses
SAMPLED_IDS_PAD=8, so the compiled token buffer is[B, 8]even though each request supplies[B, 1]. The new path initializes padded columns once, writes token IDs only to column 0, and tests the real padded ABI. Block-table rows are cleared and rewritten whenever page IDs change, including when requests move into previously padded rows.Validation
python -m pytest tests/test_batching.py tests/test_parallel.py tests/test_deepseek_v4.py— 55 passedgit diff --checkchecks passedtask_20260716_210028_3707328928)