A high-throughput and memory-efficient inference and serving engine for LLMs
[PD][Bugfix] Fix KV Cache sharing with HMA (#44629)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jun 9, 2026 at 15:10 UTC
6690a0c4de37ac965475a6ea90ad4271de77a767
3 files changed
+38
-12
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+2
@@ -35,6 +35,8 @@ sw_attn_configs=(
35
"GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192" # SW model
36
"GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
37
"GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
38
+ # Gemma4: SW + cross-layer KV sharing
39
+ "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-4-E2B-it VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
40
)
41
42
# Select config array based on DP_EP env var
tests/v1/kv_connector/nixl_integration/test_accuracy.py
+28
-11
@@ -24,6 +24,7 @@ EXPECTED_VALUES = {
24
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8": 0.84,
25
"ibm-granite/granite-4.0-h-tiny": 0.77,
26
"Qwen/Qwen3.5-0.8B": 0.33,
27
+ "google/gemma-4-E2B-it": 0.485,
28
}
29
30
SIMPLE_PROMPT = (
@@ -49,17 +50,33 @@ def test_accuracy():
50
"""Run the end to end accuracy test."""
51
run_simple_prompt()
52
52
- model_args = (
53
- f"model={MODEL_NAME},"
54
- f"base_url={BASE_URL}/completions,"
55
- f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False"
56
- )
57
-
58
- results = lm_eval.simple_evaluate(
59
- model="local-completions",
60
- model_args=model_args,
61
- tasks=TASK,
62
- )
53
+ if "gemma-4" in MODEL_NAME:
54
+ # Gemma4 is quite sensible to having a chat template applied, so we evaluate
55
+ # on chat completions.
56
+ model_args = (
57
+ f"model={MODEL_NAME},"
58
+ f"base_url={BASE_URL}/chat/completions,"
59
+ f"num_concurrent={NUM_CONCURRENT},"
60
+ "tokenizer_backend=huggingface"
61
+ )
62
+ results = lm_eval.simple_evaluate(
63
+ model="local-chat-completions",
64
+ model_args=model_args,
65
+ tasks=TASK,
66
+ num_fewshot=5,
67
+ apply_chat_template=True,
68
+ )
69
+ else:
70
+ model_args = (
71
+ f"model={MODEL_NAME},"
72
+ f"base_url={BASE_URL}/completions,"
73
+ f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False"
74
+ )
75
+ results = lm_eval.simple_evaluate(
76
+ model="local-completions",
77
+ model_args=model_args,
78
+ tasks=TASK,
79
+ )
80
81
measured_value = results["results"][TASK][FILTER]
82
expected_value = EXPECTED_VALUES.get(MODEL_NAME)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+8
-1
@@ -850,7 +850,14 @@ class NixlConnectorWorker:
850
# However, physical page_size may differ when kernel requires a specific
851
# block size. This leads to SSM and FA layers having different num_blocks.
852
# `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this.
853
- layer_spec = self._layer_specs[layer_name]
853
+ layer_spec = self._layer_specs.get(layer_name)
854
+ if layer_spec is None:
855
+ logger.debug(
856
+ "Skipping layer %s as no KVCache spec is present. "
857
+ "This is likely because the layer is sharing its KV cache",
858
+ layer_name,
859
+ )
860
+ continue
861
if isinstance(layer_spec, UniformTypeKVCacheSpecs):
862
# MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs
863
layer_spec = layer_spec.kv_cache_specs[layer_name]