perf(runtime): avoid Kimi MLA projection layout copies - #950
Open
Max191 wants to merge 1 commit into
Open
Conversation
Run the MLA query and value projections directly from their supported transposed activation layouts into the caller-provided strided outputs. This uses torch.bmm's out parameter to remove the activation materialization, temporary result, and final copy while preserving the projection contract. Performance evidence On physical MI350X GPUs 4-7, five upstream/main runs at b28e18d were compared with three candidate runs. Each run served eight successful requests at concurrency eight with exactly 8,192 input and 8,192 output tokens, graph capture including the 8,192-token prefill graph, profiler disabled, greedy sampling, seed 20260724, and zero client warmups. EAGLE3 used one draft step and verify width two. The target and draft revisions were 42c6da365fdd32ae482ee7eb8bdc563da31c3404 and 27f31052f73704614af4abe6a65e766663959f29. Median TTFT moved from 4405.845 to 4302.860 ms (+2.337%), TPOT from 27.375 to 26.454 ms (+3.364%), E2EL from 228633.556 to 220785.828 ms (+3.432%), and output throughput from 280.960 to 291.261 tok/s (+3.666%). Candidate TPOT ranged from 26.420 to 26.511 ms and output throughput from 290.374 to 291.265 tok/s; the corresponding baseline ranges were 27.249 to 27.430 ms and 280.569 to 281.944 tok/s. Every run completed 8/8 requests. Generated text, per-request response-event vectors, the 33,016-event total, and the 98.473467% acceptance proxy matched the baseline exactly. Pre-commit and diff checks passed. Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
Collaborator
|
Hi @Max191 please fix the conflicts thanks |
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
2 participants
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.
Runs the MLA query and value projections directly from their supported transposed activation layouts into the caller-provided strided outputs. This removes the activation materialization and final copy. This also removes an amd-specific specialization that proves to not be necessary based on benchmark performance.
Performance evidence
On 4 MI350X GPUs (attn-tp4 moe-tp4), five baseline Kimi 2.5 runs were compared with three candidate runs. Each run served eight successful requests at concurrency eight with exactly 8,192 input and 8,192 output tokens. EAGLE3 used one draft step and verify width two.
Median TTFT moved from 4405.845 to 4302.860 ms (+2.337%), TPOT from 27.375 to 26.454 ms (+3.364%), E2EL from 228633.556 to 220785.828 ms (+3.432%), and output throughput from 280.960 to 291.261 tok/s (+3.666%). Candidate TPOT ranged from 26.420 to 26.511 ms and output throughput from 290.374 to 291.265 tok/s; the corresponding baseline ranges were 27.249 to 27.430 ms and 280.569 to 281.944 tok/s.
Every run completed 8/8 requests. Generated text, per-request response-event vectors, the 33,016-event total, and the 98.473467% acceptance proxy matched the baseline exactly.