Skip to content

feat(cambricon): add reshape and cache flash provider - #938

Draft
baominghelly wants to merge 1 commit into
masterfrom
feat/cambricon-reshape-and-cache-flash
Draft

feat(cambricon): add reshape and cache flash provider#938
baominghelly wants to merge 1 commit into
masterfrom
feat/cambricon-reshape-and-cache-flash

Conversation

@baominghelly

Copy link
Copy Markdown
Contributor

Summary

  • Add a Cambricon native ReshapeAndCacheFlash provider in src/native/cambricon/ops/reshape_and_cache_flash/.
  • Scatter token-major key/value rows into paged flash-attention caches using slot_mapping, including negative-slot skips and stride-aware NHD/HND cache layouts.
  • Support FP16, BF16, and FP32 on the caller-provided Cambricon stream.

Motivation

ReshapeAndCacheFlash did not have a Cambricon native provider, so Cambricon inference paths could not populate paged KV caches through the current InfiniOps interface.

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

Targeted Cambricon configure, build, and install (this is not the repository-wide smoke build):

cmake -S /workspace/src/InfiniOps-llama-reshape-cache-flash \
  -B /workspace/build/infiniops-reshape-and-cache-flash-pr \
  -DCMAKE_BUILD_TYPE=Release \
  -DWITH_CAMBRICON=ON \
  -DGENERATE_PYTHON_BINDINGS=ON \
  -DINFINI_OPS_OPS=reshape_and_cache_flash \
  -DINFINI_RT_ROOT=/workspace/install/infinirt-llama-trace \
  -DCMAKE_INSTALL_PREFIX=/workspace/build/infiniops-reshape-and-cache-flash-pr-site/infini
cmake --build /workspace/build/infiniops-reshape-and-cache-flash-pr \
  --target ops --parallel 32
cmake --install /workspace/build/infiniops-reshape-and-cache-flash-pr

[ 50%] Built target infiniops
[100%] Built target ops

Targeted Cambricon test:

PYTHONPATH=/workspace/build/infiniops-reshape-and-cache-flash-pr-site \
  /torch/venv3/pytorch/bin/python3.10 -m pytest -q \
  tests/test_reshape_and_cache_flash.py --devices cambricon

........................                                                 [100%]
24 passed in 0.20s

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 Yes targeted ops build passed 24 targeted tests passed; repository smoke/full suite not run
Moore No N/A - not affected N/A - not affected
Ascend No N/A - not affected N/A - not affected
Full `pytest` output (optional)
........................                                                 [100%]
24 passed in 0.20s

Benchmark / Performance Impact

N/A - this adds a previously unavailable Cambricon provider; no performance benchmark was run.

Notes for Reviewers

  • Each MLU logical task handles one or more (token, head) pairs and writes one contiguous head_size vector into the cache slot selected by slot_mapping.
  • Negative slots are skipped. The caller contract must provide in-range non-negative slots and avoid duplicate concurrent destinations.
  • Tensor strides are used for token/head input addressing and block/page/head cache addressing, covering padded sources plus logical NHD and HND cache layouts.
  • The current public interface only accepts kv_cache_dtype == "auto", so cache dtypes match key/value dtypes and k_scale / v_scale do not participate in the copy.
  • The implementation is correctness-oriented direct device copying; NRAM/vectorized performance optimization is intentionally deferred.
  • Existing cross-platform tests already cover the provider, so this PR adds no test-only changes.
  • The repository smoke build and full suite were not run; the PR remains draft pending broader validation.

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.

1 participant