From 38ca7a899cd3a9aa6f1e948b1b63f4fae7a685ce Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Tue, 15 Sep 2026 16:27:11 -0400 Subject: [PATCH] [Bugfix][HiSparse] Preserve per-layer offsets in KV cache bindings (#57027) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- .buildkite/test_areas/lm_eval.yaml | 18 +++++++++++++++++ .../gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml | 20 +++++++++++++++++++ tests/evals/gsm8k/configs/models-hisparse.txt | 1 + vllm/v1/hisparse/binding.py | 15 ++++++++------ 4 files changed, 48 insertions(+), 6 deletions(-) create mode 100644 tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml create mode 100644 tests/evals/gsm8k/configs/models-hisparse.txt diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index d9afefc82e..ec45f7e69e 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -116,6 +116,24 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt +- label: ":nvidia: (B200) LM Eval HiSparse Nightly" + key: lm-eval-hisparse-nightly-4xb200 + timeout_in_minutes: 60 + device: b200-k8s + num_devices: 4 + optional: true + source_file_dependencies: + - .buildkite/test_areas/lm_eval.yaml + - tests/evals/gsm8k/ + - vllm/v1/hisparse/ + - vllm/distributed/kv_transfer/kv_connector/v1/hisparse/ + - vllm/v1/attention/backends/mla/ + - vllm/model_executor/layers/attention/sparse_mla_attention.py + - vllm/v1/worker/gpu/ + - csrc/libtorch_stable/hisparse_kernels.cu + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hisparse.txt + - label: ":nvidia: (B200) LM Eval PCP" key: lm-eval-pcp-4xb200 timeout_in_minutes: 360 diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml new file mode 100644 index 0000000000..0098719259 --- /dev/null +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-HiSparse.yaml @@ -0,0 +1,20 @@ +model_name: "nvidia/GLM-5.2-NVFP4" +accuracy_threshold: 0.90 +tolerance: 0.0 +num_questions: 1319 +num_fewshot: 5 +max_tokens: 256 +max_concurrency: 32 +request_timeout_seconds: 1800 +startup_max_wait_seconds: 1200 +server_args: >- + --tensor-parallel-size 4 + --max-model-len 8192 + --max-num-seqs 32 + --max-num-batched-tokens 4096 + --enforce-eager + --kv-cache-dtype bfloat16 + --attention-config '{"hisparse_config":{}}' + --kv-transfer-config '{"kv_connector":"HiSparseConnector","kv_role":"kv_both","kv_connector_extra_config":{"host_pool_gib":32}}' +env: + VLLM_USE_V2_MODEL_RUNNER: "1" diff --git a/tests/evals/gsm8k/configs/models-hisparse.txt b/tests/evals/gsm8k/configs/models-hisparse.txt new file mode 100644 index 0000000000..53ff181f4f --- /dev/null +++ b/tests/evals/gsm8k/configs/models-hisparse.txt @@ -0,0 +1 @@ +GLM-5.2-NVFP4-HiSparse.yaml diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index c23abc9cb4..716211fbb3 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -205,9 +205,12 @@ def bind_hisparse_kv_caches( """Bind existing cache storage and block tables; return the bound handles.""" assert host_pool.registered is not None tensor_configs = { - name: tensor_config + name: ( + tensor_config, + tensor_config.offset + layer_index * tensor_config.layer_stride, + ) for tensor_config in kv_cache_config.kv_cache_tensors - for name in tensor_config.layers + for layer_index, name in enumerate(tensor_config.layers) } resident_source_index = 0 for group_id, group in enumerate(kv_cache_config.kv_cache_groups): @@ -216,12 +219,12 @@ def bind_hisparse_kv_caches( for cache_name in group.layer_names: assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX) layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)] - tensor_config = tensor_configs[cache_name] + tensor_config, byte_offset = tensor_configs[cache_name] assert not tensor_config.host_resident cache_handle = _get_hisparse_cache(forward_context, layer_name) cache_handle.bind_cache( kv_caches[cache_name], - byte_offset=tensor_config.offset, + byte_offset=byte_offset, block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size, @@ -249,11 +252,11 @@ def bind_hisparse_kv_caches( raise RuntimeError("HiSparse hot tensors must share one GPU backing.") layer_name = cache_name[: -len(HISPARSE_HOT_SUFFIX)] cache_handle = _get_hisparse_cache(forward_context, layer_name) - tensor_config = tensor_configs[cache_name] + tensor_config, byte_offset = tensor_configs[cache_name] assert not tensor_config.host_resident cache_handle.runtime.bind_hot_cache( raw_tensor, - byte_offset=tensor_config.offset, + byte_offset=byte_offset, block_stride=tensor_config.block_stride, num_blocks=kv_cache_config.num_blocks, block_size=group.kv_cache_spec.block_size,