From 5dfe256bb9cc26df1f602898d6350c36345b5011 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 29 Jun 2026 16:07:48 +0800 Subject: [PATCH] refactor: rename quant policy to kv cache dtype --- .github/scripts/eval_chat_config.py | 4 +- .../benchmark/test_apiserver_performance.py | 8 +- .../benchmark/test_prefixcache_performance.py | 8 +- .../benchmark/test_throughput_performance.py | 8 +- .../interface/pipeline/test_pipeline_func.py | 2 +- autotest/tools/common_case_config.py | 60 ++++---- autotest/tools/pipeline/llm_case.py | 6 +- autotest/tools/pipeline/mllm_case.py | 8 +- autotest/utils/config_utils.py | 52 +++---- autotest/utils/constant.py | 2 +- autotest/utils/evaluate_utils.py | 8 +- benchmark/benchmark_guided.py | 4 +- benchmark/profile_pipeline_api.py | 4 +- benchmark/profile_throughput.py | 8 +- docs/en/inference/turbomind_config.md | 10 +- docs/en/quantization/kv_quant.md | 32 ++--- docs/zh_cn/inference/turbomind_config.md | 12 +- docs/zh_cn/quantization/kv_quant.md | 32 ++--- lmdeploy/cli/cli.py | 4 +- lmdeploy/cli/serve.py | 8 +- lmdeploy/cli/utils.py | 26 +--- lmdeploy/messages.py | 64 ++++++--- lmdeploy/pytorch/backends/attention.py | 4 +- .../backends/cuda/attention/default.py | 24 ++-- .../pytorch/backends/cuda/attention/fa3.py | 20 +-- .../pytorch/backends/cuda/attention/mla.py | 10 +- lmdeploy/pytorch/backends/cuda/op_backend.py | 2 +- .../backends/dlinfer/ascend/op_backend.py | 5 +- .../pytorch/backends/dlinfer/attention.py | 2 +- lmdeploy/pytorch/config.py | 7 +- lmdeploy/pytorch/engine/cache_engine.py | 112 +++++++-------- lmdeploy/pytorch/engine/config_builder.py | 2 +- lmdeploy/pytorch/engine/model_agent/agent.py | 2 +- .../pytorch/kernels/cuda/fill_kv_cache.py | 38 +++--- .../pytorch/kernels/cuda/flatten_kv_cache.py | 48 +++---- .../pytorch/kernels/cuda/pagedattention.py | 84 ++++++------ lmdeploy/pytorch/kernels/cuda/turbo_quant.py | 2 +- lmdeploy/pytorch/model_inputs.py | 12 +- lmdeploy/pytorch/nn/attention.py | 6 +- lmdeploy/turbomind/turbomind.py | 2 +- src/turbomind/engine/engine.cc | 2 +- src/turbomind/engine/engine_config.h | 2 +- .../kernels/attention/attention_params.h | 2 +- src/turbomind/kernels/attention/decoding.cu | 4 +- .../kernels/attention/kv_cache_utils_v2.cu | 16 +-- .../kernels/attention/kv_cache_utils_v2.h | 8 +- .../kernels/attention/test_attention.cu | 26 ++-- src/turbomind/models/llama/SequenceManager.cc | 4 +- src/turbomind/models/llama/SequenceManager.h | 2 +- src/turbomind/models/llama/llama_utils.h | 2 +- .../models/llama/unified_attention_layer.cc | 6 +- .../models/llama/unified_attention_layer.h | 4 +- src/turbomind/models/llama/unified_decoder.cc | 2 +- tests/pytorch/kernel/test_fa3_attention.py | 4 +- tests/pytorch/kernel/test_fill_kv_cache.py | 24 ++-- tests/pytorch/kernel/test_flatten_kv_cache.py | 38 +++--- tests/pytorch/kernel/test_paged_attention.py | 53 ++++---- tests/pytorch/kernel/test_turboquant.py | 6 +- tests/pytorch/kernel/turboquant_utils.py | 6 +- .../test_lmdeploy/test_fp8_kv_cache_dtype.py | 128 ++++++++++++++++++ .../test_lmdeploy/test_fp8_kv_cache_policy.py | 82 ----------- ...quant_policy.py => test_kv_cache_dtype.py} | 70 +++++----- 62 files changed, 654 insertions(+), 589 deletions(-) create mode 100644 tests/test_lmdeploy/test_fp8_kv_cache_dtype.py delete mode 100644 tests/test_lmdeploy/test_fp8_kv_cache_policy.py rename tests/test_lmdeploy/{test_quant_policy.py => test_kv_cache_dtype.py} (89%) diff --git a/.github/scripts/eval_chat_config.py b/.github/scripts/eval_chat_config.py index 29a8ee21d5..ec9d3e3f23 100644 --- a/.github/scripts/eval_chat_config.py +++ b/.github/scripts/eval_chat_config.py @@ -316,11 +316,11 @@ model['path'] = model['path'] + '-inner-w8a8' for model in [v for k, v in locals().items() if k.endswith('_kvint4')]: - model['engine_config']['quant_policy'] = 4 + model['engine_config']['kv_cache_dtype'] = 4 model['abbr'] = model['abbr'] + '_kvint4' for model in [v for k, v in locals().items() if k.endswith('_kvint8')]: - model['engine_config']['quant_policy'] = 8 + model['engine_config']['kv_cache_dtype'] = 8 model['abbr'] = model['abbr'] + '_kvint8' for model in [v for k, v in locals().items() if k.startswith('pytorch_')]: diff --git a/autotest/benchmark/test_apiserver_performance.py b/autotest/benchmark/test_apiserver_performance.py index 77b5210573..1d5f960184 100644 --- a/autotest/benchmark/test_apiserver_performance.py +++ b/autotest/benchmark/test_apiserver_performance.py @@ -113,7 +113,7 @@ def test_pytorch_apiserver_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-30B-A3B', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -122,7 +122,7 @@ def test_pytorch_apiserver_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-30B-A3B', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 4, + 'kv_cache_dtype': 4, 'parallel_config': { 'tp': 2 }, @@ -131,7 +131,7 @@ def test_pytorch_apiserver_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-30B-A3B', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 2 }, @@ -140,7 +140,7 @@ def test_pytorch_apiserver_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-VL-30B-A3B-Instruct', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 2 }, diff --git a/autotest/benchmark/test_prefixcache_performance.py b/autotest/benchmark/test_prefixcache_performance.py index e655cf1ca6..16f01b93c3 100644 --- a/autotest/benchmark/test_prefixcache_performance.py +++ b/autotest/benchmark/test_prefixcache_performance.py @@ -96,7 +96,7 @@ def test_pytorch_prefix_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-30B-A3B', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -105,7 +105,7 @@ def test_pytorch_prefix_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-VL-30B-A3B-Instruct', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 2 }, @@ -114,7 +114,7 @@ def test_pytorch_prefix_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-30B-A3B', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -123,7 +123,7 @@ def test_pytorch_prefix_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-VL-30B-A3B-Instruct', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 2 }, diff --git a/autotest/benchmark/test_throughput_performance.py b/autotest/benchmark/test_throughput_performance.py index 72a98a6c80..526703d4fc 100644 --- a/autotest/benchmark/test_throughput_performance.py +++ b/autotest/benchmark/test_throughput_performance.py @@ -96,7 +96,7 @@ def test_pytorch_throughput_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-30B-A3B', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -105,7 +105,7 @@ def test_pytorch_throughput_tp16(config, run_config, worker_id): 'model': 'Qwen/Qwen3-VL-30B-A3B-Instruct', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 2 }, @@ -123,7 +123,7 @@ def test_throughput_func_tp2(config, run_config, worker_id): 'model': 'meta-llama/Meta-Llama-3-1-8B-Instruct', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -132,7 +132,7 @@ def test_throughput_func_tp2(config, run_config, worker_id): 'model': 'Qwen/Qwen3-VL-8B-Instruct', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 1 }, diff --git a/autotest/interface/pipeline/test_pipeline_func.py b/autotest/interface/pipeline/test_pipeline_func.py index 85246754b8..6e4387df8a 100644 --- a/autotest/interface/pipeline/test_pipeline_func.py +++ b/autotest/interface/pipeline/test_pipeline_func.py @@ -618,7 +618,7 @@ def test_backend_config_validate_turbomind(config, model, backend, worker_id): pipeline(model_path, backend_config=backend_config) with pytest.raises(pydantic.ValidationError): - backend_config = backend(quant_policy=1) + backend_config = backend(kv_cache_dtype=1) pipeline(model_path, backend_config=backend_config) with pytest.raises(pydantic.ValidationError): diff --git a/autotest/tools/common_case_config.py b/autotest/tools/common_case_config.py index 592c5870c6..05092642a2 100644 --- a/autotest/tools/common_case_config.py +++ b/autotest/tools/common_case_config.py @@ -3,7 +3,7 @@ 'model': 'Qwen/Qwen3-30B-A3B', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -12,7 +12,7 @@ 'model': 'mistralai/Mixtral-8x7B-Instruct-v0.1', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -23,7 +23,7 @@ 'model': 'Qwen/Qwen3-0.6B', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -32,7 +32,7 @@ 'model': 'Qwen/Qwen3-0.6B-inner-4bits', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -41,7 +41,7 @@ 'model': 'Qwen/Qwen3-8B', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 1 }, @@ -52,7 +52,7 @@ 'model': 'OpenGVLab/InternVL3-8B', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -61,7 +61,7 @@ 'model': 'OpenGVLab/InternVL3-8B', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 1 }, @@ -72,7 +72,7 @@ 'model': 'OpenGVLab/InternVL3_5-30B-A3B', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -81,7 +81,7 @@ 'model': 'OpenGVLab/InternVL3_5-30B-A3B', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 2 }, @@ -92,7 +92,7 @@ 'model': 'google/gemma-2-9b-it', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 1 }, @@ -103,7 +103,7 @@ 'model': 'Qwen/Qwen3-VL-30B-A3B-Instruct', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -112,7 +112,7 @@ 'model': 'Qwen/Qwen3-VL-30B-A3B-Instruct', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 1 }, @@ -123,7 +123,7 @@ 'model': 'THUDM/cogvlm-chat-hf', 'backend': 'turbomind', 'communicator': 'cuda-ipc', - 'quant_policy': 4, + 'kv_cache_dtype': 4, 'parallel_config': { 'tp': 1 }, @@ -132,7 +132,7 @@ 'model': 'THUDM/cogvlm-chat-hf', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -143,7 +143,7 @@ 'model': 'Qwen/Qwen3-30B-A3B', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -152,7 +152,7 @@ 'model': 'OpenGVLab/InternVL3-38B', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -162,7 +162,7 @@ BASE_MODELSCOPE_CONFIG = [{ 'model': 'Qwen/Qwen2.5-7B-Instruct', 'communicator': 'cuda-ipc', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -173,7 +173,7 @@ }, { 'model': 'Qwen/Qwen2.5-7B-Instruct', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 1 }, @@ -193,7 +193,7 @@ 'model': 'meta-llama/Llama-2-7b-chat-hf', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -208,7 +208,7 @@ 'model': 'baichuan-inc/Baichuan2-13B-Chat', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -224,7 +224,7 @@ 'model': 'Qwen/Qwen3-30B-A3B', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 2 }, @@ -233,7 +233,7 @@ 'model': 'mistralai/Mixtral-8x7B-Instruct-v0.1', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -242,7 +242,7 @@ 'model': 'Qwen/Qwen3.5-35B-A3B', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -258,7 +258,7 @@ 'model': 'meta-llama/Meta-Llama-3-1-8B-Instruct', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -267,7 +267,7 @@ 'model': 'Qwen/Qwen3-0.6B', 'backend': 'pytorch', 'communicator': 'nccl', - 'quant_policy': 8, + 'kv_cache_dtype': 8, 'parallel_config': { 'tp': 1 }, @@ -277,7 +277,7 @@ BASE_TOOLCALL_TEST_LLM = [{ 'model': 'Qwen/Qwen3-8B', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -287,7 +287,7 @@ }, { 'model': 'meta-llama/Meta-Llama-3-1-70B-Instruct', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 4 }, @@ -297,7 +297,7 @@ }, { 'model': 'Qwen/Qwen3-30B-A3B', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, @@ -309,7 +309,7 @@ BASE_REASONING_TEST_LLM = [{ 'model': 'Qwen/Qwen3-VL-30B-A3B-Instruct', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 1 }, @@ -319,7 +319,7 @@ }, { 'model': 'Qwen/Qwen3-30B-A3B', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, diff --git a/autotest/tools/pipeline/llm_case.py b/autotest/tools/pipeline/llm_case.py index a68fab4a74..9c411b38f5 100644 --- a/autotest/tools/pipeline/llm_case.py +++ b/autotest/tools/pipeline/llm_case.py @@ -13,14 +13,14 @@ def run_pipeline_chat_test(model_path, run_config, cases_path, is_pr_test: bool = False): backend = run_config.get('backend') communicator = run_config.get('communicator') - quant_policy = run_config.get('quant_policy') + kv_cache_dtype = run_config.get('kv_cache_dtype') extra_params = run_config.get('extra_params', {}) parallel_config = run_config.get('parallel_config', {}) if backend == 'pytorch': - backend_config = PytorchEngineConfig(quant_policy=quant_policy) + backend_config = PytorchEngineConfig(kv_cache_dtype=kv_cache_dtype) else: - backend_config = TurbomindEngineConfig(communicator=communicator, quant_policy=quant_policy) + backend_config = TurbomindEngineConfig(communicator=communicator, kv_cache_dtype=kv_cache_dtype) # quant format model_lower = model_path.lower() diff --git a/autotest/tools/pipeline/mllm_case.py b/autotest/tools/pipeline/mllm_case.py index abe44e0ab9..43f0dd7d85 100644 --- a/autotest/tools/pipeline/mllm_case.py +++ b/autotest/tools/pipeline/mllm_case.py @@ -139,16 +139,18 @@ def _best_effort_runtime_cleanup(pipe=None, device: str | None = None) -> None: def run_pipeline_mllm_test(model_path, run_config, resource_path, is_pr_test: bool = False): backend = run_config.get('backend') communicator = run_config.get('communicator') - quant_policy = run_config.get('quant_policy') + kv_cache_dtype = run_config.get('kv_cache_dtype') extra_params = run_config.get('extra_params', {}) parallel_config = run_config.get('parallel_config', {}) if 'pytorch' == backend: - backend_config = PytorchEngineConfig(session_len=65152, quant_policy=quant_policy, cache_max_entry_count=0.6) + backend_config = PytorchEngineConfig(session_len=65152, + kv_cache_dtype=kv_cache_dtype, + cache_max_entry_count=0.6) else: backend_config = TurbomindEngineConfig(session_len=65152, communicator=communicator, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, cache_max_entry_count=0.6) # quant format diff --git a/autotest/utils/config_utils.py b/autotest/utils/config_utils.py index 00facd6964..7af0df17d9 100644 --- a/autotest/utils/config_utils.py +++ b/autotest/utils/config_utils.py @@ -380,15 +380,15 @@ def _quant_cfg_for_entry(entry: dict[str, Any]) -> dict[str, list[str]]: def _is_kvint_enabled_in_entry( backend: str, base_model: str, - quant_policy: int, + kv_cache_dtype: int, quant_cfg: dict[str, list[str]], ) -> bool: - if quant_policy == 0: + if kv_cache_dtype == 0: return True enabled = set(quant_cfg.get(backend) or []) - if quant_policy in (4, 8): - return f'kvint{quant_policy}' in enabled - if quant_policy == 42: + if kv_cache_dtype in (4, 8): + return f'kvint{kv_cache_dtype}' in enabled + if kv_cache_dtype == 42: return 'kvint42' in enabled return False @@ -427,7 +427,7 @@ def _build_run_config_entry( backend: str, communicator: str, parallel_config: dict[str, int], - quant_policy: int, + kv_cache_dtype: int, config: dict[str, Any], func_type: str, extra: dict[str, Any] | None, @@ -447,7 +447,7 @@ def _build_run_config_entry( 'model': model_id, 'backend': backend, 'communicator': communicator, - 'quant_policy': quant_policy, + 'kv_cache_dtype': kv_cache_dtype, 'parallel_config': copy.deepcopy(parallel_config), 'extra_params': merged_extra, } @@ -503,11 +503,11 @@ def _get_func_config_list_per_model( for model in models_for_quant: qcfg = quant_cfg - for quant_policy in [0, 4, 8, 42]: - if not _is_kvint_enabled_in_entry(backend, _base_model_name(model), quant_policy, qcfg): + for kv_cache_dtype in [0, 4, 8, 42]: + if not _is_kvint_enabled_in_entry(backend, _base_model_name(model), kv_cache_dtype, qcfg): continue for communicator in backend_map[backend]: - sig = (model, communicator, quant_policy, launch_extra_sig, '') + sig = (model, communicator, kv_cache_dtype, launch_extra_sig, '') if sig in seen: continue seen.add(sig) @@ -517,7 +517,7 @@ def _get_func_config_list_per_model( backend, communicator, parallel_config, - quant_policy, + kv_cache_dtype, config, func_type, extra, @@ -525,14 +525,14 @@ def _get_func_config_list_per_model( run_configs.append(run_config) # Runtime fp8 (--model-format fp8) is a separate case at - # quant_policy 0, mirroring legacy flat-yaml fp8_model_list. + # kv_cache_dtype 0, mirroring legacy flat-yaml fp8_model_list. if ( - quant_policy == 0 + kv_cache_dtype == 0 and _is_fp8_enabled_in_entry(backend, qcfg) and 'fp8' not in model.lower() and run_config.get('extra_params', {}).get('model-format') is None ): - fp8_sig = (model, communicator, quant_policy, launch_extra_sig, 'fp8') + fp8_sig = (model, communicator, kv_cache_dtype, launch_extra_sig, 'fp8') if fp8_sig not in seen: seen.add(fp8_sig) fp8_config = copy.deepcopy(run_config) @@ -561,14 +561,14 @@ def get_cli_common_param(run_config: dict[str, Any]) -> str: backend = run_config.get('backend') model = run_config.get('model') communicator = run_config.get('communicator') - quant_policy = run_config.get('quant_policy') + kv_cache_dtype = run_config.get('kv_cache_dtype') extra_params = run_config.get('extra_params', {}) parallel_config = run_config.get('parallel_config', {}) cli_params = [f'--backend {backend}', f'--communicator {communicator}'] # Optional params - if quant_policy != 0: - cli_params.append(f'--quant-policy {quant_policy}') + if kv_cache_dtype != 0: + cli_params.append(f'--kv-cache-dtype {kv_cache_dtype}') # quant format model_lower = model.lower() @@ -697,9 +697,9 @@ def get_model_list(config: dict[str, Any], return [m for m in typed_models if _base_model_name(m) in chat_bases] -def _is_kvint_model(config: dict[str, Any], backend: str, model: str, quant_policy: int) -> bool: - """Check KV quant policy support via per-model ``quantization`` blocks.""" - if quant_policy == 0: +def _is_kvint_model(config: dict[str, Any], backend: str, model: str, kv_cache_dtype: int) -> bool: + """Check KV cache dtype support via per-model ``quantization`` blocks.""" + if kv_cache_dtype == 0: return True env_key = _model_matrix_env_key(config) deps_profile = get_deps_profile_selector() @@ -710,7 +710,7 @@ def _is_kvint_model(config: dict[str, Any], backend: str, model: str, quant_poli layout = _parallel_layout(_entry_engine_config(entry)) if backend not in _normalize_entry_backends(entry, config, layout): continue - return _is_kvint_enabled_in_entry(backend, base, quant_policy, _quant_cfg_for_entry(entry)) + return _is_kvint_enabled_in_entry(backend, base, kv_cache_dtype, _quant_cfg_for_entry(entry)) return False def _base_model_name(model: str) -> str: @@ -987,7 +987,7 @@ def get_case_str_by_config(run_config: dict[str, Any], is_simple: bool = True) - model_name = run_config['model'] backend_type = run_config['backend'] communicator = run_config.get('communicator', 'nccl') - quant_policy = run_config.get('quant_policy', 0) + kv_cache_dtype = run_config.get('kv_cache_dtype', 0) parallel_config = run_config.get('parallel_config', {'tp': 1}) extra_params = run_config.get('extra_params', {}) @@ -1010,7 +1010,7 @@ def get_case_str_by_config(run_config: dict[str, Any], is_simple: bool = True) - else: extra_params_case += f'_{k}{v}'.replace('_', '-').replace('/', '-').replace('.', '-') - return f'{backend_type}_{pure_model_name}_{communicator}_{parallel_str}_{quant_policy}{extra_params_case}' + return f'{backend_type}_{pure_model_name}_{communicator}_{parallel_str}_{kv_cache_dtype}{extra_params_case}' def parse_config_by_case(case_str: str) -> dict[str, Any]: @@ -1029,9 +1029,9 @@ def parse_config_by_case(case_str: str) -> dict[str, Any]: quant_idx = i break if quant_idx is None: - raise ValueError(f'No numeric quant policy found in case string: {case_str}') + raise ValueError(f'No numeric KV cache dtype found in case string: {case_str}') - quant_policy = int(case_parts[quant_idx]) + kv_cache_dtype = int(case_parts[quant_idx]) parallel_parts = case_parts[3:quant_idx] # Convert parallel str to dict, e.g: ['tp1','dp2'] -> {'tp':1, 'dp':2} @@ -1049,5 +1049,5 @@ def parse_config_by_case(case_str: str) -> dict[str, Any]: 'model': model, 'communicator': communicator, 'parallel_config': parallel_config, - 'quant_policy': quant_policy, + 'kv_cache_dtype': kv_cache_dtype, } diff --git a/autotest/utils/constant.py b/autotest/utils/constant.py index 87e2759395..b801cc1150 100644 --- a/autotest/utils/constant.py +++ b/autotest/utils/constant.py @@ -258,7 +258,7 @@ def _deps_profile_is_legacy() -> bool: 'model': 'Qwen/Qwen2.5-32B-Instruct', 'backend': 'turbomind', 'communicator': 'nccl', - 'quant_policy': 0, + 'kv_cache_dtype': 0, 'parallel_config': { 'tp': 2 }, diff --git a/autotest/utils/evaluate_utils.py b/autotest/utils/evaluate_utils.py index cd1d8f793e..b22cc631bd 100644 --- a/autotest/utils/evaluate_utils.py +++ b/autotest/utils/evaluate_utils.py @@ -79,7 +79,7 @@ def write_to_summary(case_name, result, msg, metrics, result_dir): model = config['model'] communicator = config['communicator'] parallel_config_str = config['parallel_config'] - quant_policy = config['quant_policy'] + kv_cache_dtype = config['kv_cache_dtype'] dataset_name = [] dataset_metrics = [] @@ -92,7 +92,7 @@ def write_to_summary(case_name, result, msg, metrics, result_dir): summary_file = os.environ.get('GITHUB_STEP_SUMMARY', '') md_summary_file = f'{result_dir}/summary_{case_name}.md' - summary_line = f'| {model} | {quant_policy} | {backend} | {communicator} | {parallel_config_str} | {status} | {summary_dataset_metrics} |\n' # noqa: E501 + summary_line = f'| {model} | {kv_cache_dtype} | {backend} | {communicator} | {parallel_config_str} | {status} | {summary_dataset_metrics} |\n' # noqa: E501 write_header = not os.path.exists(md_summary_file) or os.path.getsize(md_summary_file) == 0 with open(md_summary_file, 'a') as f: @@ -100,7 +100,7 @@ def write_to_summary(case_name, result, msg, metrics, result_dir): dash_line = '-----|' * (len(metrics.keys())) f.write('## Model Evaluation Results\n') f.write( - f'| Model | QuantPolicy | Backend | Communicator | Parallel config | Status | {summary_dataset_name} |\n' # noqa + f'| Model | KVCacheDType | Backend | Communicator | Parallel config | Status | {summary_dataset_name} |\n' # noqa ) f.write(f'|-------|-------------|---------|--------------|----|--------|{dash_line}\n') f.write(summary_line) @@ -111,7 +111,7 @@ def write_to_summary(case_name, result, msg, metrics, result_dir): dash_line = '-----|' * (len(metrics.keys())) f.write('## Model Evaluation Results\n') f.write( - f'| Model | QuantPolicy | Backend | Communicator | Parallel config | Status | {summary_dataset_name} |\n' # noqa + f'| Model | KVCacheDType | Backend | Communicator | Parallel config | Status | {summary_dataset_name} |\n' # noqa ) f.write(f'|-------|-------------|---------|--------------|----|--------|{dash_line}\n') f.write(summary_line) diff --git a/benchmark/benchmark_guided.py b/benchmark/benchmark_guided.py index b3ac6d0812..faef613c01 100644 --- a/benchmark/benchmark_guided.py +++ b/benchmark/benchmark_guided.py @@ -501,7 +501,7 @@ def parse_args(): tb_group._group_actions.append(cache_block_seq_len_act) tb_group._group_actions.append(prefix_caching_act) ArgumentHelper.model_format(tb_group, default='hf') - ArgumentHelper.quant_policy(tb_group, default=0) + ArgumentHelper.kv_cache_dtype(tb_group, default=0) ArgumentHelper.num_tokens_per_iter(tb_group) ArgumentHelper.max_prefill_iters(tb_group) ArgumentHelper.communicator(tb_group) @@ -547,7 +547,7 @@ def main(): session_len=args.session_len, cache_block_seq_len=args.cache_block_seq_len, model_format=args.model_format, - quant_policy=args.quant_policy, + kv_cache_dtype=args.kv_cache_dtype, num_tokens_per_iter=args.num_tokens_per_iter, max_prefill_iters=args.max_prefill_iters, enable_prefix_caching=args.enable_prefix_caching, diff --git a/benchmark/profile_pipeline_api.py b/benchmark/profile_pipeline_api.py index e679224b5a..6e8abd28b8 100644 --- a/benchmark/profile_pipeline_api.py +++ b/benchmark/profile_pipeline_api.py @@ -281,7 +281,7 @@ def parse_args(): tb_group._group_actions.append(cache_block_seq_len_act) tb_group._group_actions.append(prefix_caching_act) ArgumentHelper.model_format(tb_group, default='hf') - ArgumentHelper.quant_policy(tb_group, default=0) + ArgumentHelper.kv_cache_dtype(tb_group, default=0) ArgumentHelper.num_tokens_per_iter(tb_group) ArgumentHelper.max_prefill_iters(tb_group) ArgumentHelper.communicator(tb_group) @@ -302,7 +302,7 @@ def main(): session_len=args.session_len, cache_block_seq_len=args.cache_block_seq_len, model_format=args.model_format, - quant_policy=args.quant_policy, + kv_cache_dtype=args.kv_cache_dtype, num_tokens_per_iter=args.num_tokens_per_iter, max_prefill_iters=args.max_prefill_iters, enable_prefix_caching=args.enable_prefix_caching, diff --git a/benchmark/profile_throughput.py b/benchmark/profile_throughput.py index 6e098bb55e..350c27d88f 100644 --- a/benchmark/profile_throughput.py +++ b/benchmark/profile_throughput.py @@ -326,7 +326,7 @@ def parse_args(): cache_count_act = ArgumentHelper.cache_max_entry_count(pt_group) cache_block_seq_len_act = ArgumentHelper.cache_block_seq_len(pt_group) prefix_caching_act = ArgumentHelper.enable_prefix_caching(pt_group) - quant_policy_act = ArgumentHelper.quant_policy(pt_group, default=0) + kv_cache_dtype_act = ArgumentHelper.kv_cache_dtype(pt_group, default=0) dtype_act = ArgumentHelper.dtype(pt_group) # turbomind engine args @@ -335,7 +335,7 @@ def parse_args(): tb_group._group_actions.append(cache_count_act) tb_group._group_actions.append(cache_block_seq_len_act) tb_group._group_actions.append(prefix_caching_act) - tb_group._group_actions.append(quant_policy_act) + tb_group._group_actions.append(kv_cache_dtype_act) tb_group._group_actions.append(dtype_act) ArgumentHelper.dp(tb_group) @@ -362,7 +362,7 @@ def main(): cache_max_entry_count=args.cache_max_entry_count, cache_block_seq_len=args.cache_block_seq_len, model_format=args.model_format, - quant_policy=args.quant_policy, + kv_cache_dtype=args.kv_cache_dtype, num_tokens_per_iter=args.num_tokens_per_iter, max_prefill_iters=args.max_prefill_iters, async_=args.async_, @@ -379,7 +379,7 @@ def main(): device_type=args.device, eager_mode=args.eager_mode, enable_prefix_caching=args.enable_prefix_caching, - quant_policy=args.quant_policy, + kv_cache_dtype=args.kv_cache_dtype, dtype=args.dtype, distributed_executor_backend=args.distributed_executor_backend, dllm_block_length=args.dllm_block_length, diff --git a/docs/en/inference/turbomind_config.md b/docs/en/inference/turbomind_config.md index 5f549fb0c7..a79ceb6f0d 100644 --- a/docs/en/inference/turbomind_config.md +++ b/docs/en/inference/turbomind_config.md @@ -34,7 +34,7 @@ cache_max_entry_count = 0.5 cache_block_seq_len = 128 cache_chunk_size = 1 enable_prefix_caching = false -quant_policy = 0 +kv_cache_dtype = 0 max_position_embeddings = 2048 rope_scaling_factor = 0.0 use_logn_attn = 0 @@ -107,8 +107,8 @@ Since k/v block is the smallest granularity for reuse in prefix caching, if the ### kv quantization and inference switch -- `quant_policy=4` means 4bit k/v quantization and inference -- `quant_policy=8` indicates 8bit k/v quantization and inference +- `kv_cache_dtype='int4'` means 4bit k/v quantization and inference +- `kv_cache_dtype='int8'` indicates 8bit k/v quantization and inference Please refer to [kv quant](../quantization/kv_quant.md) for detailed guide. @@ -152,7 +152,7 @@ step_length = 1 cache_max_entry_count = 48 cache_chunk_size = 1 use_context_fmha = 1 -quant_policy = 0 +kv_cache_dtype = 0 max_position_embeddings = 2048 use_dynamic_ntk = 0 use_logn_attn = 0 @@ -198,7 +198,7 @@ TurboMind allocates k/v cache memory based on `session_len`, `cache_chunk_size`, ### kv int8 switch -When initiating 8bit k/v inference, change `quant_policy = 4` and `use_context_fmha = 0`. Please refer to [kv int8](../quantization/kv_quant.md) for a guide. +When initiating 8bit k/v inference, change `kv_cache_dtype = 'int8'` and `use_context_fmha = 0`. Please refer to [kv int8](../quantization/kv_quant.md) for a guide. ### long context switch diff --git a/docs/en/quantization/kv_quant.md b/docs/en/quantization/kv_quant.md index eda4841c4a..bf45b5da1b 100644 --- a/docs/en/quantization/kv_quant.md +++ b/docs/en/quantization/kv_quant.md @@ -57,15 +57,15 @@ TurboQuant achieves efficient compression through two key steps: Tested on H200 with Qwen3-30B-A3B-Base model and ShareGPT dataset: -| Metric | Baseline (quant_policy=0) | TurboQuant (quant_policy=42) | Change | -| ------------------ | ------------------------- | ---------------------------- | ---------- | -| Input throughput | 2368.8 tok/s | 2195.8 tok/s | -7.3% | -| Output throughput | 2186.7 tok/s | 2027.0 tok/s | -7.3% | -| Request throughput | 10.74 req/s | 9.96 req/s | -7.3% | -| Mean E2E latency | 5.888s | 6.348s | +7.8% | -| Mean TTFT | 1.139s | 1.235s | +8.4% | -| Mean TPOT | 0.024s | 0.026s | +8.3% | -| Mean ITL | 0.059s | 0.059s | ~unchanged | +| Metric | Baseline (kv_cache_dtype=0) | TurboQuant (kv_cache_dtype=42) | Change | +| ------------------ | --------------------------- | ------------------------------ | ---------- | +| Input throughput | 2368.8 tok/s | 2195.8 tok/s | -7.3% | +| Output throughput | 2186.7 tok/s | 2027.0 tok/s | -7.3% | +| Request throughput | 10.74 req/s | 9.96 req/s | -7.3% | +| Mean E2E latency | 5.888s | 6.348s | +7.8% | +| Mean TTFT | 1.139s | 1.235s | +8.4% | +| Mean TPOT | 0.024s | 0.026s | +8.3% | +| Mean ITL | 0.059s | 0.059s | ~unchanged | **Test configuration**: GPU: H200, Model: Qwen3-30B-A3B-Base, Dataset: ShareGPT, Concurrency: 64, Requests: 5000 @@ -97,15 +97,15 @@ pip install lmdeploy ## Usage -Applying kv quantization and inference via LMDeploy is quite straightforward. Simply set the `quant_policy` parameter. +Applying kv quantization and inference via LMDeploy is quite straightforward. Simply set the `kv_cache_dtype` parameter. -**LMDeploy specifies that `quant_policy=4` stands for 4-bit kv, `quant_policy=8` indicates 8-bit kv, and `quant_policy=42` indicates TurboQuant.** +**LMDeploy specifies that `kv_cache_dtype='int4'` stands for 4-bit kv, `kv_cache_dtype='int8'` indicates 8-bit kv, and `kv_cache_dtype='turbo_quant'` indicates TurboQuant. The numeric values `4`, `8`, and `42` are also accepted.** ### Offline inference ```python from lmdeploy import pipeline, TurbomindEngineConfig -engine_config = TurbomindEngineConfig(quant_policy=8) +engine_config = TurbomindEngineConfig(kv_cache_dtype='int8') pipe = pipeline("internlm/internlm2_5-7b-chat", backend_config=engine_config) response = pipe(["Hi, pls intro yourself", "Shanghai is"]) print(response) @@ -114,19 +114,19 @@ print(response) ### Serving ```shell -lmdeploy serve api_server internlm/internlm2_5-7b-chat --quant-policy 8 +lmdeploy serve api_server internlm/internlm2_5-7b-chat --kv-cache-dtype int8 ``` ### TurboQuant -TurboQuant uses `quant_policy=42`, **PytorchEngine only**: +TurboQuant uses `kv_cache_dtype='turbo_quant'`, **PytorchEngine only**: ```python from lmdeploy import pipeline, PytorchEngineConfig engine_config = PytorchEngineConfig( tp=1, cache_max_entry_count=0.8, - quant_policy=42 # TurboQuant: K=4bit QJL4 + V=2bit MSE + kv_cache_dtype='turbo_quant' # TurboQuant: K=4bit QJL4 + V=2bit MSE ) pipe = pipeline("Qwen/Qwen3-8B", backend_config=engine_config) response = pipe.infer("Hello, how are you?", max_new_tokens=30) @@ -147,7 +147,7 @@ We apply kv quantization of LMDeploy to several LLM models and utilize OpenCompa | race-middle | 9a54b6 | accuracy | 41.57 | 41.78 | 41.23 | 88.93 | 88.93 | 88.93 | 92.76 | 92.83 | 92.55 | 87.33 | 87.26 | 86.28 | | race-high | 9a54b6 | accuracy | 39.65 | 39.77 | 40.77 | 85.33 | 85.31 | 84.62 | 90.51 | 90.42 | 90.42 | 82.53 | 82.59 | 82.02 | -For detailed evaluation methods, please refer to [this](../benchmark/evaluate_with_opencompass.md) guide. Remember to pass `quant_policy` to the inference engine in the config file. +For detailed evaluation methods, please refer to [this](../benchmark/evaluate_with_opencompass.md) guide. Remember to pass `kv_cache_dtype` to the inference engine in the config file. ## Performance diff --git a/docs/zh_cn/inference/turbomind_config.md b/docs/zh_cn/inference/turbomind_config.md index 1033ecaa6d..e408491775 100644 --- a/docs/zh_cn/inference/turbomind_config.md +++ b/docs/zh_cn/inference/turbomind_config.md @@ -34,7 +34,7 @@ cache_max_entry_count = 0.5 cache_block_seq_len = 128 cache_chunk_size = 1 enable_prefix_caching = false -quant_policy = 0 +kv_cache_dtype = 0 max_position_embeddings = 2048 rope_scaling_factor = 0.0 use_logn_attn = 0 @@ -109,10 +109,10 @@ cache_block_seq_len * num_layer * kv_head_num * size_per_head * 2 * sizeof(kv_da ### kv 量化推理开关 -`quant_policy`是 kv 量化和推理开关。 +`kv_cache_dtype`是 kv 量化和推理开关。 -- `quant_policy=4` 代表 4bit k/v 量化和推理 -- `quant_policy=8` 代表 8bit k/v 量化和推理 +- `kv_cache_dtype='int4'` 代表 4bit k/v 量化和推理 +- `kv_cache_dtype='int8'` 代表 8bit k/v 量化和推理 具体使用方法，请参考 [kv quant](../quantization/kv_quant.md) 部署文档 @@ -156,7 +156,7 @@ step_length = 1 cache_max_entry_count = 48 cache_chunk_size = 1 use_context_fmha = 1 -quant_policy = 0 +kv_cache_dtype = 0 max_position_embeddings = 2048 use_dynamic_ntk = 0 use_logn_attn = 0 @@ -202,7 +202,7 @@ TurboMind 根据 `session_len`、 `cache_chunk_size` 和 `cache_max_entry_count` ### kv int8 开关 -当启动 8bit k/v 推理时，需要修改参数 `quant_policy` 和 `use_context_fmha`。详细内容请查阅 [kv int8](../quantization/kv_quant.md) 部署文档。 +当启动 8bit k/v 推理时，需要修改参数 `kv_cache_dtype` 和 `use_context_fmha`。详细内容请查阅 [kv int8](../quantization/kv_quant.md) 部署文档。 ### 外推能力开关 diff --git a/docs/zh_cn/quantization/kv_quant.md b/docs/zh_cn/quantization/kv_quant.md index 2651402af2..f1be207ba4 100644 --- a/docs/zh_cn/quantization/kv_quant.md +++ b/docs/zh_cn/quantization/kv_quant.md @@ -57,15 +57,15 @@ TurboQuant 通过两个关键步骤实现高效压缩： 在 H200 上使用 Qwen3-30B-A3B-Base 模型、ShareGPT 数据集进行测试： -| 指标 | Baseline (quant_policy=0) | TurboQuant (quant_policy=42) | 变化 | -| -------------- | ------------------------- | ---------------------------- | ----- | -| 输入吞吐 | 2368.8 tok/s | 2195.8 tok/s | -7.3% | -| 输出吞吐 | 2186.7 tok/s | 2027.0 tok/s | -7.3% | -| 请求吞吐 | 10.74 req/s | 9.96 req/s | -7.3% | -| 平均端到端延迟 | 5.888s | 6.348s | +7.8% | -| 平均 TTFT | 1.139s | 1.235s | +8.4% | -| 平均 TPOT | 0.024s | 0.026s | +8.3% | -| 平均 ITL | 0.059s | 0.059s | 持平 | +| 指标 | Baseline (kv_cache_dtype=0) | TurboQuant (kv_cache_dtype=42) | 变化 | +| -------------- | --------------------------- | ------------------------------ | ----- | +| 输入吞吐 | 2368.8 tok/s | 2195.8 tok/s | -7.3% | +| 输出吞吐 | 2186.7 tok/s | 2027.0 tok/s | -7.3% | +| 请求吞吐 | 10.74 req/s | 9.96 req/s | -7.3% | +| 平均端到端延迟 | 5.888s | 6.348s | +7.8% | +| 平均 TTFT | 1.139s | 1.235s | +8.4% | +| 平均 TPOT | 0.024s | 0.026s | +8.3% | +| 平均 ITL | 0.059s | 0.059s | 持平 | **测试配置**：GPU: H200, 模型: Qwen3-30B-A3B-Base, 数据集: ShareGPT, 并发: 64, 请求数: 5000 @@ -97,15 +97,15 @@ pip install lmdeploy ## 应用示例 -通过 LMDeploy 应用 kv 量化非常简单，只需要设定 `quant_policy` 参数。 +通过 LMDeploy 应用 kv 量化非常简单，只需要设定 `kv_cache_dtype` 参数。 -**LMDeploy 规定 `quant_policy=4` 表示 kv int4 量化，`quant_policy=8` 表示 kv int8 量化，`quant_policy=42` 表示 TurboQuant 量化。** +**LMDeploy 规定 `kv_cache_dtype='int4'` 表示 kv int4 量化，`kv_cache_dtype='int8'` 表示 kv int8 量化，`kv_cache_dtype='turbo_quant'` 表示 TurboQuant 量化。也可以继续使用数值 `4`、`8` 和 `42`。** ### 离线推理 ```python from lmdeploy import pipeline, TurbomindEngineConfig -engine_config = TurbomindEngineConfig(quant_policy=8) +engine_config = TurbomindEngineConfig(kv_cache_dtype='int8') pipe = pipeline("internlm/internlm2_5-7b-chat", backend_config=engine_config) response = pipe(["Hi, pls intro yourself", "Shanghai is"]) print(response) @@ -114,19 +114,19 @@ print(response) ### 推理服务 ```shell -lmdeploy serve api_server internlm/internlm2_5-7b-chat --quant-policy 8 +lmdeploy serve api_server internlm/internlm2_5-7b-chat --kv-cache-dtype int8 ``` ### TurboQuant 量化 -TurboQuant 量化使用 `quant_policy=42`，**仅支持 PytorchEngine**： +TurboQuant 量化使用 `kv_cache_dtype='turbo_quant'`，**仅支持 PytorchEngine**： ```python from lmdeploy import pipeline, PytorchEngineConfig engine_config = PytorchEngineConfig( tp=1, cache_max_entry_count=0.8, - quant_policy=42 # TurboQuant: K=4bit QJL4 + V=2bit MSE + kv_cache_dtype='turbo_quant' # TurboQuant: K=4bit QJL4 + V=2bit MSE ) pipe = pipeline("Qwen/Qwen3-8B", backend_config=engine_config) response = pipe.infer("Hello, how are you?", max_new_tokens=30) @@ -147,7 +147,7 @@ print(response.text) | race-middle | 9a54b6 | accuracy | 41.57 | 41.78 | 41.23 | 88.93 | 88.93 | 88.93 | 92.76 | 92.83 | 92.55 | 87.33 | 87.26 | 86.28 | | race-high | 9a54b6 | accuracy | 39.65 | 39.77 | 40.77 | 85.33 | 85.31 | 84.62 | 90.51 | 90.42 | 90.42 | 82.53 | 82.59 | 82.02 | -具体的评测方式可以参考[这份指南](../benchmark/evaluate_with_opencompass.md)。评测时，请在config文件中，为推理引擎添加 `quant_policy` 参数。 +具体的评测方式可以参考[这份指南](../benchmark/evaluate_with_opencompass.md)。评测时，请在config文件中，为推理引擎添加 `kv_cache_dtype` 参数。 ## 推理效率 diff --git a/lmdeploy/cli/cli.py b/lmdeploy/cli/cli.py index 26c412bffb..4c466d6589 100644 --- a/lmdeploy/cli/cli.py +++ b/lmdeploy/cli/cli.py @@ -62,7 +62,7 @@ def add_parser_chat(): session_len_act = ArgumentHelper.session_len(pt_group) cache_max_entry_act = ArgumentHelper.cache_max_entry_count(pt_group) prefix_caching_act = ArgumentHelper.enable_prefix_caching(pt_group) - quant_policy = ArgumentHelper.quant_policy(pt_group) + kv_cache_dtype = ArgumentHelper.kv_cache_dtype(pt_group) # turbomind args tb_group = parser.add_argument_group('TurboMind engine arguments') @@ -72,7 +72,7 @@ def add_parser_chat(): tb_group._group_actions.append(session_len_act) tb_group._group_actions.append(cache_max_entry_act) tb_group._group_actions.append(prefix_caching_act) - tb_group._group_actions.append(quant_policy) + tb_group._group_actions.append(kv_cache_dtype) ArgumentHelper.model_format(tb_group) ArgumentHelper.rope_scaling_factor(tb_group) ArgumentHelper.communicator(tb_group) diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index 01ac1d44f1..5e2725d772 100644 --- a/lmdeploy/cli/serve.py +++ b/lmdeploy/cli/serve.py @@ -120,7 +120,7 @@ def add_parser_api_server(): cache_block_seq_len_act = ArgumentHelper.cache_block_seq_len(pt_group) prefix_caching_act = ArgumentHelper.enable_prefix_caching(pt_group) max_prefill_token_num_act = ArgumentHelper.max_prefill_token_num(pt_group) - quant_policy = ArgumentHelper.quant_policy(pt_group) + kv_cache_dtype = ArgumentHelper.kv_cache_dtype(pt_group) model_format = ArgumentHelper.model_format(pt_group) hf_overrides = ArgumentHelper.hf_overrides(pt_group) disable_metrics = ArgumentHelper.disable_metrics(pt_group) @@ -146,7 +146,7 @@ def add_parser_api_server(): tb_group._group_actions.append(cache_block_seq_len_act) tb_group._group_actions.append(prefix_caching_act) tb_group._group_actions.append(max_prefill_token_num_act) - tb_group._group_actions.append(quant_policy) + tb_group._group_actions.append(kv_cache_dtype) tb_group._group_actions.append(model_format) tb_group._group_actions.append(num_nodes_act) tb_group._group_actions.append(node_rank_act) @@ -239,7 +239,7 @@ def api_server(args): prefix_cache_state_budget=args.prefix_cache_state_budget, prefix_cache_decode_state_interval=args.prefix_cache_decode_state_interval, device_type=args.device, - quant_policy=args.quant_policy, + kv_cache_dtype=args.kv_cache_dtype, eager_mode=args.eager_mode, max_prefill_token_num=args.max_prefill_token_num, cudagraph_capture_batch_sizes=args.cudagraph_capture_batch_sizes, @@ -271,7 +271,7 @@ def api_server(args): max_batch_size=max_batch_size, session_len=args.session_len, model_format=args.model_format, - quant_policy=args.quant_policy, + kv_cache_dtype=args.kv_cache_dtype, rope_scaling_factor=args.rope_scaling_factor, cache_max_entry_count=args.cache_max_entry_count, cache_block_seq_len=args.cache_block_seq_len, diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 1dda62a7e8..1de6b9d08d 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -264,27 +264,15 @@ def max_batch_size(parser): 'automatically set it according to the device') @staticmethod - def quant_policy(parser, default: int = 0): - """Add argument quant_policy to parser.""" + def kv_cache_dtype(parser, default: int = 0): + """Add argument kv_cache_dtype to parser.""" - from lmdeploy.messages import QuantPolicy + from lmdeploy.messages import normalize_kv_cache_dtype - _aliases = {p.name.lower(): p.value for p in QuantPolicy} - _aliases['fp8_e4m3'] = QuantPolicy.FP8.value - - def _parse(x): - key = x.lower() - if key in _aliases: - return _aliases[key] - v = int(x) - if v not in list(QuantPolicy): - raise ValueError(f'invalid quant_policy: {x!r}') - return v - - return parser.add_argument('--quant-policy', - type=_parse, - default=default, - help='KV cache quant policy: none/int4/int8/fp8/fp8_e5m2/' + return parser.add_argument('--kv-cache-dtype', + type=normalize_kv_cache_dtype, + default=normalize_kv_cache_dtype(default), + help='KV cache dtype: auto/int4/int8/fp8/fp8_e4m3/fp8_e5m2/' 'turbo_quant (or 0/4/8/16/17/42). fp8 defaults to fp8_e4m3.') @staticmethod diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index 9486cc62b0..d467db06ac 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -17,15 +17,43 @@ logger = get_logger('lmdeploy') -class QuantPolicy(enum.IntEnum): - """Quantization policy constants for KV cache.""" - NONE = 0 +class KVCacheDType(enum.IntEnum): + """KV cache dtype constants.""" + AUTO = 0 INT4 = 4 # 4-bit KV cache INT8 = 8 # 8-bit KV cache FP8 = 16 # FP8 KV cache (float8_e4m3fn, per-tensor scale) FP8_E5M2 = 17 # FP8 KV cache (float8_e5m2, per-tensor scale) TURBO_QUANT = 42 # TurboQuant: K=4bit QJL4 + V=2bit MSE + +_KV_CACHE_DTYPE_ALIASES = { + dtype.name.lower(): dtype + for dtype in KVCacheDType +} +_KV_CACHE_DTYPE_ALIASES['fp8_e4m3'] = KVCacheDType.FP8 + + +def normalize_kv_cache_dtype(kv_cache_dtype: KVCacheDType | int | str) -> KVCacheDType: + """Normalize KV cache dtype config values.""" + if isinstance(kv_cache_dtype, KVCacheDType): + return kv_cache_dtype + + if isinstance(kv_cache_dtype, str): + key = kv_cache_dtype.lower() + if key in _KV_CACHE_DTYPE_ALIASES: + return _KV_CACHE_DTYPE_ALIASES[key] + try: + kv_cache_dtype = int(kv_cache_dtype) + except ValueError as e: + raise ValueError(f'invalid kv_cache_dtype: {kv_cache_dtype}') from e + + try: + return KVCacheDType(kv_cache_dtype) + except ValueError as e: + raise ValueError(f'invalid kv_cache_dtype: {kv_cache_dtype}') from e + + LogitsProcessor = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] """LogitsProcessor is a function that takes a tensor of input_ids, the logits tensor for the next token, and returns a modified tensor of logits to sample @@ -205,7 +233,7 @@ def __post_init__(self): self.repetition_ngram_threshold = 0 -@pydantic_dataclass +@pydantic_dataclass(config={'extra': 'forbid'}) class TurbomindEngineConfig: """TurboMind Engine config. @@ -247,7 +275,7 @@ class TurbomindEngineConfig: a k/v block, default to 64 enable_prefix_caching: enable cache prompts for block reuse, default to False - quant_policy: default to 0. For TurboMind, when k/v is quantized + kv_cache_dtype: default to 0. For TurboMind, when k/v is quantized into int4 or int8, set it to 4 or 8, respectively rope_scaling_factor: scaling factor used for dynamic ntk, default to 0. TurboMind follows the implementation of transformer @@ -297,7 +325,7 @@ class TurbomindEngineConfig: cache_chunk_size: int = -1 cache_block_seq_len: int = 64 enable_prefix_caching: bool = False - quant_policy: int = 0 + kv_cache_dtype: KVCacheDType | int | str = KVCacheDType.AUTO rope_scaling_factor: float = 0.0 use_logn_attn: bool = False download_dir: str | None = None @@ -318,14 +346,11 @@ def __post_init__(self): assert self.dtype in ['auto', 'float16', 'bfloat16'] assert self.tp >= 1, 'tp must be a positive integer' assert self.cache_max_entry_count > 0, 'invalid cache_max_entry_count' - try: - self.quant_policy = QuantPolicy(self.quant_policy) - except ValueError as e: - raise ValueError(f'invalid quant_policy: {self.quant_policy}') from e - assert self.quant_policy not in ( - QuantPolicy.FP8, - QuantPolicy.FP8_E5M2, - ), 'invalid quant_policy for TurboMind, FP8 quantization is not supported' + self.kv_cache_dtype = normalize_kv_cache_dtype(self.kv_cache_dtype) + assert self.kv_cache_dtype not in ( + KVCacheDType.FP8, + KVCacheDType.FP8_E5M2, + ), 'invalid kv_cache_dtype for TurboMind, FP8 quantization is not supported' assert self.rope_scaling_factor >= 0, 'invalid rope_scaling_factor' assert self.max_prefill_token_num >= 0, \ 'invalid max_prefill_token_num' @@ -391,7 +416,7 @@ class PytorchEngineConfig: revision: The specific model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. - quant_policy: default to 0. When k/v is quantized into int4, + kv_cache_dtype: default to 0. When k/v is quantized into int4, int8, fp8, or fp8_e5m2, set it to 4, 8, 16, or 17, respectively distributed_executor_backend: backend of distributed backend, @@ -449,7 +474,7 @@ class PytorchEngineConfig: custom_module_map: dict[str, str] = None download_dir: str = None revision: str = None - quant_policy: QuantPolicy = QuantPolicy.NONE + kv_cache_dtype: KVCacheDType | int | str = KVCacheDType.AUTO distributed_executor_backend: str = None empty_init: bool = False enable_microbatch: bool = False @@ -490,10 +515,7 @@ def __post_init__(self): assert self.num_gpu_blocks >= 0, 'invalid num_gpu_blocks' assert self.prefix_cache_state_budget >= 0, 'invalid prefix_cache_state_budget' assert self.prefix_cache_decode_state_interval >= 0, 'invalid prefix_cache_decode_state_interval' - try: - self.quant_policy = QuantPolicy(self.quant_policy) - except ValueError as e: - raise ValueError(f'invalid quant_policy: {self.quant_policy}') from e + self.kv_cache_dtype = normalize_kv_cache_dtype(self.kv_cache_dtype) assert self.device_type in ['cuda', 'ascend', 'maca', 'camb'], (f'invalid device_type: {self.device_type}') assert self.kernel_block_size >= 16 and \ (self.kernel_block_size & (self.kernel_block_size - 1)) == 0, \ @@ -506,7 +528,7 @@ def __post_init__(self): if self.prefix_cache_decode_state_interval > 0: assert self.prefix_cache_decode_state_interval % self.block_size == 0, ( 'prefix_cache_decode_state_interval must be a multiple of block_size') - if self.quant_policy > 0 and self.device_type not in ['cuda', 'ascend']: + if self.kv_cache_dtype > 0 and self.device_type not in ['cuda', 'ascend']: assert False, \ 'kv cache quantization only works for CUDA and ASCEND.' if self.device_type == 'camb' and self.block_size != 16: diff --git a/lmdeploy/pytorch/backends/attention.py b/lmdeploy/pytorch/backends/attention.py index 29b9f27290..b7d74e16d7 100644 --- a/lmdeploy/pytorch/backends/attention.py +++ b/lmdeploy/pytorch/backends/attention.py @@ -6,7 +6,7 @@ import torch -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType @dataclass @@ -20,7 +20,7 @@ class AttentionMetadata: fill_seqlens: torch.Tensor = None cu_seqlens_q: torch.Tensor = None cu_seqlens_k: torch.Tensor = None - quant_policy: QuantPolicy = QuantPolicy.NONE + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO T = TypeVar('T', bound=AttentionMetadata) diff --git a/lmdeploy/pytorch/backends/cuda/attention/default.py b/lmdeploy/pytorch/backends/cuda/attention/default.py index f53d988c63..68a562c283 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/default.py +++ b/lmdeploy/pytorch/backends/cuda/attention/default.py @@ -3,7 +3,7 @@ import torch -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from lmdeploy.pytorch.backends.attention import AttentionImpl, AttentionMetadata from lmdeploy.utils import get_logger @@ -24,7 +24,7 @@ class TritonAttentionMetadata(AttentionMetadata): q_seqlens: Length of each query sequence [batch_size]. kv_start_loc: Start location of each KV sequence [batch_size]. kv_seqlens: Length of each KV sequence [batch_size]. - quant_policy: Quantization policy (0=none, 4=int4, 8=int8, 16/17=per-tensor fp8). + kv_cache_dtype: Quantization policy (0=none, 4=int4, 8=int8, 16/17=per-tensor fp8). kv_flatten_size: Total size of flattened KV cache. tile_scheduler_metadata: Scheduler metadata for Flash MLA. num_splits: Number of splits for Flash MLA. @@ -40,7 +40,7 @@ class TritonAttentionMetadata(AttentionMetadata): q_seqlens: torch.Tensor = None kv_start_loc: torch.Tensor = None kv_seqlens: torch.Tensor = None - quant_policy: QuantPolicy = QuantPolicy.NONE + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO kv_flatten_size: int = None # flash mla tile_scheduler_metadata: torch.Tensor = None @@ -153,7 +153,7 @@ def _fill_kv_cache_impl( """Fill kv cache.""" kv_seqlens = attn_metadata.kv_seqlens block_offsets = attn_metadata.block_offsets - quant_policy = attn_metadata.quant_policy + kv_cache_dtype = attn_metadata.kv_cache_dtype # fill seqlen args fill_seqlens, fill_max_q_seqlen, fill_q_start_loc = self._get_fill_meta( @@ -175,7 +175,7 @@ def _fill_kv_cache_impl( block_offsets=block_offsets, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, ) def _forward_decoding( @@ -205,7 +205,7 @@ def _forward_decoding( Attention output tensor. """ block_offsets = attn_metadata.block_offsets - quant_policy = attn_metadata.quant_policy + kv_cache_dtype = attn_metadata.kv_cache_dtype attn_output = self.paged_attention_fwd( query, @@ -221,7 +221,7 @@ def _forward_decoding( # custom args sinks=learnable_sink, alibi_slopes=self.alibi_slopes, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, ) @@ -257,7 +257,7 @@ def _forward_prefill( kv_start_loc = attn_metadata.kv_start_loc kv_seqlens = attn_metadata.kv_seqlens kv_flatten_size = attn_metadata.kv_flatten_size - quant_policy = attn_metadata.quant_policy + kv_cache_dtype = attn_metadata.kv_cache_dtype # Prepare flattened KV cache BLOCK_BS = k_cache.size(1) @@ -275,13 +275,13 @@ def _forward_prefill( out_dtype=query.dtype, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, flatten_kv_layout=kv_layout, ) - # For quant_policy==QuantPolicy.TURBO_QUANT, flattened K/V are in rotated domain. + # For kv_cache_dtype==KVCacheDType.TURBO_QUANT, flattened K/V are in rotated domain. # Rotate Q to match, and inverse-rotate output afterwards. - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: from lmdeploy.pytorch.kernels.cuda.turbo_quant import ( hadamard_rotate, hadamard_rotate_inv, @@ -308,7 +308,7 @@ def _forward_prefill( ) # Inverse-rotate output back to original domain - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: attn_output = hadamard_rotate_inv(attn_output) return attn_output diff --git a/lmdeploy/pytorch/backends/cuda/attention/fa3.py b/lmdeploy/pytorch/backends/cuda/attention/fa3.py index 84aeefa189..efd66512e9 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/fa3.py +++ b/lmdeploy/pytorch/backends/cuda/attention/fa3.py @@ -1,7 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from lmdeploy.utils import get_logger from .default import TritonAttentionImpl, TritonAttentionMetadata, _cdiv @@ -103,13 +103,13 @@ def _decoding_speculative( Returns: Attention output tensor. """ - quant_policy = attn_metadata.quant_policy + kv_cache_dtype = attn_metadata.kv_cache_dtype # TurboQuant stores packed uint8 data in cache, which FA3's native # flash_attn_with_kvcache cannot dequantize directly. - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: raise NotImplementedError( - 'quant_policy=QuantPolicy.TURBO_QUANT is not supported with ' + 'kv_cache_dtype=KVCacheDType.TURBO_QUANT is not supported with ' 'FA3 speculative decoding (max_q_seqlen > 1). ' 'FA3 speculative decoding accesses raw KV cache directly ' 'and cannot dequantize TurboQuant packed data. ' @@ -165,7 +165,7 @@ def _decoding_standard( Attention output tensor. """ block_offsets = attn_metadata.block_offsets - quant_policy = attn_metadata.quant_policy + kv_cache_dtype = attn_metadata.kv_cache_dtype attn_output = self.paged_attention_fwd( query, @@ -183,7 +183,7 @@ def _decoding_standard( # custom args k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, ) return attn_output @@ -258,7 +258,7 @@ def _forward_prefill( kv_start_loc = attn_metadata.kv_start_loc kv_seqlens = attn_metadata.kv_seqlens kv_flatten_size = attn_metadata.kv_flatten_size - quant_policy = attn_metadata.quant_policy + kv_cache_dtype = attn_metadata.kv_cache_dtype # Flatten KV cache for varlen attention block_size = k_cache.size(1) @@ -273,7 +273,7 @@ def _forward_prefill( out_dtype=query.dtype, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, flatten_kv_layout='shd', ) @@ -281,7 +281,7 @@ def _forward_prefill( # For TurboQuant, flattened K/V are in rotated domain. # Rotate Q to match, and inverse-rotate output afterwards. - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: from lmdeploy.pytorch.kernels.cuda.turbo_quant import ( hadamard_rotate, hadamard_rotate_inv, @@ -303,7 +303,7 @@ def _forward_prefill( ) # Inverse-rotate output back to original domain - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: attn_output = hadamard_rotate_inv(attn_output) return attn_output diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index 2fc4fdb298..81ce3661f0 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -4,7 +4,7 @@ import torch -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from lmdeploy.utils import get_logger from .default import TritonAttentionImpl, TritonAttentionMetadata @@ -327,7 +327,7 @@ def run_flatten_kv_cache(self, kv_seqlens = attn_metadata.kv_seqlens block_offsets = attn_metadata.block_offsets kv_flatten_size = attn_metadata.kv_flatten_size - quant_policy = attn_metadata.quant_policy + kv_cache_dtype = attn_metadata.kv_cache_dtype is_fp8_kvcache = k_cache.dtype == torch.float8_e4m3fn BLOCK_BS = k_cache.size(1) @@ -361,7 +361,7 @@ def run_flatten_kv_cache(self, out_dtype=out_dtype, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, flatten_kv_layout=flatten_kv_layout, ) @@ -405,8 +405,8 @@ def _fill_kv_cache_impl(self, block_offsets = attn_metadata.block_offsets kv_seqlens = attn_metadata.kv_seqlens - quant_policy = attn_metadata.quant_policy - assert quant_policy == QuantPolicy.NONE + kv_cache_dtype = attn_metadata.kv_cache_dtype + assert kv_cache_dtype == KVCacheDType.AUTO # fill seqlen args fill_seqlens, fill_max_q_seqlen, fill_q_start_loc = self._get_fill_meta( diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index 5608bac35f..68f98447ab 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -224,7 +224,7 @@ def update_step_context(cls, step_context): kv_start_loc=kv_start_loc, kv_seqlens=kv_seqlens, kv_flatten_size=kv_flatten_size, - quant_policy=step_context.kv_quant_policy, + kv_cache_dtype=step_context.kv_cache_dtype, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_kv_seqlen=step_context.max_kv_seqlen, diff --git a/lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py b/lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py index 8cc41f5df8..277bb12a90 100644 --- a/lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py +++ b/lmdeploy/pytorch/backends/dlinfer/ascend/op_backend.py @@ -10,6 +10,7 @@ import torch import torch.distributed as dist +from lmdeploy.messages import KVCacheDType from lmdeploy.pytorch import envs as _envs from lmdeploy.pytorch.config import BackendConfig, CacheConfig, ModelConfig from lmdeploy.pytorch.distributed import get_dist_manager @@ -364,7 +365,7 @@ def get_moe_group_name(group): max_kv_seq_len) q_seqlens_cpu = update_q_seqlens(step_context.is_decoding, is_prefill_no_cache, q_seqlens_cpu) - if not cls.enable_graph and step_context.kv_quant_policy == 8: + if not cls.enable_graph and step_context.kv_cache_dtype == KVCacheDType.INT8: record_file = os.getenv('ASCEND_QUANT_RECORD_FILE') assert record_file, 'please specify valid ASCEND_QUANT_RECORD_FILE' path = Path(record_file) @@ -403,7 +404,7 @@ def get_moe_group_name(group): is_prefill_no_cache=is_prefill_no_cache, max_q_seq_len=max_q_seq_len, max_kv_seq_len=max_kv_seq_len, - quant_policy=step_context.kv_quant_policy, + kv_cache_dtype=step_context.kv_cache_dtype, quant_meta=AscendKVQuantMeta.quant_meta, has_initial_state=has_initial_state, ) diff --git a/lmdeploy/pytorch/backends/dlinfer/attention.py b/lmdeploy/pytorch/backends/dlinfer/attention.py index a8eea27545..80567f6a2c 100644 --- a/lmdeploy/pytorch/backends/dlinfer/attention.py +++ b/lmdeploy/pytorch/backends/dlinfer/attention.py @@ -83,7 +83,7 @@ def forward( is_prefill_no_cache = attn_metadata.is_prefill_no_cache max_q_seq_len = attn_metadata.max_q_seq_len max_kv_seq_len = attn_metadata.max_kv_seq_len - quant_bits = attn_metadata.quant_policy + quant_bits = attn_metadata.kv_cache_dtype cu_seq_lens_kv = attn_metadata.cu_seq_lens_kv if attn_metadata.quant_meta is not None: diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py index 344e8b53b7..47dc023145 100644 --- a/lmdeploy/pytorch/config.py +++ b/lmdeploy/pytorch/config.py @@ -6,7 +6,7 @@ import torch -from lmdeploy.messages import PytorchEngineConfig, QuantPolicy +from lmdeploy.messages import KVCacheDType, PytorchEngineConfig, normalize_kv_cache_dtype from lmdeploy.pytorch.disagg.config import EngineRole, MigrationBackend from lmdeploy.pytorch.utils import maybe_register_config_serialize_by_value from lmdeploy.utils import get_logger, is_bf16_supported @@ -118,7 +118,7 @@ class CacheConfig: max_prefill_token_num: int = 8192 cudagraph_capture_batch_sizes: list[int] | None = None enable_prefix_caching: bool = False - quant_policy: QuantPolicy = QuantPolicy.NONE + kv_cache_dtype: KVCacheDType | int | str = KVCacheDType.AUTO device_type: str = 'cuda' num_state_caches: int = None prefix_cache_state_budget: int = 0 @@ -141,6 +141,7 @@ def __post_init__(self): self.enable_prefix_caching = False if self.kernel_block_size == -1: self.kernel_block_size = self.block_size + self.kv_cache_dtype = normalize_kv_cache_dtype(self.kv_cache_dtype) if self.prefix_cache_decode_state_interval > 0: assert self.prefix_cache_decode_state_interval % self.block_size == 0, ( 'prefix_cache_decode_state_interval must be a multiple of block_size') @@ -641,7 +642,7 @@ def from_config( max_prefill_token_num=target_cache_cfg.max_prefill_token_num, cudagraph_capture_batch_sizes=target_cache_cfg.cudagraph_capture_batch_sizes, device_type=target_cache_cfg.device_type, - quant_policy=target_cache_cfg.quant_policy, + kv_cache_dtype=target_cache_cfg.kv_cache_dtype, migration_backend=target_cache_cfg.migration_backend) obj = cls( model=model, diff --git a/lmdeploy/pytorch/engine/cache_engine.py b/lmdeploy/pytorch/engine/cache_engine.py index 6dd6ffb6d3..a7a4983e52 100644 --- a/lmdeploy/pytorch/engine/cache_engine.py +++ b/lmdeploy/pytorch/engine/cache_engine.py @@ -20,7 +20,7 @@ ) from lmdeploy.utils import get_logger -from ...messages import QuantPolicy +from ...messages import KVCacheDType from ..config import CacheConfig, ModelConfig KVCache = tuple[torch.Tensor, torch.Tensor] @@ -46,43 +46,43 @@ def __post_init__(self): self.aligned_size = round_up(self.size, self.alignment) -def _get_kv_cache_dtype(model_config: ModelConfig): - kv_cache_dtype = model_config.dtype +def _get_kv_cache_tensor_dtype(model_config: ModelConfig): + kv_cache_tensor_dtype = model_config.dtype if model_config.use_mla_fp8_cache: - kv_cache_dtype = torch.float8_e4m3fn - return kv_cache_dtype + kv_cache_tensor_dtype = torch.float8_e4m3fn + return kv_cache_tensor_dtype -_FP8_CACHE_DTYPES = { - QuantPolicy.FP8: torch.float8_e4m3fn, - QuantPolicy.FP8_E5M2: torch.float8_e5m2, +_FP8_CACHE_TENSOR_DTYPES = { + KVCacheDType.FP8: torch.float8_e4m3fn, + KVCacheDType.FP8_E5M2: torch.float8_e5m2, } -_KV_CACHE_QUANT_POLICY_DESCS = { - QuantPolicy.FP8: 'fp8_e4m3 KV cache', - QuantPolicy.FP8_E5M2: 'fp8_e5m2 KV cache', - QuantPolicy.INT4: 'int4 KV cache', - QuantPolicy.INT8: 'int8 KV cache', - QuantPolicy.TURBO_QUANT: 'TurboQuant KV cache', +_KV_CACHE_DTYPE_DESCS = { + KVCacheDType.FP8: 'fp8_e4m3 KV cache', + KVCacheDType.FP8_E5M2: 'fp8_e5m2 KV cache', + KVCacheDType.INT4: 'int4 KV cache', + KVCacheDType.INT8: 'int8 KV cache', + KVCacheDType.TURBO_QUANT: 'TurboQuant KV cache', } -def _is_fp8_quant_policy(quant_policy: QuantPolicy): - """Return whether quant policy stores KV payload as torch FP8.""" - return quant_policy in _FP8_CACHE_DTYPES +def _is_fp8_kv_cache_dtype(kv_cache_dtype: KVCacheDType): + """Return whether the KV cache payload is stored as torch FP8.""" + return kv_cache_dtype in _FP8_CACHE_TENSOR_DTYPES -def _get_fp8_cache_dtype(quant_policy: QuantPolicy): - """Get the cache tensor dtype for an FP8 KV-cache quant policy.""" +def _get_fp8_cache_tensor_dtype(kv_cache_dtype: KVCacheDType): + """Get the cache tensor dtype for an FP8 KV cache dtype.""" try: - return _FP8_CACHE_DTYPES[quant_policy] + return _FP8_CACHE_TENSOR_DTYPES[kv_cache_dtype] except KeyError as e: - raise ValueError(f'Not an FP8 quant policy: {quant_policy}') from e + raise ValueError(f'Not an FP8 KV cache dtype: {kv_cache_dtype}') from e -def _describe_kv_cache_quant_policy(quant_policy: QuantPolicy): - """Describe the active KV-cache quantization policy for logs.""" - return _KV_CACHE_QUANT_POLICY_DESCS.get(quant_policy) +def _describe_kv_cache_dtype(kv_cache_dtype: KVCacheDType): + """Describe the active KV cache dtype for logs.""" + return _KV_CACHE_DTYPE_DESCS.get(kv_cache_dtype) # 512*1 + 4*4 + 64*2 = 656 @@ -119,26 +119,26 @@ def __init__( self.block_size = cache_config.kernel_block_size self.num_layers = model_config.num_layers - self.kv_cache_dtype = _get_kv_cache_dtype(self.model_config) + self.kv_cache_tensor_dtype = _get_kv_cache_tensor_dtype(self.model_config) if self.model_config.use_mla_fp8_cache: - cache_config.quant_policy = 0 + cache_config.kv_cache_dtype = KVCacheDType.AUTO - if _is_fp8_quant_policy(cache_config.quant_policy): - self.kv_cache_dtype = _get_fp8_cache_dtype(cache_config.quant_policy) + if _is_fp8_kv_cache_dtype(cache_config.kv_cache_dtype): + self.kv_cache_tensor_dtype = _get_fp8_cache_tensor_dtype(cache_config.kv_cache_dtype) assert self.cache_config.device_type in ['cuda'], \ f'FP8 quantization is only supported on CUDA device, but got {self.cache_config.device_type}.' - elif cache_config.quant_policy > 0: + elif cache_config.kv_cache_dtype > 0: if self.cache_config.device_type in ['cuda']: - self.kv_cache_dtype = torch.uint8 + self.kv_cache_tensor_dtype = torch.uint8 elif self.cache_config.device_type in ['ascend', 'npu']: - self.kv_cache_dtype = torch.int8 + self.kv_cache_tensor_dtype = torch.int8 else: raise ValueError(f'unsupported device_type {self.cache_config.device_type}') - quant_desc = _describe_kv_cache_quant_policy(cache_config.quant_policy) - if quant_desc is not None: - logger.info('Using %s.', quant_desc) + dtype_desc = _describe_kv_cache_dtype(cache_config.kv_cache_dtype) + if dtype_desc is not None: + logger.info('Using %s.', dtype_desc) # Initialize the cache. self.local_gpu_cache = self.allocate_gpu_cache() @@ -183,7 +183,7 @@ def _get_key_block_shape_impl(cls, block_size: int, head_size: int, world_size: int = 1, - quant_policy: QuantPolicy = QuantPolicy.NONE): + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO): """Get single block shape.""" attn_backend = get_backend() dtype = model_config.dtype @@ -199,9 +199,9 @@ def _get_key_block_shape_impl(cls, return (block_size, num_heads, MLA_FP8_HEAD_DIM) # pack head_dim to uint8 (4-bit) - if quant_policy == QuantPolicy.INT4 or quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.INT4 or kv_cache_dtype == KVCacheDType.TURBO_QUANT: assert head_size % 2 == 0, \ - f'head_size: {head_size}, quant_policy: {quant_policy}' + f'head_size: {head_size}, kv_cache_dtype: {kv_cache_dtype}' head_size = head_size // 2 return attn_backend.get_k_block_shape(block_size, num_heads, head_size, dtype) @@ -211,7 +211,7 @@ def _get_value_block_shape_impl(cls, block_size: int, head_size: int, world_size: int = 1, - quant_policy: QuantPolicy = QuantPolicy.NONE): + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO): """Get single block shape.""" attn_backend = get_backend() dtype = model_config.dtype @@ -227,13 +227,13 @@ def _get_value_block_shape_impl(cls, # flash mla shared key and value return (block_size, num_heads, 0) - if quant_policy == QuantPolicy.TURBO_QUANT: # pack head_dim to uint8 (2-bit for V cache) + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: # pack head_dim to uint8 (2-bit for V cache) assert head_size % 4 == 0, \ - f'head_size: {head_size}, quant_policy: {quant_policy}' + f'head_size: {head_size}, kv_cache_dtype: {kv_cache_dtype}' head_size = head_size // 4 - elif quant_policy == QuantPolicy.INT4: # pack head_dim to uint8 (4-bit) + elif kv_cache_dtype == KVCacheDType.INT4: # pack head_dim to uint8 (4-bit) assert head_size % 2 == 0, \ - f'head_size: {head_size}, quant_policy: {quant_policy}' + f'head_size: {head_size}, kv_cache_dtype: {kv_cache_dtype}' head_size = head_size // 2 return attn_backend.get_v_block_shape(block_size, num_heads, head_size, dtype) @@ -249,13 +249,13 @@ def get_k_cache_desc(cls, model_config: ModelConfig, cache_config: CacheConfig, block_size=cache_config.kernel_block_size, head_size=head_size, world_size=world_size, - quant_policy=cache_config.quant_policy, + kv_cache_dtype=cache_config.kv_cache_dtype, ) shape = list(shape) - dtype = _get_kv_cache_dtype(model_config) - if _is_fp8_quant_policy(cache_config.quant_policy): - dtype = _get_fp8_cache_dtype(cache_config.quant_policy) - elif cache_config.quant_policy in (QuantPolicy.INT4, QuantPolicy.INT8, QuantPolicy.TURBO_QUANT): + dtype = _get_kv_cache_tensor_dtype(model_config) + if _is_fp8_kv_cache_dtype(cache_config.kv_cache_dtype): + dtype = _get_fp8_cache_tensor_dtype(cache_config.kv_cache_dtype) + elif cache_config.kv_cache_dtype in (KVCacheDType.INT4, KVCacheDType.INT8, KVCacheDType.TURBO_QUANT): dtype = torch.uint8 return CacheDesc(shape=shape, dtype=dtype) @@ -270,13 +270,13 @@ def get_v_cache_desc(cls, model_config: ModelConfig, cache_config: CacheConfig, block_size=cache_config.kernel_block_size, head_size=head_size, world_size=world_size, - quant_policy=cache_config.quant_policy, + kv_cache_dtype=cache_config.kv_cache_dtype, ) shape = list(shape) - dtype = _get_kv_cache_dtype(model_config) - if _is_fp8_quant_policy(cache_config.quant_policy): - dtype = _get_fp8_cache_dtype(cache_config.quant_policy) - elif cache_config.quant_policy in (QuantPolicy.INT4, QuantPolicy.INT8, QuantPolicy.TURBO_QUANT): + dtype = _get_kv_cache_tensor_dtype(model_config) + if _is_fp8_kv_cache_dtype(cache_config.kv_cache_dtype): + dtype = _get_fp8_cache_tensor_dtype(cache_config.kv_cache_dtype) + elif cache_config.kv_cache_dtype in (KVCacheDType.INT4, KVCacheDType.INT8, KVCacheDType.TURBO_QUANT): dtype = torch.uint8 return CacheDesc(shape=shape, dtype=dtype) @@ -284,17 +284,17 @@ def get_v_cache_desc(cls, model_config: ModelConfig, cache_config: CacheConfig, def get_quant_cache_descs(cls, k_cache_desc: CacheDesc, v_cache_desc: CacheDesc, model_config: ModelConfig, cache_config: CacheConfig): """Get quant cache descs.""" - if cache_config.quant_policy == QuantPolicy.NONE: + if cache_config.kv_cache_dtype == KVCacheDType.AUTO: return [] - if _is_fp8_quant_policy(cache_config.quant_policy): + if _is_fp8_kv_cache_dtype(cache_config.kv_cache_dtype): # Regular FP8 KV cache uses fixed scalar scales from Attention, not # per-token scale/zero cache tensors. return [] dtype = model_config.dtype - # For quant_policy==QuantPolicy.TURBO_QUANT, K uses 4-bit quantization (has MSE norm and QJL norm), + # For kv_cache_dtype==KVCacheDType.TURBO_QUANT, K uses 4-bit quantization (has MSE norm and QJL norm), # V uses 2-bit quantization (only has MSE norm) - if cache_config.quant_policy == QuantPolicy.TURBO_QUANT: + if cache_config.kv_cache_dtype == KVCacheDType.TURBO_QUANT: key_scale_zero_shape = k_cache_desc.shape[:-1] + [2] val_scale_zero_shape = v_cache_desc.shape[:-1] + [1] else: diff --git a/lmdeploy/pytorch/engine/config_builder.py b/lmdeploy/pytorch/engine/config_builder.py index b35d8f51ed..c58c51fdcf 100644 --- a/lmdeploy/pytorch/engine/config_builder.py +++ b/lmdeploy/pytorch/engine/config_builder.py @@ -77,7 +77,7 @@ def build_cache_config(engine_config: PytorchEngineConfig): enable_prefix_caching=engine_config.enable_prefix_caching, prefix_cache_state_budget=engine_config.prefix_cache_state_budget, prefix_cache_decode_state_interval=engine_config.prefix_cache_decode_state_interval, - quant_policy=engine_config.quant_policy, + kv_cache_dtype=engine_config.kv_cache_dtype, device_type=engine_config.device_type, migration_backend=engine_config.migration_backend, role=engine_config.role, diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py index f9c28e4159..dae5c59341 100644 --- a/lmdeploy/pytorch/engine/model_agent/agent.py +++ b/lmdeploy/pytorch/engine/model_agent/agent.py @@ -171,7 +171,7 @@ def model_forward( cache_config=cache_engine.cache_config, kv_caches=cache_engine.gpu_cache, state_caches=state_cache_engine.state_caches, - kv_quant_policy=cache_engine.cache_config.quant_policy, + kv_cache_dtype=cache_engine.cache_config.kv_cache_dtype, ) with ctx_mgr.context(context): diff --git a/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py b/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py index df36fcbba0..4d716d2fd5 100644 --- a/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py +++ b/lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py @@ -6,17 +6,17 @@ import triton.language as tl from torch import Tensor -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from .turbo_quant import get_lloyd_max_codebook, hadamard_rotate -# Triton-compatible quantization policy constants +# Triton-compatible KV cache dtype constants # Python Enum cannot be used in Triton kernels, so we define these as module-level # constants which Triton will inline at compile time. -Q_POLICY_NONE = tl.constexpr(0) -Q_POLICY_INT4 = tl.constexpr(4) -Q_POLICY_INT8 = tl.constexpr(8) -Q_POLICY_TURBO = tl.constexpr(42) +KV_CACHE_DTYPE_NONE = tl.constexpr(0) +KV_CACHE_DTYPE_INT4 = tl.constexpr(4) +KV_CACHE_DTYPE_INT8 = tl.constexpr(8) +KV_CACHE_DTYPE_TURBO = tl.constexpr(42) @triton.jit @@ -538,11 +538,11 @@ def _fill_page_quant( stride_szh: tl.constexpr, stride_szd: tl.constexpr, BLOCK_D: tl.constexpr, - quant_policy: tl.constexpr, + kv_cache_dtype: tl.constexpr, is_value: tl.constexpr, ): """Fill page.""" - if quant_policy == Q_POLICY_INT8: + if kv_cache_dtype == KV_CACHE_DTYPE_INT8: return _fill_page_quant_int8(state_ptr, cache_ptr, scales_zeros_ptr, @@ -564,7 +564,7 @@ def _fill_page_quant( stride_szh=stride_szh, stride_szd=stride_szd, BLOCK_D=BLOCK_D) - elif quant_policy == Q_POLICY_INT4: + elif kv_cache_dtype == KV_CACHE_DTYPE_INT4: return _fill_page_quant_int4(state_ptr, cache_ptr, scales_zeros_ptr, @@ -586,7 +586,7 @@ def _fill_page_quant( stride_szh=stride_szh, stride_szd=stride_szd, BLOCK_D=BLOCK_D) - elif quant_policy == Q_POLICY_TURBO: + elif kv_cache_dtype == KV_CACHE_DTYPE_TURBO: if is_value: return _fill_page_quant_turbo_int2(state_ptr, cache_ptr, @@ -635,7 +635,7 @@ def _fill_page_quant( stride_szd=stride_szd, BLOCK_D=BLOCK_D) else: - tl.static_assert(False, 'Unsupported quant policy') + tl.static_assert(False, 'Unsupported KV cache dtype') @triton.jit @@ -679,7 +679,7 @@ def _fill_kv_cache_quant_kernel( stride_vszb: tl.constexpr, stride_vszh: tl.constexpr, stride_vszd: tl.constexpr, - quant_policy: tl.constexpr, + kv_cache_dtype: tl.constexpr, stride_boff, BLOCK: tl.constexpr, BLOCK_D: tl.constexpr, @@ -752,7 +752,7 @@ def _fill_kv_cache_quant_kernel( stride_szh=stride_kszh, stride_szd=stride_kszd, BLOCK_D=BLOCK_D, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, is_value=False) if BLOCK_DV > 0: @@ -779,7 +779,7 @@ def _fill_kv_cache_quant_kernel( stride_szh=stride_vszh, stride_szd=stride_vszd, BLOCK_D=BLOCK_DV, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, is_value=True) @@ -794,7 +794,7 @@ def fill_kv_cache(k_states: Tensor, block_offsets: Tensor, k_scales_zeros: Tensor = None, v_scales_zeros: Tensor = None, - quant_policy: QuantPolicy = QuantPolicy.NONE, + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO, kv_layout: str = 'bshd'): """Fill key/value state to cache for paged attention. @@ -838,7 +838,7 @@ def fill_kv_cache(k_states: Tensor, v_centroids = torch.empty((1,), device=k_states.device, dtype=torch.float32) v_boundaries = torch.empty((1,), device=k_states.device, dtype=torch.float32) - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: raw_k_dim = k_states.size(-1) if raw_k_dim & (raw_k_dim - 1) != 0: raise ValueError(f'TurboQuant K requires power-of-2 raw dim, got {raw_k_dim}') @@ -875,7 +875,7 @@ def fill_kv_cache(k_states: Tensor, grid = (num_heads, max_num_blocks, batch_size) is_decoding = max_num_blocks == 1 - if quant_policy == QuantPolicy.NONE: + if kv_cache_dtype == KVCacheDType.AUTO: _fill_kv_cache_kernel[grid]( k_states, v_states, @@ -909,7 +909,7 @@ def fill_kv_cache(k_states: Tensor, num_warps=4, num_stages=3, ) - elif quant_policy in (QuantPolicy.FP8, QuantPolicy.FP8_E5M2): + elif kv_cache_dtype in (KVCacheDType.FP8, KVCacheDType.FP8_E5M2): assert k_scales_zeros is not None, 'FP8 KV cache requires k scale.' assert v_scales_zeros is not None, 'FP8 KV cache requires v scale.' finfo = torch.finfo(k_caches.dtype) @@ -991,7 +991,7 @@ def fill_kv_cache(k_states: Tensor, stride_vszb=v_scales_zeros.stride(s_dim), stride_vszh=v_scales_zeros.stride(h_dim), stride_vszd=v_scales_zeros.stride(d_dim), - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, stride_boff=block_offsets.stride(0), BLOCK=BLOCK, BLOCK_D=BLOCK_D, diff --git a/lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py b/lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py index 42ca2b9998..e43873c506 100644 --- a/lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py +++ b/lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py @@ -5,17 +5,17 @@ import triton.language as tl from torch import Tensor -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from .turbo_quant import get_lloyd_max_codebook -# Triton-compatible quantization policy constants +# Triton-compatible KV cache dtype constants # Python Enum cannot be used in Triton kernels, so we define these as module-level # constants which Triton will inline at compile time. -Q_POLICY_NONE = tl.constexpr(0) -Q_POLICY_INT4 = tl.constexpr(4) -Q_POLICY_INT8 = tl.constexpr(8) -Q_POLICY_TURBO = tl.constexpr(42) +KV_CACHE_DTYPE_NONE = tl.constexpr(0) +KV_CACHE_DTYPE_INT4 = tl.constexpr(4) +KV_CACHE_DTYPE_INT8 = tl.constexpr(8) +KV_CACHE_DTYPE_TURBO = tl.constexpr(42) @triton.jit @@ -263,7 +263,7 @@ def _flatten_kv_cache_quant( stride_vos: tl.constexpr, stride_vod: tl.constexpr, stride_boff, - quant_policy: tl.constexpr, + kv_cache_dtype: tl.constexpr, OUT_SIZE, HEAD_DIM_K: tl.constexpr, HEAD_DIM_V: tl.constexpr, @@ -310,12 +310,12 @@ def _flatten_kv_cache_quant( b_off = tl.load(block_offsets_ptr + batch_id * stride_boff + page_id) b_off = b_off.to(tl.int64) offs_bs = tl.arange(0, BLOCK_BS) - if quant_policy == Q_POLICY_INT4: + if kv_cache_dtype == KV_CACHE_DTYPE_INT4: HALF_HDK: tl.constexpr = HEAD_DIM_K // 2 HALF_HDV: tl.constexpr = HEAD_DIM_V // 2 offs_dk = tl.arange(0, BLOCK_DK) % HALF_HDK offs_dv = tl.arange(0, BLOCK_DV) % HALF_HDV - elif quant_policy == Q_POLICY_TURBO: + elif kv_cache_dtype == KV_CACHE_DTYPE_TURBO: # K is QJL4 packed in int4 => packed dim = HEAD_DIM_K // 2 # V is TurboQuant MSE int2 => packed dim = HEAD_DIM_V // 4 HALF_HDK: tl.constexpr = HEAD_DIM_K // 2 @@ -348,10 +348,10 @@ def _flatten_kv_cache_quant( # K path # ----------------------- kc = tl.load(kc_ptrs) - if quant_policy == Q_POLICY_INT4 or quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_INT4 or kv_cache_dtype == KV_CACHE_DTYPE_TURBO: kc = _dequant_int4(kc, HEAD_DIM_K, BLOCK_DK) - if quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_TURBO: # QJL4: # low 3bit = mse idx # high 1bit = qjl sign @@ -375,12 +375,12 @@ def _flatten_kv_cache_quant( # V path # ----------------------- vc = tl.load(vc_ptrs) - if quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_TURBO: vc = _dequant_int2(vc, HEAD_DIM_V, BLOCK_DV) - elif quant_policy == Q_POLICY_INT4: + elif kv_cache_dtype == KV_CACHE_DTYPE_INT4: vc = _dequant_int4(vc, HEAD_DIM_V, BLOCK_DV) - if quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_TURBO: # V is TurboQuant MSE int2, meta only stores norm vs = tl.load(vsz_ptrs) vq = tl.load(v_codebook_ptr + vc.to(tl.int32)) @@ -403,7 +403,7 @@ def flatten_kv_cache(k_caches: Tensor, out_dtype: torch.dtype = None, k_scales_zeros: Tensor = None, v_scales_zeros: Tensor = None, - quant_policy: QuantPolicy = QuantPolicy.NONE, + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO, kv_layout: str = 'bshd', flatten_kv_layout: str = 'hsd'): """Recover paged KV cache to contiguous KV cache. @@ -422,7 +422,7 @@ def flatten_kv_cache(k_caches: Tensor, raise RuntimeError('Unsupported layout.') if out_dtype is None: - if quant_policy in (QuantPolicy.FP8, QuantPolicy.FP8_E5M2, QuantPolicy.TURBO_QUANT): + if kv_cache_dtype in (KVCacheDType.FP8, KVCacheDType.FP8_E5M2, KVCacheDType.TURBO_QUANT): out_dtype = torch.float16 else: out_dtype = k_caches.dtype @@ -437,10 +437,10 @@ def flatten_kv_cache(k_caches: Tensor, num_heads = k_caches.size(h_dim) k_head_dim = k_caches.size(d_dim) v_head_dim = v_caches.size(d_dim) - if quant_policy == QuantPolicy.INT4: + if kv_cache_dtype == KVCacheDType.INT4: k_head_dim *= 2 v_head_dim *= 2 - elif quant_policy == QuantPolicy.TURBO_QUANT: + elif kv_cache_dtype == KVCacheDType.TURBO_QUANT: k_head_dim *= 2 # K packed int4 => raw dim *2 v_head_dim *= 4 # V packed int2 => raw dim *4 BLOCK_DK = triton.next_power_of_2(k_head_dim) @@ -449,7 +449,7 @@ def flatten_kv_cache(k_caches: Tensor, shared_kv = k_caches.data_ptr() == v_caches.data_ptr() and v_head_dim < k_head_dim if flatten_kv_layout == 'hsd': k_states = k_caches.new_empty(num_heads, out_size, k_head_dim, dtype=out_dtype) - if quant_policy == QuantPolicy.NONE and shared_kv: + if kv_cache_dtype == KVCacheDType.AUTO and shared_kv: v_states = k_states[..., :v_head_dim] v_head_dim = 0 else: @@ -460,7 +460,7 @@ def flatten_kv_cache(k_caches: Tensor, stride_vos = v_states.stride(1) elif flatten_kv_layout == 'shd': k_states = k_caches.new_empty(out_size, num_heads, k_head_dim, dtype=out_dtype) - if quant_policy == QuantPolicy.NONE and shared_kv: + if kv_cache_dtype == KVCacheDType.AUTO and shared_kv: v_states = k_states[..., :v_head_dim] v_head_dim = 0 else: @@ -475,7 +475,7 @@ def flatten_kv_cache(k_caches: Tensor, # The extra batch zero-fills the padded tail. Current callers pad to the # next block plus one guard block, so at most two pages are needed. grid = (max(num_blocks, 2), batch_size + 1, num_heads) - if quant_policy == QuantPolicy.NONE: + if kv_cache_dtype == KVCacheDType.AUTO: _flatten_kv_cache[grid]( k_caches, v_caches, @@ -506,7 +506,7 @@ def flatten_kv_cache(k_caches: Tensor, BLOCK_DK=BLOCK_DK, BLOCK_DV=BLOCK_DV, ) - elif quant_policy in (QuantPolicy.FP8, QuantPolicy.FP8_E5M2): + elif kv_cache_dtype in (KVCacheDType.FP8, KVCacheDType.FP8_E5M2): assert k_scales_zeros is not None, 'FP8 KV cache requires k scale.' assert v_scales_zeros is not None, 'FP8 KV cache requires v scale.' _flatten_kv_cache_fp8_scalar[grid]( @@ -542,7 +542,7 @@ def flatten_kv_cache(k_caches: Tensor, BLOCK_DV=BLOCK_DV, ) else: - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: # K = QJL4 => 3bit centroid codebook k_codebook, _ = get_lloyd_max_codebook(k_head_dim, bits=3, device=k_caches.device) # V = TurboQuant MSE int2 => 2bit centroid codebook @@ -585,7 +585,7 @@ def flatten_kv_cache(k_caches: Tensor, stride_vos=stride_vos, stride_vod=v_states.stride(2), stride_boff=block_offsets.stride(0), - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, OUT_SIZE=out_size, HEAD_DIM_K=k_head_dim, HEAD_DIM_V=v_head_dim, diff --git a/lmdeploy/pytorch/kernels/cuda/pagedattention.py b/lmdeploy/pytorch/kernels/cuda/pagedattention.py index 5f4f797733..e575923707 100644 --- a/lmdeploy/pytorch/kernels/cuda/pagedattention.py +++ b/lmdeploy/pytorch/kernels/cuda/pagedattention.py @@ -9,7 +9,7 @@ from packaging import version from torch import Tensor -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from lmdeploy.utils import get_logger from .turbo_quant import hadamard_rotate @@ -17,14 +17,14 @@ logger = get_logger('lmdeploy') -# Triton-compatible quantization policy constants +# Triton-compatible KV cache dtype constants # Python Enum cannot be used in Triton kernels, so we define these as module-level # constants which Triton will inline at compile time. -Q_POLICY_NONE = tl.constexpr(0) -Q_POLICY_INT4 = tl.constexpr(4) -Q_POLICY_INT8 = tl.constexpr(8) -Q_POLICY_FP8 = tl.constexpr(16) -Q_POLICY_TURBO = tl.constexpr(42) +KV_CACHE_DTYPE_NONE = tl.constexpr(0) +KV_CACHE_DTYPE_INT4 = tl.constexpr(4) +KV_CACHE_DTYPE_INT8 = tl.constexpr(8) +KV_CACHE_DTYPE_FP8 = tl.constexpr(16) +KV_CACHE_DTYPE_TURBO = tl.constexpr(42) TRITON_VERSION = version.parse(triton.__version__) VERSION_300 = version.parse('3.0.0') @@ -300,7 +300,7 @@ def _fwd_grouped_split_quant_kernel( stride_vszbs: tl.constexpr, stride_vszh: tl.constexpr, stride_vszd: tl.constexpr, - quant_policy: tl.constexpr, + kv_cache_dtype: tl.constexpr, stride_ok: tl.constexpr, stride_obs: tl.constexpr, stride_oh: tl.constexpr, @@ -384,7 +384,7 @@ def _fwd_grouped_split_quant_kernel( # initialize pointer to m and l m_i = tl.zeros([BLOCK_H], dtype=tl.float32) - float('inf') l_i = tl.zeros([BLOCK_H], dtype=tl.float32) - if quant_policy == Q_POLICY_INT4 or quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_INT4 or kv_cache_dtype == KV_CACHE_DTYPE_TURBO: packed_k_dim: tl.constexpr = head_size // 2 # K: raw dim -> packed dim (two halves packed into one byte) @@ -403,7 +403,7 @@ def _fwd_grouped_split_quant_kernel( + packed_offs_dk1[:, None] * stride_kd + offs_n[None, :] * stride_kbs) - if quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_TURBO: # V: packed dim = head_size_v, raw dim = head_size_v * 4 raw_offs_dv = tl.arange(0, BLOCK_DV * 4) packed_offs_dv = raw_offs_dv % head_size_v @@ -415,7 +415,7 @@ def _fwd_grouped_split_quant_kernel( offs_dv = raw_offs_dv acc = tl.zeros([BLOCK_H, BLOCK_DV * 4], dtype=tl.float32) else: - # quant_policy == Q_POLICY_INT4, V is 4-bit, packed dim = head_size_v, raw dim = head_size_v * 2 + # kv_cache_dtype == KV_CACHE_DTYPE_INT4, V is 4-bit, packed dim = head_size_v, raw dim = head_size_v * 2 raw_offs_dv = tl.arange(0, BLOCK_DV * 2) packed_offs_dv = raw_offs_dv % head_size_v shift_vd = (raw_offs_dv // head_size_v) * 4 @@ -449,10 +449,10 @@ def _fwd_grouped_split_quant_kernel( # -- compute qk ---- # k = tl.load(k_ptrs + b_offset * stride_kp) k = tl.load(k_ptr + off_k + b_offset * stride_kp) - if quant_policy == Q_POLICY_INT4 or quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_INT4 or kv_cache_dtype == KV_CACHE_DTYPE_TURBO: k = (k >> shift_kd) & 0x0F - if quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_TURBO: kmse_norm = tl.load(ksz_ptrs + b_offset * stride_kszp) kqjl_norm = tl.load(ksz_ptrs + b_offset * stride_kszp + stride_kszd) @@ -460,7 +460,7 @@ def _fwd_grouped_split_quant_kernel( k_cent = _k4v2_k_centroid((k & 0x7), head_size) k_sign = ((k >> 3) & 0x1).to(tl.float32) * 2.0 - 1.0 k = (kmse_norm * (k_cent + kqjl_norm * k_sign)).to(q.dtype) - elif quant_policy == Q_POLICY_FP8: + elif kv_cache_dtype == KV_CACHE_DTYPE_FP8: ks = tl.load(KScalesZeros).to(tl.float32) k = k.to(q.dtype) else: @@ -470,36 +470,36 @@ def _fwd_grouped_split_quant_kernel( if BLOCK_DMODEL1 != 0: k1 = tl.load(k_ptr + off_k1 + b_offset * stride_kp) - if quant_policy == Q_POLICY_INT4 or quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_INT4 or kv_cache_dtype == KV_CACHE_DTYPE_TURBO: k1 = (k1 >> shift_k1d) & 0x0F - if quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_TURBO: kmse_norm = tl.load(ksz_ptrs + b_offset * stride_kszp) kqjl_norm = tl.load(ksz_ptrs + b_offset * stride_kszp + stride_kszd) k1_cent = _k4v2_k_centroid((k1 & 0x7), head_size) k1_sign = ((k1 >> 3) & 0x1).to(tl.float32) * 2.0 - 1.0 k1 = (kmse_norm * (k1_cent + kqjl_norm * k1_sign)).to(q.dtype) - elif quant_policy == Q_POLICY_FP8: + elif kv_cache_dtype == KV_CACHE_DTYPE_FP8: k1 = k1.to(q.dtype) else: k1 = ((k1 - kz) * ks).to(q.dtype) # -- load / dequant v ---- - if quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_TURBO: v = tl.load(v_ptr + off_v + b_offset * stride_vp) v = (v >> shift_vd[None, :]) & 0x03 - elif quant_policy == Q_POLICY_INT4: + elif kv_cache_dtype == KV_CACHE_DTYPE_INT4: v = tl.load(v_ptr + off_v + b_offset * stride_vp) v = (v >> shift_vd[None, :]) & 0x0F else: v = tl.load(v_ptr + off_v + b_offset * stride_vp) - if quant_policy == Q_POLICY_TURBO: + if kv_cache_dtype == KV_CACHE_DTYPE_TURBO: vs = tl.load(vsz_ptrs + b_offset * stride_vszp) v = _k4v2_v_centroid(v, head_size_v) v = (v * vs).to(q.dtype) - elif quant_policy == Q_POLICY_FP8: + elif kv_cache_dtype == KV_CACHE_DTYPE_FP8: vs = tl.load(VScalesZeros).to(tl.float32) v = v.to(q.dtype) else: @@ -512,7 +512,7 @@ def _fwd_grouped_split_quant_kernel( qk += tl.dot(q, k) if BLOCK_DMODEL1 != 0: qk += tl.dot(q1, k1) - if quant_policy == Q_POLICY_FP8: + if kv_cache_dtype == KV_CACHE_DTYPE_FP8: qk *= ks qk *= sm_scale if logit_softcapping > 0.0: @@ -547,7 +547,7 @@ def _fwd_grouped_split_quant_kernel( acc = acc * alpha[:, None] # update acc - if quant_policy == Q_POLICY_FP8: + if kv_cache_dtype == KV_CACHE_DTYPE_FP8: p = p * vs p, v = _convert_pv(p, v) acc += tl.dot(p, v) @@ -561,9 +561,9 @@ def _fwd_grouped_split_quant_kernel( offs_dv[None, :] * stride_od) tl.store(acc_out_ptr + off_acc, acc, mask=mask_h[:, None] & mask_dv[None, :]) - if quant_policy == Q_POLICY_INT4: + if kv_cache_dtype == KV_CACHE_DTYPE_INT4: off_meta = (cur_batch * stride_obs + split_k_id * stride_ok + cur_head * stride_oh + head_size_v * 2) - elif quant_policy == Q_POLICY_TURBO: + elif kv_cache_dtype == KV_CACHE_DTYPE_TURBO: off_meta = (cur_batch * stride_obs + split_k_id * stride_ok + cur_head * stride_oh + head_size_v * 4) else: off_meta = (cur_batch * stride_obs + split_k_id * stride_ok + cur_head * stride_oh + head_size_v) @@ -770,7 +770,7 @@ def flash_attn_with_kvcache( alibi_slopes: Tensor = None, k_scales_zeros: Tensor = None, v_scales_zeros: Tensor = None, - quant_policy: QuantPolicy = QuantPolicy.NONE, + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO, sinks: Tensor = None, kv_layout: str = 'bshd', ): @@ -807,8 +807,8 @@ def flash_attn_with_kvcache( shared_kv = k_cache.data_ptr() == v_cache.data_ptr() # quant42 K/V have different semantics and meta shape, should not share buffer - if quant_policy == QuantPolicy.TURBO_QUANT: - assert not shared_kv, 'quant_policy==42 does not support shared_kv' + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: + assert not shared_kv, 'kv_cache_dtype==42 does not support shared_kv' def _get_block_d(Lk): """Get block d.""" @@ -823,11 +823,11 @@ def _get_block_d(Lk): # shape constraints Lq, Lk, Lv = q.shape[-1], k_cache.shape[d_dim], v_cache.shape[d_dim] - if quant_policy == QuantPolicy.INT4 or quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.INT4 or kv_cache_dtype == KVCacheDType.TURBO_QUANT: # K uses 4-bit: Lq == Lk * 2 - # For quant_policy==QuantPolicy.TURBO_QUANT, V uses 2-bit: raw V dim == Lv * 4 + # For kv_cache_dtype==KVCacheDType.TURBO_QUANT, V uses 2-bit: raw V dim == Lv * 4 assert Lq == Lk * 2 - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: o = q.new_empty(q.shape[:-1] + (Lv * 4, )) else: o = q.new_empty(q.shape[:-1] + (Lv * 2, )) @@ -835,7 +835,7 @@ def _get_block_d(Lk): assert Lq == Lk o = q.new_empty(q.shape[:-1] + (Lv, )) - # quant_policy == QuantPolicy.TURBO_QUANT: interpret as + # kv_cache_dtype == KVCacheDType.TURBO_QUANT: interpret as # - K: QJL4 = 3bit MSE centroid + 1bit QJL sign # - V: TurboQuant MSE int2 # Implementation: @@ -843,7 +843,7 @@ def _get_block_d(Lk): # - K dequant as mse_norm * (centroid[idx3] + qjl_norm * sign) # - V dequant as norm * centroid[idx2] # - output inverse-rotated because V is still rotated before caching - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: real_k_dim = Lq real_v_dim = Lv * 4 if real_k_dim & (real_k_dim - 1) != 0: @@ -891,10 +891,10 @@ def _get_block_d(Lk): else: num_warps, num_stages = _kernel_meta_sm9x(BLOCK_DMODEL, BLOCK_H) - is_fp8_scalar = quant_policy in (QuantPolicy.FP8, QuantPolicy.FP8_E5M2) + is_fp8_scalar = kv_cache_dtype in (KVCacheDType.FP8, KVCacheDType.FP8_E5M2) SPLIT_K = _get_split_k(q.device.index, grid_1, batch, num_warps) - if quant_policy == QuantPolicy.INT4 or quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.INT4 or kv_cache_dtype == KVCacheDType.TURBO_QUANT: acc = q.new_empty(num_tokens, head, SPLIT_K, o.shape[-1] + 2, dtype=torch.float32) else: acc = q.new_empty(num_tokens, head, SPLIT_K, Lv + 2, dtype=torch.float32) @@ -905,7 +905,7 @@ def _get_block_d(Lk): batch, ) - if quant_policy != QuantPolicy.NONE: + if kv_cache_dtype != KVCacheDType.AUTO: if is_fp8_scalar: assert k_scales_zeros is not None, 'FP8 KV cache requires k scale.' assert v_scales_zeros is not None, 'FP8 KV cache requires v scale.' @@ -913,7 +913,7 @@ def _get_block_d(Lk): v_scales_arg = v_scales_zeros stride_kszp = stride_kszbs = stride_kszh = stride_kszd = 0 stride_vszp = stride_vszbs = stride_vszh = stride_vszd = 0 - triton_quant_policy = QuantPolicy.FP8 + triton_kv_cache_dtype = KVCacheDType.FP8 else: k_scales_arg = k_scales_zeros v_scales_arg = v_scales_zeros @@ -925,7 +925,7 @@ def _get_block_d(Lk): stride_vszbs = v_scales_zeros.stride(s_dim) stride_vszh = v_scales_zeros.stride(h_dim) stride_vszd = v_scales_zeros.stride(d_dim) - triton_quant_policy = quant_policy + triton_kv_cache_dtype = kv_cache_dtype _fwd_grouped_split_quant_kernel[grid](q, k_cache, v_cache, @@ -955,7 +955,7 @@ def _get_block_d(Lk): stride_vszbs=stride_vszbs, stride_vszh=stride_vszh, stride_vszd=stride_vszd, - quant_policy=triton_quant_policy, + kv_cache_dtype=triton_kv_cache_dtype, stride_ok=acc.stride(-2), stride_obs=acc.stride(-4), stride_oh=acc.stride(-3), @@ -1020,14 +1020,14 @@ def _get_block_d(Lk): num_warps = 2 grid = (head, num_tokens) - if quant_policy == QuantPolicy.INT4: + if kv_cache_dtype == KVCacheDType.INT4: Lv *= 2 BLOCK_DV *= 2 - elif quant_policy == QuantPolicy.TURBO_QUANT: + elif kv_cache_dtype == KVCacheDType.TURBO_QUANT: Lv *= 4 BLOCK_DV *= 4 - if quant_policy == QuantPolicy.TURBO_QUANT: + if kv_cache_dtype == KVCacheDType.TURBO_QUANT: LOG2_DV = int(math.log2(BLOCK_DV)) _fused_reduce_hadamard_kernel[grid](acc, o, diff --git a/lmdeploy/pytorch/kernels/cuda/turbo_quant.py b/lmdeploy/pytorch/kernels/cuda/turbo_quant.py index 1f6fddae9a..41fa24ab8f 100644 --- a/lmdeploy/pytorch/kernels/cuda/turbo_quant.py +++ b/lmdeploy/pytorch/kernels/cuda/turbo_quant.py @@ -2,7 +2,7 @@ """TurboQuant quantization utilities. This module provides: -- Hadamard transform (orthogonal rotation) for quant_policy==QuantPolicy.TURBO_QUANT +- Hadamard transform (orthogonal rotation) for kv_cache_dtype==KVCacheDType.TURBO_QUANT - Lloyd-Max codebook for 2-bit (V cache) and 3-bit (K cache) quantization """ import logging diff --git a/lmdeploy/pytorch/model_inputs.py b/lmdeploy/pytorch/model_inputs.py index 270bc34b0e..236a658bbe 100644 --- a/lmdeploy/pytorch/model_inputs.py +++ b/lmdeploy/pytorch/model_inputs.py @@ -10,7 +10,7 @@ # from torch import distributed as dist import lmdeploy.pytorch.distributed as dist -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from lmdeploy.pytorch.backends import get_backend from lmdeploy.pytorch.config import CacheConfig, DLLMConfig, ModelConfig, QuantizationConfig from lmdeploy.pytorch.multimodal.data_type import MultiModalData @@ -306,7 +306,7 @@ class StepContext: input_multimodals: list[MultiModalData] | None = None vision_inputs: VisionModelInputs | None = None attn_metadata: Any = None - kv_quant_policy: QuantPolicy = QuantPolicy.NONE + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO model_metas: list[dict[str, Any]] | None = None dp_meta: DPMeta | None = None enable_microbatch: bool = False @@ -334,7 +334,7 @@ def new( cache_config: CacheConfig, kv_caches: list | None = None, state_caches: list | None = None, - kv_quant_policy: QuantPolicy = QuantPolicy.NONE, + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO, ): """Build step context. @@ -382,7 +382,7 @@ def new( max_kv_seqlen=inputs.max_kv_seqlen, local_adapter_ids=inputs.local_adapter_ids, vision_inputs=inputs.vision_inputs, - kv_quant_policy=kv_quant_policy, + kv_cache_dtype=kv_cache_dtype, model_metas=inputs.model_metas, dp_meta=inputs.dp_meta, enable_microbatch=inputs.enable_microbatch, @@ -478,7 +478,7 @@ def build_context( cache_config: CacheConfig, kv_caches: list | None = None, state_caches: list | None = None, - kv_quant_policy: QuantPolicy = QuantPolicy.NONE, + kv_cache_dtype: KVCacheDType = KVCacheDType.AUTO, ): """Build context.""" return StepContext.new( @@ -487,7 +487,7 @@ def build_context( cache_config, kv_caches, state_caches, - kv_quant_policy, + kv_cache_dtype, ) diff --git a/lmdeploy/pytorch/nn/attention.py b/lmdeploy/pytorch/nn/attention.py index 39db348c9e..4776878989 100644 --- a/lmdeploy/pytorch/nn/attention.py +++ b/lmdeploy/pytorch/nn/attention.py @@ -2,7 +2,7 @@ import torch from torch import nn -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from lmdeploy.pytorch.distributed import get_tp_world_rank from ..backends import OpType, get_backend @@ -107,8 +107,8 @@ def forward( """forward.""" self._lazy_init(query.device) - quant_policy = attn_metadata.quant_policy - if quant_policy in (QuantPolicy.FP8, QuantPolicy.FP8_E5M2): + kv_cache_dtype = attn_metadata.kv_cache_dtype + if kv_cache_dtype in (KVCacheDType.FP8, KVCacheDType.FP8_E5M2): # Reuse the scale/zero arguments as scalar-scale channels for FP8. if self.k_scale.device != query.device: self.k_scale = self.k_scale.to(device=query.device, non_blocking=True) diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py index 2df03cf340..fb3fbdac18 100644 --- a/lmdeploy/turbomind/turbomind.py +++ b/lmdeploy/turbomind/turbomind.py @@ -231,7 +231,7 @@ def _from_hf(self, model_path: str, engine_config: TurbomindEngineConfig, ec = _tm.EngineConfig() ec.data_type = dtype_map[engine_config.dtype] ec.cache_block_seq_len = engine_config.cache_block_seq_len - ec.quant_policy = engine_config.quant_policy + ec.kv_cache_dtype = int(engine_config.kv_cache_dtype) ec.max_batch_size = engine_config.max_batch_size ec.max_prefill_token_num = engine_config.max_prefill_token_num ec.session_len = engine_config.session_len diff --git a/src/turbomind/engine/engine.cc b/src/turbomind/engine/engine.cc index de0eed861d..de469a0042 100644 --- a/src/turbomind/engine/engine.cc +++ b/src/turbomind/engine/engine.cc @@ -257,7 +257,7 @@ void Engine::Impl::CreateSequenceManager() weights_.kv_head_num / param_.attn_tp_size, weights_.num_layer, weights_.layer_types, - param_.quant_policy, + param_.kv_cache_dtype, weights_.data_type, weights_.data_type, // runtime_dtype = data_type linear_key_head_dim, diff --git a/src/turbomind/engine/engine_config.h b/src/turbomind/engine/engine_config.h index 2c0381e9c4..996d1cd80d 100644 --- a/src/turbomind/engine/engine_config.h +++ b/src/turbomind/engine/engine_config.h @@ -14,7 +14,7 @@ struct EngineConfig { #define ENGINE_FIELDS(X) \ X(DataType, data_type) \ X(int, cache_block_seq_len, 0) \ - X(int, quant_policy, 0) \ + X(int, kv_cache_dtype, 0) \ X(int, tune_layer_num, 1) \ X(int, max_batch_size, 0) \ X(int, max_prefill_token_num, 0) \ diff --git a/src/turbomind/kernels/attention/attention_params.h b/src/turbomind/kernels/attention/attention_params.h index e9696327f9..655e4df5d3 100644 --- a/src/turbomind/kernels/attention/attention_params.h +++ b/src/turbomind/kernels/attention/attention_params.h @@ -75,7 +75,7 @@ struct AttentionParams { bool use_logn_attn; int max_position_embeddings; - int quant_policy; + int kv_cache_dtype; int max_split_k; int* split_cnt; diff --git a/src/turbomind/kernels/attention/decoding.cu b/src/turbomind/kernels/attention/decoding.cu index ed53b322f2..7df70efaa1 100644 --- a/src/turbomind/kernels/attention/decoding.cu +++ b/src/turbomind/kernels/attention/decoding.cu @@ -14,8 +14,8 @@ void dispatchDecoding(const AttentionParams& params) { using namespace attention; - const bool is_kv_int8 = params.quant_policy & QuantPolicy::kCacheKVInt8; - const bool is_kv_int4 = params.quant_policy & QuantPolicy::kCacheKVInt4; + const bool is_kv_int8 = params.kv_cache_dtype & KVCacheDType::kCacheKVInt8; + const bool is_kv_int4 = params.kv_cache_dtype & KVCacheDType::kCacheKVInt4; const int query_group_sz = params.num_heads / params.num_kv_heads; TM_CHECK(!(is_kv_int4 && is_kv_int8)); diff --git a/src/turbomind/kernels/attention/kv_cache_utils_v2.cu b/src/turbomind/kernels/attention/kv_cache_utils_v2.cu index 14d64caed2..938a9ea6ef 100644 --- a/src/turbomind/kernels/attention/kv_cache_utils_v2.cu +++ b/src/turbomind/kernels/attention/kv_cache_utils_v2.cu @@ -223,7 +223,7 @@ void invokeProcessKV_v2(char** blocks, int head_num, int head_dim, int batch_size, - int quant_policy, + int kv_cache_dtype, cudaStream_t stream) { @@ -282,10 +282,10 @@ void invokeProcessKV_v2(char** blocks, TM_UNREACHABLE; }; - if (quant_policy & QuantPolicy::kCacheKVInt8) { + if (kv_cache_dtype & KVCacheDType::kCacheKVInt8) { dispatch(uint8_t{}); } - else if (quant_policy & QuantPolicy::kCacheKVInt4) { + else if (kv_cache_dtype & KVCacheDType::kCacheKVInt4) { dispatch(uint4_t{}); } else { @@ -317,7 +317,7 @@ void invokeProcessKV_v2(char** blocks, int head_num, \ int head_dim, \ int batch_size, \ - int quant_policy, \ + int kv_cache_dtype, \ cudaStream_t stream); INSTANTIATE_invokeProcessKV_v2(half); @@ -474,7 +474,7 @@ void invokeFlattenKV_v2(T* k, int head_num, int head_dim, int batch_size, - int quant_policy, + int kv_cache_dtype, cudaStream_t stream) { @@ -530,10 +530,10 @@ void invokeFlattenKV_v2(T* k, TM_UNREACHABLE; }; - if (quant_policy & QuantPolicy::kCacheKVInt8) { + if (kv_cache_dtype & KVCacheDType::kCacheKVInt8) { dispatch(uint8_t{}); } - else if (quant_policy & QuantPolicy::kCacheKVInt4) { + else if (kv_cache_dtype & KVCacheDType::kCacheKVInt4) { dispatch(uint4_t{}); } else { @@ -562,7 +562,7 @@ void invokeFlattenKV_v2(T* k, int head_num, \ int head_dim, \ int batch_size, \ - int quant_policy, \ + int kv_cache_dtype, \ cudaStream_t stream); INSTANTIATE_invokeFlattenKV_v2(half); diff --git a/src/turbomind/kernels/attention/kv_cache_utils_v2.h b/src/turbomind/kernels/attention/kv_cache_utils_v2.h index ca5259dcea..28f5f2e279 100644 --- a/src/turbomind/kernels/attention/kv_cache_utils_v2.h +++ b/src/turbomind/kernels/attention/kv_cache_utils_v2.h @@ -29,7 +29,7 @@ void invokeProcessKV_v2(char** blocks, int head_num, int head_dim, int batch_size, - int quant_policy, + int kv_cache_dtype, cudaStream_t stream = {}); template @@ -56,7 +56,7 @@ void invokeProcessKV_v2_(const AttentionParams& params) params.num_kv_heads, params.size_per_head, params.batch_size, - params.quant_policy, + params.kv_cache_dtype, params.stream); } @@ -79,7 +79,7 @@ void invokeFlattenKV_v2(T* k, int head_num, int head_dim, int batch_size, - int quant_policy, + int kv_cache_dtype, cudaStream_t stream = {}); /// TODO: remove `sum_k_len` @@ -105,7 +105,7 @@ void invokeFlattenKV_v2_(const AttentionParams& params, int sum_k_len) params.num_kv_heads, params.size_per_head, params.batch_size, - params.quant_policy, + params.kv_cache_dtype, params.stream); } diff --git a/src/turbomind/kernels/attention/test_attention.cu b/src/turbomind/kernels/attention/test_attention.cu index 147821c293..ae266c8b53 100644 --- a/src/turbomind/kernels/attention/test_attention.cu +++ b/src/turbomind/kernels/attention/test_attention.cu @@ -79,7 +79,7 @@ void TestBlocks(const thrust::universal_vector& k_cache, // [B, H, S, const size_t block_seq_len, const size_t batch_size, const int rope_dim, - int quant_policy) + int kv_cache_dtype) { const size_t seq_len = k_cache.size() / (head_dim * head_num * batch_size); const size_t n_blocks = (seq_len + block_seq_len - 1) / block_seq_len; @@ -165,7 +165,7 @@ void TestBlocks(const thrust::universal_vector& k_cache, // [B, H, S, head_num, head_dim, batch_size, - quant_policy)); + kv_cache_dtype)); } thrust::universal_vector kv_cache_2(kv_cache.size()); @@ -191,7 +191,7 @@ void TestBlocks(const thrust::universal_vector& k_cache, // [B, H, S, head_num, head_dim, batch_size, - quant_policy)); + kv_cache_dtype)); } cudaDeviceSynchronize(); @@ -280,19 +280,19 @@ int test_attention() constexpr size_t kSequenceLen = 0; constexpr int kMaxSplitK = 1; - constexpr int kBlockSz = 64; + constexpr int kBlockSz = 64; #endif #if KV_INT8 - using Tkv = uint8_t; - constexpr int kQuantPolicy = QuantPolicy::kCacheKVInt8; + using Tkv = uint8_t; + constexpr int kKVCacheDType = KVCacheDType::kCacheKVInt8; #elif KV_INT4 - using Tkv = uint4_t; - constexpr int kQuantPolicy = QuantPolicy::kCacheKVInt4; + using Tkv = uint4_t; + constexpr int kKVCacheDType = KVCacheDType::kCacheKVInt4; #else - using Tkv = T; - constexpr int kQuantPolicy = 0; + using Tkv = T; + constexpr int kKVCacheDType = 0; #endif static_assert(KvHeadNum > 0); @@ -377,7 +377,7 @@ int test_attention() kBlockSz, kBatchSize, kRoPEDim, - kQuantPolicy); + kKVCacheDType); thrust::universal_vector output_ref = output; thrust::universal_vector k_cache_ref_ptrs(kBatchSize); @@ -430,7 +430,7 @@ int test_attention() int(2 * kBatchSize * kContextLen * kHeadDim), int(kBatchSize * kContextLen * kHeadDim)}; - params.quant_policy = kQuantPolicy; + params.kv_cache_dtype = kKVCacheDType; params.finished = finished.data().get(); params.rope_theta = rope_base.data().get(); @@ -578,7 +578,7 @@ int test_attention() KvHeadNum, kHeadDim, kBatchSize, - kQuantPolicy)); + kKVCacheDType)); cudaDeviceSynchronize(); const size_t effective_context_len = Causal ? std::min(kContextLen, (size_t)params.window_size) : kContextLen; diff --git a/src/turbomind/models/llama/SequenceManager.cc b/src/turbomind/models/llama/SequenceManager.cc index 1693699555..79d503882f 100644 --- a/src/turbomind/models/llama/SequenceManager.cc +++ b/src/turbomind/models/llama/SequenceManager.cc @@ -35,7 +35,7 @@ SequenceManager::SequenceManager(int head_dim, int kv_head_num, int num_layer, const std::vector& layer_types, - int quant_policy, + int kv_cache_dtype, DataType data_type, DataType runtime_dtype, int linear_key_head_dim, @@ -102,7 +102,7 @@ SequenceManager::SequenceManager(int head_dim, } const int dbits = byte_size(runtime_dtype, 8); - const int elem_bits = quant_policy ? quant_policy : dbits; + const int elem_bits = kv_cache_dtype ? kv_cache_dtype : dbits; BlockConfig block_config{ head_dim, diff --git a/src/turbomind/models/llama/SequenceManager.h b/src/turbomind/models/llama/SequenceManager.h index 73ec7e71c2..e39c008b0d 100644 --- a/src/turbomind/models/llama/SequenceManager.h +++ b/src/turbomind/models/llama/SequenceManager.h @@ -98,7 +98,7 @@ class SequenceManager { int kv_head_num, int num_layer, const std::vector& layer_types, - int quant_policy, + int kv_cache_dtype, DataType data_type, DataType runtime_dtype, int linear_key_head_dim, diff --git a/src/turbomind/models/llama/llama_utils.h b/src/turbomind/models/llama/llama_utils.h index 75503673da..153b3d3e71 100644 --- a/src/turbomind/models/llama/llama_utils.h +++ b/src/turbomind/models/llama/llama_utils.h @@ -9,7 +9,7 @@ namespace turbomind { -enum QuantPolicy +enum KVCacheDType { kNone = 0x00, // reserve 0x01 and 0x02 for backward compatibility diff --git a/src/turbomind/models/llama/unified_attention_layer.cc b/src/turbomind/models/llama/unified_attention_layer.cc index cc0e529eb4..57e5b871b7 100644 --- a/src/turbomind/models/llama/unified_attention_layer.cc +++ b/src/turbomind/models/llama/unified_attention_layer.cc @@ -92,7 +92,7 @@ UnifiedAttentionLayer::~UnifiedAttentionLayer() aux_stream_ = {}; } -UnifiedAttentionLayer::UnifiedAttentionLayer(int quant_policy, +UnifiedAttentionLayer::UnifiedAttentionLayer(int kv_cache_dtype, const std::vector& layer_types, int layer_num, std::vector attn_weights, @@ -100,7 +100,7 @@ UnifiedAttentionLayer::UnifiedAttentionLayer(int quant const Context& ctx, int phases, bool init): - quant_policy_{quant_policy}, + kv_cache_dtype_{kv_cache_dtype}, rope_{attn_weights[0]->rope}, engine_param_{engine}, cp_fn_ctx_{ctx.comm.d_comm, ctx.comm.d_cp_group}, @@ -567,7 +567,7 @@ Tensor UnifiedAttentionLayer::core_attention(Tensor& qkv, const ForwardParam& p, params.arch = arch_; params.stream = stream; - params.quant_policy = quant_policy_; + params.kv_cache_dtype = kv_cache_dtype_; return params; }; diff --git a/src/turbomind/models/llama/unified_attention_layer.h b/src/turbomind/models/llama/unified_attention_layer.h index 79c20d3115..b9942b7ff2 100644 --- a/src/turbomind/models/llama/unified_attention_layer.h +++ b/src/turbomind/models/llama/unified_attention_layer.h @@ -55,7 +55,7 @@ class UnifiedAttentionLayer { ~UnifiedAttentionLayer(); - UnifiedAttentionLayer(int quant_policy, + UnifiedAttentionLayer(int kv_cache_dtype, const std::vector& layer_types, int layer_num, std::vector attn_weights, @@ -80,7 +80,7 @@ class UnifiedAttentionLayer { void qk_norm(Tensor& qkv, const WeightType& weights); private: - const int quant_policy_; + const int kv_cache_dtype_; const core::RopeConfig rope_; const EngineParam engine_param_; const Context& context_; diff --git a/src/turbomind/models/llama/unified_decoder.cc b/src/turbomind/models/llama/unified_decoder.cc index 2562156634..aad06baaf0 100644 --- a/src/turbomind/models/llama/unified_decoder.cc +++ b/src/turbomind/models/llama/unified_decoder.cc @@ -67,7 +67,7 @@ UnifiedDecoder::UnifiedDecoder(const EngineParam& engine, } } - attn_layer_ = std::make_unique(engine.quant_policy, + attn_layer_ = std::make_unique(engine.kv_cache_dtype, model_weight.layer_types, model_weight.num_layer, attn_weights, diff --git a/tests/pytorch/kernel/test_fa3_attention.py b/tests/pytorch/kernel/test_fa3_attention.py index 9f2526699b..71087664b2 100644 --- a/tests/pytorch/kernel/test_fa3_attention.py +++ b/tests/pytorch/kernel/test_fa3_attention.py @@ -1,7 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import torch -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from lmdeploy.pytorch.backends.cuda.attention.default import TritonAttentionMetadata from lmdeploy.pytorch.backends.cuda.attention.fa3 import FA3Impl @@ -18,7 +18,7 @@ def _make_prefill_metadata(q_seqlens, block_offsets): q_seqlens=q_seqlens, kv_start_loc=cu_seqlens[:-1], kv_seqlens=q_seqlens, - quant_policy=QuantPolicy.NONE, + kv_cache_dtype=KVCacheDType.AUTO, kv_flatten_size=int(q_seqlens.sum().item()), cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens.clone(), diff --git a/tests/pytorch/kernel/test_fill_kv_cache.py b/tests/pytorch/kernel/test_fill_kv_cache.py index 00751e2d8a..8a4cf9565b 100644 --- a/tests/pytorch/kernel/test_fill_kv_cache.py +++ b/tests/pytorch/kernel/test_fill_kv_cache.py @@ -3,7 +3,7 @@ import pytest import torch -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType # Import common TurboQuant utilities from turboquant_utils from .turboquant_utils import ( @@ -287,7 +287,7 @@ def test_fill_kv_cache(self, k_states, v_states, k_caches, v_caches, k_scales_ze class TestFillKVCacheInt42(TestFillKVCacheInt4): - """quant_policy == QuantPolicy.TURBO_QUANT: + """kv_cache_dtype == KVCacheDType.TURBO_QUANT: - K: QJL4 = 3bit MSE + 1bit QJL - V: TurboQuant MSE int2 @@ -395,7 +395,7 @@ def test_fill_kv_cache(self, k_states, v_states, k_caches, v_caches, k_scales_ze block_offsets, k_scales_zeros, v_scales_zeros, - QuantPolicy.TURBO_QUANT, + KVCacheDType.TURBO_QUANT, ) torch.testing.assert_close(k_caches, gt[0]) @@ -418,8 +418,8 @@ def test_qjl4_reference_sanity(self, head_dim): assert cos > 0.80, f'QJL4 reference cosine too low: {cos}' def test_fill_kv_cache_quant42_vs_python_reference(self): - """Test fill_kv_cache with quant_policy=QuantPolicy.TURBO_QUANT against - Python reference. + """Test fill_kv_cache with kv_cache_dtype=KVCacheDType.TURBO_QUANT + against Python reference. This test verifies that the fill_kv_cache kernel produces the same quantized output as the Python reference implementation. @@ -486,7 +486,7 @@ def test_fill_kv_cache_quant42_vs_python_reference(self): block_offsets, k_scales_zeros=blocked_ksz, v_scales_zeros=blocked_vsz, - quant_policy=QuantPolicy.TURBO_QUANT, + kv_cache_dtype=KVCacheDType.TURBO_QUANT, ) # Python reference quantization - only for the last token (the one being written) @@ -730,8 +730,8 @@ def fp8_dtype(self): yield torch.float8_e4m3fn @pytest.fixture - def quant_policy(self): - yield QuantPolicy.FP8 + def kv_cache_dtype(self): + yield KVCacheDType.FP8 @pytest.fixture def head_dim(self, request): @@ -788,7 +788,7 @@ def gt(self, k_states, v_states, k_caches, v_caches, seq_lens, history_lens, blo ], indirect=True) def test_fill_kv_cache(self, k_states, v_states, k_caches, v_caches, block_offsets, q_start_loc, q_seq_length, - kv_seq_length, max_q_seq_length, gt, quant_policy, fp8_dtype): + kv_seq_length, max_q_seq_length, gt, kv_cache_dtype, fp8_dtype): from lmdeploy.pytorch.kernels.cuda.fill_kv_cache import fill_kv_cache gt_k, gt_v, k_scale, v_scale = gt fill_kv_cache(k_states, @@ -800,7 +800,7 @@ def test_fill_kv_cache(self, k_states, v_states, k_caches, v_caches, block_offse kv_seq_length, max_q_seq_length, block_offsets, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, k_scales_zeros=k_scale, v_scales_zeros=v_scale) _assert_fp8_cache_close(k_caches, gt_k, fp8_dtype, k_scale) @@ -815,5 +815,5 @@ def fp8_dtype(self): yield torch.float8_e5m2 @pytest.fixture - def quant_policy(self): - yield QuantPolicy.FP8_E5M2 + def kv_cache_dtype(self): + yield KVCacheDType.FP8_E5M2 diff --git a/tests/pytorch/kernel/test_flatten_kv_cache.py b/tests/pytorch/kernel/test_flatten_kv_cache.py index 9a3eaac60f..c46185ae99 100644 --- a/tests/pytorch/kernel/test_flatten_kv_cache.py +++ b/tests/pytorch/kernel/test_flatten_kv_cache.py @@ -1,7 +1,7 @@ import pytest import torch -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType # Import common TurboQuant utilities from turboquant_utils from .turboquant_utils import ( @@ -154,7 +154,7 @@ def test_flatten_kv_cache(self, k_quant, v_quant, kv_seqlens, block_offsets, out out_dtype=out_dtype, k_scales_zeros=k_sz, v_scales_zeros=v_sz, - quant_policy=nbits) + kv_cache_dtype=nbits) torch.testing.assert_close(k_states, gt[0], atol=atol, rtol=rtol) torch.testing.assert_close(v_states, gt[1], atol=atol, rtol=rtol) @@ -285,7 +285,7 @@ def test_flatten_kv_cache_quant_zeroes_padded_tail(flatten_kv_layout, nbits): out_dtype=torch.float16, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, - quant_policy=nbits, + kv_cache_dtype=nbits, flatten_kv_layout=flatten_kv_layout) if flatten_kv_layout == 'hsd': @@ -342,11 +342,11 @@ def fp8_dtype(self): yield torch.float8_e4m3fn @pytest.fixture - def quant_policy(self): - yield QuantPolicy.FP8 + def kv_cache_dtype(self): + yield KVCacheDType.FP8 def test_flatten_kv_cache(self, k_caches, v_caches, kv_lens, kv_seqlens, block_offsets, block_size, num_heads, - out_size, head_dim, out_dtype, fp8_dtype, quant_policy): + out_size, head_dim, out_dtype, fp8_dtype, kv_cache_dtype): from lmdeploy.pytorch.kernels.cuda.flatten_kv_cache import flatten_kv_cache k_caches_fp8, k_scale, k_dequant = quant_fp8_scalar(k_caches, fp8_dtype, scale=0.25) @@ -362,13 +362,13 @@ def test_flatten_kv_cache(self, k_caches, v_caches, kv_lens, kv_seqlens, block_o out_dtype=out_dtype, k_scales_zeros=k_scale, v_scales_zeros=v_scale, - quant_policy=quant_policy) + kv_cache_dtype=kv_cache_dtype) torch.testing.assert_close(k_states, gt[0], atol=1e-3, rtol=1e-5) torch.testing.assert_close(v_states, gt[1], atol=1e-3, rtol=1e-5) def test_flatten_kv_cache_requires_scale(self, k_caches, v_caches, kv_seqlens, block_offsets, out_size, out_dtype, - fp8_dtype, quant_policy): + fp8_dtype, kv_cache_dtype): from lmdeploy.pytorch.kernels.cuda.flatten_kv_cache import flatten_kv_cache k_caches_fp8 = k_caches.to(fp8_dtype) @@ -381,11 +381,11 @@ def test_flatten_kv_cache_requires_scale(self, k_caches, v_caches, kv_seqlens, b block_offsets, out_size=out_size, out_dtype=out_dtype, - quant_policy=quant_policy) + kv_cache_dtype=kv_cache_dtype) @pytest.mark.parametrize('flatten_kv_layout', ['hsd', 'shd']) def test_flatten_kv_cache_zeroes_padded_tail(self, k_caches, v_caches, kv_seqlens, block_offsets, block_size, - out_size, out_dtype, fp8_dtype, quant_policy, flatten_kv_layout): + out_size, out_dtype, fp8_dtype, kv_cache_dtype, flatten_kv_layout): from lmdeploy.pytorch.kernels.cuda.flatten_kv_cache import flatten_kv_cache padded_out_size = _div_up(out_size, block_size) * block_size + block_size @@ -400,7 +400,7 @@ def test_flatten_kv_cache_zeroes_padded_tail(self, k_caches, v_caches, kv_seqlen out_dtype=out_dtype, k_scales_zeros=k_scale, v_scales_zeros=v_scale, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, flatten_kv_layout=flatten_kv_layout) if flatten_kv_layout == 'hsd': @@ -418,8 +418,8 @@ def fp8_dtype(self): yield torch.float8_e5m2 @pytest.fixture - def quant_policy(self): - yield QuantPolicy.FP8_E5M2 + def kv_cache_dtype(self): + yield KVCacheDType.FP8_E5M2 @pytest.mark.skipif(torch.cuda.get_device_capability()[0] < 9, reason='require device with cc>=9.0') @@ -514,14 +514,14 @@ def test_flatten_kv_cache_zeroes_padded_tail(self, k_cache_mla, kv_seqlens, bloc # ============================================================================= -# Tests for quant_policy=QuantPolicy.TURBO_QUANT (TurboQuant) flatten_kv_cache +# Tests for kv_cache_dtype=KVCacheDType.TURBO_QUANT (TurboQuant) flatten_kv_cache # ============================================================================= class TestFlattenKVCacheQuant42: - """Test flatten_kv_cache with quant_policy=QuantPolicy.TURBO_QUANT + """Test flatten_kv_cache with kv_cache_dtype=KVCacheDType.TURBO_QUANT (TurboQuant). - quant_policy=QuantPolicy.TURBO_QUANT uses: + kv_cache_dtype=KVCacheDType.TURBO_QUANT uses: - K: QJL4 (3bit MSE + 1bit QJL), stored in rotate domain - V: TurboQuant MSE int2, stored in rotate domain @@ -621,7 +621,7 @@ def out_dtype(self): def test_flatten_kv_cache_quant42(self, k_caches, v_caches, kv_seqlens, block_offsets, k_scales_zeros, v_scales_zeros, out_dtype, head_dim, head_dim_v, num_heads): - """Test flatten_kv_cache with quant_policy=QuantPolicy.TURBO_QUANT. + """Test flatten_kv_cache with kv_cache_dtype=KVCacheDType.TURBO_QUANT. This test verifies that: 1. The flatten function runs without error @@ -634,7 +634,7 @@ def test_flatten_kv_cache_quant42(self, k_caches, v_caches, kv_seqlens, block_of ) - # Run flatten with quant_policy=QuantPolicy.TURBO_QUANT + # Run flatten with kv_cache_dtype=KVCacheDType.TURBO_QUANT k_states, v_states = flatten_kv_cache( k_caches, v_caches, @@ -642,7 +642,7 @@ def test_flatten_kv_cache_quant42(self, k_caches, v_caches, kv_seqlens, block_of block_offsets, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, - quant_policy=QuantPolicy.TURBO_QUANT, + kv_cache_dtype=KVCacheDType.TURBO_QUANT, kv_layout='bshd', flatten_kv_layout='shd', out_dtype=out_dtype, diff --git a/tests/pytorch/kernel/test_paged_attention.py b/tests/pytorch/kernel/test_paged_attention.py index 5742c9b16f..d0482b3b5d 100644 --- a/tests/pytorch/kernel/test_paged_attention.py +++ b/tests/pytorch/kernel/test_paged_attention.py @@ -3,7 +3,7 @@ import pytest import torch -from lmdeploy.messages import QuantPolicy +from lmdeploy.messages import KVCacheDType from lmdeploy.pytorch.kernels.cuda.turbo_quant import ( hadamard_rotate, hadamard_rotate_inv, @@ -502,7 +502,7 @@ def test_paged_attention(self, conti_q, blocked_kv, block_offsets, kv_seqlens, c blocked_v, k_scales_zeros=blocked_ksz, v_scales_zeros=blocked_vsz, - quant_policy=nbits, + kv_cache_dtype=nbits, page_table=block_offsets, cache_seqlens=kv_seqlens) if nbits == 4: @@ -527,7 +527,7 @@ def test_window_attention(self, conti_q, blocked_kv, block_offsets, kv_seqlens, blocked_v, k_scales_zeros=blocked_ksz, v_scales_zeros=blocked_vsz, - quant_policy=nbits, + kv_cache_dtype=nbits, page_table=block_offsets, cache_seqlens=kv_seqlens, window_size=win_size) @@ -555,9 +555,9 @@ def fp8_dtype(self): yield torch.float8_e4m3fn @pytest.fixture - def quant_policy(self): - from lmdeploy.messages import QuantPolicy - yield QuantPolicy.FP8 + def kv_cache_dtype(self): + from lmdeploy.messages import KVCacheDType + yield KVCacheDType.FP8 @pytest.fixture def blocked_kv(self, batched_kv, seq_lens, history_lens, block_offsets, block_size, num_heads_k, feat_dim, @@ -576,7 +576,7 @@ def gt(self, batched_q, blocked_kv, mask): @pytest.mark.parametrize(['num_heads_q', 'num_heads_k'], [(8, 2), (2, 2)], indirect=True) @pytest.mark.parametrize('history_lens', [(50, 40, 30, 20)], indirect=True) @pytest.mark.parametrize('block_size', [16], indirect=True) - def test_paged_attention(self, conti_q, blocked_kv, block_offsets, kv_seqlens, conti_gt, quant_policy): + def test_paged_attention(self, conti_q, blocked_kv, block_offsets, kv_seqlens, conti_gt, kv_cache_dtype): from lmdeploy.pytorch.kernels.cuda import flash_attn_with_kvcache blocked_k, blocked_v, k_scale, v_scale, _, _ = blocked_kv @@ -586,7 +586,7 @@ def test_paged_attention(self, conti_q, blocked_kv, block_offsets, kv_seqlens, c blocked_v, k_scales_zeros=k_scale, v_scales_zeros=v_scale, - quant_policy=quant_policy, + kv_cache_dtype=kv_cache_dtype, page_table=block_offsets, cache_seqlens=kv_seqlens) torch.testing.assert_close(out, conti_gt, atol=1e-3, rtol=1e-5) @@ -599,9 +599,9 @@ def fp8_dtype(self): yield torch.float8_e5m2 @pytest.fixture - def quant_policy(self): - from lmdeploy.messages import QuantPolicy - yield QuantPolicy.FP8_E5M2 + def kv_cache_dtype(self): + from lmdeploy.messages import KVCacheDType + yield KVCacheDType.FP8_E5M2 class TestPagedAttentionBlockDecoding(TestPagedAttentionBase): @@ -676,7 +676,7 @@ def test_paged_attention(self, conti_q, blocked_kv, block_offsets, kv_seqlens, l # ============================================================================= -# quant_policy=QuantPolicy.TURBO_QUANT Tests (TurboQuant: K=QJL4, V=TurboQuant MSE int2) +# kv_cache_dtype=KVCacheDType.TURBO_QUANT Tests (TurboQuant: K=QJL4, V=TurboQuant MSE int2) # ============================================================================= def _make_blocked_cache_quant42(batched_k, @@ -688,7 +688,8 @@ def _make_blocked_cache_quant42(batched_k, num_heads_k, feat_dim, feat_dim_v): - """Create full blocked KV cache with quant_policy=QuantPolicy.TURBO_QUANT. + """Create full blocked KV cache with + kv_cache_dtype=KVCacheDType.TURBO_QUANT. This matches the semantics of the standard paged attention tests: the cache already contains the full KV sequence of length @@ -787,10 +788,10 @@ def _recover_kv_from_blocked_cache(blocked_k, class TestPagedAttentionQuant42(TestPagedAttentionBase): - """Test quant_policy=QuantPolicy.TURBO_QUANT (TurboQuant) attention kernel - numerical correctness. + """Test kv_cache_dtype=KVCacheDType.TURBO_QUANT (TurboQuant) attention + kernel numerical correctness. - quant_policy=QuantPolicy.TURBO_QUANT uses: + kv_cache_dtype=KVCacheDType.TURBO_QUANT uses: - K: QJL4 (3bit MSE + 1bit QJL) - V: TurboQuant MSE int2 @@ -949,7 +950,8 @@ def conti_gt(self, gt, seq_lens): @pytest.mark.parametrize('history_lens', [(8, 4, 2, 1)], indirect=True) @pytest.mark.parametrize('block_size', [16], indirect=True) def test_paged_attention(self, conti_q, blocked_kv, block_offsets, kv_seqlens, conti_gt): - """Test paged attention with quant_policy=QuantPolicy.TURBO_QUANT.""" + """Test paged attention with + kv_cache_dtype=KVCacheDType.TURBO_QUANT.""" from lmdeploy.pytorch.kernels.cuda import flash_attn_with_kvcache blocked_k, blocked_v, blocked_ksz, blocked_vsz = blocked_kv @@ -959,7 +961,7 @@ def test_paged_attention(self, conti_q, blocked_kv, block_offsets, kv_seqlens, c blocked_v, k_scales_zeros=blocked_ksz, v_scales_zeros=blocked_vsz, - quant_policy=QuantPolicy.TURBO_QUANT, + kv_cache_dtype=KVCacheDType.TURBO_QUANT, page_table=block_offsets, cache_seqlens=kv_seqlens, ) @@ -971,12 +973,13 @@ def test_paged_attention(self, conti_q, blocked_kv, block_offsets, kv_seqlens, c class TestPagedAttentionFP16vsQuant42(TestPagedAttentionBase): - """Compare FP16 vs quant_policy=QuantPolicy.TURBO_QUANT attention outputs. + """Compare FP16 vs kv_cache_dtype=KVCacheDType.TURBO_QUANT attention + outputs. - This test verifies that quant_policy=QuantPolicy.TURBO_QUANT (TurboQuant) produces numerically + This test verifies that kv_cache_dtype=KVCacheDType.TURBO_QUANT (TurboQuant) produces numerically reasonable results compared to FP16 baseline. - quant_policy=QuantPolicy.TURBO_QUANT uses: + kv_cache_dtype=KVCacheDType.TURBO_QUANT uses: - K: QJL4 (3bit MSE + 1bit QJL) - V: TurboQuant MSE int2 """ @@ -1072,7 +1075,7 @@ def blocked_kv_fp16(self, batched_kv, seq_lens, history_lens, block_offsets, blo @pytest.fixture def blocked_kv_quant42(self, batched_kv, seq_lens, history_lens, block_offsets, block_size, num_heads_k, feat_dim, feat_dim_v): - """Build quant_policy=QuantPolicy.TURBO_QUANT blocked KV cache.""" + """Build kv_cache_dtype=KVCacheDType.TURBO_QUANT blocked KV cache.""" batched_k, batched_v = batched_kv yield _make_blocked_cache_quant42(batched_k, batched_v, seq_lens, history_lens, block_offsets, block_size, num_heads_k, feat_dim, feat_dim_v) @@ -1089,13 +1092,13 @@ def out_fp16(self, conti_q, blocked_kv_fp16, block_offsets, kv_seqlens): blocked_v, page_table=block_offsets, cache_seqlens=kv_seqlens, - quant_policy=QuantPolicy.NONE, + kv_cache_dtype=KVCacheDType.AUTO, ) yield out @pytest.fixture def out_quant42(self, conti_q, blocked_kv_quant42, block_offsets, kv_seqlens): - """Run attention with quant_policy=QuantPolicy.TURBO_QUANT cache.""" + """Run attention with kv_cache_dtype=KVCacheDType.TURBO_QUANT cache.""" from lmdeploy.pytorch.kernels.cuda import flash_attn_with_kvcache blocked_k, blocked_v, blocked_ksz, blocked_vsz = blocked_kv_quant42 @@ -1105,7 +1108,7 @@ def out_quant42(self, conti_q, blocked_kv_quant42, block_offsets, kv_seqlens): blocked_v, k_scales_zeros=blocked_ksz, v_scales_zeros=blocked_vsz, - quant_policy=QuantPolicy.TURBO_QUANT, + kv_cache_dtype=KVCacheDType.TURBO_QUANT, page_table=block_offsets, cache_seqlens=kv_seqlens, ) diff --git a/tests/pytorch/kernel/test_turboquant.py b/tests/pytorch/kernel/test_turboquant.py index 742bf5c6c8..92f5ca431a 100644 --- a/tests/pytorch/kernel/test_turboquant.py +++ b/tests/pytorch/kernel/test_turboquant.py @@ -1,7 +1,7 @@ -"""Tests for TurboQuant (quant_policy=QuantPolicy.TURBO_QUANT). +"""Tests for TurboQuant (kv_cache_dtype=KVCacheDType.TURBO_QUANT). This module contains kernel-level tests for TurboQuant MSE quantization, -which is used by quant_policy=QuantPolicy.TURBO_QUANT (K=QJL4, V=2bit mixed precision). +which is used by kv_cache_dtype=KVCacheDType.TURBO_QUANT (K=QJL4, V=2bit mixed precision). TurboQuant is a quantization method that: - Uses Lloyd-Max algorithm for optimal quantization @@ -34,7 +34,7 @@ class TestTurboQuantMSE: """Verify TurboQuant MSE quantization-dequantization correctness. - These tests verify the core TurboQuant MSE algorithm used by quant_policy=QuantPolicy.TURBO_QUANT. + These tests verify the core TurboQuant MSE algorithm used by kv_cache_dtype=KVCacheDType.TURBO_QUANT. """ @pytest.fixture diff --git a/tests/pytorch/kernel/turboquant_utils.py b/tests/pytorch/kernel/turboquant_utils.py index c302974787..a2491ea338 100644 --- a/tests/pytorch/kernel/turboquant_utils.py +++ b/tests/pytorch/kernel/turboquant_utils.py @@ -1,8 +1,8 @@ -"""Common test utilities for TurboQuant (quant_policy=QuantPolicy.TURBO_QUANT) -kernel tests. +"""Common test utilities for TurboQuant +(kv_cache_dtype=KVCacheDType.TURBO_QUANT) kernel tests. This module contains shared helper functions for testing TurboQuant quantization, -which is used by quant_policy=QuantPolicy.TURBO_QUANT (K=QJL4, V=2bit mixed precision). +which is used by kv_cache_dtype=KVCacheDType.TURBO_QUANT (K=QJL4, V=2bit mixed precision). TurboQuant is a quantization method that: - Uses Lloyd-Max algorithm for optimal quantization diff --git a/tests/test_lmdeploy/test_fp8_kv_cache_dtype.py b/tests/test_lmdeploy/test_fp8_kv_cache_dtype.py new file mode 100644 index 0000000000..c12ac339ff --- /dev/null +++ b/tests/test_lmdeploy/test_fp8_kv_cache_dtype.py @@ -0,0 +1,128 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +from types import SimpleNamespace + +import pytest +import torch +from pydantic_core import ValidationError + +from lmdeploy.cli.utils import ArgumentHelper +from lmdeploy.messages import KVCacheDType, PytorchEngineConfig, TurbomindEngineConfig +from lmdeploy.pytorch.config import CacheConfig +from lmdeploy.pytorch.engine.cache_engine import ( + CacheDesc, + CacheEngine, + _describe_kv_cache_dtype, + _get_fp8_cache_tensor_dtype, +) + + +@pytest.mark.parametrize( + ('value', 'expected'), + [ + ('auto', KVCacheDType.AUTO), + ('0', KVCacheDType.AUTO), + ('int4', KVCacheDType.INT4), + ('4', KVCacheDType.INT4), + ('int8', KVCacheDType.INT8), + ('8', KVCacheDType.INT8), + ('fp8', KVCacheDType.FP8), + ('fp8_e4m3', KVCacheDType.FP8), + ('16', KVCacheDType.FP8), + ('fp8_e5m2', KVCacheDType.FP8_E5M2), + ('17', KVCacheDType.FP8_E5M2), + ('turbo_quant', KVCacheDType.TURBO_QUANT), + ('42', KVCacheDType.TURBO_QUANT), + ], +) +def test_kv_cache_dtype_aliases(value, expected): + parser = argparse.ArgumentParser() + ArgumentHelper.kv_cache_dtype(parser) + + assert parser.parse_args(['--kv-cache-dtype', value]).kv_cache_dtype == expected + + +@pytest.mark.parametrize('value', ['none', 'bad']) +def test_kv_cache_dtype_rejects_invalid_aliases(value): + parser = argparse.ArgumentParser() + ArgumentHelper.kv_cache_dtype(parser) + + with pytest.raises(SystemExit): + parser.parse_args(['--kv-cache-dtype', value]) + + +def test_removed_cli_flag_is_rejected(): + parser = argparse.ArgumentParser() + ArgumentHelper.kv_cache_dtype(parser) + old_flag = '--' + 'quant' + '-policy' + + with pytest.raises(SystemExit): + parser.parse_args([old_flag, 'fp8']) + + +def test_kv_cache_dtype_uses_configured_default(): + parser = argparse.ArgumentParser() + ArgumentHelper.kv_cache_dtype(parser, default=KVCacheDType.FP8.value) + + assert parser.parse_args([]).kv_cache_dtype == KVCacheDType.FP8 + + +def test_pytorch_config_accepts_fp8_kv_cache_dtypes(): + config = PytorchEngineConfig(kv_cache_dtype=KVCacheDType.FP8_E5M2) + + assert config.kv_cache_dtype == KVCacheDType.FP8_E5M2 + + +@pytest.mark.parametrize('value', [KVCacheDType.FP8.value, 'fp8', 'fp8_e4m3']) +def test_pytorch_config_normalizes_kv_cache_dtype_value(value): + config = PytorchEngineConfig(kv_cache_dtype=value) + + assert config.kv_cache_dtype == KVCacheDType.FP8 + + +def test_pytorch_config_rejects_invalid_kv_cache_dtype(): + with pytest.raises(ValueError, match='invalid kv_cache_dtype: 99'): + PytorchEngineConfig(kv_cache_dtype=99) + + +def test_pytorch_config_rejects_removed_keyword(): + old_kw = 'quant' + '_policy' + with pytest.raises(TypeError): + PytorchEngineConfig(**{old_kw: KVCacheDType.FP8}) + + +@pytest.mark.parametrize('kv_cache_dtype', [KVCacheDType.FP8, KVCacheDType.FP8_E5M2]) +def test_turbomind_config_rejects_fp8_kv_cache_dtypes(kv_cache_dtype): + with pytest.raises(ValidationError, match='invalid kv_cache_dtype'): + TurbomindEngineConfig(kv_cache_dtype=kv_cache_dtype) + + +def test_turbomind_config_rejects_removed_keyword(): + old_kw = 'quant' + '_policy' + with pytest.raises(ValidationError): + TurbomindEngineConfig(**{old_kw: KVCacheDType.INT8}) + + +def test_fp8_kv_cache_dtype_mapping(): + assert _get_fp8_cache_tensor_dtype(KVCacheDType.FP8) is torch.float8_e4m3fn + assert _get_fp8_cache_tensor_dtype(KVCacheDType.FP8_E5M2) is torch.float8_e5m2 + + +def test_fp8_kv_cache_log_description(): + assert _describe_kv_cache_dtype(KVCacheDType.FP8) == 'fp8_e4m3 KV cache' + assert _describe_kv_cache_dtype(KVCacheDType.FP8_E5M2) == 'fp8_e5m2 KV cache' + assert _describe_kv_cache_dtype(KVCacheDType.AUTO) is None + + +def test_fp8_quant_cache_descs_are_empty(): + model_config = SimpleNamespace(dtype=torch.float16) + k_desc = CacheDesc(shape=[4, 16, 2, 128], dtype=torch.float8_e4m3fn) + v_desc = CacheDesc(shape=[4, 16, 2, 128], dtype=torch.float8_e4m3fn) + + normal_cache_config = CacheConfig(max_batches=1, + block_size=16, + num_cpu_blocks=0, + num_gpu_blocks=1, + kv_cache_dtype=KVCacheDType.FP8) + + assert CacheEngine.get_quant_cache_descs(k_desc, v_desc, model_config, normal_cache_config) == [] diff --git a/tests/test_lmdeploy/test_fp8_kv_cache_policy.py b/tests/test_lmdeploy/test_fp8_kv_cache_policy.py deleted file mode 100644 index 38989c2bc9..0000000000 --- a/tests/test_lmdeploy/test_fp8_kv_cache_policy.py +++ /dev/null @@ -1,82 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. -import argparse -from types import SimpleNamespace - -import pytest -import torch -from pydantic_core import ValidationError - -from lmdeploy.cli.utils import ArgumentHelper -from lmdeploy.messages import PytorchEngineConfig, QuantPolicy, TurbomindEngineConfig -from lmdeploy.pytorch.config import CacheConfig -from lmdeploy.pytorch.engine.cache_engine import ( - CacheDesc, - CacheEngine, - _describe_kv_cache_quant_policy, - _get_fp8_cache_dtype, -) - - -def test_quant_policy_fp8_aliases(): - parser = argparse.ArgumentParser() - ArgumentHelper.quant_policy(parser) - - assert parser.parse_args(['--quant-policy', 'fp8']).quant_policy == QuantPolicy.FP8 - assert parser.parse_args(['--quant-policy', 'fp8_e4m3']).quant_policy == QuantPolicy.FP8 - assert parser.parse_args(['--quant-policy', 'fp8_e5m2']).quant_policy == QuantPolicy.FP8_E5M2 - assert parser.parse_args(['--quant-policy', '17']).quant_policy == QuantPolicy.FP8_E5M2 - - -def test_quant_policy_uses_configured_default(): - parser = argparse.ArgumentParser() - ArgumentHelper.quant_policy(parser, default=QuantPolicy.FP8.value) - - assert parser.parse_args([]).quant_policy == QuantPolicy.FP8.value - - -def test_pytorch_config_accepts_fp8_quant_policies(): - config = PytorchEngineConfig(quant_policy=QuantPolicy.FP8_E5M2) - - assert config.quant_policy == QuantPolicy.FP8_E5M2 - - -def test_pytorch_config_normalizes_quant_policy_value(): - config = PytorchEngineConfig(quant_policy=QuantPolicy.FP8.value) - - assert config.quant_policy == QuantPolicy.FP8 - - -def test_pytorch_config_rejects_invalid_quant_policy(): - with pytest.raises(ValueError, match='invalid quant_policy: 99'): - PytorchEngineConfig(quant_policy=99) - - -@pytest.mark.parametrize('quant_policy', [QuantPolicy.FP8, QuantPolicy.FP8_E5M2]) -def test_turbomind_config_rejects_fp8_quant_policies(quant_policy): - with pytest.raises(ValidationError, match='invalid quant_policy'): - TurbomindEngineConfig(quant_policy=quant_policy) - - -def test_fp8_kv_cache_dtype_mapping(): - assert _get_fp8_cache_dtype(QuantPolicy.FP8) is torch.float8_e4m3fn - assert _get_fp8_cache_dtype(QuantPolicy.FP8_E5M2) is torch.float8_e5m2 - - -def test_fp8_kv_cache_log_description(): - assert _describe_kv_cache_quant_policy(QuantPolicy.FP8) == 'fp8_e4m3 KV cache' - assert _describe_kv_cache_quant_policy(QuantPolicy.FP8_E5M2) == 'fp8_e5m2 KV cache' - assert _describe_kv_cache_quant_policy(QuantPolicy.NONE) is None - - -def test_fp8_quant_cache_descs_are_empty(): - model_config = SimpleNamespace(dtype=torch.float16) - k_desc = CacheDesc(shape=[4, 16, 2, 128], dtype=torch.float8_e4m3fn) - v_desc = CacheDesc(shape=[4, 16, 2, 128], dtype=torch.float8_e4m3fn) - - normal_cache_config = CacheConfig(max_batches=1, - block_size=16, - num_cpu_blocks=0, - num_gpu_blocks=1, - quant_policy=QuantPolicy.FP8) - - assert CacheEngine.get_quant_cache_descs(k_desc, v_desc, model_config, normal_cache_config) == [] diff --git a/tests/test_lmdeploy/test_quant_policy.py b/tests/test_lmdeploy/test_kv_cache_dtype.py similarity index 89% rename from tests/test_lmdeploy/test_quant_policy.py rename to tests/test_lmdeploy/test_kv_cache_dtype.py index 57ad7c5a16..9049b4dfe9 100644 --- a/tests/test_lmdeploy/test_quant_policy.py +++ b/tests/test_lmdeploy/test_kv_cache_dtype.py @@ -1,8 +1,8 @@ -"""Test quant_policy=QuantPolicy.TURBO_QUANT (K=4bit, V=2bit mixed precision) -for PytorchEngine. +"""Test kv_cache_dtype=KVCacheDType.TURBO_QUANT (K=4bit, V=2bit mixed +precision) for PytorchEngine. -This module tests both functional correctness and accuracy of quant_policy=QuantPolicy.TURBO_QUANT against a non- -quantized (quant_policy=QuantPolicy.NONE) baseline. +This module tests both functional correctness and accuracy of kv_cache_dtype=KVCacheDType.TURBO_QUANT against a non- +quantized (kv_cache_dtype=KVCacheDType.AUTO) baseline. """ import gc @@ -11,10 +11,10 @@ import torch from lmdeploy import GenerationConfig, PytorchEngineConfig, pipeline -from lmdeploy.messages import QuantPolicy, Response +from lmdeploy.messages import KVCacheDType, Response -# Use smaller model to avoid OOM when running both quant_policy=QuantPolicy.NONE -# and quant_policy=QuantPolicy.TURBO_QUANT +# Use smaller model to avoid OOM when running both kv_cache_dtype=KVCacheDType.AUTO +# and kv_cache_dtype=KVCacheDType.TURBO_QUANT MODEL_ID = 'Qwen/Qwen3-8B' @@ -42,7 +42,7 @@ def pipe_no_quant(model_id): engine_config = PytorchEngineConfig( tp=1, cache_max_entry_count=0.05, - quant_policy=QuantPolicy.NONE, # No quantization + kv_cache_dtype=KVCacheDType.AUTO, # No quantization ) pipe = pipeline(model_id, backend_config=engine_config, log_level='INFO') yield pipe @@ -56,7 +56,7 @@ def pipe_no_quant(model_id): @pytest.fixture(scope='class') def pipe_quant_42(model_id): - """Create pipeline with quant_policy=QuantPolicy.TURBO_QUANT. + """Create pipeline with kv_cache_dtype=KVCacheDType.TURBO_QUANT. This fixture has class scope so large model instances are released before later FP8 accuracy tests allocate their own pipelines. @@ -64,7 +64,7 @@ def pipe_quant_42(model_id): engine_config = PytorchEngineConfig( tp=1, cache_max_entry_count=0.05, - quant_policy=QuantPolicy.TURBO_QUANT, # K=4bit, V=2bit mixed precision + kv_cache_dtype=KVCacheDType.TURBO_QUANT, # K=4bit, V=2bit mixed precision ) pipe = pipeline(model_id, backend_config=engine_config, log_level='INFO') yield pipe @@ -77,11 +77,11 @@ def pipe_quant_42(model_id): # ============================================================================= -# Basic Functional Tests (quant_policy=QuantPolicy.TURBO_QUANT only) +# Basic Functional Tests (kv_cache_dtype=KVCacheDType.TURBO_QUANT only) # ============================================================================= -class TestQuantPolicy42Basic: - """Basic functional tests for quant_policy=QuantPolicy.TURBO_QUANT. +class TestKVCacheDType42Basic: + """Basic functional tests for kv_cache_dtype=KVCacheDType.TURBO_QUANT. These tests verify that the quantized model can perform basic inference without errors. They test single prompt, batch prompts, and generation config. @@ -89,11 +89,11 @@ class TestQuantPolicy42Basic: @pytest.fixture(scope='class') def pipe(self): - """Create pipeline with quant_policy=QuantPolicy.TURBO_QUANT.""" + """Create pipeline with kv_cache_dtype=KVCacheDType.TURBO_QUANT.""" engine_config = PytorchEngineConfig( tp=1, cache_max_entry_count=0.1, - quant_policy=QuantPolicy.TURBO_QUANT, + kv_cache_dtype=KVCacheDType.TURBO_QUANT, ) pipe = pipeline(MODEL_ID, backend_config=engine_config, log_level='INFO') yield pipe @@ -105,7 +105,7 @@ def pipe(self): def test_infer_single_prompt(self, pipe): """Test single prompt inference with - quant_policy=QuantPolicy.TURBO_QUANT.""" + kv_cache_dtype=KVCacheDType.TURBO_QUANT.""" prompt = 'Hello, how are you?' response = pipe.infer(prompt, max_new_tokens=30) @@ -115,7 +115,8 @@ def test_infer_single_prompt(self, pipe): assert len(response.text.strip()) > 0 def test_infer_batch_prompts(self, pipe): - """Test batch inference with quant_policy=QuantPolicy.TURBO_QUANT.""" + """Test batch inference with + kv_cache_dtype=KVCacheDType.TURBO_QUANT.""" prompts = ['What is AI?', 'Hello!'] responses = pipe.infer(prompts, max_new_tokens=20) @@ -136,15 +137,15 @@ def test_infer_with_generation_config(self, pipe): # ============================================================================= -# Accuracy Tests (quant_policy=QuantPolicy.NONE vs quant_policy=QuantPolicy.TURBO_QUANT) +# Accuracy Tests (kv_cache_dtype=KVCacheDType.AUTO vs kv_cache_dtype=KVCacheDType.TURBO_QUANT) # ============================================================================= -class TestQuantPolicy42Accuracy: - """Accuracy tests comparing quant_policy=QuantPolicy.TURBO_QUANT against +class TestKVCacheDType42Accuracy: + """Accuracy tests comparing kv_cache_dtype=KVCacheDType.TURBO_QUANT against non-quantized baseline. - These tests verify the numerical accuracy/precision of quant_policy=QuantPolicy.TURBO_QUANT - (K=4bit, V=2bit mixed precision) by comparing against quant_policy=QuantPolicy.NONE. + These tests verify the numerical accuracy/precision of kv_cache_dtype=KVCacheDType.TURBO_QUANT + (K=4bit, V=2bit mixed precision) by comparing against kv_cache_dtype=KVCacheDType.AUTO. Error thresholds are relaxed due to aggressive quantization: - MAE < 0.1 on logits @@ -296,22 +297,22 @@ def test_logprobs_sanity(self, pipe_no_quant, pipe_quant_42): # ============================================================================= -# FP8 Tests (QuantPolicy.FP8) +# FP8 Tests (KVCacheDType.FP8) # ============================================================================= @pytest.mark.skipif(_e4m3_fp8_unsupported_on_pre_sm90(), reason='Triton float8_e4m3fn conversion requires device with cc>=9.0') -class TestQuantPolicyFP8Basic: - """Basic functional tests for quant_policy=QuantPolicy.FP8.""" +class TestKVCacheDTypeFP8Basic: + """Basic functional tests for kv_cache_dtype=KVCacheDType.FP8.""" @pytest.fixture(scope='class') def pipe(self): - """Create pipeline with quant_policy=QuantPolicy.FP8.""" + """Create pipeline with kv_cache_dtype=KVCacheDType.FP8.""" engine_config = PytorchEngineConfig( tp=1, cache_max_entry_count=0.1, - quant_policy=QuantPolicy.FP8, + kv_cache_dtype=KVCacheDType.FP8, ) pipe = pipeline(MODEL_ID, backend_config=engine_config, log_level='INFO') yield pipe @@ -322,7 +323,8 @@ def pipe(self): torch.cuda.empty_cache() def test_infer_single_prompt(self, pipe): - """Test single prompt inference with quant_policy=QuantPolicy.FP8.""" + """Test single prompt inference with + kv_cache_dtype=KVCacheDType.FP8.""" prompt = 'Hello, how are you?' response = pipe.infer(prompt, max_new_tokens=30) @@ -332,7 +334,7 @@ def test_infer_single_prompt(self, pipe): assert len(response.text.strip()) > 0 def test_infer_batch_prompts(self, pipe): - """Test batch inference with quant_policy=QuantPolicy.FP8.""" + """Test batch inference with kv_cache_dtype=KVCacheDType.FP8.""" prompts = ['What is AI?', 'Hello!'] responses = pipe.infer(prompts, max_new_tokens=20) @@ -354,15 +356,15 @@ def test_infer_with_generation_config(self, pipe): @pytest.mark.skipif(_e4m3_fp8_unsupported_on_pre_sm90(), reason='Triton float8_e4m3fn conversion requires device with cc>=9.0') -class TestQuantPolicyFP8Accuracy: - """Accuracy tests comparing quant_policy=QuantPolicy.FP8 against non- +class TestKVCacheDTypeFP8Accuracy: + """Accuracy tests comparing kv_cache_dtype=KVCacheDType.FP8 against non- quantized baseline. FP8 (float8_e4m3fn, per-tensor scale) is more precise than 4-bit TurboQuant, so thresholds are tighter: MAE < 0.05, Max AE < 0.3. Uses class-scoped fixtures to avoid holding three models in GPU memory simultaneously - when running the full test suite alongside TestQuantPolicy42Accuracy. + when running the full test suite alongside TestKVCacheDType42Accuracy. """ @pytest.fixture(scope='class') @@ -371,7 +373,7 @@ def pipe_no_quant(self): engine_config = PytorchEngineConfig( tp=1, cache_max_entry_count=0.05, - quant_policy=QuantPolicy.NONE, + kv_cache_dtype=KVCacheDType.AUTO, ) pipe = pipeline(MODEL_ID, backend_config=engine_config, log_level='INFO') yield pipe @@ -387,7 +389,7 @@ def pipe_quant_fp8(self): engine_config = PytorchEngineConfig( tp=1, cache_max_entry_count=0.05, - quant_policy=QuantPolicy.FP8, + kv_cache_dtype=KVCacheDType.FP8, ) pipe = pipeline(MODEL_ID, backend_config=engine_config, log_level='INFO') yield pipe