389e493055
* CUTLASS 3.8 Release * update * Update README.md * Revert "Update README.md" This reverts commit b353e36fe83e0815f99b44e46c0c95494c44726b. * update * update --------- Co-authored-by: Haicheng Wu <57973641+hwu36@users.noreply.github.com> Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
24 lines
1.8 KiB
Markdown
24 lines
1.8 KiB
Markdown
# FMHA for Blackwell: Forward
|
|
|
|
This sample provides code for fused multi-head attention forward, context, or generation phase.
|
|
It supports HeadDims of 32, 64, and 128, and fp8, fp16, and bf16 input data types.
|
|
|
|
For forward or context usage, use an M-blocking (Seqlen-Q) of 256 and an N-blocking (Seqlen-K) of 128.
|
|
For generation usage, use an M-blocking (Num-Groups) of 128 (although the limit is currently 32 for actual Num-Groups), and a N-blocking (Seqlen-K) of 64, 128 or 256.
|
|
|
|
Context loads are done via TMA, whereas generation usage utilized `cp.async` and is thus more amenable to complex load patterns.
|
|
|
|
For variable sequence lenght, the code requires a batch of valid (but never used) padding memory ahead of the first input batch. This is achieved with least overhead by leaving one batch free and then arranging QKV consecutively.
|
|
|
|
The approach of this implementation is to reuse the selection logic of the collective gemm builder and recombine the result into an FMHA kernel.
|
|
The kernel and collective layer are then formulated to be fmha-specific.
|
|
The design assigns two tiles to each threadblock, and pingpongs between them in terms of matrix-matrix multiplication and softmax.
|
|
|
|
The example builds four binaries, showcasing the context and generation usage for fp8 and fp16.
|
|
For detailed information on how to invoke them, check out either the tests in `CMakeLists.txt` or the `--help` for them.
|
|
|
|
To modify the code for fusions, `collective/fmha_fusion.hpp` provides the easiest customization point.
|
|
The `apply_mask` function is called with the accumulator of the first GEMM and the logical positions of those elements.
|
|
It is well-suited for applying masks or activations.
|
|
More complex fusions that require memory loads would require modifying the mainloop collective to orchestrate the load via TMA.
|