Cute Dsl Fmhav2 Runner#

class CuteDslFMHAV2Runner#

Runner for the CuTe DSL FMHA-v2 kernels.

Dense LLM kernels consume separate BSND Q/K/V tensors, while native-paged LLM kernels consume BSND Q/O and the Edge-LLM NHD paged KV pool. The special padding and vision-block kernels consume dense, separate Q/K/V tensors, while ViT kernels consume packed, separate Q/K/V tensors. The runner remains separate from CuteDslFMHARunner because the optimized SM100/101/110 family has a distinct ABI.

Public Functions

CuteDslFMHAV2Runner(
int32_t numQHeads,
int32_t numKVHeads,
int32_t headDim,
int32_t batchSize = 0,
int32_t seqLenQ = 0,
int32_t kvSeqLen = 0,
bool useSmallD64 = true
)#
~CuteDslFMHAV2Runner() = default#
CuteDslFMHAV2Runner(CuteDslFMHAV2Runner const&) = delete#
CuteDslFMHAV2Runner &operator=(CuteDslFMHAV2Runner const&) = delete#
bool preflightLlm(
cudaStream_t stream,
int32_t slidingWindowSize = INT_MAX
)#

Ensures the exact causal or sliding-causal LLM variant selected by run() is loaded.

bool preflightPaged(
cudaStream_t stream,
int32_t slidingWindowSize = INT_MAX
)#

Ensures the exact native-paged causal or sliding-causal LLM variant selected by runPaged() is loaded.

bool preflightPadding(cudaStream_t stream)#

Ensures the dense non-causal padding variant selected by runPadding() is loaded.

bool preflightVisionBlock(cudaStream_t stream)#

Ensures the vision-block variant selected by runVisionBlock() is loaded.

bool preflightViT(cudaStream_t stream)#

Ensures the exact packed ViT variant selected by run() is loaded.

bool run(
void const *qPtr,
void const *kPtr,
void const *vPtr,
void *oPtr,
int32_t const *cuKVSeqLens,
cudaStream_t stream,
float attentionScale,
int32_t slidingWindowSize = INT_MAX
)#

Runs causal or sliding-causal LLM context attention over dense FP16 K/V.

bool runPaged(
void const *qPtr,
void const *pagedKVPoolPtr,
int32_t const *kvCachePageList,
void *oPtr,
int32_t const *cuQSeqLens,
int32_t const *cuKVSeqLens,
int32_t numFlatPages,
int32_t maxPagesPerSeq,
int32_t tokensPerPage,
cudaStream_t stream,
float attentionScale,
int32_t slidingWindowSize = INT_MAX
)#

Runs FP16 causal or sliding-causal attention directly against an FP16 NHD paged KV pool.

bool runPadding(
void const *qPtr,
void const *kPtr,
void const *vPtr,
void *oPtr,
int32_t const *cuQSeqLens,
int32_t const *cuKVSeqLens,
cudaStream_t stream,
float attentionScale
)#

Runs dense non-causal padded context attention with independent logical Q/KV lengths.

bool runVisionBlock(
void const *qPtr,
void const *kPtr,
void const *vPtr,
void *oPtr,
int32_t const *cuKVSeqLens,
int32_t const *blockBegin,
int32_t const *blockEnd,
cudaStream_t stream,
float attentionScale,
int32_t slidingWindowSize
)#

Runs Gemma4 vision-block attention: sliding-causal OR same-image-block.

bool run(
void const *qPtr,
void const *kPtr,
void const *vPtr,
void *oPtr,
int32_t const *cuSeqLens,
int32_t totalSeqLen,
int32_t maxSeqLen,
int32_t batchSize,
cudaStream_t stream,
float attentionScale
)#

Runs packed varlen, bidirectional ViT attention.

Public Static Functions

static bool canImplement(
int32_t numQHeads,
int32_t numKVHeads,
int32_t headSize,
int32_t smVersion,
nvinfer1::DataType dataType,
CuteDslFMHAV2MaskType maskType
)#

Returns whether the target AOT family covers this dense context shape.

static bool canImplementPaged(
int32_t numQHeads,
int32_t numKVHeads,
int32_t headSize,
int32_t smVersion,
nvinfer1::DataType dataType,
CuteDslFMHAV2MaskType maskType
)#

Returns whether the target AOT family covers native FP16 Edge-LLM paged KV for this context shape.

static bool canImplementViT(
int32_t headSize,
int32_t smVersion,
nvinfer1::DataType dataType
)#

Returns whether the target AOT family covers this packed ViT shape.