Repository navigation
Add CuTeDSL MXFP8 K-groups scale rearrange kernel - #4604
Closed
alexsamardzic wants to merge 4 commits into
Closed
alexsamardzic wants to merge 4 commits into
alexsamardzic wants to merge 4 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4604
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (5 Unrelated Failures)As of commit 3aa99d6 with merge base 3818749 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following jobs failed but was present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This was referenced Jul 24, 2026
Collaborator
Author
|
PR replaces the existing Triton MXFP8 K-groups scale rearrange kernel with a CuTeDSL implementation. It includes pytest coverage against a plain PyTorch reference and a validate/benchmark script for comparing correctness and performance against the current Triton path. To test: pytest -q test/prototype/moe_training/test_kernels.py::test_cutedsl_mx_block_rearrange_2d_K_groupsBenchmarking results, for M,K derived from some realistic models
Benchmarking command to produce the table above{
for s in \
2048,4096 4096,4096 7168,4096 8192,4096 14336,4096 28672,4096 \
2048,8192 4096,8192 7168,8192 8192,8192 14336,8192 28672,8192; do
IFS=, read -r r c <<< "$s"
PYTHONPATH=$PWD python scripts/validate_k_groups_cutedsl.py \
--rows $r --cols $c --groups 8 --multiple-of 1 \
--cuda-graph-bench --graph-iters 1000
done
} 2>/dev/null | perl -ne '
BEGIN {
print "| MxK | Triton us | CuTeDSL us | CuTeDSL speedup |\n";
print "|---:|---:|---:|---:|\n";
}
if (/shape=\((\d+),\s*(\d+)\)/) {
$shape = "$1x$2";
$triton_ok = undef;
$cutedsl_ok = undef;
}
elsif (/triton_equal=(True|False)/) {
$triton_ok = $1;
die "triton_equal=False for $shape\n" if $triton_ok ne "True";
}
elsif (/cutedsl_equal=(True|False)/) {
$cutedsl_ok = $1;
die "cutedsl_equal=False for $shape\n" if $cutedsl_ok ne "True";
}
elsif (/triton_us=([\d.]+)/) { $triton = $1; }
elsif (/cutedsl_us=([\d.]+)/) {
$cutedsl = $1;
die "missing correctness lines for $shape\n"
unless defined($triton_ok) && defined($cutedsl_ok);
printf "| %s | %.2f | %.2f | %.2f |\n",
$shape, $triton, $cutedsl, $triton / $cutedsl;
}
'Benchmarking results, for M,K being powers of two
|
Collaborator
Author
|
Replaced by #4704. |
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.
Stack from ghstack (oldest at bottom):