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 preflightPagedRagged(

cudaStream_t stream,

int32_t slidingWindowSize = INT_MAX

)

Ensures the packed-Q/O native-paged causal or sliding-causal variant selected by runPagedRagged() 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 runPagedRagged(

void const *qPtr,

void const *pagedKVPoolPtr,

int32_t const *kvCachePageList,

void *oPtr,

int32_t const *cuQSeqLens,

int32_t const *cuKVSeqLens,

int32_t totalQSeqLen,

int32_t maxQSeqLen,

int32_t numFlatPages,

int32_t maxPagesPerSeq,

int32_t tokensPerPage,

cudaStream_t stream,

float attentionScale,

int32_t slidingWindowSize = INT_MAX

)

Runs causal or sliding-causal attention over packed Q/O and an FP16 NHD paged KV pool.

cuQSeqLens and cuKVSeqLens are monotonic prefix arrays of length batchSize + 1. cuQSeqLens starts at zero and ends at totalQSeqLen. cuKVSeqLens contains the complete logical KV history, including prior cached tokens. maxQSeqLen is the maximum adjacent difference in cuQSeqLens, not the packed buffer allocation capacity.

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 canImplementPagedRagged(

int32_t numQHeads,

int32_t numKVHeads,

int32_t headSize,

int32_t smVersion,

nvinfer1::DataType dataType,

CuteDslFMHAV2MaskType maskType

)

Returns whether the target AOT family covers packed-Q/O native paged attention.

static bool canImplementViT(

int32_t headSize,

int32_t smVersion,

nvinfer1::DataType dataType

)

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