Conversation
There was a problem hiding this comment.
Pull request overview
This pull request introduces the TensorScatter operator (opset 24) to the CPU execution provider. The operator performs scatter operations on tensor caches, commonly used in LLM inference for KV-cache updates. It supports two modes: linear (sequential writes) and circular (with wrap-around), with configurable axis and optional write indices.
Changes:
- Added kernel registration for
TensorScatteroperator in opset 24 - Implemented operator logic supporting linear and circular scatter modes with axis normalization and shape validation
- Added comprehensive unit tests covering 3D/4D tensors, both modes, float16 support, and optional write indices
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| onnxruntime/core/providers/cpu/cpu_execution_provider.cc | Registers the TensorScatter kernel for opset 24 in the CPU execution provider |
| onnxruntime/core/providers/cpu/llm/tensorscatter.h | Defines the TensorScatter class interface with axis and mode attributes |
| onnxruntime/core/providers/cpu/llm/tensorscatter.cc | Implements the scatter operation logic for both linear and circular modes with proper memory handling |
| onnxruntime/test/providers/cpu/llm/tensorscatter_op_test.cc | Provides comprehensive test coverage for various tensor shapes, modes, and data types |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 7 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This pull request introduces a new ONNX operator kernel,
TensorScatter, to the CPU execution provider, along with its implementation and comprehensive unit tests. The new operator supports both linear and circular scatter modes, handles various tensor shapes and types, and allows for optional write indices. The changes are grouped below by theme.Operator kernel registration:
TensorScatteroperator for opset 24 in the CPU execution provider, ensuring it is available for inference and training. [1] [2]Operator implementation:
TensorScatterkernel implementation intensorscatter.cc, supporting both linear and circular modes, handling axis normalization, shape validation, and efficient memory copying for scatter operations.TensorScatterclass interface intensorscatter.h, including constructor and compute method, with attributes for axis and mode.Testing:
TensorScatteroperator intensorscatter_op_test.cc, covering linear and circular modes, explicit axis handling, float16 support, and scenarios with and without write indices.