diff --git a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp index ca7e920ce7fa..2f4c33ae0a21 100644 --- a/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp +++ b/cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2025, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2020-2026, NVIDIA CORPORATION. All rights reserved. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -146,6 +146,8 @@ void FusedMHARunnerV2::setupKernelParams(MHARunnerParams runnerParams) { // Packed QKV input layout, [B, S, H * D + H_kv * D + H_kv * Dv]. mKernelParams.qkv_ptr = runnerParams.qkvPtr; + // Packed-QKV kernels also read the cumulative KV lengths. + mKernelParams.cu_kv_seqlens = reinterpret_cast(runnerParams.cuQSeqLenPtr); mKernelParams.q_stride_in_bytes = mKernelParams.k_stride_in_bytes = mKernelParams.v_stride_in_bytes = get_size_in_bytes(mFixedParams.numQHeads * mFixedParams.headSize + mFixedParams.numKvHeads * mFixedParams.headSize diff --git a/docs/source/torch/features/quantization.md b/docs/source/torch/features/quantization.md index 47cc745165b4..3ec8478dfb3a 100644 --- a/docs/source/torch/features/quantization.md +++ b/docs/source/torch/features/quantization.md @@ -16,3 +16,51 @@ git clone https://github.com/NVIDIA/Model-Optimizer.git cd Model-Optimizer/examples/llm_ptq scripts/huggingface_example.sh --model --quant fp8 --export_fmt hf ``` + +## INT8 KV cache + +Dense decoder-only models with FP16 or BF16 weights can store their KV cache in +INT8 using the `TRTLLM` attention backend. The KV tensors use one byte per element, +half the storage of an FP16/BF16 cache with the same number of elements. Allocator +granularity can prevent small pools from reserving exactly half the GPU memory. +Model weights and other runtime allocations are unchanged. + +The checkpoint must contain calibrated, positive scalar dequantization scales +`model.layers..self_attn.k_proj.k_scale` and +`model.layers..self_attn.v_proj.v_scale` for each attention layer (using the +model's corresponding projection paths). The native kernels share one FP32 scale +between K and V, so the loader uses the larger of the two values. Values are stored +as `clamp(round(x / scale), -128, 127)` and read back as `value * scale`. Full +checkpoint loading fails if either scale is missing or invalid. Weight-only +partial reloads preserve the existing scales; scale updates must supply both K +and V scales together. Large differences between the K and V ranges can cause +substantial accuracy loss with a shared scale. Calibrate for INT8's range, +including the effect of rotary +position embeddings on K; FP8 or FP4 scales cannot be reused unchanged. +`TRTLLM_LOAD_KV_SCALES` must remain enabled (`1`, the default) for INT8; +disabling calibrated scale loading raises an error. + +```python +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import KvCacheConfig + +with LLM( + model="/path/to/checkpoint-with-int8-kv-scales", + backend="pytorch", + dtype="bfloat16", + attn_backend="TRTLLM", + enable_chunked_prefill=False, + kv_cache_config=KvCacheConfig(dtype="int8", enable_block_reuse=False), +) as llm: + output = llm.generate("Explain the role of a KV cache.", SamplingParams(max_tokens=64)) +``` + +This path supports full prefill followed by single-token decoding with paged KV +storage. It does not support cached/chunked prefill, prefix block reuse, +speculative decoding, MLA, sparse or cross attention, context parallelism, +hybrid attention/state-space models, disaggregated serving, KV connectors, +or quantized projection weights. These +combinations are rejected because they require different kernel paths. Select +`max_num_tokens` large enough to hold the longest complete prompt when chunked +prefill is disabled. Validate model quality with representative calibration and +evaluation data before using the quantized cache in a deployment. diff --git a/tensorrt_llm/_torch/attention/attention.py b/tensorrt_llm/_torch/attention/attention.py index 16f928721f58..7dc24137fe6c 100644 --- a/tensorrt_llm/_torch/attention/attention.py +++ b/tensorrt_llm/_torch/attention/attention.py @@ -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)): + 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, diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 5a5715b1d273..2792873a3a91 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -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: + 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." + ) + 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. diff --git a/tensorrt_llm/_torch/attention/backends/utils.py b/tensorrt_llm/_torch/attention/backends/utils.py index 67dbfe38eb6b..0a51c3085fcc 100644 --- a/tensorrt_llm/_torch/attention/backends/utils.py +++ b/tensorrt_llm/_torch/attention/backends/utils.py @@ -82,6 +82,10 @@ def create_attention( if attention_chunk_size is not None and backend_name.upper() != "TRTLLM": raise ValueError( f"Backend {backend_name} does not support chunked attention.") + if (quant_config is not None + and quant_config.layer_quant_mode.has_int8_kv_cache() + and backend_name.upper() != "TRTLLM"): + raise ValueError("INT8 KV cache requires the TRTLLM attention backend.") attn_cls = get_attention_backend(backend_name, sparse_params=sparse_params) if is_mla_enable: diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 1286c619f524..3444518dbcd1 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -640,6 +640,7 @@ def load_weights_fused_qkv_linear( module: Linear, weights: List[Dict], allow_partial_loading: bool = False) -> None: + module._load_kv_cache_scales(weights, allow_partial_loading) q_weight, k_weight, v_weight = load_weights_fused_qkv_helper( module, weights, allow_partial_loading=allow_partial_loading) if not allow_partial_loading: @@ -655,24 +656,6 @@ def load_weights_fused_qkv_linear( copy_weight_shard(module.weight, weight, shard_offset, shard_size) - if hasattr(module, "kv_scales") and os.environ.get( - "TRTLLM_LOAD_KV_SCALES", "1") == "1": - k_scales = [ - w["k_scale"][...].reshape([]) for w in weights if "k_scale" in w - ] - v_scales = [ - w["v_scale"][...].reshape([]) for w in weights if "v_scale" in w - ] - if k_scales: - assert v_scales, "k_scale and v_scale must be loaded together" - copy_weight( - module.kv_scales, - torch.tensor( - [1.0, max(k_scales).item(), - max(v_scales).item()], - dtype=torch.float32)) - module.inv_kv_scales.data = 1.0 / module.kv_scales - def load_weights_fused_gate_up_linear( self, module: Linear, @@ -3919,6 +3902,18 @@ def create_weights(self): self.out_features, self.has_bias, self.dtype) + if (self.quant_config is not None + and self.quant_config.layer_quant_mode.has_int8_kv_cache() + and self.weights_loading_config.weight_mode + == WeightMode.FUSED_QKV_LINEAR): + # INT8 kernels share one per-layer scale between K and V. + self.kv_cache_scaling_factor = Parameter(torch.ones( + 1, dtype=torch.float32), + requires_grad=False) + self.inv_kv_cache_scaling_factor = Parameter(torch.ones( + 1, dtype=torch.float32), + requires_grad=False) + self._weights_created = True self._weights_transformed = False @@ -4114,6 +4109,71 @@ def load_weights(self, weight_mode, allow_partial_loading=allow_partial_loading) + def _load_kv_cache_scales(self, weights: list[dict], + allow_partial_loading: bool) -> None: + """Load unquantized fused-QKV cache scales at one checkpoint hook. + + INT8 requires calibrated paired scalars; FP4 retains its optional + per-K/V scales and the existing environment opt-out. INT8 rejects that + opt-out rather than silently running with uncalibrated unity scales. + Weight-only partial reloads leave the existing scale parameters intact. + """ + int8_cache = (self.quant_config is not None and + self.quant_config.layer_quant_mode.has_int8_kv_cache()) + if not int8_cache and not hasattr(self, "kv_scales"): + return + if os.environ.get("TRTLLM_LOAD_KV_SCALES", "1") != "1": + if int8_cache: + raise ValueError( + "INT8 KV cache requires TRTLLM_LOAD_KV_SCALES=1 " + "to load calibrated scales.") + return + scales = { + name: [w[name][...] for w in weights if name in w] + for name in ("k_scale", "v_scale") + } + if not int8_cache: + # Preserve the existing optional FP4 scale-loading contract. + if scales["k_scale"]: + assert scales[ + "v_scale"], "k_scale and v_scale must be loaded together" + k_scale = max(scale.reshape([]) + for scale in scales["k_scale"]).item() + v_scale = max(scale.reshape([]) + for scale in scales["v_scale"]).item() + scale = self.kv_scales.new_tensor([1.0, k_scale, v_scale]) + copy_weight(self.kv_scales, scale) + copy_weight(self.inv_kv_scales, scale.reciprocal()) + return + if allow_partial_loading and not any(scales.values()): + return + if not all(scales.values()): + raise ValueError( + "INT8 KV cache requires calibrated k_scale and v_scale " + "checkpoint tensors; both scales must be loaded together.") + values = [] + for tensors in scales.values(): + for scale in tensors: + if scale.numel() != 1: + raise ValueError( + "INT8 KV cache scales must be scalar tensors.") + value = scale.item() + if not math.isfinite(value) or value <= 0: + raise ValueError( + "INT8 KV cache scales must be finite and positive.") + values.append(value) + # The largest dequantization scale covers both calibrated ranges. + # Copy into stable parameters so weight reloads preserve graph pointers. + scale = self.kv_cache_scaling_factor.new_tensor([max(values)]) + inverse = scale.reciprocal() + if not (torch.isfinite(scale).all() and torch.isfinite(inverse).all() + and (scale > 0).all() and (inverse > 0).all()): + raise ValueError( + "INT8 KV cache scale and reciprocal must be finite and " + "positive in float32.") + copy_weight(self.kv_cache_scaling_factor, scale) + copy_weight(self.inv_kv_cache_scaling_factor, inverse) + def process_weights_after_loading(self): self.quant_method.process_weights_after_loading(self) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7e31ad99cff7..bda3fcf43294 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -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(): + 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( ): diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 97d117448bb9..d2c06da7fe93 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -453,7 +453,9 @@ def _get_static_cache_size_layer_components( cache_size_per_token = kv_factor * head_dim quant_config = model_config.quant_config - 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_fp8_kv_cache() or quant_config.quant_mode.has_int8_kv_cache() + ): layer_size = cache_size_per_token elif quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache(): layer_size = math.ceil(cache_size_per_token / 2) + math.ceil(cache_size_per_token / 16) @@ -4004,6 +4006,7 @@ def get_cache_bytes_per_token(self) -> int: def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): if self.dtype not in ( + DataType.INT8, DataType.FP8, DataType.HALF, DataType.BF16, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index a7ee276f387b..b8305abafcff 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -55,12 +55,13 @@ resolve_ssm_cache_dtype) _KV_CACHE_MAP = { + "int8": QuantAlgo.INT8.value, "fp8": QuantAlgo.FP8.value, "fp8_ds_mla": QuantAlgo.FP8.value, "nvfp4": QuantAlgo.NVFP4.value, "auto": "auto" } -_VALID_KV_CACHE_DTYPES = ("fp8", "fp8_ds_mla", "nvfp4", "auto") +_VALID_KV_CACHE_DTYPES = ("int8", "fp8", "fp8_ds_mla", "nvfp4", "auto") # Dense models do not consume the MoE backend or MoE mapping dimensions. Their # runtime topology remains bounded by the general and attention dimensions. @@ -298,6 +299,8 @@ def _get_random_min_max(dtype: torch.dtype) -> Tuple[int, int]: ".inv_input_scale", ".kv_scales", ".inv_kv_scales", + ".kv_cache_scaling_factor", + ".inv_kv_cache_scaling_factor", ".alpha", ".scalar_alpha", ) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 6077f62d1ae8..5e8b464e70b5 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -1589,8 +1589,9 @@ def get_cache_size_per_token(model_config: ModelConfigPython, mem_per_token = kv_factor * num_attention_layers * head_dim # The data type bytes. quant_config = model_config.quant_config - 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_fp8_kv_cache() + or quant_config.quant_mode.has_int8_kv_cache()): mem_per_token *= 1 elif quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache( ): @@ -1615,8 +1616,8 @@ def get_cache_bytes_per_token(self): cache_size_per_token = self.kv_factor * sum( self.num_kv_heads_per_layer) * self.head_dim - if self.dtype not in (DataType.FP8, DataType.HALF, DataType.BF16, - DataType.FLOAT, DataType.NVFP4): + if self.dtype not in (DataType.INT8, DataType.FP8, DataType.HALF, + DataType.BF16, DataType.FLOAT, DataType.NVFP4): raise ValueError(f'Cannot support {self.dtype} KV cache.') cache_size_bytes_per_token = get_size_in_bytes(cache_size_per_token, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 57843f4caf91..4a47c73e6711 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4190,14 +4190,14 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description= "The data type for the KV cache. 'auto' (default) leaves the checkpoint's " "own KV-cache quantization metadata untouched (quant_config.kv_cache_quant_algo " - "is inherited as-is); 'fp8', 'fp8_ds_mla', or 'nvfp4' override it explicitly. " + "is inherited as-is); 'int8', 'fp8', 'fp8_ds_mla', or 'nvfp4' override it explicitly. " "'fp8_ds_mla' selects the packed FP8 cache used by sparse MLA on SM90/SM120/SM121. " "Resolved at " "LLM-construction time, including when set via trtllm-serve " "--extra_llm_api_options.", telemetry=TelemetryField.categorical("auto", "float16", "bfloat16", - "float32", "fp8", "fp8_ds_mla", - "nvfp4")) + "float32", "int8", "fp8", + "fp8_ds_mla", "nvfp4")) # This is a pure python field, not a pybind field. It is only for the Pytorch backend. mamba_ssm_cache_dtype: Literal[ @@ -6641,6 +6641,8 @@ def sync_quant_config_with_kv_cache_config_dtype(self) -> 'TorchLlmArgs': assert self.quant_config is not None if self.kv_cache_config.dtype == "auto": return self + elif self.kv_cache_config.dtype == 'int8': + self.quant_config.kv_cache_quant_algo = QuantAlgo.INT8 elif self.kv_cache_config.dtype in ('fp8', 'fp8_ds_mla'): self.quant_config.kv_cache_quant_algo = QuantAlgo.FP8 elif self.kv_cache_config.dtype == 'nvfp4': diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index c1d613bb408a..830f4700cd2b 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -755,6 +755,7 @@ "float16", "bfloat16", "float32", + "int8", "fp8", "fp8_ds_mla", "nvfp4" diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 4224796b6b6a..909427267837 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -14,6 +14,10 @@ l0_a10: backend: pytorch tests: # ------------- PyTorch tests --------------- + - unittest/_torch/attention/test_int8_kv_cache.py + - unittest/_torch/attention/test_packed_qkv_fmha.py + - unittest/_torch/test_int8_kv_scales.py + - unittest/_torch/test_int8_kv_llm.py - unittest/_torch/sampler/test_torch_sampler.py - unittest/_torch/sampler/test_penalties.py - unittest/_torch/test_tensor_lru_cache.py diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 407c04da0168..e36399d6e558 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -31,6 +31,7 @@ l0_b200: backend: pytorch tests: # ------------- PyTorch tests --------------- + - unittest/_torch/attention/test_packed_qkv_fmha.py - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_transfer_stress - unittest/others/test_lora_manager.py - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4 diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 970be6e80ff2..a39b45cce964 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -51,6 +51,7 @@ l0_cpu: - unittest/_torch/multimodal - unittest/_torch/ray_orchestrator/single_gpu/test_cache_transceiver_comm.py - unittest/_torch/speculative/hw_agnostic + - unittest/_torch/test_int8_kv_config.py - unittest/_torch/test_mmap_utils.py - unittest/_torch/test_model_config.py - unittest/_torch/test_utils.py diff --git a/tests/unittest/_torch/attention/test_int8_kv_cache.py b/tests/unittest/_torch/attention/test_int8_kv_cache.py new file mode 100644 index 000000000000..040c51deb67b --- /dev/null +++ b/tests/unittest/_torch/attention/test_int8_kv_cache.py @@ -0,0 +1,402 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exercise INT8 cache writes and reads through the native attention backend.""" + +from dataclasses import replace + +import pytest +import torch +from backend_case import ( + BackendCase, + _assert_cache_contains_new_tokens, + _build_kv_cache_manager, + generate_inputs, +) + +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +_PROMPT_LENS = [63, 31] +_DECODE_STEPS = 3 + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _quantize(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + # Native INT8 conversion uses round-to-nearest-even with signed saturation. + """Match native signed INT8 rounding and saturation for exact cache assertions.""" + return (x.float() / scale).round().clamp(-128, 127).to(torch.int8) + + +def _metadata( + case: BackendCase, manager: KVCacheManager, prompt_lens: list[int] +) -> TrtllmAttentionMetadata: + """Prepare actual cache-manager metadata for a context or decode batch.""" + metadata = TrtllmAttentionMetadata( + num_contexts=case.num_contexts, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=case.num_cached_tokens + ), + seq_lens=torch.tensor(case.seq_lens, dtype=torch.int32), + max_num_requests=case.num_seqs, + max_num_tokens=case.max_num_tokens, + kv_cache_manager=manager, + request_ids=list(range(case.num_seqs)), + prompt_lens=prompt_lens, + kv_layout="HND", + ) + metadata.prepare() + assert not metadata.use_paged_context_fmha + return metadata + + +def _backend(case: BackendCase) -> TrtllmAttention: + """Construct the production TRTLLM backend with INT8 KV quantization enabled.""" + return create_attention( + "TRTLLM", + layer_idx=0, + num_heads=case.num_heads, + num_kv_heads=case.num_kv_heads, + head_dim=case.head_dim, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.INT8), + ) + + +def _forward( + attention: TrtllmAttention, + metadata: TrtllmAttentionMetadata, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: torch.Tensor, +) -> torch.Tensor: + """Pass packed QKV and explicit reciprocal scales through native attention.""" + forward_args = AttentionForwardArgs( + attention_mask=PredefinedAttentionMask.CAUSAL, + kv_scale_orig_quant=scale.reciprocal(), + kv_scale_quant_orig=scale, + ) + qkv = torch.cat((q, k, v), dim=-1) + result = attention.forward(qkv, None, None, metadata, forward_args=forward_args) + if isinstance(result, tuple): + result = result[0] + return result[: q.shape[0]] + + +def _reference( + case: BackendCase, + q: torch.Tensor, + keys: list[torch.Tensor], + values: list[torch.Tensor], +) -> torch.Tensor: + """Compute causal attention in float32 over the explicit logical cache.""" + outputs = [] + offset = 0 + repeats = case.num_heads // case.num_kv_heads + for q_len, cached_len, key, value in zip( + case.seq_lens, case.num_cached_tokens, keys, values, strict=True + ): + query = q[offset : offset + q_len].view(q_len, case.num_heads, case.head_dim) + key = key.repeat_interleave(repeats, dim=1) + value = value.repeat_interleave(repeats, dim=1) + scores = query.float().transpose(0, 1) @ key.float().permute(1, 2, 0) + scores *= case.head_dim**-0.5 + q_positions = torch.arange(q_len, device=q.device) + cached_len + kv_positions = torch.arange(key.shape[0], device=q.device) + scores.masked_fill_(kv_positions[None, :] > q_positions[:, None], -torch.inf) + output = scores.softmax(dim=-1) @ value.float().transpose(0, 1) + outputs.append(output.transpose(0, 1).reshape(q_len, -1)) + offset += q_len + return torch.cat(outputs) + + +def _assert_cache( + case: BackendCase, + manager: KVCacheManager, + keys: list[torch.Tensor], + values: list[torch.Tensor], +) -> None: + """Compare the complete logical K/V cache against expected integer contents.""" + _assert_cache_contains_new_tokens( + manager, + 0, + list(range(case.num_seqs)), + case.token_nums, + [0] * case.num_seqs, + [torch.stack((key, value)) for key, value in zip(keys, values, strict=True)], + kv_layout="HND", + cache_kind="kv", + ) + + +@pytest.mark.parametrize("use_v2", [False, True], ids=["v1", "v2"]) +@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) +@pytest.mark.parametrize("num_kv_heads", [4, 2], ids=["mha", "gqa"]) +@pytest.mark.parametrize("scale_value", [1 / 32, 1 / 16], ids=["scale32", "scale16"]) +@torch.inference_mode() +def test_int8_kv_cache_prefill_and_decode( + dtype: str, num_kv_heads: int, scale_value: float, use_v2: bool +) -> None: + """Verify real cache writes, saturation, and reads across a page boundary.""" + case = BackendCase( + num_heads=4, + num_kv_heads=num_kv_heads, + head_dim=128, + seq_lens=_PROMPT_LENS, + num_cached_tokens=[0, 0], + num_contexts=2, + dtype=dtype, + page_size=64, + use_kv_cache_manager_v2=use_v2, + ) + allocation = replace(case, seq_lens=[n + _DECODE_STEPS for n in _PROMPT_LENS]) + manager = _build_kv_cache_manager(allocation, "TRTLLM", torch.int8) + reference_manager = _build_kv_cache_manager(allocation, "TRTLLM", case.compute_dtype) + try: + request_ids = list(range(case.num_seqs)) + manager.add_dummy_requests(request_ids, allocation.token_nums) + cache = manager.get_buffers(0, kv_layout="HND") + reference_cache = reference_manager.get_buffers(0, kv_layout="HND") + assert cache.dtype == torch.int8 + assert cache.shape[1:] == reference_cache.shape[1:] + assert cache[0].nbytes * 2 == reference_cache[0].nbytes + if not use_v2: + assert cache.shape == reference_cache.shape + assert cache.nbytes * 2 == reference_cache.nbytes + # V2 rounds the backing allocation to an arena size, so the same + # minimum-size arena can expose more INT8 pages than FP16/BF16 pages. + cache.zero_() + + attention = _backend(case) + scale = torch.tensor([scale_value], dtype=torch.float32, device="cuda") + inputs = generate_inputs(case, seed=17) + q, k, v = inputs["q"], inputs["new_k"], inputs["new_v"] + # Cover both saturation directions and ties-to-even in the cache writes. + k[0, :4] = torch.tensor([20, -20, 2.5 * scale_value, 3.5 * scale_value]) + v[0, :4] = torch.tensor([-20, 20, -2.5 * scale_value, -3.5 * scale_value]) + keys = [x.view(-1, num_kv_heads, case.head_dim) for x in k.split(case.seq_lens)] + values = [x.view(-1, num_kv_heads, case.head_dim) for x in v.split(case.seq_lens)] + + metadata = _metadata(case, manager, _PROMPT_LENS) + actual = _forward(attention, metadata, q, k, v, scale) + expected = _reference(case, q, keys, values) + atol, rtol = (0.04, 0.01) if dtype == "bfloat16" else (0.015, 0.005) + torch.testing.assert_close(actual.float(), expected, atol=atol, rtol=rtol) + + quantized_keys = [_quantize(key, scale) for key in keys] + quantized_values = [_quantize(value, scale) for value in values] + _assert_cache(case, manager, quantized_keys, quantized_values) + + for step in range(_DECODE_STEPS): + decode = replace( + case, + seq_lens=[1, 1], + num_cached_tokens=[n + step for n in _PROMPT_LENS], + num_contexts=0, + ) + inputs = generate_inputs(decode, seed=31 + step) + q = inputs["q"] + # Put the new token on the quantization grid: MMHA may use the + # unquantized new K/V while XQA reads it back from the cache. + k = (_quantize(inputs["new_k"], scale).float() * scale).to(case.compute_dtype) + v = (_quantize(inputs["new_v"], scale).float() * scale).to(case.compute_dtype) + for index in request_ids: + new_key = _quantize(k[index].view(1, num_kv_heads, case.head_dim), scale) + new_value = _quantize(v[index].view(1, num_kv_heads, case.head_dim), scale) + quantized_keys[index] = torch.cat((quantized_keys[index], new_key)) + quantized_values[index] = torch.cat((quantized_values[index], new_value)) + + metadata = _metadata(decode, manager, _PROMPT_LENS) + actual = _forward(attention, metadata, q, k, v, scale) + expected = _reference( + decode, + q, + [key.float() * scale for key in quantized_keys], + [value.float() * scale for value in quantized_values], + ) + torch.testing.assert_close(actual.float(), expected, atol=atol, rtol=rtol) + _assert_cache(decode, manager, quantized_keys, quantized_values) + finally: + reference_manager.shutdown() + manager.shutdown() + + +@pytest.mark.parametrize("cached_tokens,paged_context", [(7, False), (0, True), (7, True)]) +@torch.inference_mode() +def test_int8_kv_cache_rejects_cached_context(cached_tokens: int, paged_context: bool) -> None: + """Cached context must fail with either packed or paged-context metadata.""" + case = BackendCase( + num_heads=4, + num_kv_heads=2, + head_dim=128, + seq_lens=[2], + num_cached_tokens=[cached_tokens], + num_contexts=1, + page_size=64, + ) + manager = _build_kv_cache_manager(case, "TRTLLM", torch.int8) + try: + manager.add_dummy_requests([0], case.token_nums) + manager.get_buffers(0, kv_layout="HND").zero_() + attention = _backend(case) + metadata = _metadata(case, manager, case.token_nums) + metadata.use_paged_context_fmha = paged_context + inputs = generate_inputs(case, seed=23) + scale = torch.tensor([1 / 32], dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="INT8"): + _forward(attention, metadata, inputs["q"], inputs["new_k"], inputs["new_v"], scale) + finally: + manager.shutdown() + + +@torch.inference_mode() +def test_int8_kv_mixed_prefill_and_decode() -> None: + # Cached tokens are valid for the decode request in a mixed batch. + """Cached decode tokens remain valid beside an uncached context request.""" + case = BackendCase( + num_heads=4, + num_kv_heads=2, + head_dim=128, + seq_lens=[5, 1], + num_cached_tokens=[0, 7], + num_contexts=1, + page_size=64, + ) + manager = _build_kv_cache_manager(case, "TRTLLM", torch.int8) + try: + manager.add_dummy_requests([0, 1], case.token_nums) + manager.get_buffers(0, kv_layout="HND").zero_() + scale = torch.tensor([1 / 32], dtype=torch.float32, device="cuda") + inputs = generate_inputs(case, seed=42) + q = inputs["q"] + k = (_quantize(inputs["new_k"], scale).float() * scale).to(case.compute_dtype) + v = (_quantize(inputs["new_v"], scale).float() * scale).to(case.compute_dtype) + keys = [ + k[:5].view(5, 2, 128), + torch.cat((torch.zeros(7, 2, 128, device="cuda"), k[5:].view(1, 2, 128))), + ] + values = [ + v[:5].view(5, 2, 128), + torch.cat((torch.zeros(7, 2, 128, device="cuda"), v[5:].view(1, 2, 128))), + ] + actual = _forward(_backend(case), _metadata(case, manager, [5, 7]), q, k, v, scale) + torch.testing.assert_close( + actual.float(), _reference(case, q, keys, values), atol=0.015, rtol=0.005 + ) + _assert_cache( + case, + manager, + [_quantize(x, scale) for x in keys], + [_quantize(x, scale) for x in values], + ) + finally: + manager.shutdown() + + +@pytest.mark.parametrize("scale_kind", ["missing", "cpu", "float16", "vector"]) +@pytest.mark.parametrize("scale_field", ["kv_scale_orig_quant", "kv_scale_quant_orig"]) +@pytest.mark.parametrize("after_warmup", [False, True]) +@torch.inference_mode() +def test_int8_kv_rejects_invalid_scale_tensor( + scale_kind: str, scale_field: str, after_warmup: bool +) -> None: + """Validate each scale independently, including replacements after warmup.""" + case = BackendCase( + num_heads=4, + num_kv_heads=2, + head_dim=128, + seq_lens=[2], + num_cached_tokens=[0], + num_contexts=1, + page_size=64, + ) + manager = _build_kv_cache_manager(case, "TRTLLM", torch.int8) + try: + manager.add_dummy_requests([0], case.token_nums) + inputs = generate_inputs(case, seed=29) + attention = _backend(case) + metadata = _metadata(case, manager, case.token_nums) + scale = torch.tensor([1 / 32], dtype=torch.float32, device="cuda") + forward_args = AttentionForwardArgs( + attention_mask=PredefinedAttentionMask.CAUSAL, + kv_scale_orig_quant=scale.reciprocal(), + kv_scale_quant_orig=scale, + ) + qkv = torch.cat((inputs["q"], inputs["new_k"], inputs["new_v"]), dim=-1) + if after_warmup: + attention.forward(qkv, None, None, metadata, forward_args=forward_args) + invalid_scale = ( + None + if scale_kind == "missing" + else torch.ones( + 2 if scale_kind == "vector" else 1, + dtype=torch.float16 if scale_kind == "float16" else torch.float32, + device="cpu" if scale_kind == "cpu" else "cuda", + ) + ) + setattr(forward_args, scale_field, invalid_scale) + with pytest.raises(ValueError, match="scalar float32 KV scales"): + attention.forward(qkv, None, None, metadata, forward_args=forward_args) + finally: + manager.shutdown() + + +@pytest.mark.parametrize( + ("invalid", "error"), + [ + ("float32", "FP16 or BF16"), + ("missing_cache", "active KV cache"), + ("inactive_cache", "active KV cache"), + ("cross_attention", "cross-attention"), + ("helix", "context parallelism"), + ], +) +@pytest.mark.parametrize("after_warmup", [False, True]) +@torch.inference_mode() +def test_int8_kv_rejects_unsupported_forward_inputs( + invalid: str, error: str, after_warmup: bool +) -> None: + """Each validation fails in the real backend before native attention runs.""" + case = BackendCase( + num_heads=4, + num_kv_heads=2, + head_dim=128, + seq_lens=[2], + num_cached_tokens=[0], + num_contexts=1, + page_size=64, + ) + manager = _build_kv_cache_manager(case, "TRTLLM", torch.int8) + try: + manager.add_dummy_requests([0], case.token_nums) + metadata = _metadata(case, manager, case.token_nums) + inputs = generate_inputs(case, seed=29) + q, k, v = inputs["q"], inputs["new_k"], inputs["new_v"] + attention = _backend(case) + scale = torch.tensor([1 / 32], dtype=torch.float32, device="cuda") + if after_warmup: + _forward(attention, metadata, q, k, v, scale) + if invalid == "float32": + q, k, v = q.float(), k.float(), v.float() + elif invalid == "missing_cache": + metadata.kv_cache_params = None + elif invalid == "inactive_cache": + metadata.kv_cache_params.use_cache = False + elif invalid == "cross_attention": + metadata.seq_lens_kv = metadata.seq_lens.clone() + assert metadata.is_cross and not metadata.enable_helix + else: + metadata.enable_helix = True + assert not metadata.is_cross + with pytest.raises(ValueError, match=error): + _forward(attention, metadata, q, k, v, scale) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/attention/test_packed_qkv_fmha.py b/tests/unittest/_torch/attention/test_packed_qkv_fmha.py new file mode 100644 index 000000000000..2c02d76612b7 --- /dev/null +++ b/tests/unittest/_torch/attention/test_packed_qkv_fmha.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Regress packed-QKV context FMHA independently of KV-cache quantization.""" + +import pytest +import torch +from backend_case import BackendCase, generate_inputs, run_backend + +from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.mark.parametrize("dtype", ["float16", "bfloat16"]) +@pytest.mark.parametrize("num_kv_heads", [4, 2], ids=["mha", "gqa"]) +@pytest.mark.parametrize("prompt_lens", [[31, 7], [63, 31]], ids=["short", "page_boundary"]) +@torch.inference_mode() +def test_packed_qkv_fmha_context_lengths( + dtype: str, num_kv_heads: int, prompt_lens: list[int], monkeypatch: pytest.MonkeyPatch +) -> None: + """Unequal prompt lengths reach native packed FMHA with full-precision KV.""" + monkeypatch.setenv("TLLM_FMHA_LIBS", "fallback") + original_forward = FallbackFmha.forward + launches = [] + + def checked_forward(self, q, k, v, metadata, forward_args): + """Ensure this regression cannot silently select a different FMHA path.""" + assert k is None and v is None + assert metadata.num_contexts == 2 + assert not metadata.use_paged_context_fmha + assert not self.attn.quant_config.layer_quant_mode.has_kv_cache_quant() + launches.append(metadata.num_contexts) + return original_forward(self, q, k, v, metadata, forward_args) + + monkeypatch.setattr(FallbackFmha, "forward", checked_forward) + case = BackendCase( + num_heads=4, + num_kv_heads=num_kv_heads, + head_dim=128, + seq_lens=prompt_lens, + num_cached_tokens=[0, 0], + num_contexts=2, + dtype=dtype, + page_size=64, + ) + inputs = generate_inputs(case, seed=47) + expected = run_backend(case, "VANILLA", inputs, kv_dtype=case.compute_dtype, kv_layout="NHD") + actual = run_backend(case, "TRTLLM", inputs, kv_dtype=case.compute_dtype, kv_layout="HND") + assert launches + atol, rtol = (0.04, 0.01) if dtype == "bfloat16" else (0.015, 0.005) + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) diff --git a/tests/unittest/_torch/test_int8_kv_config.py b/tests/unittest/_torch/test_int8_kv_config.py new file mode 100644 index 000000000000..6abccfb39e19 --- /dev/null +++ b/tests/unittest/_torch/test_int8_kv_config.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +import pytest +import torch +from transformers import LlamaConfig + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.pyexecutor._util import CacheCost, _create_kv_cache_manager +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.model_loader import ( + initialize_dummy_weights, + validate_and_set_kv_cache_quant, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, TorchLlmArgs +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +pytestmark = pytest.mark.cpu_only + + +def _model_config(kv_quant: QuantAlgo | None = QuantAlgo.INT8) -> ModelConfig: + """Build a dense two-layer configuration for CPU cache-sizing checks.""" + return ModelConfig( + pretrained_config=LlamaConfig( + hidden_size=256, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + num_hidden_layers=2, + vocab_size=512, + ), + quant_config=QuantConfig(kv_cache_quant_algo=kv_quant), + ) + + +@pytest.mark.parametrize("checkpoint_quant", [None, QuantAlgo.FP8, QuantAlgo.INT8]) +def test_int8_kv_explicit_override_keeps_layers_in_sync( + checkpoint_quant: QuantAlgo | None, monkeypatch: pytest.MonkeyPatch +) -> None: + """An explicit INT8 override must update both global and per-layer quantization.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + config = _model_config(checkpoint_quant) + config.quant_config_dict = { + "model.layers.0.self_attn": QuantConfig(kv_cache_quant_algo=checkpoint_quant), + "model.layers.1.self_attn": QuantConfig(kv_cache_quant_algo=checkpoint_quant), + } + validate_and_set_kv_cache_quant(config, "int8") + for quant in [config.quant_config, *config.quant_config_dict.values()]: + assert quant.layer_quant_mode.has_int8_kv_cache() + assert not quant.layer_quant_mode.has_fp8_kv_cache() + + +def test_int8_kv_auto_preserves_checkpoint(monkeypatch: pytest.MonkeyPatch) -> None: + """Automatic cache selection must preserve checkpoint INT8 metadata.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + config = _model_config() + validate_and_set_kv_cache_quant(config, "auto") + assert config.quant_config.layer_quant_mode.has_int8_kv_cache() + + +@pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("manager_type", [KVCacheManager, KVCacheManagerV2]) +def test_int8_kv_estimation_uses_half_the_bf16_bytes(tp_size: int, manager_type: type) -> None: + """Both managers must budget one byte per INT8 element before allocation.""" + mapping = Mapping(world_size=tp_size, tp_size=tp_size) + expected_bytes = 2 * 2 * (2 // tp_size) * 64 + options = {"tokens_per_block": 32} if manager_type is KVCacheManagerV2 else {} + int8_bytes = CacheCost.from_raw( + manager_type.get_cache_size_per_token(_model_config(), mapping, **options) + ).slope + bf16_bytes = CacheCost.from_raw( + manager_type.get_cache_size_per_token(_model_config(None), mapping, **options) + ).slope + assert int8_bytes == expected_bytes + assert bf16_bytes == 2 * int8_bytes + + +@pytest.mark.parametrize("manager_type", [KVCacheManager, KVCacheManagerV2]) +def test_int8_kv_pool_accounting(manager_type: type) -> None: + # Sizing is CPU-only; constructor allocation is covered by GPU integration tests. + """Runtime pool accounting must match static INT8/BF16 element sizes.""" + manager = manager_type.__new__(manager_type) + manager.dtype = DataType.INT8 + manager.kv_factor = 2 + manager.kv_cache_type = CacheType.SELF + manager.num_local_layers = 2 + manager.num_kv_heads_per_layer = [2, 2] + manager.total_num_kv_heads_per_layer = [2, 2] + manager.head_dim = 64 + manager.head_dim_per_layer = [64, 64] + assert manager.get_cache_bytes_per_token() == 512 + manager.dtype = DataType.BF16 + assert manager.get_cache_bytes_per_token() == 1024 + + +class _CaptureCacheManager: + def __init__(self, *args, **kwargs) -> None: + self.dtype = kwargs["dtype"] + + +def _create_test_manager( + config: ModelConfig, + *, + reuse: bool = False, + speculative: bool = False, + chunked_prefill: bool = False, + is_disagg: bool = False, + kv_connector: bool = False, +) -> _CaptureCacheManager: + """Exercise the real cache factory without allocating GPU cache pages.""" + engine = SimpleNamespace( + model=SimpleNamespace(model_config=config), + dtype=torch.bfloat16, + is_draft_model=False, + attn_runtime_features=SimpleNamespace(chunked_prefill=chunked_prefill), + ) + return _create_kv_cache_manager( + model_engine=engine, + kv_cache_manager_cls=_CaptureCacheManager, + mapping=Mapping(), + kv_cache_config=KvCacheConfig(enable_block_reuse=reuse), + tokens_per_block=32, + max_seq_len=128, + max_batch_size=2, + spec_config=SimpleNamespace() if speculative else None, + sparse_attention_config=None, + max_num_tokens=128, + max_beam_width=1, + kv_connector_manager=SimpleNamespace() if kv_connector else None, + is_disagg=is_disagg, + ) + + +def test_int8_kv_factory_selects_int8_pool() -> None: + """The real factory must pass INT8 to the cache-manager constructor.""" + assert _create_test_manager(_model_config()).dtype == DataType.INT8 + + +@pytest.mark.parametrize( + ("options", "error"), + [ + ({"reuse": True}, "enable_block_reuse=False"), + ({"speculative": True}, "speculative decoding"), + ({"chunked_prefill": True}, "enable_chunked_prefill=False"), + ({"is_disagg": True}, "disaggregated serving"), + ({"kv_connector": True}, "KV connectors"), + ], +) +def test_int8_kv_rejects_paged_context_features(options: dict, error: str) -> None: + """Unsupported cache population paths must fail before pool allocation.""" + with pytest.raises(ValueError, match=error): + _create_test_manager(_model_config(), **options) + + +@pytest.mark.parametrize("model_kind", ["hybrid", "encoder_decoder"]) +def test_int8_kv_rejects_unsupported_model_families(model_kind: str) -> None: + """Hybrid and encoder-decoder models must not enter the dense INT8 path.""" + config = _model_config() + if model_kind == "hybrid": + config.pretrained_config.hybrid_override_pattern = "M*" + else: + config.is_encoder_decoder = True + with pytest.raises(ValueError, match="dense decoder-only"): + _create_test_manager(config) + + +def test_dummy_weights_preserve_int8_kv_calibration() -> None: + """Dummy-weight initialization must not randomize cache calibration parameters.""" + projection = torch.nn.Module() + projection.register_parameter( + "kv_cache_scaling_factor", torch.nn.Parameter(torch.ones(1), requires_grad=False) + ) + projection.register_parameter( + "inv_kv_cache_scaling_factor", torch.nn.Parameter(torch.ones(1), requires_grad=False) + ) + projection.register_parameter("weight", torch.nn.Parameter(torch.zeros(8))) + model = torch.nn.Module() + model.add_module("qkv_proj", projection) + initialize_dummy_weights(model) + torch.testing.assert_close(projection.kv_cache_scaling_factor, torch.ones(1)) + torch.testing.assert_close(projection.inv_kv_cache_scaling_factor, torch.ones(1)) + assert torch.count_nonzero(projection.weight) > 0 + + +def test_int8_kv_public_args_sync_quantization() -> None: + """Public INT8 cache configuration must replace an earlier FP8 selection.""" + args = TorchLlmArgs.model_construct( + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8), + kv_cache_config=KvCacheConfig(dtype="int8", enable_block_reuse=False), + ) + args.sync_quant_config_with_kv_cache_config_dtype() + assert args.quant_config.layer_quant_mode.has_int8_kv_cache() + assert not args.quant_config.layer_quant_mode.has_fp8_kv_cache() diff --git a/tests/unittest/_torch/test_int8_kv_llm.py b/tests/unittest/_torch/test_int8_kv_llm.py new file mode 100644 index 000000000000..a8297fd1c35d --- /dev/null +++ b/tests/unittest/_torch/test_int8_kv_llm.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path + +import pytest +import torch +from safetensors.torch import save_file +from transformers import LlamaConfig, LlamaForCausalLM + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import DecodeCudaGraphConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +@pytest.mark.parametrize("use_v2", [False, True], ids=["v1", "v2"]) +def test_int8_kv_checkpoint_generation_with_cuda_graph(tmp_path: Path, use_v2: bool) -> None: + """Explicit INT8 generation is deterministic across eager and CUDA Graph runs.""" + config = LlamaConfig( + vocab_size=256, + hidden_size=256, + intermediate_size=512, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=256, + dtype="float16", + bos_token_id=1, + eos_token_id=2, + pad_token_id=0, + ) + config.architectures = ["LlamaForCausalLM"] + config.save_pretrained(tmp_path) + with torch.random.fork_rng(devices=[]): + torch.manual_seed(20260909) + model = LlamaForCausalLM(config).half() + weights = model.state_dict() + # Known test scales for this small random checkpoint; no downloaded model. + for layer in range(2): + prefix = f"model.layers.{layer}.self_attn" + weights[f"{prefix}.k_proj.k_scale"] = torch.tensor([1 / 64], dtype=torch.float32) + weights[f"{prefix}.v_proj.v_scale"] = torch.tensor([1 / 32], dtype=torch.float32) + save_file(weights, str(tmp_path / "model.safetensors"), metadata={"format": "pt"}) + del weights, model + + results = [] + resolved_dtypes = [] + for use_cuda_graph in (False, True): + with LLM( + model=tmp_path, + backend="pytorch", + skip_tokenizer_init=True, + dtype="float16", + attn_backend="TRTLLM", + max_batch_size=2, + max_num_tokens=128, + max_seq_len=256, + cuda_graph_config=DecodeCudaGraphConfig(batch_sizes=[1, 2]) if use_cuda_graph else None, + enable_chunked_prefill=False, + disable_overlap_scheduler=True, + kv_cache_config=KvCacheConfig( + dtype="int8", + enable_block_reuse=False, + free_gpu_memory_fraction=0.001, + use_kv_cache_manager_v2=use_v2, + ), + ) as llm: + resolved_dtypes.append(llm.args.quant_config.kv_cache_quant_algo) + assert llm.args.kv_cache_config.dtype == "int8" + assert llm.args.kv_cache_config.use_kv_cache_manager_v2 == use_v2 + assert (llm.args.cuda_graph_config is not None) == use_cuda_graph + assert resolved_dtypes[-1] == QuantAlgo.INT8 + outputs = llm.generate( + [[1] + list(range(10, 72)), [1] + list(range(100, 130))], + SamplingParams(end_id=2, pad_id=0, max_tokens=8, temperature=0, ignore_eos=True), + ) + ids = [output.outputs[0].token_ids for output in outputs] + assert len(ids) == 2 and all(len(tokens) == 8 for tokens in ids) + results.append(ids) + assert resolved_dtypes == [QuantAlgo.INT8, QuantAlgo.INT8] + assert results[0] == results[1] diff --git a/tests/unittest/_torch/test_int8_kv_scales.py b/tests/unittest/_torch/test_int8_kv_scales.py new file mode 100644 index 000000000000..263db758867c --- /dev/null +++ b/tests/unittest/_torch/test_int8_kv_scales.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from tensorrt_llm._torch.attention.attention import Attention +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.linear import Linear, WeightMode, WeightsLoadingConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _qkv(kv_quant: QuantAlgo = QuantAlgo.INT8) -> Linear: + """Construct an unquantized fused QKV projection with quantized KV storage.""" + with torch.device("cuda"): + return Linear( + 4, + 12, + bias=False, + dtype=torch.float16, + reduce_output=False, + quant_config=QuantConfig(kv_cache_quant_algo=kv_quant), + weights_loading_config=WeightsLoadingConfig(weight_mode=WeightMode.FUSED_QKV_LINEAR), + fused_weight_shard_indices_mapping={"q": (0, 4), "k": (4, 4), "v": (8, 4)}, + ) + + +def _weights(k_scale=0.03125, v_scale=0.0625) -> list[dict]: + """Create separate Q/K/V checkpoint shards with optional calibration tensors.""" + weights = [{"weight": torch.full((4, 4), float(i + 1), dtype=torch.float16)} for i in range(3)] + if k_scale is not None: + weights[1]["k_scale"] = torch.as_tensor(k_scale, dtype=torch.float64) + if v_scale is not None: + weights[2]["v_scale"] = torch.as_tensor(v_scale, dtype=torch.float64) + return weights + + +def test_int8_kv_load_and_partial_reload() -> None: + """Calibration and weights reload independently without moving scale pointers.""" + qkv = _qkv() + pointers = (qkv.kv_cache_scaling_factor.data_ptr(), qkv.inv_kv_cache_scaling_factor.data_ptr()) + qkv.load_weights(_weights()) + assert qkv.kv_cache_scaling_factor.dtype == torch.float32 + assert qkv.kv_cache_scaling_factor.item() == 0.0625 + assert qkv.inv_kv_cache_scaling_factor.item() == 16 + torch.testing.assert_close(qkv.weight, torch.cat([w["weight"] for w in _weights()]).cuda()) + qkv.load_weights( + [{}, {"weight": torch.full((4, 4), 7, dtype=torch.float16)}, {}], allow_partial_loading=True + ) + assert qkv.kv_cache_scaling_factor.item() == 0.0625 + assert qkv.inv_kv_cache_scaling_factor.item() == 16 + assert (qkv.weight[4:8] == 7).all() + qkv.load_weights( + [{}, {"k_scale": torch.tensor(0.125)}, {"v_scale": torch.tensor(0.25)}], + allow_partial_loading=True, + ) + assert qkv.kv_cache_scaling_factor.item() == 0.25 + assert qkv.inv_kv_cache_scaling_factor.item() == 4 + assert pointers == ( + qkv.kv_cache_scaling_factor.data_ptr(), + qkv.inv_kv_cache_scaling_factor.data_ptr(), + ) + + +@pytest.mark.parametrize( + "bad_scale", [0.0, -1.0, float("nan"), float("inf"), 1e-45, 1e40, [0.1, 0.2]] +) +def test_int8_kv_rejects_invalid_scale_without_changing_loaded_scales(bad_scale) -> None: + """Invalid calibration must leave previously loaded scale values intact.""" + qkv = _qkv() + qkv.load_weights(_weights()) + with pytest.raises(ValueError, match="INT8 KV cache"): + qkv.load_weights(_weights(bad_scale, bad_scale)) + assert qkv.kv_cache_scaling_factor.item() == 0.0625 + assert qkv.inv_kv_cache_scaling_factor.item() == 16 + + +@pytest.mark.parametrize("k_scale,v_scale", [(None, None), (None, 0.1), (0.1, None)]) +def test_int8_kv_requires_paired_checkpoint_scales(k_scale, v_scale) -> None: + """A full checkpoint must include both calibrated K and V scales.""" + with pytest.raises(ValueError, match="both scales must be loaded together"): + _qkv().load_weights(_weights(k_scale, v_scale)) + + +def test_int8_kv_partial_reload_rejects_unpaired_scale() -> None: + """Partial scale updates must supply K and V together.""" + qkv = _qkv() + qkv.load_weights(_weights()) + with pytest.raises(ValueError, match="both scales must be loaded together"): + qkv.load_weights([{}, {"k_scale": torch.tensor(0.125)}, {}], allow_partial_loading=True) + assert qkv.kv_cache_scaling_factor.item() == 0.0625 + + +@pytest.mark.parametrize("backend", ["VANILLA", "FLASHINFER"]) +def test_int8_kv_rejects_unsupported_backend(backend: str) -> None: + """INT8 cache configuration must fail before selecting another attention backend.""" + with pytest.raises(ValueError, match="INT8 KV cache requires the TRTLLM"): + create_attention( + backend, + layer_idx=0, + num_heads=4, + num_kv_heads=2, + head_dim=128, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.INT8), + ) + + +def test_int8_kv_rejects_quantized_projections() -> None: + """INT8 KV cache must not silently enable untested quantized-weight combinations.""" + with pytest.raises(ValueError, match="unquantized FP16/BF16 projections"): + create_attention( + "TRTLLM", + layer_idx=0, + num_heads=4, + num_kv_heads=2, + head_dim=128, + quant_config=QuantConfig(quant_algo=QuantAlgo.FP8, kv_cache_quant_algo=QuantAlgo.INT8), + ) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_int8_kv_scale_loading_environment(enabled: bool, monkeypatch: pytest.MonkeyPatch) -> None: + """INT8 must not silently use unity calibration when scale loading is disabled.""" + monkeypatch.setenv("TRTLLM_LOAD_KV_SCALES", str(int(enabled))) + qkv = _qkv() + if enabled: + qkv.load_weights(_weights()) + assert qkv.kv_cache_scaling_factor.item() == 0.0625 + else: + with pytest.raises(ValueError, match="TRTLLM_LOAD_KV_SCALES=1"): + qkv.load_weights(_weights()) + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_fp4_kv_scale_hook_preserves_optional_loading( + enabled: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """The shared checkpoint hook retains FP4 opt-out and stable reload buffers.""" + monkeypatch.setenv("TRTLLM_LOAD_KV_SCALES", str(int(enabled))) + qkv = _qkv(QuantAlgo.NVFP4) + pointers = (qkv.kv_scales.data_ptr(), qkv.inv_kv_scales.data_ptr()) + qkv.load_weights(_weights(None, None)) + torch.testing.assert_close(qkv.kv_scales, torch.ones_like(qkv.kv_scales)) + qkv.load_weights(_weights()) + expected = qkv.kv_scales.new_tensor([1.0, 0.03125, 0.0625] if enabled else [1.0] * 3) + torch.testing.assert_close(qkv.kv_scales, expected) + torch.testing.assert_close(qkv.inv_kv_scales, expected.reciprocal()) + qkv.load_weights([{}, {"weight": torch.zeros(4, 4)}, {}], allow_partial_loading=True) + torch.testing.assert_close(qkv.kv_scales, expected) + qkv.load_weights( + [{}, {"k_scale": torch.tensor(0.125)}, {"v_scale": torch.tensor(0.25)}], + allow_partial_loading=True, + ) + expected = qkv.kv_scales.new_tensor([1.0, 0.125, 0.25] if enabled else [1.0] * 3) + torch.testing.assert_close(qkv.kv_scales, expected) + torch.testing.assert_close(qkv.inv_kv_scales, expected.reciprocal()) + assert pointers == (qkv.kv_scales.data_ptr(), qkv.inv_kv_scales.data_ptr()) + + +def test_int8_kv_attention_construction_rejects_context_parallelism() -> None: + """Reject Helix CP from the actual Attention constructor before any forward.""" + config = ModelConfig( + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.INT8), + mapping=Mapping(world_size=2, cp_size=2, cp_config={"cp_type": "HELIX"}), + attn_backend="TRTLLM", + ) + with torch.device("cuda"), pytest.raises(ValueError, match="without context parallelism"): + Attention( + hidden_size=512, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=128, + bias=False, + dtype=torch.float16, + config=config, + layer_idx=0, + reduce_output=False, + )