Repository navigation
[perf](kt-kernel): AVX2 MXFP4 decodes each weight group once for the last 2-3 rows of an expert - #2253
Open
05yuki wants to merge 1 commit into
Open
[perf](kt-kernel): AVX2 MXFP4 decodes each weight group once for the last 2-3 rows of an expert#225305yuki wants to merge 1 commit into
05yuki wants to merge 1 commit into
Conversation
…last 2-3 rows of an expert The group-32 fast path runs an expert's rows in 4-row blocks and the rest one at a time, decoding every 32-value weight group again for each of those rows. With 2, 3, 6 or 7 rows on an expert (a small decode batch, a speculative verify, the few rows of a short prefill) the remaining 2-3 rows now share one decode a group. Each row keeps the single-row loop's arithmetic, so every output is bit-identical to before. KT_MXFP4_REM_ROWS=0 restores the row-at-a-time loop. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01C2UBZdnrDZSBh8yQz8rc9g
This branch has not been deployed
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
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.
The AVX2 MXFP4 group-32 fast path runs an expert's rows in 4-row blocks and the rest one at a time. Each of those remaining rows decodes every 32-value weight group again (PSHUFB, unpack), which is what bounds this kernel, so 2 rows on an expert cost 1.5x one row and 3 rows 2.1x. Small decode batches, speculative verify steps and the few rows a short prefill puts on each expert land exactly there.
operators/avx2/mxfp4-moe.hpp: when 2 or 3 rows remain after the 4-row blocks (m = 2, 3, 6, 7), each group is decoded once and applied to all of them. Each row keeps the single-row loop's own arithmetic (one product chain a group, the scale folded into two sums by group parity, groups in order), so every output is bit-identical to before.KT_MXFP4_REM_ROWS=0restores the row-at-a-time loop.Measured on one socket of a dual EPYC 7452 (Zen 2, 32 threads), DeepSeek-V4.1-Flash's TP2 rank shapes (hidden 5120, intermediate 1152, 384 experts) with synthetic weights, 4 experts a call, every row on the same 4 experts, median of 2 x 200 calls (kt-mxfp4-decode-bw.py
--same):Outputs: the hash over 64 calls is the same before and after at every row count above, and each row of the 2- and 3-row calls is byte-identical to that row computed alone.
🤖 Generated with Claude Code
https://claude.ai/code/session_01C2UBZdnrDZSBh8yQz8rc9g