-
Notifications
You must be signed in to change notification settings - Fork 2.8k
[None][feat] Add calibrated INT8 KV cache to PyTorch backend #18953
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -699,6 +699,13 @@ def __init__( | |
| self.create_weights() | ||
|
|
||
| def create_weights(self): | ||
| if (self.quant_config is not None | ||
| and self.quant_config.layer_quant_mode.has_int8_kv_cache() | ||
| and (self.attn_backend.upper() != "TRTLLM" | ||
| or self.mapping.cp_size > 1)): | ||
|
Comment on lines
+704
to
+705
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win Add construction-time coverage for context-parallelism rejection.
🤖 Prompt for AI Agents |
||
| raise ValueError( | ||
| "INT8 KV cache requires TRTLLM attention without context parallelism." | ||
| ) | ||
| # self.attn has no weights but has states that are related to quant_config, | ||
| # which could be modified after __init__ | ||
| self.attn.update_quant_config(self.quant_config) | ||
|
|
@@ -895,6 +902,10 @@ def _attn_impl( | |
| ): | ||
| kv_scale_orig_quant = self.qkv_proj.inv_kv_scales | ||
| kv_scale_quant_orig = self.qkv_proj.kv_scales | ||
| elif (self.quant_config is not None | ||
| and self.quant_config.layer_quant_mode.has_int8_kv_cache()): | ||
| kv_scale_orig_quant = self.qkv_proj.inv_kv_cache_scaling_factor | ||
| kv_scale_quant_orig = self.qkv_proj.kv_cache_scaling_factor | ||
|
|
||
| attn_output = self.attn.forward( | ||
| q, | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -1531,6 +1531,18 @@ def __init__( | |
| def update_quant_config(self, new_quant_config: Optional[QuantConfig]): | ||
| self.quant_config = new_quant_config or QuantConfig() | ||
| self.quant_mode = int(self.quant_config.layer_quant_mode) | ||
| self.has_int8_kv_cache = self.quant_config.layer_quant_mode.has_int8_kv_cache( | ||
| ) | ||
| if (self.has_int8_kv_cache | ||
| and self.quant_config.layer_quant_mode.has_any_quant( | ||
| exclude_kv_cache=True)): | ||
| raise ValueError( | ||
| "INT8 KV cache currently requires unquantized FP16/BF16 projections." | ||
| ) | ||
| if self.has_int8_kv_cache and (self.is_mla_enable | ||
| or self.sparse_params is not None): | ||
| raise ValueError( | ||
| "INT8 KV cache does not support MLA or sparse attention.") | ||
|
|
||
| self.has_fp8_qdq = self.has_fp8_kv_cache = self.has_nvfp4 = False | ||
| if self.quant_config is not None: | ||
|
|
@@ -1819,6 +1831,32 @@ def forward( | |
| metadata, | ||
| TrtllmAttentionMetadata, | ||
| ) | ||
| if self.has_int8_kv_cache: | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This validation block runs on every forward — per layer, per step in eager mode — but most of the checks are invariants of the configuration (q dtype, scale dtype/device/numel, use_cache), not of the batch. Only |
||
| if q.dtype not in (torch.float16, torch.bfloat16): | ||
| raise ValueError( | ||
| "INT8 KV cache requires FP16 or BF16 attention inputs.") | ||
| if metadata.kv_cache_params is None or not metadata.kv_cache_params.use_cache: | ||
| raise ValueError("INT8 KV cache requires an active KV cache.") | ||
| # Decode-only steps have no context prefixes to inspect. Tensor and | ||
| # cache guards remain per-call: backend callers may replace them. | ||
| cached_context = metadata.num_contexts > 0 and any( | ||
| n > 0 for n in metadata.kv_cache_params. | ||
| num_cached_tokens_per_seq[:metadata.num_contexts]) | ||
| if metadata.is_cross or metadata.enable_helix: | ||
| raise ValueError( | ||
| "INT8 KV cache does not support cross-attention or context parallelism." | ||
| ) | ||
|
Comment on lines
+1835
to
+1848
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win Add coverage for all INT8 validation branches.
🤖 Prompt for AI Agents |
||
| if metadata.use_paged_context_fmha or cached_context: | ||
| raise ValueError( | ||
| "INT8 KV cache does not support paged context attention or cached prefill." | ||
| ) | ||
| for scale in (forward_args.kv_scale_orig_quant, | ||
| forward_args.kv_scale_quant_orig): | ||
| if (scale is None or scale.dtype != torch.float32 | ||
| or scale.device != q.device or scale.numel() != 1): | ||
| raise ValueError( | ||
| "INT8 KV cache requires scalar float32 KV scales " | ||
| "on the attention input device.") | ||
| # Cross-attention uses the THOP path; the trtllm-gen backend API does | ||
| # not carry encoder K/V tensors yet. | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2465,7 +2465,29 @@ def _create_kv_cache_manager( | |
| # use cache_layer_idx to read from the target layer's cache slot via | ||
| # Gemma4Attention. No layer_mask exclusion needed here. | ||
|
|
||
| if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache(): | ||
| if quant_config is not None and quant_config.quant_mode.has_int8_kv_cache(): | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The rejection list covers block reuse, spec decode, and chunked prefill, but not disaggregated serving: |
||
| if is_disagg or kv_connector_manager is not None: | ||
| raise ValueError( | ||
| "INT8 KV cache does not support disaggregated serving or KV connectors." | ||
| ) | ||
| if is_hybrid_linear(config) or _model_config.is_encoder_decoder: | ||
| raise ValueError( | ||
| "INT8 KV cache currently supports dense decoder-only models.") | ||
| if kv_cache_config.enable_block_reuse: | ||
| raise ValueError( | ||
| "INT8 KV cache requires kv_cache_config.enable_block_reuse=False; " | ||
| "paged context attention does not support INT8 KV cache.") | ||
| if spec_config is not None: | ||
| raise ValueError( | ||
| "INT8 KV cache does not support speculative decoding.") | ||
| if (model_engine is not None | ||
| and model_engine.attn_runtime_features.chunked_prefill): | ||
| raise ValueError( | ||
| "INT8 KV cache requires enable_chunked_prefill=False; " | ||
| "paged context attention does not support INT8 KV cache.") | ||
| kv_cache_dtype = tensorrt_llm.bindings.DataType.INT8 | ||
| elif quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache( | ||
| ): | ||
| kv_cache_dtype = tensorrt_llm.bindings.DataType.FP8 | ||
| elif quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache( | ||
| ): | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This initializes
cu_kv_seqlensfor every PACKED_QKV launch — all FP16/BF16 context attention, not just INT8 KV. The aliasing tocuQSeqLenPtris correct for this layout (kv_len == q_len), and setting a previously uninitialized pointer is strictly safer, but it's a standalone bugfix to a shared kernel path buried inside a feature PR. Consider splitting it into its own PR with a dedicated test so it can land, be bisected, and be reverted independently of the INT8 feature; at minimum call it out for the FMHA owners and confirm multi-arch CI covers it.