Mila::Dnn::Compute::IPackedKvInference Struct Reference

KV-cache inference interface for packed-QKV MHA backends. More...

Inheritance diagram for Mila::Dnn::Compute::IPackedKvInference:

Public Member Functions
virtual void	decode (const ITensor &input, ITensor &output, dim_t position)=0
	Process a single autoregressive token against the KV cache.
virtual void	prefill (const ITensor &qkv, ITensor &output)=0
	Populate the KV cache from a packed QKV sequence and compute output.
Public Member Functions inherited from Mila::Dnn::Compute::IKvCacheLifecycle
virtual void	initializeKvCache (dim_t batch_size, dim_t max_sequence_length)=0
	Allocate the KV cache for a given batch size and maximum sequence length.
virtual void	resetKvCache ()=0
	Reset the KV cache to an empty state, preserving the allocation.
virtual bool	rewindKvCache (dim_t position)=0
	Rewind the logical cache fill position without touching device K/V buffer contents, so positions [0, position) can be reused by a subsequent partial prefill (PromptCaching.md).

Detailed Description

KV-cache inference interface for packed-QKV MHA backends.

Implemented by GPT-style MHA backends (e.g. CudaMultiHeadAttentionOp). Uses fused QKV input throughout – Q, K, and V are concatenated along the feature axis and split internally by the backend kernel.

Position is implicit: GPT-style MHA always begins prefill at position 0. Absolute positional encoding is handled upstream by Lpe, not inside attention.

Two-phase inference protocol: prefill – populate the KV cache from the full prompt sequence. decode – process one autoregressive token against the accumulated cache.

Member Function Documentation

decode()

virtual void Mila::Dnn::Compute::IPackedKvInference::decode ( const ITensor & input, ITensor & output, dim_t position )

pure virtual

Process a single autoregressive token against the KV cache.

Parameters

input	Packed QKV single-token input [B, 1, 3 * embedding_dim].
output	Pre-allocated output [B, 1, embedding_dim].
position	Zero-based absolute sequence position into the KV cache.

prefill()

virtual void Mila::Dnn::Compute::IPackedKvInference::prefill ( const ITensor & qkv, ITensor & output )

pure virtual

Populate the KV cache from a packed QKV sequence and compute output.

Parameters

qkv	Packed QKV input [B, T, 3 * embedding_dim].
output	Pre-allocated attention output [B, T, embedding_dim].

---

The documentation for this struct was generated from the following file:

• Mila/Src/Dnn/Compute/Operations/IPackedKvInference.ixx