feat(moe): consume prepared stage1 activation scales - #4762
Conversation
ab5cabe to
d8c367b
Compare
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
There was a problem hiding this comment.
Pull request overview
This PR updates fused_moe_2stages to optionally consume a caller-prepared stage-1 activation (already quantized) and its corresponding sorted E8M0 scale layout, avoiding redundant quantization when the selected kernel requires prequant=True.
Changes:
- Add
_is_prepared_stage1_input()gate to detect when a prepared stage-1 activation/scale pair can be consumed directly. - Update
fused_moe_2stagesto bypass the A8W4 quantize+sort path when the prepared-input gate is satisfied. - Add a focused unit test for the prepared-input gating logic.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
aiter/fused_moe.py |
Adds prepared-input gate and uses it to skip redundant stage-1 quantization for prequant kernels. |
op_tests/test_fused_moe_prequant.py |
Adds unit coverage for the prepared-input gating behavior. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| assert ( | ||
| _is_prepared_stage1_input(metadata, hidden_states, q_dtype, scale) is expected | ||
| ) |
There was a problem hiding this comment.
Fixed in 37fb420. The test matrix now rejects uint8/FP32 scales and a wrong 1x32 column layout, in addition to the existing prequant/dtype/no-scale negative cases. The positive case uses the exact E8M0 dtype.
| q_dtype_a: torch.dtype | None, | ||
| a1_scale: torch.Tensor | None, | ||
| ) -> bool: | ||
| return metadata.prequant and hidden_states.dtype == q_dtype_a and a1_scale is not None |
There was a problem hiding this comment.
Fixed in 37fb420. The fast path is now limited to contiguous 2D MXFP8 activations plus same-device contiguous 2D E8M0 scales whose second dimension exactly matches the 1x32 activation layout. BF16/FP16 and unrelated scale dtypes/layouts no longer satisfy the gate.
Allow a caller to pass an already-quantized MXFP8 stage-1 activation and its sorted E8M0 scale layout to fused_moe_2stages. This avoids launching fused_dynamic_mxfp8_quant_moe_sort a second time when a producer already owns both outputs. The fast path is deliberately restricted to contiguous 2D FP8 activations and same-device contiguous E8M0 scales with the expected 1x32 column shape. BF16/FP16 inputs, unrelated scale dtypes or layouts, non-prequant kernels, and callers without scales retain the existing path. Keep the generic consumer contract separate from the Kimi-K3 exact-M3 producer so each boundary can be reviewed and validated independently. Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com>
d8c367b to
37fb420
Compare
The prepared Stage-1 fast path previously inferred a route-sorted, SBM32-preshuffled E8M0 layout from dtype, contiguity, and column count alone. An undersized or token-major scale tensor could therefore enter the consumer and produce invalid reads or incorrect output. Require the explicit preshuffled layout tag and enough scale rows for both sorted_ids and the metadata block extent. Unsupported layouts fail closed to the existing quantization path. Add negative coverage for missing and wrong tags, undersized and token-major scales, non-vector sorted IDs, and metadata-driven extents. Validation: 15 unit tests; Ruff; Black; compileall; diff-check. Signed-off-by: Yanyuan Qin <yanyuan.qin@amd.com>
Why
fused_moe_2stagesalways quantizes an A8W4 stage-1 input when the selectedkernel declares
prequant=True. That prevents a caller-owned producer frompassing an already-quantized activation and its consumer-ready, route-sorted
E8M0 scale layout: the same activation is quantized and materialized again.
The prepared path must also fail closed. Dtype, contiguity, and scale column
count alone do not prove that a scale tensor has the route-SBM32 preshuffled
layout or enough rows for the selected metadata. Accepting an undersized or
token-major tensor could otherwise cause invalid reads or incorrect output.
What changed
Consume the caller-provided activation directly only when the complete
prepared-stage1 contract agrees:
q_dtype_a;mx_e8m0_route_sbm32_preshuffled_v1layout tag;Missing/wrong layout tags, undersized or token-major scales, BF16 callers,
non-prequant kernels, and callers without scales retain the existing path.
This PR deliberately contains no model-specific selector or producer.
Tests
python3 -m py_compile \ aiter/fused_moe.py \ op_tests/test_fused_moe_prequant.py ruff check \ aiter/fused_moe.py \ op_tests/test_fused_moe_prequant.py pytest -q op_tests/test_fused_moe_prequant.py # 15 passed git diff --check origin/main...HEADThe negative cases cover missing and wrong tags, undersized and token-major
scales, non-vector sorted IDs, and metadata-driven extents. On gfx950, the
consumer was also exercised with a prepared FP8 activation and sorted E8M0
scales through the production two-stage A8W4 kernels. Output and repeated HIP
Graph replay were bit-exact against the existing path.
This PR makes no standalone performance claim; the producer that removes the
redundant work is intentionally split into a follow-up.