Skip to content

fix(hicache): handle tensor device buffers - #42785

Closed
AnkitNakhawa wants to merge 1 commit into
sgl-project:mainfrom
AnkitNakhawa:fix/npu-hicache-tensor-buffer
Closed

AnkitNakhawa wants to merge 1 commit into
sgl-project:mainfrom
AnkitNakhawa:fix/npu-hicache-tensor-buffer

Conversation

@AnkitNakhawa

@AnkitNakhawa AnkitNakhawa commented Oct 6, 2026 •

Copy link
Copy Markdown

Motivation

Fixes #42672.

NPU HiCache without FIA exposes k_buffer and v_buffer as multi-layer tensors. MHATokenToKVPoolHost used getattr(..., None) or () before iterating them, which evaluates the tensor as a boolean and raises RuntimeError: Boolean value of Tensor with more than one value is ambiguous during host-pool construction.

Modifications

  • Add a small iterator that handles missing buffers explicitly without boolean evaluation.
  • Preserve existing list/tuple behavior and let multi-layer tensors yield zero-copy per-layer views.
  • Add a CPU regression test that constructs an MHA host pool from stacked multi-layer tensor buffers.

Accuracy Tests

Not applicable. This changes buffer discovery during HiCache initialization and does not change model computation or outputs.

Speed Tests and Profiling

Not applicable. The change runs only during host-pool construction.

Testing

  • pre-commit run --files python/sglang/srt/mem_cache/pool_host/mha.py test/registered/unit/mem_cache/test_mem_pool_host.py
  • python -m py_compile python/sglang/srt/mem_cache/pool_host/mha.py test/registered/unit/mem_cache/test_mem_pool_host.py
  • Direct CPU assertions with PyTorch 2.14 for tensor, list, missing, and None buffer representations

The registered test was not collected locally because the minimal macOS environment does not contain the full SGLang development dependency set. Ascend NPU end-to-end validation was not run locally and should be covered by NPU CI.

Checklist

  • Format code with pre-commit.
  • Add focused regression coverage.
  • Documentation update is not required for this bug fix.
  • Accuracy and speed benchmarks are not applicable.
  • Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ❌ Run #38449365184
Latest PR Test (Extra): ❌ Run #38449365223
Latest PR Test (AMD ROCm 10): ❌ Run #38449365306

@OrangeRedeng

Copy link
Copy Markdown
Collaborator

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Oct 9, 2026
@ping1jing2

Copy link
Copy Markdown
Collaborator

/rerun-failed-ci

@AnkitNakhawa
AnkitNakhawa force-pushed the fix/npu-hicache-tensor-buffer branch from c83d581 to f2f028a Compare October 13, 2026 18:45
@vlserov

vlserov commented Oct 14, 2026

Copy link
Copy Markdown
Contributor

Original issue has been fixed with PR #43711. Can be withdrawn

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

hicache Hierarchical Caching for SGLang run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] NPU HiCache: MHATokenToKVPoolHost crashes when the device k_buffer is a single tensor (since #40326)

4 participants