Util Kernels#

void trt_edgellm::kernel::calCuQCuKVSeqLens( rt::Tensor const &inputSeqLen, rt::Tensor const &kvCacheStartIndices, rt::Tensor &cuQSeqLens, rt::Tensor &cuKVSeqLens, cudaStream_t stream )#

Build backend-neutral sequence metadata for ragged paged context attention.

cuQSeqLens is the exclusive prefix sum of the active input lengths. cuKVSeqLens is the exclusive prefix sum of kvCacheStartIndices[b] + inputSeqLen[b], so it includes the complete logical KV history for chunked prefill. The outputs can be passed unchanged to either the optimized Blackwell or FMHA-v2 ragged paged backend.

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::calSWAChunkedPrefillMetadata( rt::Tensor const &inputSeqLen, rt::Tensor const &kvCacheStartIndices, rt::Tensor &cuQSeqLens, rt::Tensor &cuKVSeqLens, rt::Tensor &kvCacheEndIdxs, rt::Tensor &paddedCuKVSeqLens, int32_t runtimeSeqLen, int32_t slidingWindowSize, cudaStream_t stream )#

Compute sequence metadata for paged SWA chunked prefill.

The temporary KV source contains the previous resident window followed by the current chunk. The KV prefix sums therefore use min(kvCacheStartIndices[b], slidingWindowSize) + inputSeqLen[b]. kvCacheEndIdxs uses the padded runtime sequence length so the generic RoPE/write kernel assigns position start + tokenOffset to every row. paddedCuKVSeqLens uses the same resident prefix plus runtimeSeqLen, which preserves the causal offset when ragged chunks are padded to the maximum query length.

void trt_edgellm::kernel::assemblePagedSWAChunkedPrefillFMHAKV( rt::Tensor const &swaPool, rt::Tensor const &swaPageTable, rt::Tensor const &k, rt::Tensor const &v, rt::Tensor const &inputSeqLen, rt::Tensor const &kvCacheStartIndices, rt::Tensor &kWorkspace, rt::Tensor &vWorkspace, int32_t slidingWindowSize, cudaStream_t stream )#

Assemble split FP16 K/V for SWA chunked prefill from a paged resident window plus the current chunk.

For each batch row, logical tokens [start - min(start, W), start) are gathered through swaPageTable; newly roped K/V from the current chunk are appended directly. The outputs have shape [B, W + S, Hkv, D], are padded with zeros, and are consumed with cu-seqlens from calSWAChunkedPrefillMetadata(). The persistent pool remains bounded independently of the maximum sequence length.

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 vision-block overlay prefill kernels.

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::launchScatterActiveRows( void const *source, void *destination, int32_t const *stateIndices, int32_t activeRows, int32_t residentRows, size_t rowBytes, cudaStream_t stream )#: Scatter contiguous active rows into a resident-slot tensor using a device-side row map. Invalid resident rows are ignored.

void trt_edgellm::kernel::launchGatherTokenAlignedRope( float const *source, float *output, int32_t const *positions, int32_t const *queryStartOffsets, int32_t const *queryLengths, int32_t const *stateIndices, int32_t numTokens, int32_t numSequences, int32_t sourceRows, int32_t cacheCapacity, int32_t rotaryDim, cudaStream_t stream )#

Gather token-aligned RoPE rows from a shared cache or a resident-slot cache.

sourceRows == 1 broadcasts the standard RoPE cache. Otherwise sequence i resolves its source row through stateIndices[i]. Only rows in the query intervals are gathered; padding remains zero.