Update MLA decode example - #3357
Open
Observer007 wants to merge 2 commits into
Open
Conversation
Port the 4.6 MLA FP8 decode schedule to the public CuTeDSL example. The kernel now uses two softmax warp groups, splits QK and PV MMA issue across separate warps, and keeps the balanced 152/152/56 register split that avoids local spills in the cu12/cu13 generated code while preserving the cu13 performance path. Update the MLA static scheduler helper for split_pv plumbing and the corrected non-persistent coordinate decode required by the folded/fanned-out schedule. Example validated on B200 with nvidia-cutlass-dsl 4.6.0a0, CUDA_VERSION 13.3: py_compile for mla_decode_fp8.py and mla_helpers.py; b=4 s_k=1024 persistent reference check passed; b=296 s_q=1 s_k=8192 h=128 l=512 r=64 persistent skip-ref perf averaged 424.65 us over 20 iterations at 802 MHz core clock.
Port the 4.6 Blackwell MLA FP16 decode updates onto the public CuTeDSL example. This brings over the fold_sq/spec-decoding path, the seq_len_q-aware QK tiler K selection, reduced load_kv staging for multi-query decode, and the corrected causal-mask tile count. Public-tree cleanup keeps the example import path under cute.blackwell and removes internal release-marker comments plus the sm107-only branch from the Blackwell example. Validated with nvidia-cutlass-dsl 4.6.0a0 public wheels on B200 locked at 802 MHz: py_compile; FP16 b=4 s_k=1024 persistent reference check on CUDA 12.9 and 13.3; FP16 b=296 s_q=1 s_k=8192 persistent skip-ref perf was 751.56 us on CUDA 12.9 and 751.36 us on CUDA 13.3. Rechecked FP8 b=296 persistent on the same branch: 438.68 us on CUDA 12.9 and 424.80 us on CUDA 13.3.
|
This PR has been labeled |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Port the 4.6 Blackwell MLA decode updates to the public CuTeDSL examples.
This includes:
152/152/56register split.fold_sq,seq_len_q-aware QK tiler K selection, reducedload_kv_stagefor multi-querydecode, and corrected causal-mask tile count.
split_pvplumbing and fix non-persistent coordinate decoding.cute.blackwell...imports and remove internal release-marker comments / SM107-only branches from Blackwell examples.Validation
Validated with public
nvidia-cutlass-dsl==4.6.0a0wheels on B200, core clock locked at 802 MHz.CUDA 12.9
_cutlass_ir.cu12loaded:CUDA_VERSION DSLCudaVersion(major=12, minor=9)py_compilepassed for MLA FP8/FP16 files.b=4, s_q=1, s_k=1024, persistent, reference check passed.b=296, s_q=1, s_k=8192, h=128, l=512, r=64, persistent751.56 usb=296configuration, persistent438.68 usCUDA 13.3
_cutlass_ir.cu13loaded:CUDA_VERSION DSLCudaVersion(major=13, minor=3)b=4, s_q=1, s_k=1024, persistent, reference check passed.b=296, s_q=1, s_k=8192, h=128, l=512, r=64, persistent751.36 usb=296configuration, persistent424.80 us