perf(qwen3): skip redundant decode d2h/host copies in greedy path - #71
perf(qwen3): skip redundant decode d2h/host copies in greedy path#71lyfne123 wants to merge 2 commits into
Conversation
The fused Qwen3-14B decode kernel does device-side embedding and greedy
sampling, so several per-step host copies are pure waste. Eliminate three,
all host-only (no kernel change, no recompile):
1. Greedy logits d2h. In the device-greedy path the token id is produced
inside the kernel (decode_sampled_ids_buffer), so host code never reads
the [BATCH, VOCAB] f32 logits. Point the kernel's logits `out` slot at a
reusable worker-resident DeviceTensor instead of the shared host buffer;
the runtime only d2h-copies host shared-memory args, so the ~9.7 MB/step
copy (16 x 152064 x 4B) disappears. Non-greedy still uses the host buffer.
2. next_hidden d2h. Decode never returns next_hidden (run_decode always
passes None to _integrated_sample_result), so its Out buffer is scratch.
Route it to a device-resident buffer unconditionally, dropping a
~160 KB/step copy.
3. Dead decode hidden copy. _prepare_decode_inputs built a [B, hidden]
bf16 buffer from batch.hidden_states (a zeros placeholder for
device-embedding executors) that no caller reads and the kernel does not
consume. Stop building/returning it.
DecodeResult.hidden_states / .logits become Optional; serving_worker reads
host logits lazily (only on the non-greedy fallback), so None is safe.
Verified on NPU (a2a3, greedy temperature=0): correct output
("The capital of France is" -> "Paris ... Washington"), 23 decode steps
through the new path, no errors. Run-to-run token divergence was confirmed
to be inherent NPU non-determinism (the same unmodified binary diverges
identically across devices at a near-tie argmax), not caused by this change.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
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:
📝 WalkthroughWalkthroughChangesDevice-greedy decode flow
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ServingWorker
participant Qwen314BModelRunner
participant DecodeKernel
participant DecodeScratch
ServingWorker->>Qwen314BModelRunner: run_decode with greedy sampling enabled
Qwen314BModelRunner->>DecodeScratch: get cached logits target
Qwen314BModelRunner->>DecodeKernel: run decode with logits_override
DecodeKernel->>DecodeScratch: write device-resident logits
Qwen314BModelRunner-->>ServingWorker: return DecodeResult with logits=None
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 by keeping the decode hidden states and logits device-resident when device-side greedy sampling is enabled, avoiding unnecessary device-to-host (d2h) copies. This is achieved by introducing reusable worker-resident scratch buffers for logits and next_hidden states, making hidden_states and logits optional in DecodeResult, and moving the host-side logits slicing logic into _sample_result_row. Feedback on the changes highlights a potential AttributeError in _sample_result_row if result.logits is None when falling back to host-side sampling, and suggests adding a defensive check to handle this case robustly.
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.
| logits = result.logits | ||
| logits_row = logits[row_idx] if logits.dim() > 1 else logits | ||
| return self.sampler.sample(logits_row, params) |
There was a problem hiding this comment.
If allow_device_sampled is True but sampled is unexpectedly None (or in any other scenario where result.logits is None), accessing logits.dim() will raise an AttributeError: 'NoneType' object has no attribute 'dim'. Adding a defensive check here to ensure logits is not None before accessing its attributes makes the sampling fallback path much more robust and provides a clearer error message.
logits = result.logits
if logits is None:
raise RuntimeError(
"Both sampled_token_ids and logits are None. Cannot perform sampling."
)
logits_row = logits[row_idx] if logits.dim() > 1 else logits
return self.sampler.sample(logits_row, params)…gression) The decode d2h/host-copy removal is numerically neutral: the kernel's greedy argmax reads identical logit values whether the logits buffer is host- or device-resident, next_hidden is write-only, and the removed decode hidden copy was never consumed. Local reruns pass chunked-prefill 3/3 on this branch and 2/2 on main; the prior CI failure was an inherent chunked-vs-full-prefill near-tie flip on the first token. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
c85cabe to
94e8975
Compare
|
和 #57 重复了 |
2 participants
Summary
The fused Qwen3-14B decode kernel does device-side embedding and greedy sampling, so several per-step host copies are pure waste. This removes three, all host-only — no kernel change, no recompile. All Qwen entry points default to
temperature=0(greedy device sampling), so this is the hot path in normal runs.16 × 152064 × 4B ≈ 9.7 MBnext_hiddend2h16 × 5120 × 2B ≈ 160 KBhiddencopy[16, 5120]alloc + copy + floatHow
The runtime only device→host copies host shared-memory args; worker-resident
DeviceTensorargs stay on device. So:decode_sampled_ids_buffer(device), and host code never reads the full[BATCH, VOCAB]logits. Point the kernel's logitsoutslot at a reusable worker-residentDeviceTensor(_decode_device_scratch) instead of the shared host buffer. The kernel's internal argmax stays device-local; the d2h disappears. Non-greedy still uses the host buffer and returns real logits.run_decodealways passesNoneto_integrated_sample_result), so itsOutbuffer is scratch. Route it to a device-resident buffer unconditionally._prepare_decode_inputsbuilt a[B, hidden]bf16 buffer frombatch.hidden_states(a zeros placeholder for device-embedding executors) that the kernel does not consume and no caller reads. Stop building/returning it.DecodeResult.hidden_states/.logitsbecomeOptional;serving_workerreads host logits lazily (only on the non-greedy fallback), soNoneis safe.engine._sample_batch_rowsalready short-circuits on device-sampled ids before touching logits.Verification (NPU, a2a3, greedy
temperature=0)"The capital of France is"→"Paris … Washington", 24 tokens over 23 decode steps through the new path, no errors.next_hiddenis write-only.Scope
🤖 Generated with Claude Code