[Triton] Clean up MoE elementwise kernels - #5259
Draft
vgokhale wants to merge 3 commits into
Draft
Conversation
…uant
The MoE quant path had four kernels that reimplemented quant kernels living
in aiter/ops/triton/quant. Fold them together.
_mxfp4_quant_kernel: deleted. It already shared _mxfp4_quant_op with
_dynamic_mxfp4_quant_kernel and was otherwise a strict subset of it (no int64
stride guard, no persistent loop, no cache modifier), so mxfp4_quant() is now a
wrapper over dynamic_mxfp4_quant. dynamic_mxfp4_quant gains optional out
params, mirroring dynamic_mxfp8_quant's `scale`, because the a4w4 GEMM needs
row-major scales while dynamic_mxfp4_quant allocates a transposed view -- and
its stride(0)==1 would flip the X_SCALES_TDM heuristic in moe_gemm_a4w4. Also
drops an eager x.to(float32) that materialised the whole activation in fp32
before the launch; the kernel casts on load.
_downcast_to_mxfp: keeps both scale schemes behind a new POW2_SCALE constexpr.
False (default, unchanged) is amax/dtype_max with the exponent rounded per
DEQUANT_SCALE_ROUNDING_MODE; True is the even_round scheme shared with
_mxfp4_quant_op/_mxfp8_quant_op. They are not interchangeable -- they disagree
on ~12% of fp4 block scales and ~0.3% of fp8 ones -- so the flag stays. With
POW2_SCALE=True the output is now bit-identical to dynamic_mxfp4_quant (fp4)
and dynamic_mxfp8_quant (e4m3). No perf difference between the two (0.98-1.09x
across four shapes): the kernel is memory bound, so the log2/exp2 pair costs
nothing and the name describes the policy, not a speed tier. _mxfp8_quant_op
gains a defaulted LOG2_DTYPE_MAX so the fast path also covers e5m2, whose 15
its hardcoded 8 got wrong.
_downcast_to_static_fp8: deleted, merged into
_static_per_tensor_quant_fp8_i8_kernel behind FAST_CONVERT (reciprocal-multiply
vs exact division). The two forms are not equivalent: at scale=448.0 a dense
bf16 sweep puts them 1 fp8 ulp apart on 64 elements, so both survive. Here the
name is literal -- the reciprocal form is 15-20% faster. Two side effects:
- The merged kernel takes its output dtype from the pointer rather than
hardcoding tl.float8e4nv, so gfx942 no longer converts fn->fnuz on store
(the buffer is allocated fnuz there). This changes gfx942 output; it was
double-rounding through the wrong exponent bias before.
- The old MoE tile indexed with offs_m[None, :] / offs_n[:, None], putting
the stride-1 axis first where Triton cannot prove stride_x_n == 1.
Correcting that makes the MoE path 1.07-1.62x faster. Against the old
row-per-program kernel: 0.94x worst case, 1.05-2.38x elsewhere, and it can
now serve wide rows that did not fit in registers at all.
The shared launcher views both tensors as 2D instead of trusting qx.stride(0),
because per_tensor_quant_triton in aiter/ops/quant.py passes an N-D output
alongside a 2-D input -- the old kernel got away with using x's row stride for
the output. No test covered that.
upcast_from_mxfp: moved to op_tests/triton_tests/utils/mxfp_ref.py. It is test
scaffolding -- it builds bf16 references so the torch comparison sees the same
values the kernel reads -- and had no callers under aiter/.
Verified on gfx950: the exact-division and reciprocal paths are bit-identical
to the kernels they replace across 9,431,040 elements x 307 scales, int8
included; mxfp4_quant is bit-identical across 6 shapes x 3 dtypes; the N-D
output path matches for 3-D and 4-D, fp8 and int8, static and dynamic.
op_tests/triton_tests/{moe,quant} plus test_activation.py: 6897 passed,
6396 skipped.
Comment-only follow-up to a8968e2. The two constexpr flags each carried a long docstring explaining both of their branches; move each down to a short comment at the branch it actually governs (POW2_SCALE in _compute_mx_quant_and_scale, FAST_CONVERT in _static_per_tensor_quant_fp8_i8_kernel) and drop the rest. Same for the wrapper-level prose on mxfp4_quant, downcast_to_static_fp8, downcast_to_mxfp, _static_per_tensor_quant_launch and _mxfp8_quant_op's LOG2_DTYPE_MAX note. No functional change. Reverified on gfx950: both static-fp8 paths still bit-identical to the kernels they replaced over 9,431,040 elements x 307 scales, mxfp4_quant bit-identical, the N-D output path still correct for 3-D and 4-D across fp8/int8 and static/dynamic, and op_tests/triton_tests/{moe,quant} plus test_activation.py unchanged at 6897 passed / 6396 skipped.
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
The dedup left three names one underscore apart: the kernel
(_static_per_tensor_quant_fp8_i8_kernel), the public wrapper
(static_per_tensor_quant_fp8_i8) and a private launcher between them holding
the tiling heuristic and the FAST_CONVERT choice. That middle tier is not a
layer this tree has anywhere else -- kernels live in _triton_kernels/, host
wrappers in the op module, and nothing sits in between -- so remove it.
fast_convert becomes a keyword argument on static_per_tensor_quant_fp8_i8,
defaulting to True so the existing positional callers (aiter/ops/quant.py) are
unaffected, and the tiling and launch move into the wrapper body alongside
every other op in that file. dynamic_per_tensor_quant_fp8_i8 and the MoE
downcast_to_static_fp8 now both go through the public wrapper. It reads better
as a public knob anyway: reciprocal-multiply vs exact division is a real
accuracy/speed choice for a caller to make.
Side effect: the dynamic path emits a second INFO log line, since it now calls
the public wrapper rather than the kernel. That is gated behind
AITER_TRITON_LOG_LEVEL (default WARNING) and is arguably accurate -- the
dynamic op does perform a static quant internally.
No behavior change otherwise. Reverified on gfx950: both scale-application
paths still bit-identical to the kernels they replaced over 9,431,040 elements
x 307 scales, int8 included; the N-D output path still correct for 3-D and 4-D
across fp8/int8 and static/dynamic; op_tests/triton_tests/{moe,quant} plus
test_activation.py at 6897 passed / 6396 skipped.
vgokhale
changed the title
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.
No description provided.