[cuda backend] reduce memory consumption on gemma4_31b by running embedding in int8#20351
Merged
Conversation
TurboQuantKVCache.update wrote the compressed cache via slice index-assignment (self.k_packed[:, :, input_pos] = ...), which lowers to index_put_ and breaks CUDA-graph capture of the decode step. Switch the 4 cache writes to index_copy_(2, input_pos, ...): a static scatter along the position dim, matching the model's flat global KV cache (Gemma4KVCache.update). Nibble packing is on the last dim (D/2); index_copy_ along dim 2 (positions) leaves the packed layout untouched. MLX overrides update(); qwen shares this path (test_turboquant 5/5).
Builds on the kv_len decode clamp. After that clamp, prefill still scanned ALL KV blocks in [0, chunk_end] for every query row, paying the full causal upper-triangle. Add a per-tile causal upper bound to tq4_sdpa's loop, gated behind a new mask_is_causal op arg (threaded through register_fake, both autotune wrappers, launcher, and body): loop_end = min(kv_len, (kv_len - Lq) + max(seq_pos) + 1). seq_pos is the query-row index; the (kv_len - Lq) offset converts it to an absolute KV position, so it is correct for chunked prefill. gemma's full-attention (global) layers pass mask_is_causal=True; sliding layers are bf16 (untouched) and qwen does not pass the flag (defaults False). Decode-safe: at Lq=1 loop_end == kv_len, so decode is byte-identical. Skipped blocks are fully masked, so prefill is bit-identical to no-skip, just faster.
…ontext) decode/prefill) The fused TQ4 attention kernel iterated the full pre-allocated KV buffer (max_seq_len) every decode/prefill step, masking empty/future blocks only after computing them, making the 10 global layers O(max_seq_len) regardless of actual context. Add an optional kv_len GPU int32 scalar to triton::tq4_sdpa that bounds the inner loop to the valid/filled length. gemma's call site passes kv_len = input_pos[0] + T (decode: pos+1; prefill chunk: chunk_end). kv_len is read on-device via tl.load (no .item(), no host sync) so the bound updates across CUDA-graph replays. When kv_len is None, HAS_KV_LEN is False and the loop runs over full Lk, so the shared qwen path is byte-identical.
…V (cover the 128k decode/prefill path)
…only') TQ4 KV cache long-context is now supported on the CUDA backend, not just MLX. Update the README section: retitle to CUDA + MLX, add a CUDA 128k example, and state that long context requires BOTH --turboquant and a larger --max-seq-len (raising --max-seq-len alone keeps a bf16 KV cache that does not fit at 128k). Note the ~27 GiB runtime peak fits a 32 GB card.
…t export memory Decode the tied GGUF token_embd/lm_head weight to a gatherable int8 IntxUnpackedToInt8Tensor for the embedding instead of dequantizing the whole ~1.4B-element weight to bf16. Halves the embedding's host + GPU-constant footprint (2 -> 1 B/elem). The decode is lossless for Q6_K; for Q4_K the int8 embedding matches the precision of the tied lm_head and the other q4_k linears (same bf16 zero-point reconstruction), so it adds no error beyond the model's existing 4-bit level. lm_head keeps its packed format (Q6_K -> CudaDp4aPlanarInt6Tensor from the same int8 decode; Q4_K -> CudaCoalescedInt4Tensor via _resolve_tied_lm_head). MLX and non-Q4_K/Q6_K paths unchanged.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20351
Note: Links to docs will display an error until the docs builds have been completed. ❌ 3 New Failures, 1 Cancelled Job, 4 Unrelated FailuresAs of commit e322831 with merge base 63b4c4d ( NEW FAILURES - The following jobs have failed:
CANCELLED JOB - The following job was cancelled. Please retry:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
metascroy
reviewed
Jun 17, 2026
| def test_load_converts_weights(self): | ||
| """GGUF -> CUDA: Q4_K -> CudaCoalescedInt4Tensor, Q6_K -> CudaDp4aPlanarInt6Tensor, | ||
| embedding bf16.""" | ||
| embedding int8 (gatherable).""" |
Contributor
There was a problem hiding this comment.
Is cuda using int8 embedding + lm_head tied now?
Before it used int8 embedding + int4 lm_head? How does that change perf?
Contributor
Author
There was a problem hiding this comment.
no it is still int4 lm_head. This PR only switches the GGUF embedding from bf16 to int8 (it was previously dequantized to bf16); lm_head is unchanged.
metascroy
approved these changes
Jun 17, 2026
mergennachin
approved these changes
Jun 18, 2026
mergennachin
changed the title
This PR needs a
|
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
3 participants
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.
Decode the tied GGUF token_embd/lm_head weight to a gatherable int8 IntxUnpackedToInt8Tensor for the embedding instead of dequantizing the whole ~1.4B-element weight to bf16 to reduce the memory consumption during exporation
Halves the embedding's host + GPU-constant footprint (2 -> 1 B/elem) and reduce about 1.5 gb vram memory.
The decode is lossless for Q6_K; for Q4_K the int8 embedding matches the precision of the tied lm_head and the other q4_k linears (same bf16 zero-point reconstruction), so it adds no error beyond the model's existing 4-bit level. Also create test case to verify it.
lm_head keeps its packed format (Q6_K -> CudaDp4aPlanarInt6Tensor from the same int8 decode; Q4_K -> CudaCoalescedInt4Tensor via _resolve_tied_lm_head). MLX and non-Q4_K/Q6_K paths unchanged.
Memory consumption: 28.3 GB with context length == 128k + turboquant