Add kv_cache_host_memory_bytes as a configurable runtime setting (#…

…1303) * client side validation for non fp8 kv cache and fp8 context fmha * enable chunked context as default * add kv_cache_host_memory_bytes to trt-llm runtime config * fix test * bump pyproject version
basetenlabs · Jan 9, 2025 · 5f835d0 · 5f835d0
1 parent c07100b
commit 5f835d0
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 3 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.58.post1"
+version = "0.9.59.rc0"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

diff --git a/truss/base/trt_llm_config.py b/truss/base/trt_llm_config.py
@@ -81,6 +81,7 @@ class TrussSpecDecMode(str, Enum):
 
 class TrussTRTLLMRuntimeConfiguration(BaseModel):
     kv_cache_free_gpu_mem_fraction: float = 0.9
+    kv_cache_host_memory_bytes: Optional[int] = None
     enable_chunked_context: bool = True
     batch_scheduler_policy: TrussTRTLLMBatchSchedulerPolicy = (
         TrussTRTLLMBatchSchedulerPolicy.GUARANTEED_NO_EVICT

diff --git a/truss/tests/conftest.py b/truss/tests/conftest.py
@@ -409,8 +409,8 @@ def modify_handle(h: TrussHandle):
                 },
                 "runtime": {
                     "kv_cache_free_gpu_mem_fraction": 0.9,
-                    "enabled_chunked_context": False,
-                    "num_draft_tokens": None,
+                    "kv_cache_host_memory_bytes": 1000,
+                    "enabled_chunked_context": True,
                     "batch_scheduler_policy": TrussTRTLLMBatchSchedulerPolicy.GUARANTEED_NO_EVICT.value,
                 },
             }

diff --git a/truss/tests/trt_llm/test_trt_llm_config.py b/truss/tests/trt_llm/test_trt_llm_config.py
@@ -17,6 +17,7 @@ def test_trt_llm_configuration_init_and_migrate_deprecated_runtime_fields(
     trt_llm_config = TRTLLMConfiguration(**deprecated_trtllm_config["trt_llm"])
     assert trt_llm_config.runtime.model_dump() == {
         "kv_cache_free_gpu_mem_fraction": 0.1,
+        "kv_cache_host_memory_bytes": None,
         "enable_chunked_context": True,
         "batch_scheduler_policy": TrussTRTLLMBatchSchedulerPolicy.MAX_UTILIZATION.value,
         "request_default_max_tokens": 10,
@@ -32,6 +33,7 @@ def test_trt_llm_configuration_init_and_migrate_deprecated_runtime_fields_existi
     )
     assert trt_llm_config.runtime.model_dump() == {
         "kv_cache_free_gpu_mem_fraction": 0.1,
+        "kv_cache_host_memory_bytes": None,
         "enable_chunked_context": True,
         "batch_scheduler_policy": TrussTRTLLMBatchSchedulerPolicy.MAX_UTILIZATION.value,
         "request_default_max_tokens": 10,