Decoder XQA Runner#

class DecoderXQARunner#

Decoder XQA (eXtended Query Attention) kernel runner.

Public Functions

DecoderXQARunner(
nvinfer1::DataType const dataType,
nvinfer1::DataType const kvDataType,
int32_t batchSize,
int32_t numQHeads,
int32_t numKvHeads,
int32_t headSize,
int32_t smVersion
) noexcept#

Constructor for DecoderXQARunner.

Parameters:

  • dataType – [in] Data type for computation

  • kvDataType – [in] KV cache data type

  • batchSize – [in] Batch size

  • numQHeads – [in] Number of query heads

  • numKvHeads – [in] Number of key-value heads

  • headSize – [in] Head dimension size

  • smVersion – [in] CUDA SM version

DecoderXQARunner() noexcept = default#
~DecoderXQARunner() noexcept = default#
void dispatchXQAKernel(
XQALaunchParams &params,
cudaStream_t const &stream
)#

Dispatch XQA kernel and compute the attention result.

Parameters:

  • params – [inout] Launch parameters for XQA kernel

  • stream – [in] CUDA stream for kernel execution

Throws:

  • std::runtime_error – if device pointers are invalid or no available kernel available

  • std::runtime_error – if a CUDA driver error occurs

void dispatchSpecDecodeXQAKernel(
XQALaunchParams &params,
cudaStream_t const &stream
)#

Dispatch spec-decode XQA kernel for tree attention.

Parameters:

  • params – [inout] Launch parameters for XQA kernel

  • stream – [in] CUDA stream for kernel execution

Throws:

  • std::runtime_error – if device pointers are invalid or no available kernel available

  • std::runtime_error – if a CUDA driver error occurs

XQALaunchParams initXQAParams() noexcept#

Initialize XQA parameters with MHA and hardware configuration.

The XQA parameter can be used by prepareToRun() to query kernel to dispatch. Device pointer shall be setup by caller to dispatch XQA kernel.

Returns:

Initialized XQA launch parameters

Public Static Functions

static bool loadDecodeXQAKernelFromCubin(
XQAJitKey const &key,
void const *cubinData,
size_t cubinSize
)#

Load one decoder XQA kernel from a JIT-generated module image.

Parameters:

  • key – [in] JIT key that describes the kernel variant

  • cubinData – [in] Cubin or PTX bytes returned by NVRTC

  • cubinSize – [in] Number of bytes in cubinData

Throws:

std::runtime_error – if a CUDA driver error occurs

Returns:

True if the kernel loaded successfully and was inserted into the runtime cache

struct XQALaunchParams#

Launch parameters for XQA (eXtended Query Attention) kernel.

Public Members

void *output = nullptr#

Device memory pointers to launch XQA kernel.

Output tensor

void const *qInputPtr = nullptr#

Query input pointer.

KVCache kvCache#

KV cache structure.

float kScale = 1.0f#

K dequant scale (quantized -> original), host scalar.

float vScale = 1.0f#

V dequant scale (quantized -> original), host scalar.

uint32_t slidingWinSize = 0#

Sliding window size (0 = no sliding window)

int32_t *semaphores = nullptr#

Semaphores for synchronization.

void *scratch = nullptr#

Scratch memory.

void *treeAttnMask = nullptr#

Unique device memory pointer for spec-decode tree attention.

Tree attention mask

int32_t *qCuSeqLen = nullptr#

Cumulative query sequence lengths.

float const *attentionSinks = nullptr#

Attention sinks parameter.

int32_t numQheads = 0#

MHA parameters to locate a kernel to launch.

Number of query heads

int32_t numKVheads = 0#

Number of key-value heads.

int32_t headSize = 0#

Head dimension size.

int32_t batchSize = 0#

Batch size.

int32_t qSeqLen = 0#

Parameters for spec-decode tree attention.

Query sequence length

float attentionScale = 1.0F#

Absolute QK^T multiplier.

int32_t headGroupSize = 0#

Head group size.

nvinfer1::DataType dataType#

I/O data type of the kernel.

nvinfer1::DataType kvDataType#

KV cache data type.