Enhance KVCacheManager to maintain adjusted max attention window sizes. Introduced an adjusted dictionary to track window size mappings and updated the logic to reflect these changes in the max attention window vector. Updated unit tests to validate the new behavior and ensure expected outputs for various memory configurations.

Signed-off-by: qixiang-99 <[email protected]>

Author: qixiang-99

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -670,6 +670,7 @@ def calculate_cache_size_per_token(layers: Set[int]) -> int:
 
         accum_max_tokens = 0
         prev_window_size = 0
+        adjusted_dict = {}
         adjusted_max_attention_window_vec = max_attention_window_vec.copy()
 
         for window_size in sorted(window_size_to_layers):
@@ -712,12 +713,14 @@ def calculate_cache_size_per_token(layers: Set[int]) -> int:
 
             if accum_max_tokens not in adjusted_window_size_to_layers:
                 adjusted_window_size_to_layers[accum_max_tokens] = layers.copy()
-                # also update adjusted_max_attention_window_vec
-                for i, v in enumerate(adjusted_max_attention_window_vec):
-                    if v == window_size:
-                        adjusted_max_attention_window_vec[i] = accum_max_tokens
             else:
                 adjusted_window_size_to_layers[accum_max_tokens].extend(layers)
+            adjusted_dict[window_size] = accum_max_tokens
+            # also update adjusted_max_attention_window_vec
+            adjusted_max_attention_window_vec = [
+                adjusted_dict.get(v, v)
+                for v in adjusted_max_attention_window_vec
+            ]
 
             remaining_layers -= set(layers)
             prev_window_size = window_size

tests/unittest/_torch/test_resource_manager.py

@@ -433,6 +433,7 @@ def test_adjust_window_sizes_for_vswa(self):
             200: [4, 5, 6],
             7000: [7, 8],
         }
+        max_attention_window_vec = [100] * 4 + [200] * 3 + [7000] * 2
 
         model_config = self.MockModelConfig()
         model_config.num_attention_heads = 2
@@ -460,6 +461,7 @@ def test_adjust_window_sizes_for_vswa(self):
                     100: [0, 1, 2, 3],
                     130: [4, 5, 6, 7, 8],
                 },
+                [100] * 4 + [130] * 5,
                 None,
                 "limited_memory_clamped_windows"),
             (
@@ -471,6 +473,7 @@ def test_adjust_window_sizes_for_vswa(self):
                     200: [4, 5, 6],
                     1017: [7, 8],
                 },
+                [100] * 4 + [200] * 3 + [1017] * 2,
                 None,
                 "less_limited_memory_clamped_windows"),
             (
@@ -482,6 +485,7 @@ def test_adjust_window_sizes_for_vswa(self):
                     200: [4, 5, 6],
                     7000: [7, 8],
                 },
+                [100] * 4 + [200] * 3 + [7000] * 2,
                 None,
                 "sufficient_memory_no_clamping"),
             (
@@ -490,6 +494,7 @@ def test_adjust_window_sizes_for_vswa(self):
                 {
                     51: [0, 1, 2, 3, 4, 5, 6, 7, 8],
                 },
+                [51] * 9,
                 None,
                 "very_limited_memory_all_clamped"),
             (
@@ -501,15 +506,17 @@ def test_adjust_window_sizes_for_vswa(self):
                     100: [0, 1, 2, 3],
                     134: [4, 5, 6, 7, 8],
                 },
+                [100] * 4 + [134] * 5,
                 134,
                 "less_limited_memory_but_clamped_by_max_tokens"),
         ]
 
-        for memory_bytes, expected_window_sizes, max_tokens, description in test_cases:
+        for memory_bytes, expected_window_sizes, expected_max_attention_window_vec, max_tokens, description in test_cases:
             with self.subTest(case=description, memory_bytes=memory_bytes):
                 kv_cache_config = tllm.KvCacheConfig(max_tokens=max_tokens)
-                adjusted = KVCacheManager.adjust_window_sizes_for_vswa(
+                adjusted, adjusted_max_attention_window_vec = KVCacheManager.adjust_window_sizes_for_vswa(
                     window_size_to_layers=window_size_to_layers,
+                    max_attention_window_vec=max_attention_window_vec,
                     model_config=model_config,
                     kv_cache_config=kv_cache_config,
                     pool_memory_bytes=memory_bytes,
@@ -524,6 +531,13 @@ def test_adjust_window_sizes_for_vswa(self):
                     f"Memory bytes: {memory_bytes}\n"
                     f"Actual: {adjusted}\n"
                     f"Expected: {expected_window_sizes}")
+                self.assertEqual(
+                    adjusted_max_attention_window_vec,
+                    expected_max_attention_window_vec,
+                    f"Test case '{description}' failed.\n"
+                    f"Memory bytes: {memory_bytes}\n"
+                    f"Actual: {adjusted_max_attention_window_vec}\n"
+                    f"Expected: {expected_max_attention_window_vec}")
 
 
 if __name__ == "__main__":