[WIP] feat(npu): support MXFP8 datatype for A2E operator by wangyibo1005 · Pull Request #444 · vllm-project/afd-plugin

wangyibo1005 · 2026-10-12T03:53:40Z

PLEASE FILL IN THE PR DESCRIPTION AND MAKE SURE THE CHECKLIST ITEMS HAVE BEEN CONSIDERED.

Purpose

Issue

Related issue(s): #
Closing keyword, only if fully resolved: Closes #

Scope

In scope:
Out of scope:

Implementation Notes

Test Plan

Test Result

Docs Impact

Files updated:
If none, reason:

Essential PR Checklist

Purpose is clear and linked to public context when possible.
Scope is bounded.
Compatibility with vLLM v0.30.0 is considered.
No changes are made to the vLLM source checkout.
Plugin-owned classes or explicit dotted class paths are preferred over monkey patches.
Any compat shim or monkey patch is isolated, idempotent, version-guarded, documented, and tested.
Imports remain CPU-safe; CUDA-heavy work is delayed or GPU-gated.
Validation evidence is included, including skipped GPU tests when applicable.
Documentation impact is stated.

Add support for MXFP8 input to the A2E operator on A5. Signed-off-by: wangyibo1005 <2633333316@qq.com> Co-Authored-By: Claude Code <noreply@anthropic.com>

The attention side quantizes the FFN activation before the transport, so the payload halves. The FFN side receives FP8 plus an E8M0 block scale and restores the model dtype, because the router and the shared experts can only consume a model-dtype activation. The connector rejects mode 4 when the operator or the patched prepare stage is unavailable, rather than failing later inside the kernel. Depends on the A2E MXFP8 operator support in vllm-project#444.

Record the quant mode, the measured gain, the operator requirement, and the concurrency precondition that decides whether the gain shows up at all. The transport section builds on the text in vllm-project#444; rebase it after that lands.

The attention side quantizes the FFN activation before the transport, so the payload halves. The FFN side receives FP8 plus an E8M0 block scale and restores the model dtype, because the router and the shared experts can only consume a model-dtype activation. The connector rejects mode 4 when the operator or the patched prepare stage is unavailable, rather than failing later inside the kernel. Depends on the A2E MXFP8 operator support in vllm-project#444. Signed-off-by: 宋元盟 <songyuanmeng1@huawei.com>

Record the quant mode, the measured gain, the operator requirement, and the concurrency precondition that decides whether the gain shows up at all. The transport section builds on the text in vllm-project#444; rebase it after that lands. Signed-off-by: 宋元盟 <songyuanmeng1@huawei.com>

The attention side quantizes the FFN activation before the transport, so the payload halves. The FFN side receives FP8 plus an E8M0 block scale and restores the model dtype, because the router and the shared experts can only consume a model-dtype activation. The connector rejects mode 4 when the operator or the patched prepare stage is unavailable, rather than failing later inside the kernel. Depends on the A2E MXFP8 operator support in vllm-project#444. Signed-off-by: 宋元盟 <songyuanmeng1@huawei.com>

Record the quant mode, the measured gain, the operator requirement, and the concurrency precondition that decides whether the gain shows up at all. The transport section builds on the text in vllm-project#444; rebase it after that lands. Signed-off-by: 宋元盟 <songyuanmeng1@huawei.com>

