[feature] TQ serving: add TurboQuant prefill/decode branch alongside … - #114
Open
sunghajung6688 wants to merge 1 commit into
Open
[feature] TQ serving: add TurboQuant prefill/decode branch alongside …#114sunghajung6688 wants to merge 1 commit into
sunghajung6688 wants to merge 1 commit into
Conversation
…FP path Wire TurboQuant (TQ) KV-cache quantization onto the refactored serving (turboquant3 = clean main). TQ uses nibble-packed UINT8 KV caches (head_dim//2 bytes/row) + per-row FP32 scales and dispatches the 26-param decode_fwd_tq / prefill_fwd_tq kernels (pypto-lib turboquant @ aae38ce, B1' rotated-space). TQ decode returns raw logits, so supports_device_*=False and the engine samples on the host. Per file: - types.py: KvQuantConfig frozen dataclass + RuntimeConfig.kv_quant_config. - model_runner.py: _KvCachePool quant_k/v + k/v_scales pages (optional); init_kv_cache TQ alloc branch; close frees all 4. - qwen3_l3_dispatch.py: prefill_fwd_tq/decode_fwd_tq bindings + 26-param qwen3_prefill_tq_host / qwen3_decode_tq_host wrappers. - npu_executor.py: tq_mode; supports_device_*=not tq_mode; _compile_tq_kernels (load TQ modules via new _draft.py loader fallback, validate, build rot_matrices [num_layers*head_dim,head_dim] BF16 + Lloyd-Max codebook [1,16] FP32, compile 26-param prefill/decode, CPU sampling); page_size=128 enforced via existing _validate_supported_shape. - npu_runner.py: _CompiledKernels/_StaticKernelArgs TQ fields; tq_mode; _kv_bytes_per_page helper (TQ-aware, replaces 3 inline copies); pool-aware _prefill/_decode_kernel_args with 26-element TQ tuples; rot/codebook + dedicated [128] prefill rope through all 4 static-upload sites; decode hidden_states input buffer; _shared_l3_worker registers TQ programs. - cli/main.py + npu_generate.py: --tq flag threads tq_mode + kv_quant_config. FP path unchanged (non-TQ serving identical); 4-artifact 26-param alignment verified (host wrapper = dummy_args = runtime kernel_args for both kernels).
📝 WalkthroughWalkthroughTurboQuant mode is added to Qwen3-14B configuration and CLI flows, with quantized KV-cache allocation, TQ kernel compilation and dispatch, and mode-specific runner buffers, memory sizing, kernel arguments, and hidden-state staging. ChangesTurboQuant execution
Estimated code review effort: 5 (Critical) | ~120 minutes 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pypto_serving/model/common/runner/model_runner.py`:
- Around line 96-103: Update the TurboQuant allocation cleanup in the try/except
block within _alloc_kv_cache_with_retry so every tensor allocated before a
failure is freed: quant_k_pages, quant_v_pages, k_scales_pages, and
v_scales_pages. Ensure cleanup safely handles tensors whose allocation did not
complete, then re-raise the original exception so halve-and-retry behavior
remains unchanged.
In `@pypto_serving/model/qwen/npu_runner.py`:
- Around line 1271-1291: Add an explicit validation in _prepare_decode_inputs
requiring batch.hidden_states to be present whenever _tq_mode is enabled,
matching DeepSeek’s existing contract. Fail before constructing or passing
_DecodeKernelInputs to the TQ decode kernel, while preserving the current non-TQ
and valid TQ paths.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 64b275ca-a11d-4a95-862c-d04affbc5456
📒 Files selected for processing (8)
examples/model/qwen3_14b/npu_generate.pypypto_serving/__init__.pypypto_serving/cli/main.pypypto_serving/config/types.pypypto_serving/model/common/runner/model_runner.pypypto_serving/model/qwen/npu_executor.pypypto_serving/model/qwen/npu_runner.pypypto_serving/model/qwen/qwen3_l3_dispatch.py
Comment on lines
+96
to
+103
| quant_k_pages = self._alloc_kv_cache_tensor(quant_shape, torch.uint8) | ||
| try: | ||
| quant_v_pages = self._alloc_kv_cache_tensor(quant_shape, torch.uint8) | ||
| k_scales_pages = self._alloc_kv_cache_tensor(scale_shape, torch.float32) | ||
| v_scales_pages = self._alloc_kv_cache_tensor(scale_shape, torch.float32) | ||
| except Exception: | ||
| self._free_kv_cache_tensor(quant_k_pages) | ||
| raise |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Resource leak on TurboQuant allocation failure. The try block allocates three tensors (quant_v_pages, k_scales_pages, v_scales_pages) but the except frees only quant_k_pages. If k_scales_pages or v_scales_pages allocation raises, the already-allocated quant_v_pages/k_scales_pages leak. This path is reached during _alloc_kv_cache_with_retry's halve-and-retry-on-OOM loop, so the leaked device memory is never reclaimed and each retry starts with less free HBM — defeating the retry and possibly failing at the floor even when a smaller page count would have fit.
🛡️ Proposed fix to free all partially-allocated tensors
quant_k_pages = self._alloc_kv_cache_tensor(quant_shape, torch.uint8)
+ quant_v_pages = k_scales_pages = None
try:
quant_v_pages = self._alloc_kv_cache_tensor(quant_shape, torch.uint8)
k_scales_pages = self._alloc_kv_cache_tensor(scale_shape, torch.float32)
v_scales_pages = self._alloc_kv_cache_tensor(scale_shape, torch.float32)
except Exception:
- self._free_kv_cache_tensor(quant_k_pages)
+ for tensor in (quant_k_pages, quant_v_pages, k_scales_pages):
+ if tensor is not None:
+ self._free_kv_cache_tensor(tensor)
raise📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| quant_k_pages = self._alloc_kv_cache_tensor(quant_shape, torch.uint8) | |
| try: | |
| quant_v_pages = self._alloc_kv_cache_tensor(quant_shape, torch.uint8) | |
| k_scales_pages = self._alloc_kv_cache_tensor(scale_shape, torch.float32) | |
| v_scales_pages = self._alloc_kv_cache_tensor(scale_shape, torch.float32) | |
| except Exception: | |
| self._free_kv_cache_tensor(quant_k_pages) | |
| raise | |
| quant_k_pages = self._alloc_kv_cache_tensor(quant_shape, torch.uint8) | |
| quant_v_pages = k_scales_pages = None | |
| try: | |
| quant_v_pages = self._alloc_kv_cache_tensor(quant_shape, torch.uint8) | |
| k_scales_pages = self._alloc_kv_cache_tensor(scale_shape, torch.float32) | |
| v_scales_pages = self._alloc_kv_cache_tensor(scale_shape, torch.float32) | |
| except Exception: | |
| for tensor in (quant_k_pages, quant_v_pages, k_scales_pages): | |
| if tensor is not None: | |
| self._free_kv_cache_tensor(tensor) | |
| raise |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pypto_serving/model/common/runner/model_runner.py` around lines 96 - 103,
Update the TurboQuant allocation cleanup in the try/except block within
_alloc_kv_cache_with_retry so every tensor allocated before a failure is freed:
quant_k_pages, quant_v_pages, k_scales_pages, and v_scales_pages. Ensure cleanup
safely handles tensors whose allocation did not complete, then re-raise the
original exception so halve-and-retry behavior remains unchanged.
Comment on lines
+1271
to
+1291
| # TurboQuant decode consumes the input embedding (hidden_states) directly | ||
| # instead of embedding inside the kernel; stage the host-provided | ||
| # ``batch.hidden_states`` into the persistent device buffer (padding rows | ||
| # replicate row 0 so their KV writes stay idempotent). | ||
| hidden_arg = None | ||
| if self._tq_mode and batch.hidden_states is not None and compiled.decode_hidden_buffer is not None: | ||
| hidden_arg = compiled.decode_hidden_buffer | ||
| hidden_arg[:actual_batch].copy_(batch.hidden_states) | ||
| if actual_batch < kernel_batch: | ||
| hidden_arg[actual_batch:].copy_( | ||
| batch.hidden_states[0:1].expand(kernel_batch - actual_batch, -1) | ||
| ) | ||
|
|
||
| return _DecodeKernelInputs( | ||
| actual_batch=actual_batch, | ||
| token_ids=token_ids, | ||
| seq_lens=seq_lens, | ||
| block_table=block_table, | ||
| slot_mapping=slot_mapping, | ||
| logits=compiled.decode_logits_buffer, | ||
| hidden=hidden_arg, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find where DecodeBatch.hidden_states is populated and whether it can be None in TQ mode.
rg -nP --type=py -C3 '\bhidden_states\b' -g '!**/model/qwen/npu_runner.py' | rg -n 'DecodeBatch|hidden_states\s*='
rg -nP --type=py -C3 'class DecodeBatch\b'Repository: hw-native-sys/pypto-serving
Length of output: 165
🏁 Script executed:
#!/bin/bash
set -u
echo "Files:"
git ls-files | rg '(^|/)npu_runner\.py$|py\.py$' || true
echo
echo "Search DecodeBatch definitions/usages across repo:"
rg -n -C3 '\bDecodeBatch\b|class DecodeBatch\b|hidden_states\s*=' . || true
echo
echo "Show npu_runner relevant sections if present:"
if [ -f pypto_serving/model/qwen/npu_runner.py ]; then
wc -l pypto_serving/model/qwen/npu_runner.py
rg -n -C4 '_prepare_decode_inputs|_decode_kernel_args|_tq_mode|supports_device_embedding|hidden_states|decode_hidden_buffer' pypto_serving/model/qwen/npu_runner.py
fiRepository: hw-native-sys/pypto-serving
Length of output: 49992
🏁 Script executed:
#!/bin/bash
set -u
echo "DecodeBatch definition:"
sed -n '260,345p' pypto_serving/config/types.py | cat -n
echo
echo "Engine decode construction context:"
sed -n '330,385p' pypto_serving/serving/engine/engine.py | cat -n
echo
echo "Server decode construction context:"
sed -n '330,395p' pypto_serving/serving/server/serving_worker.py | cat -n
echo
echo "Qwen runner decode input preparation context:"
sed -n '1170,1293p' pypto_serving/model/qwen/npu_runner.py | cat -n
echo
echo "DeepSeek runner hidden states requirements context:"
sed -n '1578,1600p' pypto_serving/model/deepseek/npu_runner.py | cat -nRepository: hw-native-sys/pypto-serving
Length of output: 18166
Add an explicit TQ decode hidden_states guard. _DecodeBatch.hidden_states is optional (torch.Tensor | None) and can be None from the device-embedding worker path; the Qwen runner leaves inputs.hidden as None in TQ mode but _decode_kernel_args passes it as the TQ hidden_states parameter. Add a failing guard in _prepare_decode_inputs, matching the same contract DeepSeek enforces, before passing the batch to the TQ decode kernel.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pypto_serving/model/qwen/npu_runner.py` around lines 1271 - 1291, Add an
explicit validation in _prepare_decode_inputs requiring batch.hidden_states to
be present whenever _tq_mode is enabled, matching DeepSeek’s existing contract.
Fail before constructing or passing _DecodeKernelInputs to the TQ decode kernel,
while preserving the current non-TQ and valid TQ paths.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR: TurboQuant (TQ) KV-Cache Quantization Branch
Branch:
turboquant3→mainBase:
c217603(latest main)Scope: 8 files, +873 / −114
Summary
Wires TurboQuant (TQ) online KV-cache quantization onto the refactored serving stack as a branch alongside the existing FP (BF16) path. TQ is opt-in via
--tq; with the flag off, serving is byte-for-byte identical to today's non-TQ path.TQ replaces the default BF16 KV caches with:
head_dim // 2bytes/row), andand dispatches the 26-parameter
prefill_fwd_tq/decode_fwd_tqkernels (pypto-libturboquant@aae38ce, B1′ rotated-space). TQ decode returns raw logits, sosupports_device_*isFalsein TQ mode and the engine samples on the host.Design: TQ as a parallel branch, not a replacement
The FP path is untouched. Every TQ addition is gated behind
tq_mode/kv_quant_config:With
--tqabsent, every one of those branches takes theelseand behaves exactly as onmain.Changes per file
pypto_serving/config/types.py—KvQuantConfigfrozen dataclass +RuntimeConfig.kv_quant_config(KvQuantConfig | None = None).pypto_serving/model/common/runner/model_runner.py—_KvCachePoolgains optionalquant_k/v+k/v_scalespages;init_kv_cachegains a TQ allocation branch;closefrees all four.pypto_serving/model/qwen/qwen3_l3_dispatch.py—prefill_fwd_tq/decode_fwd_tqbindings + 26-parameterqwen3_prefill_tq_host/qwen3_decode_tq_hosthost wrappers.pypto_serving/model/qwen/npu_executor.py—tq_mode;supports_device_* = not tq_mode;_compile_tq_kernels(load TQ modules with the new loader fallback, validate shapes, buildrot_matrices [num_layers*head_dim, head_dim]BF16 + Lloyd-Max codebook[1, 16]FP32, compile the 26-param prefill/decode callables, CPU sampling).page_size = 128is enforced via the existing_validate_supported_shape.pypto_serving/model/qwen/npu_runner.py— TQ fields on_CompiledKernels/_StaticKernelArgs;tq_mode; TQ-aware_kv_bytes_per_page(replaces 3 inline copies); pool-aware_prefill/_decode_kernel_argswith 26-element TQ tuples; rot/codebook + dedicated[128]prefill rope threaded through all 4 static-upload sites; decodehidden_statesinput buffer;_shared_l3_workerregisters the TQ programs.pypto_serving/cli/main.py+examples/model/qwen3_14b/npu_generate.py—
--tq/--tq-modeflag threadstq_mode+kv_quant_config.pypto_serving/__init__.py— exportKvQuantConfig.Usage
Offline generate
python examples/model/qwen3_14b/npu_generate.py \ --model /path/to/Qwen3-14B \ --platform a2a3 \ --tq \ ... # remaining args as usualServing CLI
Drop
--tqto run the unchanged FP path.Verification
--tqoff, non-TQ serving is identical tomain.host wrapper signature ==
dummy_args== runtimekernel_args.decode_fwd_tqfixed-shape constraints (BATCH,NUM_LAYERS,VOCAB,REAL_VOCAB) are validated against runtime config in_validate_tq_shape.Coexistence with main
This branch rebases cleanly onto the latest
main(c217603), which added the DeepSeek-V4KVCacheSpec/KVCacheGroupSpec/num_speculative_tokens/kv_cache_groupsmachinery. The two features are independent and both are present in the merged result:RuntimeConfigcarrieskv_cache_groups(DeepSeek-V4 cache families) andkv_quant_config(TQ) side by side. TQ currently targets Qwen3-14B only; the DeepSeek-V4 path is unaffected.Closes #20