undo accidental rename due to copy-paste and correct missed config st…

…ructure change

app_tests/benchmark_tests/llm/sglang_benchmarks/conftest.py

@@ -64,18 +64,18 @@ def pre_process_model(request, tmp_path_factory):
 
 
 @pytest.fixture(scope="module")
-def write_config(request, model_test_dir):
+def write_config(request, pre_process_model):
     batch_sizes = request.param["batch_sizes"]
     prefix_sharing_algorithm = request.param["prefix_sharing_algorithm"]
 
     # Construct the new config filename
     config_path = (
-        model_test_dir
+        pre_process_model
         / f"{'_'.join(str(bs) for bs in batch_sizes)}_{prefix_sharing_algorithm}.json"
     )
 
     # Read the base config file
-    base_config_path = model_test_dir / "config.json"
+    base_config_path = pre_process_model / "config.json"
     with open(base_config_path, "r") as f:
         config = json.load(f)
 

shortfin/python/shortfin_apps/llm/components/config_struct.py

@@ -168,7 +168,7 @@ def paged_kv_unit_size_elements(self) -> int:
         size = 1
         size *= self.transformer_block_count
         size *= 2  # K and V cache line
-        size *= self.attn_head_count
+        size *= self.paged_kv_cache.attention_head_count_kv
         size *= self.attn_head_dim
         return size