connector_extra_config.quant_mode=4 quantizes the Attention activation to E4M3 inside the opaque A2E call, ships it with its E8M0 MX scale, and hands the pair to the experts on the FFN side without dequantizing. Mode 0 keeps the twelve-argument BF16 call, so an operator build that predates the MXFP8 arguments still runs. Construction refuses mode 4 unless the deployment can carry it: the common gate must be on Attention, both sides must have the same rank count between 1 and 32, three AIV cores are needed, and the operator probe must have said yes. The FFN receive refuses the token-ids slot in this mode, because the router and the shared experts need the model dtype. Quantizing the final wire tensor keeps the scale covering exactly the rows the operator sends; mode 4 cannot reach the padding branch because it requires equal rank counts. The scale travels with the routing, so prepare pads, gathers and shards it by the same rule as the activation rows. What consumes it fixes the layout: ALLGATHER re-lays it as (rows, cols // 2, 2) for the fused GMM, while MC2 hands the two-dimensional scale to its dispatch operator. MC2 gets a dispatcher subclass that forwards the scale and selects quant_mode=0 instead of letting upstream quantize an activation that already is; an explicit comm_quant_mode that disagrees is rejected rather than overruled. The tests drive the real prepare transforms and the real FFN entry point, with only the connector payload supplied, and pin the bytes the experts read: every scale row carries different values, so a transform that reorders, pads or re-lays it cannot keep them in place. This branch has no NPU run. The operator that accepts x_scale and quant_mode lives in vllm-project#444, so the mode is wired and reviewed but unverified on hardware. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>

Pingzii · 2026-11-11T02:58:27Z

Long-term this operator needs three fixes. The first two are confirmable from the source in this
branch; the third is an ABI concern on the integration side.

1. Make an A/F protocol mismatch fail loudly instead of reading another rank's bytes or hanging.

The F side derives the remote packet offset from its own batch_size, hidden_size and topk
(a2e/op_kernel/a2e.h, the layout is recomputed on both ends), and it decides the FP8 output dtype
from its local placeholder tensor (torch_binding.cpp, expand_x_out takes x.scalar_type()).
Each side's parameters are individually legal, so a disagreement silently reads the wrong remote
bytes. There is also a case where the FFN rank waits for flags the Attention rank never publishes:
an A rank with compute_gate=0 only publishes the activation-ready flag (WAIT_FLAG_OFFSET_3),
while an F rank with compute_gate=1 first waits for the ids and weights flags
(1 * attnToMoeRatio + index and WAIT_FLAG_OFFSET_2).

Suggested fix: validate the fixed part of the protocol at the control plane -- protocol version,
activation dtype, compute_gate, quant_mode, and the rank pairing -- and confirm on every transfer
that the rows actually received match what the receiver expects. If the operator is also meant to be
used standalone, those checks need an explicit entry point rather than living only in AFD's
connector.

2. Apply the routing-id capacity and core-count constraints to both modes, not only to MXFP8.

Both modes pull the whole id block into one UB region on the receive side
(recvExpertIdsWithMte: camCpGM2UB(..., recvBatchSize * topk * sizeof(int32_t))), but only the
quant_mode == 4 branch validates the budget. With topk=6, 8192 rows is 196,608 bytes against a
195,584-byte budget, so BF16 exceeds the region today and nothing rejects it; the first row count
that does not fit is 8150.

The core split has a matching hole. With compute_gate=1 and aiv_num=2, sendWithMte assigns one
core to the ids and one to the routing weights (blockNum - 1 and blockNum - BLOCK_IDX_USED_2), so
no core runs sendX and the activation-ready flag is never published. The cores < 3 check exists
only in the MXFP8 host branch.

Please make both modes reject these two inputs first; large-batch support can follow by giving the id
receive dedicated cores and moving to tiled copies, which is what would lift the whole-block UB limit.

3. Version the new native interface.

aclnnA2eGetWorkspaceSize extended its C parameter list while keeping the same symbol name. The
loader resolves it with GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize") -- a plain name lookup --
and then casts it to the prototype implied by the new call site (op_api_common.h,
ConvertToOpApiFunc). A mixed install (new Torch extension, old operator library) can therefore find
the old symbol, call it with the new signature, and corrupt the call. BF16 is affected by this too,
not just MXFP8. The safest long-term shape is a new symbol for the extended entry point, or an ABI
version checked at load time so a mixed install fails immediately. Python-level schema probing
(Tensor? x_scale and int quant_mode present) cannot detect this layer.

Secondary, not a blocker: A2E_GATE_UB_BYTES is defined separately in the binding, the meta
function and the host tiling; keeping one definition would stop the three from drifting. The MX scale
copy is the first one-byte DMA here, and it is safe because CpUB2GM writes through DataCopyPad per
byte, so UB alignment padding is dropped rather than written to GM; the existing hidden=65, batch=4 case already exercises a 12-byte scale, and keeping it guards that property. A==F is fine as
the first-version boundary.

On the device tests: the existing A5 case passes an empty E8M0 tensor where #506's FFN call passes
x_scale=None, and it compares only the activation and the MX scale. Extending it to the real
receiver shape and, under compute_gate=1, checking the activation, the MX scale, the expert ids and
the routing weights separately would cover what the transport actually depends on.

Both roles read their own `connector_extra_config.quant_mode`, and each side's value is individually legal, so an Attention role on 4 with an FFN role on 0 passed both local checks while the two A2E packet layouts differ in the activation offset and the element width. The FFN rank would read the peer's bytes as its own mode and hand the experts a silently wrong activation. The control payload now carries the sender's mode, and the FFN role compares it with its own in `recv_dp_metadata_list` -- the step the FFN worker runs before its first forward, hence before the first A2E call. A peer that declares no mode reads as 0, so an older Attention process is a mismatch rather than a silent downgrade. `quant_mode=4` also required a way to see a device whose operator build lacks the path: the operator registers the same fourteen-argument schema on every build and rejects mode 4 only at runtime, so the construction gate now checks for an Ascend A5 device as well. The dtype half of that build check needs a capability marker from the operator side (vllm-project#444); until it exists this remains a documented requirement rather than a check. Tests: the mixed 4/0 pair fails at the handshake, a matching pair passes, the send side declares its own mode, the codec round-trips the field and defaults a missing one to 0, and the A5 check rejects another device family. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>

Pingzii · 2026-11-11T05:14:40Z

One more ask from the #506 side, related to point 1's protocol handshake: please expose the operator build's MXFP8 capability.

The fourteen-argument schema is registered unconditionally, so a Python caller cannot tell an A5 build with E4M3/E5M2/E8M0 support from one that rejects quant_mode=4 at runtime (AFD_TARGET_A5, and the dtype guard inside the MXFP8 branch). Schema probing therefore accepts a build that cannot run the mode, and the failure only appears at the first transfer; the connector now checks the device family itself, but the dtype half has no local oracle.

A compile-time constant or a tiny exported predicate (for example afd_ascend.a2e_mxfp8_supported() -> bool) would let a caller refuse the mode at construction and would also let a mixed install report which build it is talking to.

Both roles read their own `connector_extra_config.quant_mode`, and each value is individually legal, so an Attention role on 4 with an FFN role on 0 passed both local checks while the two A2E packet layouts differ in the activation offset and the element width. The FFN rank would read the peer's bytes as its own mode and hand the experts a silently wrong activation. Two checks, in the order they bite: - `init_afd_connector` gathers every rank's mode over the AFD group and fails them all when the set is not a singleton. Every rank joins that group, so the 4/0 and 0/4 pairs both fail here -- before profile, capture and the first A2E call -- rather than leaving the Attention side waiting for an E2A the FFN side will never send. - The control payload also carries the sender's mode, and the FFN role compares it in `recv_dp_metadata_list`. That path covers a peer that started before the group handshake existed; a payload without the field reads as 0, which is the right answer only for a peer that never had mode 4. A `5041bd5` peer does have it, so the handshake above is the one that keeps that combination from being silently accepted. `quant_mode=4` also needed a way to see a device whose operator build lacks the path: the operator registers the same fourteen-argument schema on every build and rejects mode 4 only at runtime, so the construction gate now checks for an Ascend A5 device. The dtype half of that build check needs a capability marker from the operator side (vllm-project#444); until it exists this stays a documented requirement rather than a check. Tests: ranks must agree on the mode (4/0 and 0/4 both raise), the mixed pair fails at the handshake, a matching pair passes, the send side declares its own mode, the codec round-trips the field and defaults a missing one to 0, and the A5 check rejects another device family. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>

connector_extra_config.quant_mode=4 quantizes the Attention activation to E4M3 inside the opaque A2E call, ships it with its E8M0 MX scale, and hands the pair to the experts on the FFN side without dequantizing. Mode 0 keeps the twelve-argument BF16 call, so an operator build that predates the MXFP8 arguments still runs. Construction refuses mode 4 unless the deployment can carry it: the common gate must be on Attention, both sides must have the same rank count between 1 and 32, three AIV cores are needed, and the operator probe must have said yes. The FFN receive refuses the token-ids slot in this mode, because the router and the shared experts need the model dtype. Quantizing the final wire tensor keeps the scale covering exactly the rows the operator sends; mode 4 cannot reach the padding branch because it requires equal rank counts. The scale travels with the routing, so prepare pads, gathers and shards it by the same rule as the activation rows. What consumes it fixes the layout: ALLGATHER re-lays it as (rows, cols // 2, 2) for the fused GMM, while MC2 hands the two-dimensional scale to its dispatch operator. MC2 gets a dispatcher subclass that forwards the scale and selects quant_mode=0 instead of letting upstream quantize an activation that already is; an explicit comm_quant_mode that disagrees is rejected rather than overruled. The tests drive the real prepare transforms and the real FFN entry point, with only the connector payload supplied, and pin the bytes the experts read: every scale row carries different values, so a transform that reorders, pads or re-lays it cannot keep them in place. This branch has no NPU run. The operator that accepts x_scale and quant_mode lives in vllm-project#444, so the mode is wired and reviewed but unverified on hardware. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>

Both roles read their own `connector_extra_config.quant_mode`, and each value is individually legal, so an Attention role on 4 with an FFN role on 0 passed both local checks while the two A2E packet layouts differ in the activation offset and the element width. The FFN rank would read the peer's bytes as its own mode and hand the experts a silently wrong activation. Two checks, in the order they bite: - `init_afd_connector` gathers every rank's mode over the AFD group and fails them all when the set is not a singleton. Every rank joins that group, so the 4/0 and 0/4 pairs both fail here -- before profile, capture and the first A2E call -- rather than leaving the Attention side waiting for an E2A the FFN side will never send. - The control payload also carries the sender's mode, and the FFN role compares it in `recv_dp_metadata_list`. That path covers a peer that started before the group handshake existed; a payload without the field reads as 0, which is the right answer only for a peer that never had mode 4. A `5041bd5` peer does have it, so the handshake above is the one that keeps that combination from being silently accepted. `quant_mode=4` also needed a way to see a device whose operator build lacks the path: the operator registers the same fourteen-argument schema on every build and rejects mode 4 only at runtime, so the construction gate now checks for an Ascend A5 device. The dtype half of that build check needs a capability marker from the operator side (vllm-project#444); until it exists this stays a documented requirement rather than a check. Tests: ranks must agree on the mode (4/0 and 0/4 both raise), the mixed pair fails at the handshake, a matching pair passes, the send side declares its own mode, the codec round-trips the field and defaults a missing one to 0, and the A5 check rejects another device family. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>

connector_extra_config.quant_mode=4 quantizes the Attention activation to E4M3 inside the opaque A2E call, ships it with its E8M0 MX scale, and hands the pair to the experts on the FFN side without dequantizing. Mode 0 keeps the twelve-argument BF16 call, so an operator build that predates the MXFP8 arguments still runs. Construction refuses mode 4 unless the deployment can carry it: the common gate must be on Attention, both sides must have the same rank count between 1 and 32, three AIV cores are needed, and the operator probe must have said yes. The FFN receive refuses the token-ids slot in this mode, because the router and the shared experts need the model dtype. Quantizing the final wire tensor keeps the scale covering exactly the rows the operator sends; mode 4 cannot reach the padding branch because it requires equal rank counts. The scale travels with the routing, so prepare pads, gathers and shards it by the same rule as the activation rows. What consumes it fixes the layout: ALLGATHER re-lays it as (rows, cols // 2, 2) for the fused GMM, while MC2 hands the two-dimensional scale to its dispatch operator. MC2 gets a dispatcher subclass that forwards the scale and selects quant_mode=0 instead of letting upstream quantize an activation that already is; an explicit comm_quant_mode that disagrees is rejected rather than overruled. The tests drive the real prepare transforms and the real FFN entry point, with only the connector payload supplied, and pin the bytes the experts read: every scale row carries different values, so a transform that reorders, pads or re-lays it cannot keep them in place. This branch has no NPU run. The operator that accepts x_scale and quant_mode lives in vllm-project#444, so the mode is wired and reviewed but unverified on hardware. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>

Both roles read their own `connector_extra_config.quant_mode`, and each value is individually legal, so an Attention role on 4 with an FFN role on 0 passed both local checks while the two A2E packet layouts differ in the activation offset and the element width. The FFN rank would read the peer's bytes as its own mode and hand the experts a silently wrong activation. Two checks, in the order they bite: - `init_afd_connector` gathers every rank's mode over the AFD group and fails them all when the set is not a singleton. Every rank joins that group, so the 4/0 and 0/4 pairs both fail here -- before profile, capture and the first A2E call -- rather than leaving the Attention side waiting for an E2A the FFN side will never send. - The control payload also carries the sender's mode, and the FFN role compares it in `recv_dp_metadata_list`. That path covers a peer that started before the group handshake existed; a payload without the field reads as 0, which is the right answer only for a peer that never had mode 4. A `5041bd5` peer does have it, so the handshake above is the one that keeps that combination from being silently accepted. `quant_mode=4` also needed a way to see a device whose operator build lacks the path: the operator registers the same fourteen-argument schema on every build and rejects mode 4 only at runtime, so the construction gate now checks for an Ascend A5 device. The dtype half of that build check needs a capability marker from the operator side (vllm-project#444); until it exists this stays a documented requirement rather than a check. Tests: ranks must agree on the mode (4/0 and 0/4 both raise), the mixed pair fails at the handshake, a matching pair passes, the send side declares its own mode, the codec round-trips the field and defaults a missing one to 0, and the A5 check rejects another device family. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>

wangyibo1005 requested review from hsliuustc0106 and jiangkuaixue123 as code owners October 12, 2026 03:53

wangyibo1005 force-pushed the A2E-Quant branch from 5931383 to 03d93e2 Compare October 12, 2026 03:58

hsliuustc0106 added enhancement New feature or request Ascend Ascend NPU platform and related changes labels Oct 12, 2026

feat(npu): support MXFP8 datatype for A2E operator

b5c6410

Add support for MXFP8 input to the A2E operator on A5. Signed-off-by: wangyibo1005 <2633333316@qq.com> Co-Authored-By: Claude Code <noreply@anthropic.com>

wangyibo1005 force-pushed the A2E-Quant branch from 03d93e2 to b5c6410 Compare October 12, 2026 07:02

Pingzii mentioned this pull request Oct 20, 2026

[WIP] feat(npu): quantize the A2E FFN transport to MXFP8 #475

Closed

This was referenced Oct 29, 2026

feat(npu): carry a pre-quantized activation over the sync A2E transport #506

Open

[RFC]: External routing and MXFP8 pass-through for synchronous AFD on Ascend 950 (A5) #522

Open

Navigation Menu

Repository navigation

[WIP] feat(npu): support MXFP8 datatype for A2E operator - #444

[WIP] feat(npu): support MXFP8 datatype for A2E operator#444
wangyibo1005 wants to merge 1 commit into
vllm-project:mainfrom
wangyibo1005:A2E-Quant

wangyibo1005 commented Oct 12, 2026

Uh oh!

Pingzii commented Nov 11, 2026 •

edited

Loading

Uh oh!

Pingzii commented Nov 11, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

wangyibo1005 commented Oct 12, 2026

Purpose

Issue

Scope

Implementation Notes

Test Plan

Test Result

Docs Impact

Uh oh!

Pingzii commented Nov 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Pingzii commented Nov 11, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Pingzii commented Nov 11, 2026 •

edited

Loading