From 58d46fffcbd290dc7c12e997774eab0156666b85 Mon Sep 17 00:00:00 2001 From: Superjomn <328693+Superjomn@users.noreply.github.com> Date: Fri, 4 Jul 2025 15:14:31 +0800 Subject: [PATCH 1/3] init Signed-off-by: Superjomn <328693+Superjomn@users.noreply.github.com> --- examples/llm-api/quickstart_advanced.py | 2 +- tensorrt_llm/bench/benchmark/utils/general.py | 12 ++++--- .../bench/dataclasses/configuration.py | 1 - tensorrt_llm/bench/dataclasses/reporting.py | 14 ++++++-- tensorrt_llm/llmapi/llm_args.py | 14 +++++--- .../defs/accuracy/test_llm_api_pytorch.py | 36 +++++++++---------- .../test_disaggregated_single_gpu.py | 6 +--- .../defs/perf/pytorch_model_config.py | 7 ++++ .../_torch/modeling/test_modeling_deepseek.py | 4 +-- 9 files changed, 58 insertions(+), 38 deletions(-) diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 9065bd2f00fa..90d527562a18 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -149,6 +149,7 @@ def setup_llm(args): kv_cache_config = KvCacheConfig( enable_block_reuse=not args.disable_kv_cache_reuse, free_gpu_memory_fraction=args.kv_cache_fraction, + dtype=args.kv_cache_dtype, ) spec_decode_algo = args.spec_decode_algo.upper( @@ -194,7 +195,6 @@ def setup_llm(args): model=args.model_dir, backend='pytorch', disable_overlap_scheduler=args.disable_overlap_scheduler, - kv_cache_dtype=args.kv_cache_dtype, kv_cache_config=kv_cache_config, attn_backend=args.attention_backend, cuda_graph_config=cuda_graph_config, diff --git a/tensorrt_llm/bench/benchmark/utils/general.py b/tensorrt_llm/bench/benchmark/utils/general.py index 0073ea1d44f0..bc72b5e14670 100755 --- a/tensorrt_llm/bench/benchmark/utils/general.py +++ b/tensorrt_llm/bench/benchmark/utils/general.py @@ -88,12 +88,14 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, enable_chunked_prefill = params.get("enable_chunked_prefill", False) kv_cache_dtype = "auto" + kv_cache_config = {} if extra_llm_api_options: with open(extra_llm_api_options, 'r') as f: llm_args_dict = yaml.safe_load(f) - - if "kv_cache_dtype" in llm_args_dict: - kv_cache_dtype = llm_args_dict["kv_cache_dtype"] + kv_cache_config = llm_args_dict.get("kv_cache_config", { + "dtype": "auto", + }) + kv_cache_dtype = kv_cache_config.get("dtype", "auto") enable_chunked_prefill = llm_args_dict.get("enable_chunked_prefill", enable_chunked_prefill) @@ -158,9 +160,11 @@ def get_settings(params: dict, dataset_metadata: DatasetMetadata, model: str, "max_batch_size": max_batch_size } + kv_cache_config["dtype"] = kv_cache_dtype + pyt_options = { "cuda_graph_config": cuda_graph_config, - "kv_cache_dtype": kv_cache_dtype, + "kv_cache_config": kv_cache_config, } backend = params.get("backend", "pytorch") diff --git a/tensorrt_llm/bench/dataclasses/configuration.py b/tensorrt_llm/bench/dataclasses/configuration.py index 0d352fa068c2..77f80632088f 100755 --- a/tensorrt_llm/bench/dataclasses/configuration.py +++ b/tensorrt_llm/bench/dataclasses/configuration.py @@ -112,7 +112,6 @@ def get_pytorch_perf_config(self) -> PyTorchConfig: def get_autodeploy_perf_config(self) -> Dict: AutoDeployPerfConfig = dict ad_config = AutoDeployPerfConfig() - ad_config["kv_cache_dtype"] = "auto" ad_config["attn_backend"] = "flashinfer" return ad_config diff --git a/tensorrt_llm/bench/dataclasses/reporting.py b/tensorrt_llm/bench/dataclasses/reporting.py index d994000d6d08..476ff50ec2d4 100755 --- a/tensorrt_llm/bench/dataclasses/reporting.py +++ b/tensorrt_llm/bench/dataclasses/reporting.py @@ -11,6 +11,7 @@ from tensorrt_llm.bench.dataclasses.statistics import (BenchmarkStatistics, PercentileStats, RequestRecord) +from tensorrt_llm.llmapi import KvCacheConfig from tensorrt_llm.logger import Logger from tensorrt_llm.models.modeling_utils import SpeculativeDecodingMode @@ -275,8 +276,17 @@ def get_statistics_dict(self) -> Dict[str, Any]: model = self.rt_cfg.model_path or self.rt_cfg.model model_config = ModelConfig.from_pretrained(model, trust_remote_code=True) - validate_and_set_kv_cache_quant(model_config, - self.kwargs["kv_cache_dtype"]) + kv_cache_config = self.kwargs.get("kv_cache_config", + KvCacheConfig()) + if isinstance(kv_cache_config, KvCacheConfig): + kv_cache_dtype = kv_cache_config.dtype + elif isinstance(kv_cache_config, dict): + kv_cache_dtype = kv_cache_config.get("dtype", "auto") + else: + raise ValueError( + f"Invalid kv_cache_config type: {type(kv_cache_config)}.") + + validate_and_set_kv_cache_quant(model_config, kv_cache_dtype) stats_dict["engine"] |= { "backend": diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 1b385b6e8fc6..d1514ff96098 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -821,6 +821,10 @@ class KvCacheConfig(BaseModel, PybindMirror): use_uvm: bool = Field(default=False, description="Whether to use UVM for the KV cache.") + # This is a pure python field, not a pybind field. It is only for the Pytorch backend. + dtype: str = Field(default="auto", + description="The data type to use for the KV cache.") + def _to_pybind(self): return _KvCacheConfig( enable_block_reuse=self.enable_block_reuse, @@ -1737,6 +1741,11 @@ def validate_enable_build_cache(self): f"Invalid build_cache_config: {self.enable_build_cache}") return self + @model_validator(mode="after") + def validate_kv_cache_dtype(self): + assert self.kv_cache_config.dtype == "auto", "KvCacheConfig.dtype is not supported by the TensorRT backend." + return self + class LoadFormat(Enum): AUTO = 0 @@ -1810,9 +1819,6 @@ class TorchLlmArgs(BaseLlmArgs): "If true, will use the TRTLLM sampler instead of the PyTorch sampler. The TRTLLM sampler has a wide coverage of sampling strategies." ) - kv_cache_dtype: str = Field(default="auto", - description="Data type for KV cache.") - enable_iter_perf_stats: bool = Field( default=False, description="Enable iteration performance statistics.") @@ -2016,7 +2022,7 @@ def get_pytorch_backend_config(self) -> "PyTorchConfig": moe_backend=self.moe_config.backend, enable_mixed_sampler=self.enable_mixed_sampler, enable_trtllm_sampler=self.enable_trtllm_sampler, - kv_cache_dtype=self.kv_cache_dtype, + kv_cache_dtype=self.kv_cache_config.dtype, enable_iter_perf_stats=self.enable_iter_perf_stats, enable_iter_req_stats=self.enable_iter_req_stats, print_iter_log=self.print_iter_log, diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index b09422318085..e9afc3534a87 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -154,7 +154,7 @@ def test_fp8(self, fp8kv, attn_backend, torch_compile): ) if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + pytorch_config["kv_cache_config"] = KvCacheConfig(dtype="fp8") with LLM( f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct-FP8", quant_config=quant_config, @@ -192,7 +192,7 @@ def test_fp8_4gpus(self, tp_size, pp_size, fp8kv, attn_backend, ) if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + pytorch_config["kv_cache_config"] = KvCacheConfig(dtype="fp8") with LLM( f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct-FP8", tensor_parallel_size=tp_size, @@ -691,7 +691,7 @@ def test_fp8_block_scales(self, mtp, fp8kv, attention_dp, cuda_graph, quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" mtp_config = None mtp_nextn = 2 @@ -750,7 +750,7 @@ def test_cute_dsl_fp8_block_scales( quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" mtp_config = None if mtp_nextn > 0: @@ -861,7 +861,7 @@ def test_fp8_block_scales_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" mtp_config = None if mtp_nextn > 0: @@ -930,7 +930,7 @@ def test_cute_dsl_fp8_block_scales_4gpus( quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" mtp_config = None if mtp_nextn > 0: @@ -1016,21 +1016,19 @@ def test_nvfp4_4gpus_online_eplb(self, fp8kv): num_slots = 80 eplb_config = MoeLoadBalancerConfig(num_slots=num_slots, layer_updates_per_iter=2) - pytorch_backend_options = dict(cuda_graph_config=CudaGraphConfig(), - moe_config=MoeConfig( - backend="WIDEEP", - load_balancer=eplb_config)) + pytorch_config = dict(cuda_graph_config=CudaGraphConfig(), + moe_config=MoeConfig(backend="WIDEEP", + load_balancer=eplb_config)) quant_config = QuantConfig() quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_backend_options["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only", tensor_parallel_size=4, moe_expert_parallel_size=4, kv_cache_config=kv_cache_config, - **pytorch_backend_options, + **pytorch_config, enable_attention_dp=True, quant_config=quant_config) as llm: task = GSM8K(self.MODEL_NAME) @@ -1070,7 +1068,7 @@ def test_nvfp4(self, fp8kv, attention_dp, cuda_graph, overlap_scheduler, quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", kv_cache_config=kv_cache_config, @@ -1131,7 +1129,7 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", tensor_parallel_size=tp_size, @@ -1194,7 +1192,7 @@ def test_no_kv_cache_reuse(self, quant_dtype, mtp_nextn, fp8kv, quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" with LLM(model_path, kv_cache_config=kv_cache_config, @@ -1254,7 +1252,7 @@ def test_chunked_prefill(self, quant_dtype, kv_cache_reuse, fp8kv, quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" with LLM(model_path, kv_cache_config=kv_cache_config, @@ -1362,7 +1360,7 @@ def test_nvfp4_multi_gpus(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" mtp_config = None if mtp_nextn > 0: @@ -1412,7 +1410,7 @@ def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - pytorch_config["kv_cache_dtype"] = "fp8" + kv_cache_config.dtype = "fp8" mtp_config = None if mtp_nextn > 0: diff --git a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py index 0a392a575a40..540313cfdffc 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py +++ b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py @@ -110,14 +110,12 @@ def verify_disaggregated(model, generation_overlap, enable_cuda_graph, prompt, worker_pytorch_configs.append( dict( disable_overlap_scheduler=True, - kv_cache_dtype="auto", cuda_graph_config=CudaGraphConfig() if enable_cuda_graph else None)) # Generation worker worker_pytorch_configs.append( dict( disable_overlap_scheduler=not generation_overlap, - kv_cache_dtype="auto", cuda_graph_config=CudaGraphConfig() if enable_cuda_graph else None)) kv_cache_configs = [KvCacheConfig(max_tokens=2048 * 8) for _ in range(2)] @@ -233,18 +231,16 @@ def test_disaggregated_llama_context_capacity(model, enable_cuda_graph, worker_pytorch_configs.append( dict( disable_overlap_scheduler=True, - kv_cache_dtype="auto", cuda_graph_config=CudaGraphConfig() if enable_cuda_graph else None)) # Generation worker worker_pytorch_configs.append( dict( disable_overlap_scheduler=not generation_overlap, - kv_cache_dtype="auto", cuda_graph_config=CudaGraphConfig() if enable_cuda_graph else None)) kv_cache_configs = [ - KvCacheConfig(max_tokens=128, enable_block_reuse=False) + KvCacheConfig(max_tokens=128, enable_block_reuse=False, dtype="auto") for _ in range(2) ] model_names = [model_path(model) for _ in range(2)] diff --git a/tests/integration/defs/perf/pytorch_model_config.py b/tests/integration/defs/perf/pytorch_model_config.py index 4c0ef1840931..23ccd0f18411 100644 --- a/tests/integration/defs/perf/pytorch_model_config.py +++ b/tests/integration/defs/perf/pytorch_model_config.py @@ -17,6 +17,8 @@ Model pytorch yaml config for trtllm-bench perf tests """ +from tensorrt_llm.llmapi import KvCacheConfig + def recursive_update(d, u): for k, v in u.items(): @@ -186,4 +188,9 @@ def get_model_yaml_config(model_label: str, } base_config.update(lora_config) + kv_cache_config = base_config.get('kv_cache_config', KvCacheConfig()) + if 'kv_cache_dtype' in base_config: + kv_cache_config.dtype = base_config.pop('kv_cache_dtype', 'auto') + base_config.update({'kv_cache_config': kv_cache_config}) + return base_config diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseek.py b/tests/unittest/_torch/modeling/test_modeling_deepseek.py index e5cf9680bbff..ad242f6b28c8 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseek.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseek.py @@ -68,7 +68,6 @@ def test_deepseek_trtllmgen(model_name): pytorch_config = dict( disable_overlap_scheduler=True, - kv_cache_dtype="auto", attn_backend="TRTLLM", load_format="dummy", moe_config=MoeConfig(backend="TRTLLM"), @@ -89,7 +88,8 @@ def test_deepseek_trtllmgen(model_name): moe_tensor_parallel_size=-1, enable_attention_dp=False, speculative_config=spec_config, - kv_cache_config=KvCacheConfig(enable_block_reuse=False, + kv_cache_config=KvCacheConfig(dtype="auto", + enable_block_reuse=False, free_gpu_memory_fraction=0.4)) sampling_params = SamplingParams(max_tokens=20) From 459f15794f77cc6a5302ef8391df142cabf5fdac Mon Sep 17 00:00:00 2001 From: Superjomn <328693+Superjomn@users.noreply.github.com> Date: Fri, 11 Jul 2025 18:11:04 +0800 Subject: [PATCH 2/3] hide quant_config from PyT Signed-off-by: Superjomn <328693+Superjomn@users.noreply.github.com> --- tensorrt_llm/llmapi/llm_args.py | 46 ++++++-- tensorrt_llm/llmapi/llm_utils.py | 3 + .../defs/accuracy/test_llm_api_pytorch.py | 110 ++---------------- .../multi_gpu_modeling/test_deepseek.py | 1 - .../references_committed/llm.yaml | 7 -- 5 files changed, 46 insertions(+), 121 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d1514ff96098..8461833ddfb4 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1028,10 +1028,6 @@ class BaseLlmArgs(BaseModel): lora_config: Optional[LoraConfig] = Field( default=None, description="LoRA configuration for the model.") - # Quantization and calibration configurations - quant_config: Optional[QuantConfig] = Field( - default=None, description="Quantization config.", validate_default=True) - # Several options from ExecutorConfig, expanded here for less hierarchy kv_cache_config: KvCacheConfig = Field(default_factory=KvCacheConfig, description="KV cache config.") @@ -1212,13 +1208,6 @@ def validate_dtype(cls, v, info): raise RuntimeError("Pre SM 80 GPUs do not support bfloat16") return v - @field_validator("quant_config", mode='before') - @classmethod - def validate_quant_config(cls, v, info): - if v is None: - v = QuantConfig() - return v - @field_validator("gpus_per_node", mode='before') @classmethod def validate_gpus_per_node(cls, v, info): @@ -1660,6 +1649,10 @@ class TrtLlmArgs(BaseLlmArgs): calib_config: Optional[CalibConfig] = Field( default=None, description="Calibration config.", validate_default=True) + # Quantization and calibration configurations + quant_config: Optional[QuantConfig] = Field( + default=None, description="Quantization config.", validate_default=True) + embedding_parallel_mode: str = Field( default='SHARDING_ALONG_VOCAB', description="The embedding parallel mode.") @@ -1697,6 +1690,13 @@ def init_calib_config(cls, v): return CalibConfig() return v + @field_validator("quant_config", mode='before') + @classmethod + def validate_quant_config(cls, v, info): + if v is None: + v = QuantConfig() + return v + @model_validator(mode="after") def setup_embedding_parallel_mode(self): if self.embedding_parallel_mode == 'NONE': @@ -1872,6 +1872,19 @@ class TorchLlmArgs(BaseLlmArgs): 'MNNVL']] = Field(default='AUTO', description="Allreduce strategy to use.") + # PrivateVars + _quant_config: Optional[QuantConfig] = PrivateAttr(default=None) + + @property + def quant_config(self) -> QuantConfig: + if self._quant_config is None: + self._quant_config = QuantConfig() + return self._quant_config + + @quant_config.setter + def quant_config(self, value: QuantConfig): + self._quant_config = value + # TODO: remove backend later @field_validator('backend', mode='before') def init_backend(cls, v): @@ -1999,6 +2012,17 @@ def validate_cuda_graph_config(self) -> 'TorchLlmArgs': return self + @model_validator(mode='after') + def sync_quant_config_with_kv_cache_config_dtype(self) -> 'TorchLlmArgs': + assert self.quant_config is not None + if self.kv_cache_config.dtype == 'fp8': + self.quant_config.kv_cache_quant_algo = QuantAlgo.FP8 + else: + logger.warning( + f"Cannot sync quant_config.kv_cache_quant_algo with kv_cache_config.dtype of {self.kv_cache_config.dtype}, " + "please update the validator") + return self + # TODO: Remove this after the PyTorch backend is fully migrated to TorchLlmArgs from ExecutorConfig def get_pytorch_backend_config(self) -> "PyTorchConfig": from tensorrt_llm._torch.pyexecutor.config import PyTorchConfig diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index cf2bdb26c143..ea197a13b07f 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -407,6 +407,9 @@ def _update_from_hf_quant_config(self) -> bool: logger.info(f"Setting {key}={value} from HF quant config.") setattr(quant_config, key, value) + # Update the quant_config in llm_args for pytorch + self.llm_args.quant_config = quant_config + return True hf_config_path = f"{self._model_dir}/config.json" diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index e9afc3534a87..3d41fdfc27a6 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -22,7 +22,6 @@ KvCacheConfig, MoeConfig, MTPDecodingConfig, NGramDecodingConfig, SamplingParams, TorchCompileConfig) -from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization import QuantAlgo from ..conftest import (llm_models_root, parametrize_with_ids, @@ -50,7 +49,6 @@ def test_nvfp4(self): model_path = f"{llm_models_root()}/nvfp4-quantized/Meta-Llama-3.1-8B" with LLM(model_path) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) task = MMLU(self.MODEL_NAME) @@ -66,7 +64,6 @@ def test_nvfp4_streaming(self, stream_interval): with LLM(f"{llm_models_root()}/nvfp4-quantized/Meta-Llama-3.1-8B", stream_interval=stream_interval) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 assert llm.args.stream_interval == stream_interval task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm, streaming=True) @@ -142,7 +139,6 @@ def test_bfloat16_4gpus(self, tp_size, pp_size, attn_backend, @parametrize_with_ids("attn_backend", ["TRTLLM", "FLASHINFER"]) @parametrize_with_ids("fp8kv", [False, True]) def test_fp8(self, fp8kv, attn_backend, torch_compile): - quant_config = QuantConfig(QuantAlgo.FP8) torch_compile_config = TorchCompileConfig( enable_fullgraph=True) if torch_compile else None pytorch_config = dict( @@ -153,15 +149,11 @@ def test_fp8(self, fp8kv, attn_backend, torch_compile): disable_overlap_scheduler=torch_compile, ) if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 pytorch_config["kv_cache_config"] = KvCacheConfig(dtype="fp8") with LLM( f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct-FP8", - quant_config=quant_config, **pytorch_config) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -180,7 +172,6 @@ def test_fp8_4gpus(self, tp_size, pp_size, fp8kv, attn_backend, "Pipeline parallel with torch.compile is not supported yet.\n" "Issue: Unfusing flashinfer_fused_add_rmsnorm causes outputs to be " "discarded at graph breaks.") - quant_config = QuantConfig(QuantAlgo.FP8) torch_compile_config = TorchCompileConfig( enable_fullgraph=True) if torch_compile else None pytorch_config = dict( @@ -191,17 +182,13 @@ def test_fp8_4gpus(self, tp_size, pp_size, fp8kv, attn_backend, disable_overlap_scheduler=torch_compile, ) if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 pytorch_config["kv_cache_config"] = KvCacheConfig(dtype="fp8") with LLM( f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct-FP8", tensor_parallel_size=tp_size, pipeline_parallel_size=pp_size, - quant_config=quant_config, **pytorch_config) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -307,7 +294,6 @@ def test_fp8_prequantized(self): model_path = f"{llm_models_root()}/llama-3.2-models/Llama-3.2-1B-FP8" with LLM(model_path) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) @@ -329,7 +315,6 @@ def test_fp8_prequantized(self): model_path = f"{llm_models_root()}/llama-3.2-models/Llama-3.2-3B-Instruct-FP8" with LLM(model_path) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) task = MMLU(self.MODEL_NAME) @@ -372,7 +357,6 @@ def test_nvfp4_tp4(self): model_path = f"{llm_models_root()}/modelopt-hf-model-hub/Llama-3.3-70B-Instruct-fp4" with LLM(model_path, tensor_parallel_size=4) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) task.evaluate(llm) task = GSM8K(self.MODEL_NAME) @@ -554,7 +538,6 @@ def test_fp8_tp2(self): model_path = f"{llm_models_root()}/modelopt-hf-model-hub/Mixtral-8x7B-Instruct-v0.1-fp8" with LLM(model_path, tensor_parallel_size=2) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) task = MMLU(self.MODEL_NAME) @@ -566,7 +549,6 @@ def test_nvfp4_tp2(self): model_path = f"{llm_models_root()}/modelopt-hf-model-hub/Mixtral-8x7B-Instruct-v0.1-fp4" with LLM(model_path, tensor_parallel_size=2) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = CnnDailymail(self.MODEL_NAME) task.evaluate(llm) task = MMLU(self.MODEL_NAME) @@ -687,10 +669,7 @@ def test_fp8_block_scales(self, mtp, fp8kv, attention_dp, cuda_graph, torch_compile_config=torch_compile_config, ) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" mtp_config = None @@ -704,13 +683,10 @@ def test_fp8_block_scales(self, mtp, fp8kv, attention_dp, cuda_graph, with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -746,10 +722,7 @@ def test_cute_dsl_fp8_block_scales( moe_config=MoeConfig(backend="CUTEDSL"), ) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" mtp_config = None @@ -760,14 +733,11 @@ def test_cute_dsl_fp8_block_scales( f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config, ) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -808,14 +778,11 @@ def test_fp8_block_scales_cuda_graph_padding_4gpus(self, mtp_nextn, disable_overlap_scheduler=False, cuda_graph_config=CudaGraphConfig(enable_padding=True), ) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/fp8", tensor_parallel_size=4, kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES @@ -857,10 +824,7 @@ def test_fp8_block_scales_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, torch_compile_config=torch_compile_config, ) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" mtp_config = None @@ -873,13 +837,10 @@ def test_fp8_block_scales_4gpus(self, tp_size, pp_size, ep_size, mtp_nextn, moe_expert_parallel_size=ep_size, kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -926,10 +887,7 @@ def test_cute_dsl_fp8_block_scales_4gpus( moe_config=MoeConfig(backend="CUTEDSL"), ) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" mtp_config = None @@ -943,13 +901,10 @@ def test_cute_dsl_fp8_block_scales_4gpus( moe_expert_parallel_size=ep_size, kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config, ) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -1019,18 +974,17 @@ def test_nvfp4_4gpus_online_eplb(self, fp8kv): pytorch_config = dict(cuda_graph_config=CudaGraphConfig(), moe_config=MoeConfig(backend="WIDEEP", load_balancer=eplb_config)) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: kv_cache_config.dtype = "fp8" - with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only", - tensor_parallel_size=4, - moe_expert_parallel_size=4, - kv_cache_config=kv_cache_config, - **pytorch_config, - enable_attention_dp=True, - quant_config=quant_config) as llm: + with LLM( + f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only", + tensor_parallel_size=4, + moe_expert_parallel_size=4, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=True, + ) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -1064,21 +1018,15 @@ def test_nvfp4(self, fp8kv, attention_dp, cuda_graph, overlap_scheduler, if mtp_nextn > 0: mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -1125,10 +1073,7 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, if mtp_nextn > 0: mtp_config = MTPDecodingConfig(num_nextn_predict_layers=mtp_nextn) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" with LLM(f"{llm_models_root()}/DeepSeek-V3-Lite/nvfp4_moe_only_mtp", @@ -1137,12 +1082,9 @@ def test_nvfp4_4gpus(self, fp8kv, attention_dp, cuda_graph, moe_expert_parallel_size=ep_size, kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -1183,21 +1125,13 @@ def test_no_kv_cache_reuse(self, quant_dtype, mtp_nextn, fp8kv, if quant_dtype == "none": assert not fp8kv - quant_config = None else: - quant_config = QuantConfig() - if quant_dtype == "fp8": - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES - elif quant_dtype == "nvfp4": - quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" with LLM(model_path, kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config) as llm: if quant_dtype == "fp8": @@ -1205,8 +1139,6 @@ def test_no_kv_cache_reuse(self, quant_dtype, mtp_nextn, fp8kv, elif quant_dtype == "nvfp4": assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -1243,15 +1175,8 @@ def test_chunked_prefill(self, quant_dtype, kv_cache_reuse, fp8kv, if quant_dtype == "none": assert not fp8kv - quant_config = None else: - quant_config = QuantConfig() - if quant_dtype == "fp8": - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES - elif quant_dtype == "nvfp4": - quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" with LLM(model_path, @@ -1259,7 +1184,6 @@ def test_chunked_prefill(self, quant_dtype, kv_cache_reuse, fp8kv, enable_chunked_prefill=True, max_num_tokens=512, **pytorch_config, - quant_config=quant_config, enable_attention_dp=True, speculative_config=mtp_config) as llm: @@ -1268,9 +1192,6 @@ def test_chunked_prefill(self, quant_dtype, kv_cache_reuse, fp8kv, elif quant_dtype == "nvfp4": assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 - task = GSM8K(self.MODEL_NAME) task.evaluate(llm) @@ -1356,10 +1277,7 @@ def test_nvfp4_multi_gpus(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, cuda_graph_config=CudaGraphConfig() if cuda_graph else None, moe_config=MoeConfig(backend=moe_backend)) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.NVFP4 if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" mtp_config = None @@ -1372,14 +1290,11 @@ def test_nvfp4_multi_gpus(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, moe_expert_parallel_size=ep_size, kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config) as llm: assert llm.args.moe_backend == moe_backend assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -1406,10 +1321,7 @@ def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, cuda_graph_config=CudaGraphConfig() if cuda_graph else None, ) - quant_config = QuantConfig() - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES if fp8kv: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 kv_cache_config.dtype = "fp8" mtp_config = None @@ -1422,12 +1334,9 @@ def test_fp8_blockscale(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, moe_expert_parallel_size=ep_size, kv_cache_config=kv_cache_config, **pytorch_config, - quant_config=quant_config, enable_attention_dp=attention_dp, speculative_config=mtp_config) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8_BLOCK_SCALES - if fp8kv: - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) task.evaluate(llm) @@ -1517,7 +1426,6 @@ def test_fp8_prequantized(self): model_path = f"{llm_models_root()}/Llama-3.1-Nemotron-Nano-8B-v1-FP8" with LLM(model_path) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) task.evaluate(llm) task = GSM8K(self.MODEL_NAME) @@ -1569,7 +1477,6 @@ def test_fp8_prequantized(self, cuda_graph, tp_size, pp_size, ep_size): kv_cache_config=KvCacheConfig( free_gpu_memory_fraction=0.85)) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) task.evaluate(llm) task = GSM8K(self.MODEL_NAME) @@ -1602,7 +1509,6 @@ def test_reasoning_fp8_prequantized(self): kv_cache_config=kv_cache_config, max_batch_size=256) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.FP8 - assert llm.args.quant_config.kv_cache_quant_algo == QuantAlgo.FP8 task = MMLU(self.MODEL_NAME) task.evaluate(llm) task = GSM8K(self.MODEL_NAME) diff --git a/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py b/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py index 5d2a8b713746..5a38f0d07881 100644 --- a/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py +++ b/tests/unittest/_torch/multi_gpu_modeling/test_deepseek.py @@ -63,7 +63,6 @@ def test_deepseek_streaming(model_name, backend, quant, tp_size): pytorch_config = dict( disable_overlap_scheduler=True, - kv_cache_dtype="auto", attn_backend=backend, ) moe_config = MoeConfig(max_num_tokens=moe_max_num_tokens) diff --git a/tests/unittest/api_stability/references_committed/llm.yaml b/tests/unittest/api_stability/references_committed/llm.yaml index 4201b190a2fe..66fbdabfc5d9 100644 --- a/tests/unittest/api_stability/references_committed/llm.yaml +++ b/tests/unittest/api_stability/references_committed/llm.yaml @@ -57,13 +57,6 @@ methods: guided_decoding_backend: annotation: Optional[Literal["xgrammar", "llguidance"]] default: null - # Quantization and calibration - quant_config: - annotation: Optional[tensorrt_llm.models.modeling_utils.QuantConfig] - default: null - calib_config: - annotation: Optional[tensorrt_llm.llmapi.llm_utils.CalibConfig] - default: null # Speculative decoding speculative_config: annotation: Union[tensorrt_llm.llmapi.llm_args.DraftTargetDecodingConfig, tensorrt_llm.llmapi.llm_args.EagleDecodingConfig,tensorrt_llm.llmapi.llm_args.LookaheadDecodingConfig, tensorrt_llm.llmapi.llm_args.MedusaDecodingConfig, tensorrt_llm.llmapi.llm_args.MTPDecodingConfig, tensorrt_llm.llmapi.llm_args.NGramDecodingConfig, tensorrt_llm.llmapi.llm_args.UserProvidedDecodingConfig, NoneType] From cbaf2023d03f5125c8b6293240473cb7568e020b Mon Sep 17 00:00:00 2001 From: Superjomn <328693+Superjomn@users.noreply.github.com> Date: Wed, 16 Jul 2025 09:44:35 +0800 Subject: [PATCH 3/3] fix Signed-off-by: Superjomn <328693+Superjomn@users.noreply.github.com> --- tensorrt_llm/llmapi/llm_args.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 8461833ddfb4..5bb95d73a748 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2014,8 +2014,13 @@ def validate_cuda_graph_config(self) -> 'TorchLlmArgs': @model_validator(mode='after') def sync_quant_config_with_kv_cache_config_dtype(self) -> 'TorchLlmArgs': + if self.kv_cache_config is None: + return self + assert self.quant_config is not None - if self.kv_cache_config.dtype == 'fp8': + if self.kv_cache_config.dtype == "auto": + return self + elif self.kv_cache_config.dtype == 'fp8': self.quant_config.kv_cache_quant_algo = QuantAlgo.FP8 else: logger.warning(