Implementation:InternLM Lmdeploy Tensor
| Knowledge Sources | |
|---|---|
| Domains | Tensor_Operations, Core_Infrastructure |
| Last Updated | 2026-02-07 15:00 GMT |
Overview
Provides the primary multi-dimensional tensor abstraction combining a Buffer (flat storage) with a Layout (shape and strides), plus a string-keyed TensorMap container for named tensor collections.
Description
The Tensor class pairs a Buffer (1D typed memory with shared ownership) with a Layout (multi-dimensional shape and stride). Construction accepts various combinations: layout + dtype + device/allocator, buffer + layout, raw pointer + layout + dtype + device, or a buffer alone (1D). Data access is provided via typed data<T>(), untyped raw_data(), and null-safe data_or(other). Shape manipulation includes view(shape) (reshape), slice(base, shape) (sub-tensor with offset), squeeze(dim), transpose(a, b), and t() (2D transpose). Contiguity is checked via is_contiguous().
The typed subclass Tensor_<T> provides compile-time type safety with dtype enforcement.
TensorMap extends std::unordered_map<std::string, Tensor> with convenience methods: at(key) with descriptive error messages listing available keys, try_(key) returning a nullable pointer, contains(key), produce(key, value) for insertion with uniqueness check, consume(key) for move-out with removal, and try_consume(key). Both Tensor and TensorMap support binary serialization.
Free functions Copy(src, dst) and Clear(tensor) perform async CUDA memory operations, requiring contiguous layouts.
Usage
Used as the primary data container throughout TurboMind for model weights, activations, KV-cache entries, and I/O. TensorMap is used for passing named tensor collections between model layers and for weight loading.
Code Reference
Source Location
- Repository: InternLM_Lmdeploy
- File (header): src/turbomind/core/tensor.h
- File (impl): src/turbomind/core/tensor.cc
- Lines: tensor.h 1-415, tensor.cc 1-174
Signature
namespace turbomind::core {
class Tensor {
public:
Tensor() = default;
Tensor(Layout layout, DataType dtype, Device device);
Tensor(Layout layout, DataType dtype, Allocator& alloc);
Tensor(Buffer buffer, Layout layout);
Tensor(Buffer buffer);
Tensor(void* data, Layout layout, DataType dtype, Device device);
template<class T> Tensor(T* data, Layout layout, Device device);
Buffer& buffer() noexcept;
const Buffer& buffer() const noexcept;
DataType dtype() const;
Device device() const;
ssize_t size() const noexcept;
ssize_t byte_size() const noexcept;
explicit operator bool() const noexcept;
template<class T> T* data();
void* raw_data();
Tensor view(std::vector<ssize_t> shape) const;
Tensor slice(std::vector<ssize_t> base, std::vector<ssize_t> shape) const;
Tensor slice(ssize_t base, ssize_t size = 1) const;
Tensor borrow() const;
Tensor squeeze(int dim) const;
Tensor transpose(int a, int b) const;
Tensor t() const;
auto& layout() const noexcept;
auto& shape() const noexcept;
auto& stride() const noexcept;
int ndim() const noexcept;
bool is_contiguous() const noexcept;
};
template<class T> struct Tensor_ : public Tensor { /* typed wrapper */ };
class TensorMap : public std::unordered_map<std::string, Tensor> {
public:
Tensor& at(const std::string& key);
const Tensor& at(const std::string& key) const;
Tensor* try_(const std::string& key);
bool contains(const std::string& key) const;
void produce(const std::string& key, Tensor value);
Tensor try_consume(const std::string& key);
Tensor consume(const std::string& key);
};
void Copy(const Tensor& src, Ref<Tensor> dst_, const Stream& stream);
void Copy(const Tensor& src, Ref<Tensor> dst_);
void Clear(Ref<Tensor> a_, const Stream& stream);
void Clear(Ref<Tensor> a_);
Tensor empty_like(const Tensor& tensor);
Tensor empty_like(const Tensor& tensor, Device device);
Tensor empty_like(const Tensor& tensor, DataType dtype);
} // namespace turbomind::core
Import
#include "src/turbomind/core/tensor.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| layout | Layout | Yes | Shape and stride descriptor for the tensor |
| dtype | DataType | Yes | Element data type |
| device | Device | Conditional | Device placement (CPU, pinned, CUDA) |
| alloc | Allocator& | Conditional | Memory allocator (alternative to device) |
| buffer | Buffer | Conditional | Pre-existing buffer for wrapping |
| shape | vector<ssize_t> | view/slice | Target shape for reshape or sub-tensor extents |
Outputs
| Name | Type | Description |
|---|---|---|
| data() | T* or void* | Pointer to tensor data |
| size() | ssize_t | Total element count |
| byte_size() | ssize_t | Total bytes |
| layout() | const Layout& | Shape and stride descriptor |
| view() | Tensor | Reshaped tensor sharing the same buffer |
| slice() | Tensor | Sub-tensor with shared buffer |
| t() | Tensor | Transposed 2D tensor |
Usage Examples
#include "src/turbomind/core/tensor.h"
using namespace turbomind::core;
// Create a 2D tensor [batch, hidden] on GPU
Tensor hidden({32, 4096}, kFloat16, kDEVICE);
// Access data
half_t* ptr = hidden.data<half_t>();
// Reshape to 3D
Tensor reshaped = hidden.view({32, 32, 128});
// Slice first 8 rows
Tensor batch_slice = hidden.slice(0, 8);
// Use TensorMap for named tensors
TensorMap inputs;
inputs.produce("hidden_states", hidden);
inputs.produce("attention_mask", mask_tensor);
Tensor& h = inputs.at("hidden_states");