perf: add mla metadata parallel path#3465
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds a new “parallel planner” GPU kernel path for MLA v1.2 metadata generation, and refactors shared device-side helpers to reuse logic between kernels.
Changes:
- Introduce
kn_get_mla_metadata_v1_2_parallelwith a 3-phase plan (scan + parallel fill + tail fill). - Extract Qo-tiling and block-sum computation into reusable device helper functions.
- Add an env-var-controlled dispatch path to select the parallel kernel when constraints are met and LDS budget allows.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if(use_parallel && (scratch_bytes + qo_bytes + kv_bytes <= lds_size)) | ||
| { | ||
| using Traits = | ||
| MlaMetadataV12Traits<kPackedQoLenPerWg, kQoSplits, kUniSeqlenQo, true, kIsSparse>; | ||
| kn_get_mla_metadata_v1_2_parallel<Traits><<<grid, fill_block, lds_size, stream>>>(params); | ||
| } |
| const char* parallel_env = std::getenv("AITER_MLA_META_USE_PARALLEL"); | ||
| const bool parallel_wanted = (parallel_env == nullptr) || (std::atoi(parallel_env) != 0); |
| const char* parallel_env = std::getenv("AITER_MLA_META_USE_PARALLEL"); | ||
| const bool parallel_wanted = (parallel_env == nullptr) || (std::atoi(parallel_env) != 0); | ||
| const bool use_parallel = parallel_wanted && (max_seqlen_qo == 1) && !kQoSplits && | ||
| !kIsSparse && (params.page_size == 1) && | ||
| (params.qk_batch_ratio == 1); |
| const int32_t scratch_bytes = | ||
| static_cast<int32_t>(sizeof(int32_t)) * (3 + 5 * params.num_batches); | ||
| const int32_t qo_bytes = is_unique ? 0 : static_cast<int32_t>(sizeof(int32_t)) * params.num_batches; | ||
| const int32_t kv_bytes = static_cast<int32_t>(sizeof(int32_t)) * params.num_batches; |
|
Hi @valarLip @minmengdie, We have added detailed microbench and Kimi K2.5 end-to-end profiling results to this PR. Could you please take another look when you get a chance? Please let me know if there are any concerns or if any additional data would be helpful from my side. Thanks! |
|
Hi all, I saw the previous CI run failed because of formatting error and have it fixed. If there are anything needed from my side, please feel free to let me know:) |
|
rerun ci |
|
Hi @valarLip , Some of the CI failed at op_tests/test_dsv4_rotate_quant.py, which seems irrelevant to this PR's change. Would you mind taking at look at the CI results? Thanks! |
Motivation
The
fast_modeMLA decode metadata planner (kn_get_mla_metadata_v1_2, reached viaget_mla_metadata_v1) currently runs its greedy work-distribution as a single-warp serial scan over CUs.For MLA decode (
num_heads_k = 1), the scan length is the full CU count, so this metadata kernel is on the decode critical path. In Kimi K2.5 serving, this kernel can take around 100us at medium-to-high concurrency.Technical Details
This PR adds
kn_get_mla_metadata_v1_2_parallel, a closed-form two-phase parallel planner for the pure-decode fast path.The new planner produces metadata that matches the serial kernel, and the change is contained in
v1_2_device.cuh.Test Plan
We evaluated this change with both:
Test environment:
Test Result
Microbenchmark
All microbench cases produce matching metadata between the serial and parallel planners.
Geomean speedup across these microbench cases is 7.07x.
Kimi K2.5 MXFP4 E2E Profile
The table below shows the metadata kernel latency extracted from Kimi K2.5 MXFP4 end-to-end traces. All values are in ns. Lower is better.
median saving / layeris computed as(serial median - parallel median) / 62.In the end-to-end Kimi K2.5 MXFP4 profile, the parallel planner reduces the metadata kernel median latency by 3.22x to 9.66x across the tested concurrencies, corresponding to about 553.5 ns to 1,487.8 ns per layer.