Struct IKVCacheUpdateLayer 

pub struct IKVCacheUpdateLayer { /* private fields */ }

IKVCacheUpdateLayer

Layer that represents a KVCacheUpdate operation.

The KVCacheUpdate layer is used to cache the key or value tensors for the attention mechanism. K and V use separate KVCacheUpdate layers.

An IKVCacheUpdateLayer has three inputs (cache, update, writeIndices) and one output. In kLINEAR mode, for each batch element i, the layer copies the update tensor into the cache starting at position writeIndices[i]. Assuming no out-of-bounds writes occur, the operation for each sequence position s in [0, sequenceLength) is:

output[i, :, writeIndices[i] + s, :] = update[i, :, s, :]

The output performs in-place updates on the cache tensor, so they must share the same device memory address.

Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI.

Implementations

impl IKVCacheUpdateLayer

pub fn setCacheMode( self: Pin<&mut IKVCacheUpdateLayer>, cacheMode: KVCacheMode, ) -> bool

Set the mode of the KVCacheUpdate layer.

• cacheMode The mode of the KVCacheUpdate layer. For TensorRT 10.15, only kLINEAR mode is supported.

True if cache mode is set successfully, false otherwise.

pub fn getCacheMode(self: &IKVCacheUpdateLayer) -> KVCacheMode

Get the mode of the KVCacheUpdate layer.

The mode of the KVCacheUpdate layer.

pub fn setUpdateForm( self: Pin<&mut IKVCacheUpdateLayer>, form: AttentionIOForm, ) -> bool

Set the update form.

Default is kPADDED_BHND. When set to kPACKED_NHD, the update tensor shape is [totalTokens, numHeads, dimHead] instead of [batchSize, numHeads, numTokens, dimHead], and the updateLengths tensor must be provided.

• form The update form.

True if the update form is set successfully, false otherwise.

See [getUpdateForm()] See AttentionIOForm

pub fn getUpdateForm(self: &IKVCacheUpdateLayer) -> AttentionIOForm

Get the update form.

The update form.

See [setUpdateForm()] See AttentionIOForm

pub unsafe fn setUpdateLengths( self: Pin<&mut IKVCacheUpdateLayer>, lengths: *mut ITensor, ) -> bool

Set the update lengths tensor.

Only valid when the update form is kPACKED_NHD. Provides a cumulative token counts tensor with shape [batchSize + 1]. The first element should be 0 and the last element equals totalTokens. The number of tokens for batch i is lengths[i + 1] - lengths[i]. Must be set when update form is kPACKED_NHD.

Providing a first element that is not 0 results in undefined behavior.

• lengths A 1D tensor of type kINT32 with shape [batchSize + 1], or nullptr to clear a previously set update lengths tensor.

True if the update lengths tensor is set successfully, false otherwise.

See [getUpdateLengths()]

pub fn getUpdateLengths(self: &IKVCacheUpdateLayer) -> *mut ITensor

Get the update lengths tensor.

The update lengths tensor, or nullptr if not set.

See [setUpdateLengths()]

Trait Implementations

impl AsLayer for IKVCacheUpdateLayer

fn as_layer(&self) -> &ILayer

fn as_layer_pin_mut(&mut self) -> Pin<&mut ILayer>

impl AsLayerTyped for IKVCacheUpdateLayer

const TYPE: LayerType = LayerType::kKVCACHE_UPDATE

impl AsRef<ILayer> for IKVCacheUpdateLayer

fn as_ref(self: &IKVCacheUpdateLayer) -> &ILayer

Converts this type into a shared reference of the (usually inferred) input type.

impl ExternType for IKVCacheUpdateLayer

type Id = (n, v, i, n, f, e, r, _1, (), I, K, V, C, a, c, h, e, U, p, d, a, t, e, L, a, y, e, r)

A type-level representation of the type’s C++ namespace and type name.

type Kind = Opaque

Either cxx::kind::Opaque or cxx::kind::Trivial.

impl MakeCppStorage for IKVCacheUpdateLayer

unsafe fn allocate_uninitialized_cpp_storage() -> *mut IKVCacheUpdateLayer

Allocates heap space for this type in C++ and return a pointer to that space, but do not initialize that space (i.e. do not yet call a constructor).

unsafe fn free_uninitialized_cpp_storage(arg0: *mut IKVCacheUpdateLayer)

Frees a C++ allocation which has not yet had a constructor called.