Util Kernels#

void trt_edgellm::kernel::calCuQCuKVSeqLensAndKVEndIdxs(
rt::Tensor const &inputSeqLen,
rt::Tensor const &kvCacheStartIndices,
rt::Tensor &cuQSeqLens,
rt::Tensor &cuKVSeqLens,
rt::Tensor &kvCacheEndIdxs,
rt::OptionalOutputTensor paddedCuKVSeqLens,
int32_t const runtimeSeqLen,
cudaStream_t stream,
)#

Host-side wrapper that launches a lightweight CUDA kernel to compute prefix-sum of sequence lengths and KV cache end indices.

Note

kvCacheStartIndices is optional. If it is not provided, kvStartIndices will be assumed to be 0.

Parameters:

  • inputSeqLen – [in] int32_t tensor with shape [B]. Actual token length of each request.

  • kvCacheStartIndices – [in] int32_t tensor with shape [B]. Start index of KV cache for each request. (optional, pass in empty tensor to indicate zero start indices)

  • cuQSeqLens – [out] int32_t tensor with shape [B+1]. Exclusive prefix-sum of inputSeqLen.

  • cuKVSeqLens – [out] int32_t tensor with shape [B+1]. Exclusive prefix-sum of (kvCacheStartIndices[i] + inputSeqLen[i]). If kvCacheStartIndices is empty, this will be exclusive prefix-sum of inputSeqLen.

  • kvCacheEndIdxs – [out] int32_t tensor with shape [B]. Each element equals kvCacheStartIndices[i] + runtimeSeqLen (Here we use padding to ease later kernel launch).

  • paddedCuKVSeqLens – [out] (optional) int32_t tensor with shape [B+1]. Exclusive prefix-sum of kvCacheEndIdxs (= kvCacheStartIdx + runtimeSeqLen per batch). Pass std::nullopt to skip. Background: CuTe DSL FMHA kernel uses bottom_right_align with offset = s_k - s_q. Q is padded to runtimeSeqLen for all batches, so we must use padded KV lengths (s_k = kvCacheEndIdx per batch) to keep offset non-negative. Using actual s_k (< runtimeSeqLen for shorter batches) would produce a negative offset that masks out valid KV positions, breaking attention.

  • runtimeSeqLen – [in] Runtime sequence length (equals to the maximum of inputSeqLen).

  • stream – [in] CUDA stream used to launch the kernel.

Throws:

std::runtime_error – if tensor shapes are invalid

void trt_edgellm::kernel::launchBuildVisionPackedMask(
int32_t const *visionBlockIds,
int32_t const *contextLengths,
uint32_t *packedMask,
int32_t *cuMaskRows,
int32_t batchSize,
int32_t seqLen,
int32_t slidingWindowSize,
cudaStream_t stream,
)#

Build the FMHA_v2 CUSTOM_MASK packed mask for Gemma4 vision-block prefill.

Allowed(q, k) for q, k < contextLengths[b]: sliding-causal: k <= q and (slidingWindowSize <= 0 or k > q - slidingWindowSize), OR vision block: visionBlockIds[b][q] >= 0 and visionBlockIds[b][q] == visionBlockIds[b][k] All other positions (including rows/cols >= contextLength and padded rows/cols) have their bits cleared.

Parameters:

  • visionBlockIds – [in] int32_t [B, S]; -1 for text/audio/pad, non-negative per contiguous image run.

  • contextLengths – [in] int32_t [B]; actual token count per sequence.

  • packedMask – [out] uint32_t [getPackedMaskSizeInWords(B, S, S)] packed mask.

  • cuMaskRows – [out] int32_t [B+1]; prefix sum of padded mask rows (i * getPackedMaskRowsPerSeq(S)). Consumed by the kernels via params.cu_mask_rows.

  • batchSize – [in] Number of sequences B.

  • seqLen – [in] Padded runtime sequence length S (Q length == KV length).

  • slidingWindowSize – [in] Sliding window size counting the query itself (<= 0 means plain causal).

  • stream – [in] CUDA stream.

Throws:

std::runtime_error – on invalid arguments.

void trt_edgellm::kernel::launchBuildVisionBlockRanges(
int32_t const *visionBlockIds,
int32_t const *contextLengths,
int32_t *blockBegin,
int32_t *blockEnd,
int32_t batchSize,
int32_t seqLen,
cudaStream_t stream,
)#

Expand [B, S] vision-block IDs into per-position [blockBegin, blockEnd] interval tensors for the FFPA vision-block overlay prefill kernel.

Each contiguous run of an identical non-negative ID inside the per-batch valid prefix (contextLengths[b], clamped to seqLen) yields blockBegin = run start and blockEnd = run end for every position in the run. Text/audio positions (ID < 0) and padding positions receive the -1/-1 sentinel (empty interval). All tensors are [B, S] int32 device buffers.

void trt_edgellm::kernel::cvtKVLayoutBHSDToSplitKV(
rt::Tensor const &src,
rt::Tensor &kDst,
rt::Tensor &vDst,
rt::Tensor const &kvScaleQuantOrig,
int32_t seqLen,
cudaStream_t stream,
)#

Converts KV cache layout from [B, 2, H, S, D] into separate K and V tensors of shape [B, S, H, D].

Splits the interleaved KV source into two independent FP16 output tensors, applying FP8 dequantization when the source is FP8. Used in the chunked-prefill path so that the SEPARATE_Q_K_V FMHA kernels receive separate K and V pointers.

Parameters:

  • src – [in] Source tensor with shape [B, 2, H, S, D].

  • kDst – [out] Destination K tensor with shape [B, S, H, D] (FP16).

  • vDst – [out] Destination V tensor with shape [B, S, H, D] (FP16).

  • kvScaleQuantOrig – [in] Optional packed dequant scale tensor for FP8 KV cache (shape [2], float). Layout: [kScaleQuantOrig, vScaleQuantOrig]. Pass an empty tensor for FP16 src.

  • stream – [in] CUDA stream to launch the kernel on.

Throws:

std::runtime_error – if tensor shapes or data types are invalid.