Repository navigation
Conversation
191c4f2 to
4d2d890
Compare
jiangkuaixue123
left a comment
There was a problem hiding this comment.
[P2] Reject unsupported MXFP8 rank counts at connector construction (afd_plugin/connectors/npu/camp2p.py:554). With quant_mode=4, a 4A2F deployment can pass the AIV and schema checks and proceed to the first A2E call. The operator supplied by #444 explicitly requires equal positive Attention and expert rank counts, at most 32 each. The asymmetric padding path at line 1405 cannot satisfy that contract. Please enforce the rank constraint in the mode-4 validator, document it beside the 4A2F example, and cover asymmetric rejection and symmetric acceptance.
The GSM8K scores are from the same 1,319 questions. The two aggregate percentages alone cannot establish accuracy neutrality. Please report the paired counts (BF16-only correct and MXFP8-only correct) and a paired interval for the difference, or describe the observed 1.14-point drop without claiming equivalence. Static review of the incremental change and #444 operator contract. I did not reproduce device measurements.
## Background The synchronous CAMP2P path normally runs attention on the Attention role and the router, shared experts, and routed experts on the FFN role. This PR makes router/shared placement an explicit deployment option. With `compute_gate_on_attention=true`, Attention computes top-k and shared experts, while FFN consumes the external routing and runs the routed experts. ## Motivation **Improve eager execution through a different A/F work split.** The measured Attention-side placement delivers **+17.6% throughput** versus the same-commit FFN-side placement in the evaluated eager workload. **Establish the routing contract for quantized activation transport.** Once FFN receives the routing alongside the activation, it can consume an activation quantized on Attention. The MXFP8 pass-through builds on this contract in #506. The option is **off by default** and belongs in the common `additional_config["afd"]` block. ## Implementation ```text Default Attention : attention FFN : router + shared experts + routed experts compute_gate_on_attention=true Attention : attention + router + shared experts FFN : transported top-k + routed experts ``` - **Reuse the A2E routing slots.** Send expert IDs as `int32` and top-k weights as `float32` through the existing ids channel. - **Owned prepare and comm subclasses.** `AFDPrepareAndFinalizeWithAllGather` and `AFDPrepareAndFinalizeWithMC2` provide `prepare_external`, which takes the routing for the local input rows and returns the complete `MoEPrepareOutput` together with the aligned routing. AFD comm subclasses override only `_get_prepare_finalize()`, so the native selector still decides which path runs and ordinary execution keeps the upstream classes. There are no global prepare patches left. - **Explicit routing contract.** The routing is an argument and a return value; nothing is stashed on the upstream class, and the module no longer reads or clears an `afd_external_routing` attribute. - **One MoE component.** `AFDExternalRoutingMoE` builds its comm candidates during model initialization, where the config context exists, and each forward only looks up the candidate matching the native selection. It runs prepare, expert execution and finalize, switches the forward-context comm for the duration of the call and restores it on every exit path, and rejects unsupported types such as fused-MC2. - **Consume the external top-k on FFN.** The component hands the aligned routing straight to expert execution. - **Place shared experts on Attention.** Run them between activation send and routed-output receive, then combine. - **Keep protocol selection explicit.** The model configuration selects token-ID or external-routing semantics; the model itself only builds the component for the synchronous FFN role and delegates. The supported opt-in configuration is synchronous CAMP2P, DeepSeek-V4, TP=1, without DBO or ubatching. Both roles must use matching routing settings. ## Performance and accuracy Measured at `828582a` on one A5 4A4F node during the same evening session, with BF16 A2E transport and automatic communication selection: | Configuration | Eager throughput (token/s) | Graph throughput (token/s) | | --- | ---: | ---: | | Attention-side gate (`828582a`) | **192,887** | **326,400** | | FFN-side gate (`828582a`) | 163,997 | 331,455.5 | | Clean baseline (`bcc0126`) | 164,543 | 333,407 | | Comparison | Eager | Graph | | --- | ---: | ---: | | Attention-side vs FFN-side placement | **+17.6%** | **-1.5%** | | FFN-side placement vs clean baseline | -0.3% | -0.6% | The result supports Attention-side placement as an eager optimization and a foundation for MXFP8 pass-through. Placement remains opt-in, allowing deployments to select the appropriate A/F split for their execution mode. Graph-mode Attention-side routing scored **90.60% on the full 1,319-question GSM8K evaluation**, with normal generated output. Measurement setup: concurrency **128**, reported input **8,275** / output **32** tokens per request, `max-num-batched-tokens=3072`, `max-num-seqs=128`, and zero failed requests. Graph uses `FULL_DECODE_ONLY` with capture sizes `[8,16,32,64,128]`; eager uses `--enforce-eager`. Each cell has two runs from the same firmware/session. These are workload-specific measurements with small run counts. ### Re-measured on the target stack at `e2c7075` The refactored implementation was re-measured on the stack this PR declares: vLLM `0.32.0` @ `5df63a7`, vLLM-Ascend `bd42eea0`, BF16 A2E transport, Attention-side gate, automatic communication selection, TP=1, eager (no CUDA/ACL graph), 4A4F on one A5 node. | Run | E2E (ms) | Prefill (token/s) | Output (token/s) | Total (token/s) | Failed | | --- | ---: | ---: | ---: | ---: | ---: | | 13:28 | 5,126.2 | 1,614.2 | 728.59 | 189,136 | 0 / 128 | | 13:34 | 5,146.5 | 1,607.9 | 726.69 | 188,643 | 0 / 128 | The two runs agree within 0.4%, so they are recorded as a repeat baseline rather than a single sample. GSM8K on the same configuration scored **89.92%** on the full 1,319-question evaluation. These numbers describe this configuration only: there is no same-stack reference implementation to compare against, so they are not evidence of a performance change from the refactor. The pre-refactor table above was measured on the previous stack and stays attributed to `828582a`. ### Placement re-measured on the target stack (`e2c7075`) Same node, same stack (vLLM `0.32.0` @ `5df63a7`, vLLM-Ascend `bd42eea0`), same workload (128 concurrency, 8,275 input / 32 output tokens per request), zero failed requests. The clean baseline is upstream main at `b093863`, which keeps the gate on the FFN side. | Mode | Clean baseline `b093863` (total token/s) | Attention-side gate `e2c7075` (total token/s) | Delta | | --- | ---: | ---: | ---: | | Eager | 165,345 | 189,136 / 188,643 | **+14.2%** (mean) | | Graph (`FULL_DECODE_ONLY`) | 364,627 / 362,563 | 354,731 / 353,972 | **-2.5%** (mean) | On this workload, Attention-side placement averages about **+14.2% higher total throughput in eager** and about **2.5% lower in graph**. That is a placement comparison across two commits on one stack; it is not a pure refactor cost, and it does not mean Attention-side placement is unsuitable for graph mode. Graph mode also raises this configuration's own throughput by **87.6%** over eager (189,136 -> 354,731 total token/s). GSM8K on the target stack scored **89.92%** in eager and **89.99%** in graph mode, both on the full 1,319-question evaluation. ## Baseline and scope Rebased onto main at `e58fdec`, which includes the merged #500 opaque routing-op fix. The retired `profile_run_eager` implementation and its wiring have been removed. #506 builds MXFP8 pass-through on this routing contract. [Incremental diff for this PR: `e58fdec...2c17cce`](Pingzii/afd-plugin@e58fdec...2c17cce) ## Verification CPU, measured at `2c17cce`: **59 passed, 2 skipped**. The focused set covers the routing-id op, the runtime installer wiring, the CAMP2P token-ID path, the routing payload transforms, and the new MoE component. The maintained CPU suite also passes. Two mutation checks were run on the refactored component before the rebase, and both produced the expected failures: - disabling the per-token state cleanup makes the MC2 finalize case fail (it observes state that should already be cleared); - disabling the comm restore makes the prepare-failure case fail (it observes the AFD comm left in place). Device evidence: the head `7430a13` starts on one A5 4A4F node with BF16 transport, `compute_gate_on_attention=true`, the CAMP2P connector, TP=1 and eager mode, and completes a generation request (`POST /v1/completions` returned 200). Accuracy and throughput were measured at `e2c7075`. The pre-refactor table above stays attributed to `828582a`. --------- Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com> Co-authored-by: 宋元盟 <songyuanmeng2@huawei.com>
4d2d890 to
86495ce
Compare
c0c42e3 to
5041bd5
Compare
|
Ready for re-review at
CPU verification on this head: full |
ad9044a to
44394d1
Compare
|
Two follow-ups are in at The handshake has a test that fails if it is deleted. Rather than adapting the The guide states the startup behavior. The two bullets that had been joined are split, the Two carry-overs from the earlier, now superseded check. The per-payload mode comparison is still in the code and will go away in the follow-up that makes the startup handshake the single protocol check. What stays regardless: construction requires an Ascend A5 device ( Testing is down to the seams that own each behaviour. Five cases are gone at The convergence follow-up will also keep the JSON round-trip test (dropping only its CPU verification at this head: connectors and model_executor 323 passed / 28 skipped; ruff check/format, typos and the changed-file mypy clean. Still unrun: the real HCCL startup handshake, mode 4 end-to-end on ALLGATHER and MC2, and the operator capability marker. |
8b87fe4 to
a33f413
Compare
2f1faa5 to
ec6409f
Compare
|
Both fixed at The bound read is now inside the guard. The fake no longer lets the preset answer for this rank. After copying the preset, it overwrites this rank's own slot with what it actually published ( Guide updated with the two missing statements: the FFN must resolve to ALLGATHER or MC2 for the largest batch Attention can send, rejected at the startup handshake on both roles before Local: connectors 45 cases in that file, connectors + model_executor 333 passed / 28 skipped; ruff check/format, typos and the changed-file mypy clean. |
7c26ad8 to
9e8b284
Compare
|
Presets reordered at Mutation-checked that the case now bites: forcing the FFN to publish a passing verdict in the ALLTOALL branch makes Local: connectors 45 cases in that file, connectors + model_executor 333 passed / 28 skipped; ruff/format and pre-commit clean. PR is still 10 commits. |
dc41c70 to
a7ab213
Compare
Three unit tests fail on Windows for reasons unrelated to what they check: - `tests/unit/compat/npu/test_ascend_config_lifecycle.py` reads a source file with `Path.read_text()` and no encoding, so the GBK locale decodes UTF-8 bytes and raises `UnicodeDecodeError`; both reads now state UTF-8. - `tests/unit/package/test_cam_async_build.py` builds its manifest file list with native separators while setuptools works on posix manifest paths, so every required entry misses; the list now uses `as_posix()`. - `tests/unit/package/test_package.py` needs installation metadata to find the `vllm.general_plugins` entry point, which a source tree does not have; it now skips with that reason, and still checks the entry point when the distribution is installed. No production code changes. With these, full `tests/unit` on a Windows CPU host with `torch` and no `vLLM` is 900 passed / 89 skipped / 0 failed (previously 8 failed). Where these come from, for review context: - the six GBK failures: `tests/unit/compat/npu/test_ascend_config_lifecycle.py` arrived in #425 (`b7561d7`) with `ast.parse(path.read_text())` and no encoding, so it is latent everywhere except a locale that cannot decode UTF-8; - the manifest failure: #539 (`f2c79d8`) replaced the posix-normalised `required` set with `set(files.allfiles)`, whose entries carry native separators, so the comparison only holds where `/` is the separator. The #361 version of the test compared normalised paths and was correct; - the entry-point failure: the test first appeared as `tests/test_package.py` in #1 (`dffdb14`, "Add Phase 1 plugin skeleton") and was later moved to `tests/unit/package/test_package.py` by `d1fd84f`; it always needed installation metadata. None of the three involves the transports in #506; the tests are older than that branch. On verification scope: this PR's checks are label, pre-commit and DCO -- none of them installs the plugin or runs pytest. The numbers above come from a local run on this branch; the entry-point test exercising an installed build is therefore *not* covered by CI, before or after this change. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>
The sync transport can only carry BF16 activations today. This adds the payload half of an opt-in MXFP8 mode: the routing carries the activation MX scale, one E8M0 exponent per 32 values and one row per token. The scale follows the same row transforms as the routing -- pad, gather and split -- because a row whose scale stays behind describes another token. Padding writes the byte 127 (2**0, the neutral exponent) through the byte view, because E8M0 stores an exponent and the collectives are not required to accept the float8 dtype. That view is what the payload fakes reject, so a transform that hands a collective the float tensor instead of its bytes fails the tests rather than passing silently. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>
An operator build that predates the MXFP8 arguments would either reject the fourteen-argument call or accept it and ignore the scale. The probe answers three ways -- the schema takes x_scale and quant_mode, it does not, or it cannot be read -- so the caller decides what an unreadable build means instead of treating it as support. The check is structural, not textual. str(torch.Type) renders Optional[Tensor] in a form that does not parse back, and a bare Tensor[] would satisfy a name test while the operator expects the optional wrapper. Inspecting the types keeps both faults visible. The probe tests build their schemas with torch._C.parse_schema, so they exercise the same rendering the production path reads. 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>
quant_mode now has two legal values, so the connector tables and the support matrix state what mode 4 carries and which deployment shapes it requires. The payload contract gains activation_scale, including the row-per-token and ceil(hidden / 32) shape the transport validates. The support matrix records that mode 4 has no device run in this repository yet, rather than implying the mode is verified because it is wired. 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>
Three cases restated behaviour another seam already pins: the two `compute_ffn_output` cases re-bound the same extracted methods that the runner case drives through the real `_ffn_forward`, and the component-level row-order case asserted the same bytes and layout as that runner case. They are gone. Three pairs with identical setup became parametrize tables, keeping every scenario: the pair contract now covers E4M3 without a scale, a scale next to BF16, E5M2, and a scale with the wrong row count in one table; MC2 dispatch covers an unset and an explicit pre-quantized comm mode; the control-payload check covers 4/0, 0/4 and 4/4. The BF16 no-scale assertion moved onto the vllm-project#503 routing case, which now states `torch.bfloat16` explicitly instead of relying on the float32 default, and the startup handshake case now also pins that `p2p` is built after the collective. Mutation re-checked after the change: deleting the runner's `activation_scale` forwarding still fails both the ALLGATHER and the MC2 implementation of `test_the_ffn_runner_carries_the_received_scale_to_the_experts`. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>
Round two of the same review: keep the seam that owns each behaviour and drop the rows that only restate it. - Schema probe: the float and int `x_scale` cases both hit the same OptionalType rejection, so they are gone. `Tensor[]` stays -- it is the one that catches an element-type-only check -- and it now parses directly, so the hand-built ListType scaffolding is gone with it. - The FFN runner case keeps MC2 only: the DP-gather case above it already pins the ALLGATHER three-dimensional layout, its row order and the padding byte. - The MC2 rejection table keeps `scale-dim`, `not-mxfp` and the `comm_quant_mode` conflict. The two dtype rows restated what the normal A2E to MoE path validates before the dispatcher is reached. - The control payload's `quant_mode` assertion folds into the existing JSON round-trip; the missing-field-reads-as-0 case stays on its own. - The startup check now runs through the real `init_afd_connector` for all four outcomes (0/0, 4/4, 4/0, 0/4): the fake collective asserts the mode that was actually published, and the refused cases assert `afd_moe` and `p2p` are never built. The direct helper call went away, so the test fails if the entry point stops calling the check. - Sending is one test now: the public `send_attn_output()` drives the registered op, and the operator boundary asserts one fourteen-argument call, the quantized activation, and the scale with its mode. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>
- The payload padding case restates `padded()`: the ALLGATHER component case walks the same call and checks the original rows and the appended 127 bytes. The gather case with distinct per-rank bytes stays, since it proves the order. - The zero-block scale case can only be reached by hand: the connector already validates the column count as `ceil(hidden / 32)` before a routing payload exists. - In the MoE pair table, `row-count` is rejected by `AFDExternalRoutingInputs` even if the component's own check is bypassed; the payload-level row test stays, and so does `no-scale`, which is what catches a runner that stops forwarding the scale. - In the MC2 dispatch table, `scale-dtype` and `activation-dtype` restate what the normal A2E to MoE path validates before a dispatcher is reached. `scale-dim`, `not-mxfp` and the `comm_quant_mode` conflict stay. No production change in this commit; the Gloo payload field is untouched. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>
The operator moves a call routing ids as one `[rows, topk]` int32 block into a single UB region, so `batch_size x topk x 4` has to fit 195,584 bytes. Nothing in the connector checked it, so `--max-num-batched-tokens 8192` with `topk=6` started cleanly and the real operator would then reject the full batch. Construction now computes the largest batch the mode can send (`8,149` tokens at `topk=6`), refuses a configuration above it, and the guide states the bound next to the other mode-4 requirements. The tests cover the boundary (8,149 accepted, 8,150 refused), a configured maximum that exceeds it, and a smaller `topk` that fits under the same batch. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>
a7ab213 to
0f059d9
Compare
A mode-4 pair can pass every construction check and still have the FFN resolve to a MoE path the transport cannot feed: the pinned selector picks ALLTOALL or FUSED_MC2 depending on batch size, EP, LoRA and MegaMoE state, and the external-routing consumer supports only ALLGATHER and MC2. The control plane is one-way, so refusing after the first control payload would leave the Attention side inside an A2E waiting for an E2A that never comes. The startup handshake grows a second step. The first collective now exchanges `[quant_mode, attention batch bound]` per rank -- only the Attention role publishes a bound, and the aggregate stays FFN first, so the FFN reads its candidate from the Attention rows alone. A mixed set fails there, both roles at once. An all-BF16 set continues without the second collective, so the ordinary deployment keeps exactly one startup collective. An all-mode-4 set asks the pinned selector, with its import, capacity lookup and call all inside the guard: a selector that raises still publishes its verdict, because a rank that exits early leaves the peer waiting. The verdict travels on the second collective and every rank acts on it -- the FFN re-raises with the original exception as its cause, the Attention side reports that the peer refused. Tests drive the real `init_afd_connector` for both roles with 1A1F and record each collective separately (`mode_gather`, `verdict_gather`), so the stop point is proved rather than inferred: mixed pairs stop at the first, an unsupported path or a raising selector stop at the second, and a mode-0 pair never reaches the second. The FFN publishes no bound, the selector is asked about Attention 1024 rather than the FFN own 512, and the failing role keeps its original cause. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>
0f059d9 to
aaa4dc1
Compare
Two ways a configured pair could still reach a working transport and then read or route wrongly: - The roles read `afd.compute_gate_on_attention` independently. With the gate on the A2E ids slot carries routing ids, with it off the same slot carries token ids, and both configurations pass construction, so a mixed pair exchanges payloads successfully: the FFN reads expert ids as token ids in one direction and token ids as routing ids (with zero weights) in the other. The startup handshake now carries the flag and fails both roles when they disagree, before any A2E call. - The selector pre-check only ran for mode 4. BF16 routing has the same ALLGATHER/MC2 contract, so a pair such as 8A8F with a bound beyond the MC2 capacity started, served small batches, and only failed at prefill when the pinned selector chose ALLTOALL. The verdict now runs whenever routing is carried -- mode 4 always, mode 0 when the gate is on -- so the pair is refused at startup. A gate-off BF16 pair still exchanges exactly one collective. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>
The regression worth pinning is not that a mismatch is detected but that neither side continues: if only the FFN rejects the gate disagreement, the Attention role walks on to the verdict collective and waits for a peer that has already left. The case now runs once per role and asserts the same stop point on both -- the group exists, the first gather happened, nothing else was created -- so a one-sided rejection fails the test instead of passing it. Signed-off-by: 宋元盟 <songyuanmeng2@huawei.com>
What changed
quant_mode=4sends a pre-quantized activation over the synchronous A2E transport. Attentionquantizes to E4M3 with
torch_npu.npu_dynamic_mx_quantinside CAMP2P's opaque send op, before thatop calls the native
a2e, and the activation plus its E8M0 MX scale travel in their own regions ofthe A2E packet; the FFN side keeps the two together through padding, gathering and sharding, then
hands both to the experts through ALLGATHER or MC2 without dequantizing.
quant_mode=0keeps the original twelve-argument call, so an operator build that predates the MXFP8arguments still runs. There is no silent fallback: a deployment that asks for mode 4 and cannot be
carried is refused at startup.
Where the change lands at this head (18 files, +2,177 / -72):
camp2p/connector.py(+473) holds the construction gate, the two startup collectives and the fourteen-argument send with the six-value receive;deepseek_v4_external_routing_moe.py(+151) holds the pair contract, the per-path scale layout and the MC2 dispatcher subclass;camp2p/external_routing.py(+74) carries the scale through pad/gather/split;compat/npu/ops.py(+67) is the operator schema probe; the remaining three files (+26) are the field and the two-hop forwardingtest_camp2p_token_ids.py(+669) covers mode 4 send/receive, construction and the startup handshake with per-collective stop points;test_deepseek_v4_external_routing_moe_cpu.py(+436) drives the real component, prepare transforms and runner forwarding; the probe, payload-transform and codec files account for the rest (+209)Deployment and dependency
The operator build must come from [WIP] feat(npu): support MXFP8 datatype for A2E operator #444 and take
x_scaleplusquant_mode. The connectorprobes the operator schema at construction, but that probe only covers the input interface -- a
build that registers the arguments and still cannot run them is a separate question, tracked by the
capability marker requested in [WIP] feat(npu): support MXFP8 datatype for A2E operator #444. Until it lands, mode 4 needs an A5 device and an operator
built with the E4M3 / E5M2 / E8M0 dtypes, as stated in the connector guide.
Upgrade the Attention and FFN roles together. The startup handshake runs for
quant_mode=0too, so it can catch a peer configured the other way; a role still on a version without the
handshake never joins it and the newer side fails at startup rather than running a mixed pair.
The handshake also exchanges
afd.compute_gate_on_attention. The A2E ids slot carries routing ids with the gate on and token ids with it off, and both configurations pass construction, so a mixed pair used to exchange payloads and then read one as the other. The disagreement fails on both roles before the first A2E call.The selector pre-check runs whenever routing is carried: always for mode 4, and for mode 0 when the gate is on. A BF16 pair such as 8A8F with a bound beyond the MC2 capacity is therefore refused at startup instead of failing at prefill when the selector chooses ALLTOALL. A gate-off BF16 pair still exchanges one collective.
The handshake exchanges every rank's mode and the batch bound the Attention role can send:
different
--max-num-batched-tokens;select_moe_comm_methodabout that same bound and refuses when itresolves to anything but ALLGATHER or MC2, which is what the external-routing consumer supports.
All of it happens before
afd_moeandp2pare created, and a selector that raises stillpublishes its verdict, so neither side waits for an answer that will not come.
On the FFN side the routing slot is required in this mode: the router and shared experts still need
the model dtype, so
recv_routingis the only supported receive.Verification
CPU only -- there is no device run behind this PR.
Local Windows run (CPU, no vLLM), full
tests/unit: 953 passed, 89 skipped, 0 failed after rebasing onto the currentmain.ruff check/format, mypy on the changed files and typos are clean.
The transport itself is pinned by: the payload byte-level tests (padding row = byte 127, per-rank
distinct bytes through the gather, scale following every row transform); the startup handshake, with
the stop point asserted per collective (mode mismatch stops after the first, an unsupported path or a
raising selector after the second, and a mode-0 pair never reaches the second); the fourteen-argument
send and the six-value receive, including the five-value build that must name the shortfall instead
of raising
IndexError; the MC2 dispatch kwargs; and the runner forwarding, mutation-checked bydeleting the forwarding line and watching both comm paths fail.
Device acceptance, once a compatible A2E build is installed:
x_scale=None, with the activation, the MX scale, the expert ids and therouting weights compared row for row;
quant_mode=4and MC2selected, verify that pre-quantized dispatch (
quant_mode=0) returns E8M0dynamicScalesOutaligned with the dispatched E4M3 activations on every valid expert row, and that GMM consumes the
aligned pair.
Performance and accuracy
Historical measurements from the development stack: single node A5 4A4F, Graph, gate on Attention,
same #444 operator build as this branch expects. They are not results of this PR's CI.
f683eb1, two runs per armAccuracy, GSM8K with 1,319 problems:
quant_mode=489.76% againstquant_mode=090.90%, adifference of -1.14 pp. It was measured on the development stack with the #444 operator build,
over the same problem set; the totals are all that was recorded.