Cute Dsl FMHA Runner#

class CuteDslFMHARunner#

Unified runner for CuTe DSL compiled FMHA kernels (Blackwell SM100/101/110).

Supports two execution modes via AOT-compiled kernel variants:

  1. LLM prefill/chunked-prefill: batched Q [B,S_q,H_q,D] + combined KV cache [B,2,H_kv,Cap,D] with causal masking and optional sliding window.

  2. ViT: packed varlen separate Q/K/V [total_S,H,D] with cu_seqlens [B+1] for ragged batching, bidirectional attention (no causal mask).

Each mode has its own kernel modules and run() overload.

Public Functions

CuteDslFMHARunner(
int32_t numQHeads,
int32_t numKVHeads,
int32_t headDim,
int32_t batchSize = 0,
int32_t seqLenQ = 0,
int32_t kvCacheCapacity = 0
)#
~CuteDslFMHARunner() = default#
CuteDslFMHARunner(CuteDslFMHARunner const&) = delete#
CuteDslFMHARunner &operator=(CuteDslFMHARunner const&) = delete#
bool preflightLlm(
cudaStream_t stream,
int32_t slidingWindowSize = INT_MAX,
bool fp8Input = false,
float skipSoftmaxThresholdLog2 = 0.0F
)#

Ensures the exact dense LLM variant selected by run() is loaded.

bool preflightPaged(
cudaStream_t stream,
int32_t slidingWindowSize = INT_MAX,
bool fp8Input = false,
bool isCausal = true,
float skipSoftmaxThresholdLog2 = 0.0F,
bool useBidirectional = false
)#

Ensures the exact paged LLM variant selected by runPaged() 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 *kvPtr,
void *oPtr,
int32_t const *cuKVSeqLens,
cudaStream_t stream,
float attentionScale,
int32_t slidingWindowSize = INT_MAX,
bool fp8Input = false,
float qScale = 1.0F,
float kScale = 1.0F,
float vScale = 1.0F,
float skipSoftmaxThresholdLog2 = 0.0F
)#

LLM FMHA: batched Q + combined KV cache with causal masking.

Output is always FP16. Selects kernel variant based on fp8Input:

  • fp8Input=false → FP16 kernels (all scales ignored)

  • fp8Input=true → FP8-input / FP16-output kernels

Parameters:

  • qPtr – Query [B, S_q, H_q, D]

  • kvPtr – Combined KV cache [B, 2, H_kv, Cap, D]

  • oPtr – Output [B, S_q, H_q, D] (always FP16)

  • cuKVSeqLens – Cumulative KV sequence lengths [B+1]

  • stream – CUDA stream

  • attentionScale – Model-defined multiplier applied to QK^T before softmax. For FP8 input, the effective softmax scale is attentionScale * qScale * kScale.

  • slidingWindowSize – Sliding window size (INT_MAX = disabled)

  • fp8Input – Whether Q/KV are FP8 E4M3

  • qScale – Q dequant scale (quant→orig), ignored when fp8Input=false

  • kScale – K dequant scale (quant→orig), ignored when fp8Input=false

  • vScale – V dequant scale (quant→orig), applied to the attention output and ignored when fp8Input=false

  • skipSoftmaxThresholdLog2 – Skip-softmax (BLASST) threshold as log2(lambda). A finite negative value (lambda in (0,1)) dispatches the skip-softmax kernel variant, which skips the P*V GEMM of KV tiles whose contribution is negligible — approximate, FP16 causal only. 0.0 (the default, log2 of the degenerate lambda = 1) disables skip, mirroring the slidingWindowSize = INT_MAX sentinel convention.

bool runPaged(
void const *qPtr,
void const *pagedKVPoolPtr,
int32_t const *kvCachePageList,
void *oPtr,
int32_t const *cuKVSeqLens,
int32_t numPages,
int32_t maxPagesPerSeq,
int32_t tokensPerPage,
nvinfer1::DataType kvDataType,
cudaStream_t stream,
float attentionScale,
int32_t slidingWindowSize = INT_MAX,
bool fp8Input = false,
float qScale = 1.0F,
float kScale = 1.0F,
float vScale = 1.0F,
bool isCausal = true,
float skipSoftmaxThresholdLog2 = 0.0F,
int32_t const *bidirectionalBlockBegin = nullptr,
int32_t const *bidirectionalBlockEnd = nullptr
)#

LLM FMHA over a paged KV cache.

Dispatches a dedicated CuTe DSL AOT variant that reads K/V directly from a paged pool using kvCachePageList. The logical descriptor shape is [numPages, H_kv, tokensPerPage, D], while the physical pool is fixed to NHD [numPages, tokensPerPage, H_kv, D]. This path maps one logical K/V TMA tile to one physical page. The current CuTe DSL variants use a K/V tile width of 128 tokens, so tokensPerPage must be 128 to avoid multi-page tile stitching or a gather workspace. Setting both block range pointers selects the FP16 D512 bidirectional-mask variant; mixed null pointers are invalid.

Parameters:

  • qPtr – Query [B, S_q, H_q, D]

  • pagedKVPoolPtr – Paged KV pool [numPages, tokensPerPage, H_kv, D]

  • kvCachePageList – Page table [B, 2, maxPagesPerSeq], K pages then V pages

  • oPtr – Output [B, S_q, H_q, D] (always FP16)

  • cuKVSeqLens – Cumulative KV sequence lengths [B+1]

  • numPages – Number of pages in the paged KV pool

  • maxPagesPerSeq – Max logical pages per sequence

  • tokensPerPage – Number of tokens per page

  • kvDataType – Paged KV cache dtype (FP16 or FP8)

  • stream – CUDA stream

  • slidingWindowSize – Sliding window size (INT_MAX = disabled)

  • fp8Input – Whether Q/KV are FP8 E4M3

  • qScale – Q dequant scale, ignored when fp8Input=false

  • kScale – K dequant scale, ignored when fp8Input=false

  • vScale – V dequant scale, ignored when fp8Input=false

  • isCausal – Whether to dispatch a causal or dense non-causal variant

  • skipSoftmaxThresholdLog2 – Skip-softmax threshold as log2(lambda), or 0 to disable

  • bidirectionalBlockBegin – Optional inclusive bidirectional-block begin positions [B, S_q]. Text/padding rows use -1; every row in a disjoint contiguous vision run must repeat that run’s begin position.

  • bidirectionalBlockEnd – Optional inclusive bidirectional-block end positions [B, S_q]. Text/padding rows use -1; every row in a disjoint contiguous vision run must repeat that run’s end position.

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
)#

ViT FMHA: packed varlen separate Q/K/V, bidirectional.

Parameters:

  • qPtr – Query [total_S, H, D]

  • kPtr – Key [total_S, H, D]

  • vPtr – Value [total_S, H, D]

  • oPtr – Output [total_S, H, D]

  • cuSeqLens – Cumulative sequence lengths [B+1]

  • totalSeqLen – Sum of all sequence lengths

  • maxSeqLen – Longest individual sequence length

  • batchSize – Number of sequences

  • stream – CUDA stream

  • attentionScale – Absolute multiplier applied to QK^T before softmax

Public Static Functions

static bool canImplement(int32_t headSize, int32_t smVersion)#
static bool canImplementViT(int32_t headSize, int32_t smVersion)#