Decoder XQA Jit Compiler

struct XQAJitKey

Key that identifies one NVRTC-compiled XQA decode kernel variant.

Public Functions

inline auto asTuple() const noexcept

inline bool operator==(XQAJitKey const &other) const noexcept

Public Members

int32_t sm = {}

nvinfer1::DataType dataType = {nvinfer1::DataType::kHALF}

nvinfer1::DataType kvDataType = {nvinfer1::DataType::kHALF}

int32_t headSize = {}

int32_t qHeadsPerKv = {}

int32_t tokensPerPage = {}

bool slidingWindow = {}

bool specDecode = {}

struct XQAJitResult

Result of compiling an XQA kernel with NVRTC.

Public Members

std::vector<uint8_t> cubin

struct XQAJitKernel

One NVRTC-compiled kernel together with the key it was compiled for.

Public Members

XQAJitKey key

std::vector<uint8_t> cubin

std::vector<uint8_t> trt_edgellm::serializeXQAJitKernels(

std::vector<XQAJitKernel> const &kernels

)

Serialize compiled XQA kernels into a self-describing byte blob.

The key is serialized field by field alongside its cubin so that a deserialized kernel is registered under the key it was actually compiled for, rather than one recomputed from whatever plugin fields happen to be present at deserialization time.

std::vector<XQAJitKernel> trt_edgellm::deserializeXQAJitKernels(

void const *data,

size_t size

)

Inverse of serializeXQAJitKernels.

Throws:

std::runtime_error – if the blob is truncated or carries an unknown format version.

bool trt_edgellm::canCompileXQAKernel(

int32_t numQHeads,

int32_t numKVHeads,

int32_t headSize,

int32_t smVersion,

nvinfer1::DataType dataType,

nvinfer1::DataType kvDataType

) noexcept

Check if a given configuration can produce a valid XQA JIT kernel.

XQAJitResult trt_edgellm::compileXQAKernel(XQAJitKey const &key)

Compile one XQA kernel variant with NVRTC.

Throws:

std::runtime_error – if NVRTC support is disabled or compilation fails.