[Feat] Add MLA Decode kernel - #1644
Open
wch0810 wants to merge 3 commits into
Open
Conversation
…sserts, and msprof profiling - Developer mode (pass_configs 4x True: AUTO_CV_COMBINE/SYNC/AUTO_SYNC/MEMORY_PLANNING) - Dual-loop + persistent grid + KV L1 double-buffer - Online softmax with negative-domain min-merge + tail mask - Layered test suite: L0 (3) + L1 (9) + L2 (3) + Boundary (5, incl 1 KNOWN-LIMITATION) - Mixed tolerance precision check (fp16 atol=2^-14, rtol=2^-9, max_abs=1e-1, ratio=0.99) - 15 robustness asserts (batch/col_indices/actual_seqlen_kv/dim mismatch) - with_compile_retry decorator for intermittent bisheng compiler failures - msprof op support (--level msprof, kernel-only Task Duration) - do_bench support (--level bench, 5 warmup + 5 repeat) - Precision: 20/20 PASS (19 PASS + 1 KNOWN-LIMITATION for fp16 max overflow) - Performance: 3764 us (msprof kernel-only), gap 1.25x vs 3036 us GPU target - Zero T.barrier_all / T.sync_grid / T.Scope / set_flag / wait_flag
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
…(GPU baseline) The golden config is batch=132 (matching l0_perf_target), and the GPU baseline is 3131 us (kernel.get_profiler().do_bench(warmup=500) on GPU source main_no_split). Previous values batch=128 and target=3036 were incorrect and caused all bench runs to use the wrong config. Verified with batch=132: - Precision: 20/20 PASS (L0 3/3 + L1 9/9 + L2 3/3 + Boundary 5 incl 1 KNOWN-LIMITATION) - do_bench: 4116.31 us (gap 1.31x vs 3131 us) - msprof kernel-only: 4028.94 us (block 0 Task Duration) - ruff format/check: passed
wch0810
changed the title
Contributor
Author
|
/re-test |
|
🔄 Re-running failed jobs Original workflow run: View details Only the failed jobs will be re-executed. |
…sion-only test - example_mla_decode.py: keep only kernel + 1 CI smoke case (inline golden) - test_mla_decode.py: own ref_mla_decode golden, precision tests only (L0/L1/L2/Boundary) - Remove bench/msprof from test file (perf testing handled separately) - All comments in English - Verified: smoke 'Test Passed!' + --level all 20/20 PASS (19 PASS + 1 KNOWN-LIMITATION)
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.
Description
Add MLA Decode (DeepSeek Multi-head Latent Attention Decode) for Ascend NPU, implemented in
example_mla_decode.py+test_mla_decode.pybased on TileLang Developer mode (pass_configs four-True). Dual-loop structure + persistent grid + KV L1 double-buffer + online softmax (negative-domain min-merge) + tail mask. Single fused kernel computing fused attention score (Q@KV^T + Q_pe@K_pe^T), softmax, and output projection (P@KV).Layered test suite (20 cases: L0 3 + L1 9 + L2 3 + Boundary 5 incl 1 KNOWN-LIMITATION) +
do_bench(end-to-end) +msprof op(kernel-only). 15 robustness asserts (batch/col_indices/actual_seqlen_kv/dim mismatch) +with_compile_retrydecorator for intermittent bisheng compiler failures.Golden Config (B=132, H=128, kv_heads=1, kv_ctx=8192, dim=512, pe_dim=64, fp16)
do_bench (end-to-end, includes host overhead)
kernel.get_profiler().do_bench(warmup=500)on GPU sourcemain_no_split)msprof bottleneck analysis (block 0 representative, 20 cores highly consistent)
Bottleneck root cause: sync-bound (Cube MTE1 stall 81.83% from KV 2.15x duplicate read + L2 read_hit 52.58%, KV 1GB >> L2 192MB) + memory-bound (Vector wait 96.76% waiting for Cube sync via workspace).
Precision (float16, precision-standard.md §4.1)