Skip to content

feat(ascend): add reshape_and_cache_flash provider - #973

Merged
voltjia merged 1 commit into
masterfrom
feat/ascend-reshape-cache
Sep 4, 2026
Merged

feat(ascend): add reshape_and_cache_flash provider#973
voltjia merged 1 commit into
masterfrom
feat/ascend-reshape-cache

Conversation

@baominghelly

@baominghelly baominghelly commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Add the Ascend reshape_and_cache_flash provider in src/native/ascend/ops/reshape_and_cache_flash/kernel.h.
  • Use ACLNN ScatterPaKvCache for contiguous non-float32 inputs and an ACLNN ScatterNdUpdate fallback for float32 or non-contiguous layouts.
  • Preserve negative slot mappings by skipping the corresponding cache updates.

Motivation

Ascend model inference needs native KV-cache population for the existing reshape_and_cache_flash interface. Without this provider, paged-attention cache updates cannot use the InfiniOps slot-0 path on Ascend.

Type of Change

  • feat — new feature / new operator / new platform
  • fix — bug fix
  • perf — performance improvement (no behavioral change)
  • refactor — code restructuring without behavior change
  • test — adding or fixing tests only
  • docs — documentation only
  • build / ci — build system or CI configuration
  • chore — tooling, formatting, or other non-code changes
  • Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

bash /workspace/ascend-slot0/configure-focused-build.sh \
  /workspace/ascend-slot0/reshape-and-cache-flash \
  /workspace/ascend-slot0/build-reshape-and-cache-flash \
  /workspace/ascend-slot0/ops-reshape-and-cache-flash.json
cmake --build /workspace/ascend-slot0/build-reshape-and-cache-flash --parallel 32
bash /workspace/ascend-slot0/basic-ops-20260901/scripts/run-focused-test.sh \
  /workspace/ascend-slot0/reshape-and-cache-flash \
  /workspace/ascend-slot0/build-reshape-and-cache-flash \
  /workspace/ascend-slot0/python-reshape-and-cache-flash \
  tests/test_reshape_and_cache_flash.py

[5/5] Linking CXX shared module src/ops.cpython-311-aarch64-linux-gnu.so
24 passed, 14 warnings in 1.43s

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA No N/A - not affected N/A - not affected
Iluvatar No N/A - not affected N/A - not affected
MetaX No N/A - not affected N/A - not affected
Cambricon No N/A - not affected N/A - not affected
Moore No N/A - not affected N/A - not affected
Ascend Yes Focused build passed on Ascend 910C with CANN 9.0.0 Complete tests/test_reshape_and_cache_flash.py: 24 passed
Full `pytest` output (optional)
collected 24 items
24 passed, 14 warnings in 1.43s

Benchmark / Performance Impact

N/A. This PR adds functional provider coverage and does not claim a performance improvement.

Notes for Reviewers

  • The fast path requires contiguous key/value/cache tensors and excludes float32.
  • The fallback copies slot mappings to host, skips negative slots, constructs block/offset indices, and applies per-token scatter updates.
  • Validation covered float16, bfloat16, and float32, both NHD/HND layouts, and contiguous/non-contiguous cases in the complete operator test file.
  • The repository-wide test suite was not run.

@baominghelly
baominghelly force-pushed the feat/ascend-reshape-cache branch from c586ab3 to c131103 Compare September 3, 2026 06:37
@baominghelly
baominghelly marked this pull request as ready for review September 3, 2026 07:07
@baominghelly
baominghelly requested review from a team and voltjia September 3, 2026 07:07
@voltjia voltjia changed the title feat(ascend): add reshape_and_cache_flash provider feat(ascend): add reshape_and_cache_flash provider Sep 4, 2026
@voltjia
voltjia merged commit 82ca5d5 into master Sep 4, 2026
12 of 20 checks passed
@voltjia
voltjia deleted the feat/ascend-reshape-cache branch September 4, 2026 02:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants