[Bugfix][HiSparse] Preserve per-layer offsets in KV cache bindings (#57027)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Lucas Wilkinson committed
Sep 15, 2026 at 16:27 UTC
38ca7a899cd3a9aa6f1e948b1b63f4fae7a685ce
4 files changed
+48
-6
.buildkite/test_areas/lm_eval.yaml
+18
@@ -116,6 +116,24 @@ steps:
116
commands:
117
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
118
119
+- label: ":nvidia: (B200) LM Eval HiSparse Nightly"
120
+ key: lm-eval-hisparse-nightly-4xb200
121
+ timeout_in_minutes: 60
122
+ device: b200-k8s
123
+ num_devices: 4
124
+ optional: true
125
+ source_file_dependencies:
126
+ - .buildkite/test_areas/lm_eval.yaml
127
+ - tests/evals/gsm8k/
128
+ - vllm/v1/hisparse/
129
+ - vllm/distributed/kv_transfer/kv_connector/v1/hisparse/
130
+ - vllm/v1/attention/backends/mla/
131
+ - vllm/model_executor/layers/attention/sparse_mla_attention.py
132
+ - vllm/v1/worker/gpu/
133
+ - csrc/libtorch_stable/hisparse_kernels.cu
134
+ commands:
135
+ - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt
136
+
137
- label: ":nvidia: (B200) LM Eval PCP"
138
key: lm-eval-pcp-4xb200
139
timeout_in_minutes: 360
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml
new
+20
@@ -0,0 +1,20 @@
1
+model_name: "nvidia/GLM-5.2-NVFP4"
2
+accuracy_threshold: 0.90
3
+tolerance: 0.0
4
+num_questions: 1319
5
+num_fewshot: 5
6
+max_tokens: 256
7
+max_concurrency: 32
8
+request_timeout_seconds: 1800
9
+startup_max_wait_seconds: 1200
10
+server_args: >-
11
+ --tensor-parallel-size 4
12
+ --max-model-len 8192
13
+ --max-num-seqs 32
14
+ --max-num-batched-tokens 4096
15
+ --enforce-eager
16
+ --kv-cache-dtype bfloat16
17
+ --attention-config '{"hisparse_config":{}}'
18
+ --kv-transfer-config '{"kv_connector":"HiSparseConnector","kv_role":"kv_both","kv_connector_extra_config":{"host_pool_gib":32}}'
19
+env:
20
+ VLLM_USE_V2_MODEL_RUNNER: "1"
tests/evals/gsm8k/configs/models-hisparse.txt
new
+1
@@ -0,0 +1 @@
1
+GLM-5.2-NVFP4-HiSparse.yaml
vllm/v1/hisparse/binding.py
+9
-6
@@ -205,9 +205,12 @@ def bind_hisparse_kv_caches(
205
"""Bind existing cache storage and block tables; return the bound handles."""
206
assert host_pool.registered is not None
207
tensor_configs = {
208
- name: tensor_config
208
+ name: (
209
+ tensor_config,
210
+ tensor_config.offset + layer_index * tensor_config.layer_stride,
211
+ )
212
for tensor_config in kv_cache_config.kv_cache_tensors
210
- for name in tensor_config.layers
213
+ for layer_index, name in enumerate(tensor_config.layers)
214
}
215
resident_source_index = 0
216
for group_id, group in enumerate(kv_cache_config.kv_cache_groups):
@@ -216,12 +219,12 @@ def bind_hisparse_kv_caches(
219
for cache_name in group.layer_names:
220
assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX)
221
layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)]
219
- tensor_config = tensor_configs[cache_name]
222
+ tensor_config, byte_offset = tensor_configs[cache_name]
223
assert not tensor_config.host_resident
224
cache_handle = _get_hisparse_cache(forward_context, layer_name)
225
cache_handle.bind_cache(
226
kv_caches[cache_name],
224
- byte_offset=tensor_config.offset,
227
+ byte_offset=byte_offset,
228
block_stride=tensor_config.block_stride,
229
num_blocks=kv_cache_config.num_blocks,
230
block_size=group.kv_cache_spec.block_size,
@@ -249,11 +252,11 @@ def bind_hisparse_kv_caches(
252
raise RuntimeError("HiSparse hot tensors must share one GPU backing.")
253
layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)]
254
cache_handle = _get_hisparse_cache(forward_context, layer_name)
252
- tensor_config = tensor_configs[cache_name]
255
+ tensor_config, byte_offset = tensor_configs[cache_name]
256
assert not tensor_config.host_resident
257
cache_handle.runtime.bind_hot_cache(
258
raw_tensor,
256
- byte_offset=tensor_config.offset,
259
+ byte_offset=byte_offset,
260
block_stride=tensor_config.block_stride,
261
num_blocks=kv_cache_config.num_blocks,
262
block_size=group.kv_cache_spec.block_size,