Cute Dsl FMHA Runner#

class CuteDslFMHARunner#

Unified runner for CuTe DSL compiled FMHA kernels (Blackwell SM100+).

Supports two execution modes via separate 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#
void 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,
bool enableSkipSoftmax = false,
)#

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

  • enableSkipSoftmax – Dispatch the skip-softmax (BLASST) kernel variant, which skips the P*V GEMM of KV tiles whose contribution is negligible (threshold baked at export). Approximate — outputs may deviate from dense by up to the calibrated accuracy gate. FP16 causal only: incompatible with fp8Input and slidingWindowSize.

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

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.

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

void 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)#
static bool loadLLMKernelModule()#
static void unloadLLMKernelModule()#
static bool loadViTKernelModule()#
static void unloadViTKernelModule()#